from huggingface_hub import hf_hub_download
model_path = hf_hub_download(repo_id="jengyang/lstm-stock-prediction-model", filename="stage2_universal_lstm_20250705_170829.keras")
scaler_path = hf_hub_download(repo_id="jengyang/lstm-stock-prediction-model", filename="stage2_scalers_20250705_170829.pkl")

import gradio as gr
import tensorflow as tf # Import TensorFlow
import numpy as np
import pandas as pd
import yfinance as yf
import joblib
from datetime import datetime
import matplotlib.pyplot as plt

# The original LSTMModel was a PyTorch model. Since we are loading a Keras model,
# we don't need to define the PyTorch LSTMModel class here.

# Load model and scaler
# Load the Keras model using tf.keras.models.load_model
model = tf.keras.models.load_model(model_path)

scaler = joblib.load(scaler_path)

seq_length = 60  # Must match your training window

def predict_stock(ticker: str, days_ahead: int = 30):
    if days_ahead < 1 or days_ahead > 90:
        return "Please choose 1-90 days.", None

    # Fetch recent data
    end_date = datetime.now().strftime("%Y-%m-%d")
    start_date = (datetime.now() - pd.Timedelta(days=seq_length + 365)).strftime("%Y-%m-%d")
    data = yf.download(ticker.upper(), start=start_date, end=end_date)

    if data.empty or len(data) < seq_length:
        return f"No enough data for {ticker}. Try a popular ticker like AAPL.", None

    close_prices = data['Close'].values.values.reshape(-1, 1)

    # Scale
    scaled_data = scaler.transform(close_prices)

    # Get last sequence
    last_sequence = scaled_data[-seq_length:].reshape(1, seq_length, 1)

    predictions = []
    current_seq = last_sequence.copy()

    for _ in range(days_ahead):
        # Predict using the Keras model
        pred = model.predict(current_seq, verbose=0)[0][0] # Keras predict returns batch, step, value
        predictions.append(pred)
        # Append prediction and slide window
        current_seq = np.append(current_seq[:, 1:, :], [[pred]], axis=1) # Update sequence for next prediction

    # Inverse scale
    pred_prices = scaler.inverse_transform(np.array(predictions).reshape(-1, 1))

    # Create plot
    future_dates = pd.date_range(start=data.index[-1] + pd.Timedelta(days=1), periods=days_ahead)
    hist_dates = data.index[-60:]  # Last 60 days historical
    hist_prices = close_prices[-60:].flatten()

    plt.figure(figsize=(12, 6))
    plt.plot(hist_dates, hist_prices, label="Historical Close")
    plt.plot(future_dates, pred_prices, label="Predicted Close", linestyle="--")
    plt.title(f"{ticker.upper()} Stock Price Prediction")
    plt.xlabel("Date")
    plt.ylabel("Price ($)")
    plt.legend()
    plt.grid()

    return f"Predicted prices for next {days_ahead} days:", plt

# Gradio Interface
iface = gr.Interface(
    fn=predict_stock,
    inputs=[
        gr.Textbox(label="Stock Ticker", placeholder="e.g., AAPL, TSLA, GOOG", value="AAPL"),
        gr.Slider(1, 90, value=30, step=1, label="Days to Predict Ahead")
    ],
    outputs=[
        gr.Textbox(label="Summary"),
        gr.Plot(label="Price Chart")
    ],
    title="Keras LSTM Stock Price Predictor", # Updated title
    description="Enter a stock ticker to forecast future closing prices using a trained Keras LSTM model.", # Updated description
    examples=[["AAPL", 30], ["TSLA", 14], ["MSFT", 60]]
)

iface.launch()