import gradio as gr
import pandas as pd
import numpy as np
from statsmodels.tsa.arima.model import ARIMA
import plotly.graph_objects as go
from tensorflow import keras
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense, Dropout
from sklearn.preprocessing import MinMaxScaler

def forecast_arima(days_ahead):
    df = pd.read_excel("Microsoft_stock_data.xlsx")
    df['Date'] = pd.to_datetime(df['Date'])
    df = df.sort_values('Date')
    
    data = df['Close'].values
    
    model = ARIMA(data, order=(1,1,1))
    fitted = model.fit()
    
    forecast = fitted.forecast(steps=int(days_ahead))
    
    fig = go.Figure()
    fig.add_trace(go.Scatter(x=df['Date'].tail(100), y=data[-100:], name='Historical', line=dict(color='blue')))
    future_dates = pd.date_range(start=df['Date'].iloc[-1], periods=int(days_ahead)+1)[1:]
    fig.add_trace(go.Scatter(x=future_dates, y=forecast, name='ARIMA Forecast', line=dict(color='red')))
    fig.update_layout(title='ARIMA Stock Price Forecast', xaxis_title='Date', yaxis_title='Price')
    
    return fig

def forecast_lstm(days_ahead):
    df = pd.read_excel("Microsoft_stock_data.xlsx")
    df['Date'] = pd.to_datetime(df['Date'])
    df = df.sort_values('Date')
    
    data = df['Close'].values.reshape(-1, 1)
    
    scaler = MinMaxScaler(feature_range=(0, 1))
    scaled_data = scaler.fit_transform(data)
    

    lookback = 60
    X_train, y_train = [], []
    for i in range(lookback, len(scaled_data)):
        X_train.append(scaled_data[i-lookback:i, 0])
        y_train.append(scaled_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))

    model = Sequential([
        LSTM(units=50, return_sequences=True, input_shape=(X_train.shape[1], 1)),
        Dropout(0.2),
        LSTM(units=50, return_sequences=False),
        Dropout(0.2),
        Dense(units=25),
        Dense(units=1)
    ])
    
    model.compile(optimizer='adam', loss='mean_squared_error')
    model.fit(X_train, y_train, batch_size=32, epochs=10, verbose=0)
    

    last_sequence = scaled_data[-lookback:]
    forecast = []
    
    for _ in range(int(days_ahead)):
        prediction = model.predict(last_sequence.reshape(1, lookback, 1), verbose=0)
        forecast.append(prediction[0, 0])
        last_sequence = np.append(last_sequence[1:], prediction)
    
    forecast = scaler.inverse_transform(np.array(forecast).reshape(-1, 1))
    
 
    fig = go.Figure()
    fig.add_trace(go.Scatter(x=df['Date'].tail(100), y=data[-100:, 0], name='Historical', line=dict(color='blue')))
    future_dates = pd.date_range(start=df['Date'].iloc[-1], periods=int(days_ahead)+1)[1:]
    fig.add_trace(go.Scatter(x=future_dates, y=forecast.flatten(), name='LSTM Forecast', line=dict(color='green')))
    fig.update_layout(title='LSTM Stock Price Forecast', xaxis_title='Date', yaxis_title='Price')
    
    return fig

def forecast_comparison(days_ahead):
    df = pd.read_excel("Microsoft_stock_data.xlsx")
    df['Date'] = pd.to_datetime(df['Date'])
    df = df.sort_values('Date')
    
    data = df['Close'].values
    

    arima_model = ARIMA(data, order=(1,1,1))
    arima_fitted = arima_model.fit()
    arima_forecast = arima_fitted.forecast(steps=int(days_ahead))
    

    scaler = MinMaxScaler(feature_range=(0, 1))
    scaled_data = scaler.fit_transform(data.reshape(-1, 1))
    
    lookback = 60
    X_train, y_train = [], []
    for i in range(lookback, len(scaled_data)):
        X_train.append(scaled_data[i-lookback:i, 0])
        y_train.append(scaled_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))
    
    lstm_model = Sequential([
        LSTM(units=50, return_sequences=True, input_shape=(X_train.shape[1], 1)),
        Dropout(0.2),
        LSTM(units=50, return_sequences=False),
        Dropout(0.2),
        Dense(units=25),
        Dense(units=1)
    ])
    
    lstm_model.compile(optimizer='adam', loss='mean_squared_error')
    lstm_model.fit(X_train, y_train, batch_size=32, epochs=10, verbose=0)
    
    last_sequence = scaled_data[-lookback:]
    lstm_forecast = []
    
    for _ in range(int(days_ahead)):
        prediction = lstm_model.predict(last_sequence.reshape(1, lookback, 1), verbose=0)
        lstm_forecast.append(prediction[0, 0])
        last_sequence = np.append(last_sequence[1:], prediction)
    
    lstm_forecast = scaler.inverse_transform(np.array(lstm_forecast).reshape(-1, 1)).flatten()
    

    fig = go.Figure()
    fig.add_trace(go.Scatter(x=df['Date'].tail(100), y=data[-100:], name='Historical', line=dict(color='blue')))
    future_dates = pd.date_range(start=df['Date'].iloc[-1], periods=int(days_ahead)+1)[1:]
    fig.add_trace(go.Scatter(x=future_dates, y=arima_forecast, name='ARIMA Forecast', line=dict(color='red', dash='dash')))
    fig.add_trace(go.Scatter(x=future_dates, y=lstm_forecast, name='LSTM Forecast', line=dict(color='green', dash='dot')))
    fig.update_layout(title='ARIMA vs LSTM: Comparison', xaxis_title='Date', yaxis_title='Price')
    
    return fig


with gr.Blocks() as demo:
    gr.Markdown("# 📈 Time Series Forecasting: ARIMA vs LSTM")
    gr.Markdown("**Microsoft Stock Price Forecasting** - Compare ARIMA and LSTM models.")
    
    days = gr.Slider(1, 90, value=30, label="Days to Forecast")
    
    with gr.Tabs():
        with gr.Tab("ARIMA Model"):
            arima_plot = gr.Plot()
            arima_btn = gr.Button("Generate ARIMA Forecast")
            arima_btn.click(forecast_arima, inputs=days, outputs=arima_plot)
     
            demo.load(forecast_arima, inputs=days, outputs=arima_plot)
        
        with gr.Tab("LSTM Model"):
            lstm_plot = gr.Plot()
            lstm_btn = gr.Button("Generate LSTM Forecast")
            lstm_btn.click(forecast_lstm, inputs=days, outputs=lstm_plot)
        
        with gr.Tab("Comparison"):
            comparison_plot = gr.Plot()
            compare_btn = gr.Button("Compare Both Models")
            compare_btn.click(forecast_comparison, inputs=days, outputs=comparison_plot)
    

    days.change(forecast_arima, inputs=days, outputs=arima_plot)

demo.launch()