import gradio as gr
import yfinance as yf
import pandas as pd
import numpy as np
from sklearn.preprocessing import MinMaxScaler
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense, GRU
from datetime import datetime, timedelta
import matplotlib.pyplot as plt

# Function to convert string to datetime object
def str_to_datetime(s):
    return datetime.strptime(s, '%Y-%m-%d')

# Function to create a windowed dataframe
def df_to_windowed_df(dataframe, first_date_str, last_date_str, n=3):
    dataframe.index = pd.to_datetime(dataframe.index)
    first_date = str_to_datetime(first_date_str)
    last_date = str_to_datetime(last_date_str)

    target_date = first_date
    dates = []
    X, Y =[],[]

    last_time = False
    while True:
        df_subset = dataframe.loc[:target_date.strftime('%Y-%m-%d')].tail(n + 1)

        if len(df_subset)!= n + 1:
            print(f'Error: Window of size {n} is too large for date {target_date}')
            return

        values = df_subset['Close'].to_numpy()
        x, y = values[:-1], values[-1]

        dates.append(target_date)
        X.append(x)
        Y.append(y)

        next_week = dataframe.loc[target_date.strftime('%Y-%m-%d'):].head(2)
        if len(next_week) < 2:
            print(f'Insufficient future data for {target_date}. Ending process.')
            break

        next_date_str = next_week.index.strftime('%Y-%m-%d')
        next_date = str_to_datetime(next_date_str)

        if last_time:
            break

        target_date = next_date

        if target_date == last_date:
            last_time = True

    ret_df = pd.DataFrame({'Target Date': dates})

    X = np.array(X)
    for i in range(n):
        ret_df[f'Target-{n-i}'] = X[:, i]

    ret_df['Target'] = Y

    return ret_df

look_back = 3
num_features = 7

# Load pre-trained models and scalers
lstm_model = Sequential()
lstm_model.add(LSTM(50, activation='relu', input_shape=(look_back, num_features)))  
lstm_model.add(Dense(1))
lstm_model.compile(optimizer='adam', loss='mse')
lstm_model.load_weights('lstm_model_weights.weights.h5')  

gru_model = Sequential()
gru_model.add(GRU(50, activation='relu', input_shape=(look_back, num_features)))  # Replace with your actual GRU architecture
gru_model.add(Dense(1))
gru_model.compile(optimizer='adam', loss='mse')
gru_model.load_weights('gru_model_weights.weights.h5') 

lstm_scaler = pd.read_pickle('lstm_scaler.pkl') 
gru_scaler = pd.read_pickle('gru_scaler.pkl')

