import streamlit as st
import yfinance as yf
import pandas as pd
import matplotlib.pyplot as plt
from statsmodels.tsa.arima.model import ARIMA
from sklearn.preprocessing import MinMaxScaler
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, LSTM, GRU
import numpy as np

# Function to load stock data using yfinance
def get_stock_data(symbol, start_date, end_date):
    stock_data = yf.download(symbol, start=start_date, end=end_date)
    return stock_data['Close']

# Function to normalize data and prepare it for LSTM/GRU
def prepare_data(data):
    scaler = MinMaxScaler(feature_range=(0, 1))
    scaled_data = scaler.fit_transform(data.values.reshape(-1, 1))
    return scaled_data, scaler

# Function to create LSTM model
def create_lstm_model(input_shape):
    model = Sequential()
    model.add(LSTM(units=50, return_sequences=True, input_shape=input_shape))
    model.add(LSTM(units=50, return_sequences=True))
    model.add(LSTM(units=50))
    model.add(Dense(units=1))
    model.compile(optimizer='adam', loss='mean_squared_error')
    return model

# Function to create GRU model
def create_gru_model(input_shape):
    model = Sequential()
    model.add(GRU(units=50, return_sequences=True, input_shape=input_shape))
    model.add(GRU(units=50, return_sequences=True))
    model.add(GRU(units=50))
    model.add(Dense(units=1))
    model.compile(optimizer='adam', loss='mean_squared_error')
    return model

# Function to fit LSTM/GRU model and make predictions
def lstm_gru_forecast(data, model_type, steps):
    scaled_data, scaler = prepare_data(data)
    input_data = scaled_data.reshape(-1, 1)

    # Split data into training and testing sets
    train_size = int(len(input_data) * 0.80)
    train_data, test_data = input_data[0:train_size, :], input_data[train_size:len(input_data), :]

    x_train, y_train = [], []
    for i in range(60, len(train_data)):
        x_train.append(train_data[i - 60:i, 0])
        y_train.append(train_data[i, 0])

    x_train, y_train = np.array(x_train), np.array(y_train)
    x_train = np.reshape(x_train, (x_train.shape[0], x_train.shape[1], 1))

    # Create and fit the model
    input_shape = (x_train.shape[1], 1)
    if model_type == 'lstm':
        model = create_lstm_model(input_shape)
    elif model_type == 'gru':
        model = create_gru_model(input_shape)

    model.fit(x_train, y_train, epochs=25, batch_size=32)

    # Make predictions
    inputs = input_data[len(input_data) - len(test_data) - 60:]
    inputs = inputs.reshape(-1, 1)
    x_test = []
    for i in range(60, len(inputs)):
        x_test.append(inputs[i - 60:i, 0])
    x_test = np.array(x_test)

    x_test = np.reshape(x_test, (x_test.shape[0], x_test.shape[1], 1))
    predicted_prices = model.predict(x_test)
    predicted_prices = scaler.inverse_transform(predicted_prices)

    # Create an index for the forecasted values
    forecast_index = pd.date_range(start=data.index[-1], periods=steps + 1, freq=data.index.freq)
    
    return pd.Series(predicted_prices.flatten(), index=forecast_index[1:])

# Function to create an ensemble forecast by averaging predictions
def ensemble_forecast(predictions_list):
    return pd.DataFrame(predictions_list).mean(axis=0)

# Function to fit ARIMA model and make predictions
def arima_forecast(data, p, d, q, steps):
    # Differencing
    for i in range(d):
        data_diff = np.diff(data)
        data = data_diff

    # Autoregressive (AR) and Moving Average (MA) components
    ar_coef = np.zeros(p) if p > 0 else []
    ma_coef = np.zeros(q) if q > 0 else []

    # Initial prediction
    predictions = list(data[:p])

    # ARIMA forecasting
    for i in range(len(data) - p):
        ar_term = sum(ar_coef[j] * data[i + p - j - 1] for j in range(p))
        ma_term = sum(ma_coef[j] * (data[i + p - j - 1] - predictions[-1]) for j in range(q))
        next_prediction = data[i + p] + ar_term + ma_term
        predictions.append(next_prediction)

        # Update coefficients using online learning (optional)
        if i + p + 1 < len(data):
            ar_coef = ar_coef + (2.0 / (i + p + 2)) * (data[i + p + 1] - next_prediction) * np.flip(data[i:i + p])
            ma_coef = ma_coef + (2.0 / (i + p + 2)) * (data[i + p + 1] - next_prediction) * np.flip(predictions[i - q + 1:i + 1])

    # Inverse differencing
    for i in range(d):
        predictions = np.cumsum([data[p - 1]] + predictions)

    return predictions[-steps:]

