import gradio as gr
import yfinance as yf
from sklearn.linear_model import LinearRegression
import plotly.graph_objs as go
import numpy as np

def train_predict_wrapper(ticker, start_date, end_date, prediction_days):
    """
    Downloads stock data, trains a linear regression model, and predicts future prices.

    Args:
      ticker: The ticker symbol of the stock.
      start_date: The start date for the data (YYYY-MM-DD format).
      end_date: The end date for the data (YYYY-MM-DD format).
      prediction_days: The number of days to predict.

    Returns:
      A plot of predicted closing prices for the next `prediction_days`.
    """
    # Download stock data
    data = yf.download(ticker, start=start_date, end=end_date)
    # Extract closing price
    data = data["Close"]

    # Prepare data for model
    data = data.reset_index()
    data['Date'] = data['Date'].map(mdates.date2num)
    X = np.array(data.index).reshape(-1, 1)
    y = data['Close'].values

    # Train linear regression model
    model = LinearRegression()
    model.fit(X[:-prediction_days], y[:-prediction_days])

    # Predict future prices
    future_indices = np.array(range(len(X), len(X) + prediction_days)).reshape(-1, 1)
    predicted_prices = model.predict(future_indices)

    # Plot
    dates = [mdates.num2date(date).strftime('%Y-%m-%d') for date in data['Date']]
    future_dates = [mdates.num2date(date).strftime('%Y-%m-%d') for date in future_indices.flatten()]
    fig = go.Figure()
    fig.add_trace(go.Scatter(x=dates, y=y, mode='lines', name='Historical Prices'))
    fig.add_trace(go.Scatter(x=future_dates, y=predicted_prices, mode='lines', name='Predicted Prices'))
    fig.update_layout(title='Stock Price Prediction', xaxis_title='Date', yaxis_title='Price')

    return fig

# Define Gradio interface
iface = gr.Interface(
    fn=train_predict_wrapper,
    inputs=[
        gr.inputs.Textbox(label="Ticker Symbol"),
        gr.inputs.Textbox(label="Start Date (YYYY-MM-DD)"),
        gr.inputs.Textbox(label="End Date (YYYY-MM-DD)"),
        gr.inputs.Slider(minimum=1, maximum=30, step=1, default=5, label="Prediction Days")
    ],
    outputs="plot"
)

# Launch the app
iface.launch()