def predict_stock_price(ticker, num_days):
    # Fetch stock data
    stock_data = yf.Ticker(ticker)
    stock_data = stock_data.history(period='5y')
    stock_data.reset_index(inplace=True)
    stock_data['Date'] = pd.to_datetime(stock_data['Date']).dt.date
    stock_data = stock_data[['Date', 'Close']]
    stock_data.index = stock_data.pop('Date')

    # Prepare data for prediction
    windowed_df = df_to_windowed_df(stock_data, '2021-03-25', '2023-11-20', n=3)
    windowed_df['MA7'] = windowed_df['Target'].rolling(window=7).mean()
    windowed_df['MA21'] = windowed_df['Target'].rolling(window=21).mean()
    windowed_df['Volatility'] = windowed_df['Target'].pct_change().rolling(window=21).std()
    windowed_df['Close-MA7_Diff'] = windowed_df['Target'] - windowed_df['MA7']
    windowed_df = windowed_df.dropna()

    X = windowed_df[['Target-3', 'Target-2', 'Target-1', 'MA7', 'MA21', 'Volatility', 'Close-MA7_Diff']]
    X = np.array(X)
    X = X.reshape(X.shape[0], X.shape[1], 1)

    # Scale the data using separate scalers
    X_lstm = lstm_scaler.transform(X.reshape(X.shape, -1)).reshape(X.shape)
    X_gru = gru_scaler.transform(X.reshape(X.shape, -1)).reshape(X.shape)

    # Predict using LSTM and GRU
    lstm_predictions = lstm_model.predict(X_lstm)
    gru_predictions = gru_model.predict(X_gru)

    # Inverse transform to get actual prices
    lstm_predictions = lstm_scaler.inverse_transform(np.concatenate([X.reshape(X.shape, -1)[:,:3], lstm_predictions], axis=1))[:, -1]
    gru_predictions = gru_scaler.inverse_transform(np.concatenate([X.reshape(X.shape, -1)[:,:3], gru_predictions], axis=1))[:, -1]

    # Future predictions
    future_dates = [windowed_df['Target Date'].iloc[-1] + timedelta(days=i) for i in range(1, num_days + 1)]
    last_window_lstm = X_lstm[-1].reshape(1, 1, X_lstm.shape)
    last_window_gru = X_gru[-1].reshape(1, 1, X_gru.shape)
    future_predictions_lstm = []
    future_predictions_gru = []

    for _ in range(num_days):
        next_price_lstm = lstm_model.predict(last_window_lstm)
        future_predictions_lstm.append(next_price_lstm)
        last_window_lstm = np.roll(last_window_lstm, -1, axis=1)
        last_window_lstm[0, -1,:3] = next_price_lstm

        next_price_gru = gru_model.predict(last_window_gru)
        future_predictions_gru.append(next_price_gru)
        last_window_gru = np.roll(last_window_gru, -1, axis=1)
        last_window_gru[0, -1,:3] = next_price_gru

    future_predictions_lstm = lstm_scaler.inverse_transform(np.concatenate([np.array(future_predictions_lstm).reshape(-1, 1)] * 4, axis=1))[:, -1]
    future_predictions_gru = gru_scaler.inverse_transform(np.concatenate([np.array(future_predictions_gru).reshape(-1, 1)] * 4, axis=1))[:, -1]

    # Create plots with confidence intervals
    plt.figure(figsize=(14, 7))
    plt.plot(windowed_df['Target Date'], windowed_df['Target'], label='Actual')
    plt.plot(windowed_df['Target Date'], lstm_predictions, label='LSTM Predictions')
    plt.plot(windowed_df['Target Date'], gru_predictions, label='GRU Predictions')

    # Calculate and plot confidence intervals (example - you'll need to adjust the logic)
    std_dev = np.std(lstm_predictions)  # Calculate standard deviation
    plt.fill_between(windowed_df['Target Date'], lstm_predictions - std_dev, lstm_predictions + std_dev, color='blue', alpha=0.2, label='LSTM Confidence Interval')

    std_dev = np.std(gru_predictions)
    plt.fill_between(windowed_df['Target Date'], gru_predictions - std_dev, gru_predictions + std_dev, color='orange', alpha=0.2, label='GRU Confidence Interval')

    plt.plot(future_dates, future_predictions_lstm, label='Future LSTM Predictions', linestyle='dashed')
    plt.plot(future_dates, future_predictions_gru, label='Future GRU Predictions', linestyle='dashed')

    # Identify and annotate potential buying/selling opportunities
    # (This is a simplified example, you'll need to implement your own logic)
    for i in range(1, len(future_predictions_lstm)):
        if future_predictions_lstm[i] > future_predictions_lstm[i-1] + std_dev:
            plt.annotate('Buy', (future_dates[i], future_predictions_lstm[i]), textcoords="offset points", xytext=(0,10), ha='center', color='green')
        elif future_predictions_lstm[i] < future_predictions_lstm[i-1] - std_dev:
            plt.annotate('Sell', (future_dates[i], future_predictions_lstm[i]), textcoords="offset points", xytext=(0,10), ha='center', color='red')

    plt.xlabel('Date')
    plt.ylabel('Closing Price')
    plt.title(f'{ticker} Stock Price Prediction')
    plt.legend()
    plt.grid(True)

    return plt

# Gradio interface
iface = gr.Interface(
    fn=predict_stock_price,
    inputs=[
        gr.Textbox(lines=1, placeholder="Enter stock ticker (e.g., AMZN)", value="AMZN"),
        gr.Slider(minimum=1, maximum=365, step=1, value=30, label="Number of days to predict")
    ],
    outputs=gr.Plot(),
    title="Stock Price Prediction with LSTM and GRU",
    description="Enter a stock ticker symbol and the number of days to predict."
)

iface.launch()