import streamlit as st
import numpy as np
import pandas as pd
import yfinance as yf
import datetime as dt
import plotly.graph_objects as go
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import SimpleRNN, LSTM, GRU, Dense, Dropout
from tensorflow.keras.optimizers import SGD, Adam
from sklearn.preprocessing import MinMaxScaler

# Defining model creation functions
def create_rnn_model(input_shape: tuple):
    model = Sequential([
        SimpleRNN(50, return_sequences=True, input_shape=input_shape, activation='tanh'),
        Dropout(0.2),
        SimpleRNN(50, return_sequences=True, activation='tanh'),
        SimpleRNN(50, activation='tanh'),
        Dense(1)
    ])
    model.compile(optimizer=SGD(learning_rate=0.01, momentum=0.9), loss='mean_squared_error')
    return model

def create_lstm_model(input_shape: tuple):
    model = Sequential([
        LSTM(50, return_sequences=True, input_shape=input_shape),
        LSTM(50),
        Dense(25),
        Dense(1)
    ])
    model.compile(optimizer='adam', loss='mean_squared_error')
    return model

def create_gru_model(input_shape: tuple):
    model = Sequential([
        GRU(50, return_sequences=True, input_shape=input_shape, activation='tanh'),
        Dropout(0.2),
        GRU(50, return_sequences=True, activation='tanh'),
        GRU(50, activation='tanh'),
        Dense(1)
    ])
    model.compile(optimizer=SGD(learning_rate=0.01, momentum=0.9), loss='mean_squared_error')
    return model

def download_data(stock, start_date, end_date):
    data = yf.download(stock, start=start_date, end=end_date)
    return data

def prepare_data(data):
    scaler = MinMaxScaler(feature_range=(0, 1))
    data_scaled = scaler.fit_transform(data.reshape(-1, 1))
    return data_scaled, scaler

def create_datasets(data_scaled, look_back=50):
    X, y = [], []
    for i in range(look_back, len(data_scaled)):
        X.append(data_scaled[i-look_back:i, 0])
        y.append(data_scaled[i, 0])
    X, y = np.array(X), np.array(y)
    X = np.reshape(X, (X.shape[0], X.shape[1], 1))
    return X, y

def plot_predictions(train_data, test_data, y_train_pred, y_test_pred, model_name):
    fig = go.Figure()
    fig.add_trace(go.Scatter(x=train_data.index, y=train_data.values.flatten(), mode='lines', name='Training Data'))
    fig.add_trace(go.Scatter(x=test_data.index, y=test_data.values.flatten(), mode='lines', name='Test Data'))
    fig.add_trace(go.Scatter(x=test_data.index[look_back:], y=y_test_pred, mode='lines', name='Predicted Data'))
    fig.update_layout(title=f'{model_name} Predictions', xaxis_title='Date', yaxis_title='Stock Price')
    st.plotly_chart(fig)

# Streamlit interface
st.title('Stock Forecasting App')

stock = st.text_input('Enter Stock Ticker', 'AAPL')
start_date = st.date_input('Start Date', dt.date(2020, 1, 1))
end_date = st.date_input('End Date', dt.date.today())
look_back = 50

if st.button('Train Models'):
    data = download_data(stock, start_date, end_date)
    train_data = data['Open'][:-look_back]
    test_data = data['Open'][-look_back:]
    train_scaled, scaler = prepare_data(train_data.values)
    test_scaled, _ = prepare_data(test_data.values)

    X_train, y_train = create_datasets(train_scaled, look_back)
    X_test, y_test = create_datasets(test_scaled, look_back)

    rnn_model = create_rnn_model((look_back, 1))
    lstm_model = create_lstm_model((look_back, 1))
    gru_model = create_gru_model((look_back, 1))

    rnn_model.fit(X_train, y_train, epochs=50, batch_size=1)
    lstm_model.fit(X_train, y_train, epochs=50, batch_size=1)
    gru_model.fit(X_train, y_train, epochs=50, batch_size=1)

    y_rnn_pred = rnn_model.predict(X_test)
    y_lstm_pred = lstm_model.predict(X_test)
    y_gru_pred = gru_model.predict(X_test)

    y_rnn_pred = scaler.inverse_transform(y_rnn_pred)
    y_lstm_pred = scaler.inverse_transform(y_lstm_pred)
    y_gru_pred = scaler.inverse_transform(y_gru_pred)

    plot_predictions(train_data, test_data, y_rnn_pred, y_lstm_pred, y_gru_pred)