# Streamlit App
# Streamlit App
def main():
    st.title("Stock Price Forecasting App")

    # Load stock data using Streamlit sidebar
    symbol = st.sidebar.text_input("Enter Stock Symbol", value='AAPL')
    start_date = st.sidebar.date_input("Select Start Date", pd.to_datetime('2021-01-01'))
    end_date = st.sidebar.date_input("Select End Date", pd.to_datetime('2022-01-01'))
    stock_prices = get_stock_data(symbol, start_date, end_date)

    # ARIMA parameters using Streamlit sliders
    p = st.sidebar.slider("AR Component (p)", min_value=0, max_value=10, value=3)
    d = st.sidebar.slider("Differencing (d)", min_value=0, max_value=5, value=0)
    q = st.sidebar.slider("MA Component (q)", min_value=0, max_value=10, value=0)
    arima_forecast_steps = st.sidebar.slider("ARIMA Forecast Steps", min_value=1, max_value=100, value=30)

    # LSTM and GRU parameters using Streamlit sliders
    lstm_gru_forecast_steps = st.sidebar.slider("LSTM/GRU Forecast Steps", min_value=1, max_value=100, value=30)

    # Custom ARIMA Forecast using Streamlit button
    if st.sidebar.button("Run Custom ARIMA Forecast"):
        arima_predictions_custom = arima_forecast(stock_prices.values, p, d, q, arima_forecast_steps)
        arima_predictions_custom = pd.Series(arima_predictions_custom, index=pd.date_range(start=stock_prices.index[-1], periods=arima_forecast_steps + 1, freq=stock_prices.index.freq))

        # Display ARIMA Forecast Plot
        st.subheader("Custom ARIMA Forecast")
        st.line_chart(pd.concat([stock_prices, arima_predictions_custom], axis=1).rename(columns={0: "ARIMA Forecast"}))

    # LSTM Forecast using Streamlit button
    if st.sidebar.button("Run LSTM Forecast"):
        lstm_predictions = lstm_gru_forecast(stock_prices, 'lstm', lstm_gru_forecast_steps)

        # Display LSTM Forecast Plot
        st.subheader("LSTM Forecast")
        st.line_chart(pd.concat([stock_prices, pd.Series(lstm_predictions, index=pd.date_range(start=stock_prices.index[-1], periods=lstm_gru_forecast_steps + 1, freq=stock_prices.index.freq))], axis=1).rename(columns={0: "LSTM Forecast"}))

    # GRU Forecast using Streamlit button
    if st.sidebar.button("Run GRU Forecast"):
        gru_predictions = lstm_gru_forecast(stock_prices, 'gru', lstm_gru_forecast_steps)

        # Display GRU Forecast Plot
        st.subheader("GRU Forecast")
        st.line_chart(pd.concat([stock_prices, pd.Series(gru_predictions, index=pd.date_range(start=stock_prices.index[-1], periods=lstm_gru_forecast_steps + 1, freq=stock_prices.index.freq))], axis=1).rename(columns={0: "GRU Forecast"}))

    # Ensemble Forecast using Streamlit button
    if st.sidebar.button("Run Ensemble Forecast"):
        ensemble_predictions = ensemble_forecast([arima_predictions_custom, lstm_predictions, gru_predictions])

        # Display Ensemble Forecast Plot
        st.subheader("Ensemble Forecast")
        st.line_chart(pd.concat([stock_prices, ensemble_predictions], axis=1).rename(columns={0: "Ensemble Forecast"}))

    # Plotting Historical Stock Prices
    st.subheader("Historical Stock Prices")
    st.line_chart(stock_prices)

if __name__ == "__main__":
    main()