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

def forecast_stock(ticker, period, future_days):
    # 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()
    
    # Fit Prophet model
    model = Prophet()
    model.fit(df)
    
    # Make future dataframe
    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='Forecast'))
    fig_forecast.add_trace(go.Scatter(x=forecast['ds'], y=forecast['yhat_upper'], name='Upper Bound', line=dict(dash='dash')))
    fig_forecast.add_trace(go.Scatter(x=forecast['ds'], y=forecast['yhat_lower'], name='Lower Bound', line=dict(dash='dash')))
    fig_forecast.update_layout(title=f'Prophet Forecast for {ticker}', xaxis_title='Date', yaxis_title='Stock Price')
    
    # Create Plotly figure for components
    fig_components = px.line(forecast, x='ds', y=['trend', 'yearly', 'weekly'])
    fig_components.update_layout(title='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")
    ],
    outputs=[
        gr.Plot(label="Forecast"),
        gr.Plot(label="Forecast Components")
    ],
    title="Stock Price Forecasting with Prophet",
    description="Select a stock, historical data period, and forecast horizon to predict future stock prices."
)

# Launch the interface
iface.launch()