import streamlit as st
import math
import numpy as np
import pandas as pd
import yfinance as yf
import datetime as dt
import plotly.graph_objects as go
import tensorflow as tf
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) -> tf.keras.Model:
    """
    Constructs a Recurrent Neural Network (RNN) model with multiple SimpleRNN layers and dropout regularization,
    followed by a Dense output layer. The model is compiled with the SGD optimizer.
    Args:
    input_shape (tuple): A tuple representing the input shape of the training data, excluding the batch size.
                         For example, (timesteps, features).
    Returns:
    tf.keras.Model: The constructed and compiled TensorFlow model.
    """
    # Initializing the RNN model
    regressor = Sequential()

    # Adding the first RNN layer with dropout regularization
    regressor.add(SimpleRNN(units=50, activation="tanh", return_sequences=True, input_shape=input_shape))
    regressor.add(Dropout(0.2))

    # Adding more RNN layers
    regressor.add(SimpleRNN(units=50, activation="tanh", return_sequences=True))
    regressor.add(SimpleRNN(units=50, activation="tanh", return_sequences=True))
    regressor.add(SimpleRNN(units=50, activation="tanh"))  # Last RNN layer does not return sequences

    # Adding the output layer with a single unit and sigmoid activation for binary outcomes
    regressor.add(Dense(units=1, activation='sigmoid'))

    # Compiling the RNN with the SGD optimizer
    regressor.compile(optimizer=SGD(learning_rate=0.01, momentum=0.9, nesterov=True),
                      loss="mean_squared_error")

    return regressor


def create_lstm_model(input_shape: tuple) -> tf.keras.Model:
    """
    Constructs a Long Short-Term Memory (LSTM) model with LSTM layers and a Dense layer,
    followed by a Dense output layer. The model is compiled with the Adam optimizer.
    Args:
    input_shape (tuple): A tuple representing the input shape of the training data, excluding the batch size.
                         For example, (timesteps, features).
    Returns:
    tf.keras.Model: The constructed and compiled TensorFlow model.
    """
    # Initializing the LSTM model
    regressorLSTM = Sequential()

    # Adding LSTM layers
    regressorLSTM.add(LSTM(50, return_sequences=True, input_shape=input_shape))
    regressorLSTM.add(LSTM(50, return_sequences=False))  # Last LSTM layer does not return sequences
    regressorLSTM.add(Dense(25))  # Additional Dense layer with 25 units

    # Adding the output layer with a single unit for output
    regressorLSTM.add(Dense(1))

    # Compiling the LSTM with the Adam optimizer
    regressorLSTM.compile(optimizer='adam', loss='mean_squared_error', metrics=["accuracy"])

    return regressorLSTM


def create_gru_model(input_shape: tuple) -> tf.keras.Model:
    """
    Constructs a Gated Recurrent Unit (GRU) model with multiple GRU layers including dropout for regularization,
    and a Dense output layer. The model is compiled with the SGD optimizer.
    Args:
    input_shape (tuple): A tuple representing the input shape of the training data, excluding the batch size.
                         For example, (timesteps, features).
    Returns:
    tf.keras.Model: The constructed and compiled TensorFlow model.
    """
    # Initializing the GRU model
    regressorGRU = Sequential()

    # Adding GRU layers with dropout regularization
    regressorGRU.add(GRU(units=50, return_sequences=True, input_shape=input_shape, activation='tanh'))
    regressorGRU.add(Dropout(0.2))

    regressorGRU.add(GRU(units=50, return_sequences=True, activation='tanh'))
    regressorGRU.add(GRU(units=50, return_sequences=True, activation='tanh'))
    regressorGRU.add(GRU(units=50, activation='tanh'))  # Last GRU layer does not return sequences

    # Adding the output layer with a relu activation
    regressorGRU.add(Dense(units=1, activation='relu'))

    # Compiling the GRU with the SGD optimizer
    regressorGRU.compile(optimizer=SGD(learning_rate=0.01, momentum=0.9, nesterov=False),
                         loss='mean_squared_error')

    return regressorGRU


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)