import gradio as gr
import pandas as pd
import numpy as np
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 LSTM, Dense

def forecast_stock(file, model_choice="ARIMA", horizon=30):
    df = pd.read_csv(file, parse_dates=["Date"])
    df = df.set_index("Date").asfreq("B")
    df["Aclose"] = df["Aclose"].interpolate()
    series = df["Aclose"].dropna()

    if model_choice == "ARIMA":
        model = ARIMA(series, order=(2,1,2))
        model_fit = model.fit()
        forecast = model_fit.forecast(steps=horizon)
    else:  # LSTM
        scaler = MinMaxScaler()
        data = scaler.fit_transform(series.values.reshape(-1,1))
        X, y = [], []
        seq_len = 60
        for i in range(len(data)-seq_len):
            X.append(data[i:i+seq_len])
            y.append(data[i+seq_len])
        X, y = np.array(X), np.array(y)

        model = Sequential([
            LSTM(50, input_shape=(seq_len,1)),
            Dense(1)
        ])
        model.compile(optimizer="adam", loss="mse")
        model.fit(X, y, epochs=3, batch_size=32, verbose=0)

        # Forecast
        last_seq = data[-seq_len:].reshape(1,seq_len,1)
        preds = []
        for _ in range(horizon):
            p = model.predict(last_seq, verbose=0)[0,0]
            preds.append(p)
            last_seq = np.append(last_seq[:,1:,:], [[[p]]], axis=1)
        forecast = scaler.inverse_transform(np.array(preds).reshape(-1,1)).flatten()

    # Plot results
    plt.figure(figsize=(10,5))
    plt.plot(series[-200:], label="History")
    plt.plot(pd.date_range(series.index[-1], periods=horizon+1, freq="B")[1:], forecast, label="Forecast")
    plt.legend()
    plt.title(f"{model_choice} Stock Forecast")
    plt.xlabel("Date")
    plt.ylabel("Price")
    plt.tight_layout()
    return plt

demo = gr.Interface(
    fn=forecast_stock,
    inputs=[
        gr.File(label="Upload Stock CSV"),
        gr.Radio(["ARIMA","LSTM"], value="ARIMA", label="Model"),
        gr.Slider(5,60,step=5,label="Forecast Horizon (days)")
    ],
    outputs=gr.Plot(label="Forecast Plot"),
    title="📈 Stock Price Forecasting",
    description="Upload stock data (CSV with Date,Aclose). Choose ARIMA or LSTM."
)

if __name__ == "__main__":
    demo.launch()