from gradio_app_builder import app
import yfinance as yf
from sklearn.linear_model import LinearRegression
import json

@app.route("/")
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 list 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 and set as index
  data = data["Close"].to_frame().set_index(data.index.values)

  # Train linear regression model
  X = data.index.values[:-prediction_days].reshape(-1, 1)
  y = data.values[:-prediction_days]
  model = LinearRegression()
  model.fit(X, y)

  # Predict future prices
  future_dates = data.index.values[-prediction_days:]
  X_future = future_dates.reshape(-1, 1)
  predicted_prices = model.predict(X_future).tolist()

  # Return predicted prices as JSON
  return json.dumps(predicted_prices)

# Launch the Gradio application
app.launch()