import gradio as gr
import pandas as pd
import yfinance as yf
import matplotlib.pyplot as plt
import pickle
import io

# Load pickled models
with open("arima_model.pkl", "rb") as f:
    arima_model = pickle.load(f)

with open("prophet_daily.pkl", "rb") as f:
    prophet_daily_model = pickle.load(f)

with open("prophet_seasonal.pkl", "rb") as f:
    prophet_seasonal_model = pickle.load(f)

def forecast_stock(ticker, model_type, days=1):
    # Fetch recent data for plotting
    df = yf.download(ticker, start="2010-01-01", interval="1d")
    df.columns = df.columns.get_level_values(0)
    df = df[['Close']].sort_index()
    df = df.asfreq('B').ffill()
    
    # Forecast
    if model_type == "ARIMA":
        forecast = arima_model.forecast(steps=days)
    elif model_type == "Prophet Daily":
        future = prophet_daily_model.make_future_dataframe(periods=days, freq='B')
        forecast_df = prophet_daily_model.predict(future)
        forecast = forecast_df['yhat'].iloc[-days:].values
    else:  # Prophet Seasonal
        future = prophet_seasonal_model.make_future_dataframe(periods=days, freq='B')
        forecast_df = prophet_seasonal_model.predict(future)
        forecast = forecast_df['yhat'].iloc[-days:].values
    
    # Plot
    plt.figure(figsize=(10,5))
    plt.plot(df.index[-50:], df['Close'].values[-50:], label='Recent Actual')
    plt.plot(pd.date_range(df.index[-1]+pd.Timedelta(days=1), periods=days, freq='B'),
             forecast, label='Forecast', marker='o')
    plt.legend()
    plt.title(f"{ticker} Stock Forecast ({model_type})")
    plot_path = "temp_plot.png"  # temporary file
    plt.savefig(plot_path)
    plt.close()
    return plot_path
# Gradio interface
ticker_input = gr.Textbox(label="Ticker Symbol", value="AAPL")
model_input = gr.Radio(["ARIMA", "Prophet Daily", "Prophet Seasonal"], label="Model")
days_input = gr.Slider(1, 30, step=1, label="Days Ahead")

gr.Interface(
    forecast_stock,
    inputs=[ticker_input, model_input, days_input],
    outputs=gr.Image(type="pil"),
    live=True,
    title="Stock Price Forecasting",
    description="Forecast next n days stock prices using ARIMA or Prophet"
).launch()