app.py

import yfinance as yf
import pandas as pd
import numpy as np
from prophet import Prophet
import plotly.graph_objs as go
import plotly.express as px
import gradio as gr
from pmdarima import auto_arima

def forecast_stock(ticker, period, future_days, use_arima):
    # Fetch data
    data = yf.Ticker(ticker)
    df = data.history(period=period)
    
    if df.empty:
        return "Could not retrieve data for the selected ticker."
    
    df = df.reset_index()
    df = df[['Date', 'Close']]
    df.columns = ['ds', 'y']
    df['ds'] = pd.to_datetime(df['ds']).dt.tz_localize(None)
    df = df.dropna()
    
    # Prophet forecast
    model = Prophet()
    model.fit(df)
    future_dates = model.make_future_dataframe(periods=future_days)
    forecast = model.predict(future_dates)
    
    # Create Plotly figure for forecast
    fig_forecast = go.Figure()
    fig_forecast.add_trace(go.Scatter(x=df['ds'], y=df['y'], name='Historical'))
    fig_forecast.add_trace(go.Scatter(x=forecast['ds'], y=forecast['yhat'], name='Prophet Forecast'))
    fig_forecast.add_trace(go.Scatter(x=forecast['ds'], y=forecast['yhat_upper'], name='Prophet Upper Bound', line=dict(dash='dash')))
    fig_forecast.add_trace(go.Scatter(x=forecast['ds'], y=forecast['yhat_lower'], name='Prophet Lower Bound', line=dict(dash='dash')))
    
    if use_arima:
        # ARIMA forecast with automatic order selection
        model_arima = auto_arima(df['y'], seasonal=False, trace=True)
        results_arima = model_arima.fit(df['y'])
        arima_forecast = results_arima.predict(n_periods=future_days)
        future_dates_arima = pd.date_range(start=df['ds'].iloc[-1] + pd.Timedelta(days=1), periods=future_days)
        
        fig_forecast.add_trace(go.Scatter(x=future_dates_arima, y=arima_forecast, name='ARIMA Forecast'))
    
    fig_forecast.update_layout(title=f'Stock Price Forecast for {ticker}', xaxis_title='Date', yaxis_title='Stock Price')
    
    # Create Plotly figure for Prophet components
    fig_components = px.line(forecast, x='ds', y=['trend', 'yearly', 'weekly'])
    fig_components.update_layout(title='Prophet Forecast Components')
    
    return fig_forecast, fig_components

# Define Gradio interface
iface = gr.Interface(
    fn=forecast_stock,
    inputs=[
        gr.Dropdown(choices=['AAPL', 'GOOGL', 'MSFT', 'AMZN'], label="Stock Ticker"),
        gr.Dropdown(choices=['1y', '2y', '5y', '10y', 'max'], label="Historical Data Period"),
        gr.Slider(minimum=30, maximum=365, step=30, label="Days to Forecast"),
        gr.Checkbox(label="Include ARIMA Forecast")
    ],
    outputs=[
        gr.Plot(label="Forecast"),
        gr.Plot(label="Prophet Forecast Components")
    ],
    title="Stock Price Forecasting with Prophet and ARIMA",
    description="Select a stock, historical data period, forecast horizon, and whether to include ARIMA forecast."
)

# Launch the interface
iface.launch()