app.py

import os
import json
import gradio as gr
import pandas as pd
import numpy as np
import torch
import joblib
from utils.retrain import LSTMModel  # ensure retrain.py has unified model loader

# ---- CONFIG ----
TICKERS = ["TSLA", "NVDA", "SPY"]  # add more tickers
HORIZONS = {
    "1d": {"days": 1},
    "1w": {"days": 5},
    "4w": {"days": 20},
    "6m": {"days": 120},
    "1y": {"days": 240},
}
MODEL_TYPES = ["Per-Ticker", "Unified"]

BASE_DIR = os.path.dirname(__file__)
DATA_PATH = os.path.join(BASE_DIR, "data")
MODELS_DIR = os.path.join(BASE_DIR, "models")
DEVICE = "cpu"

# ------------------------
# MC Dropout Prediction
# ------------------------
def enable_mc_dropout(model):
    for m in model.modules():
        if isinstance(m, torch.nn.Dropout):
            m.train()
    return model

def mc_dropout_predict(model, X_tensor, y_scaler, n_samples=200):
    model = enable_mc_dropout(model)
    preds = []
    with torch.no_grad():
        for _ in range(n_samples):
            preds.append(model(X_tensor).item())
    preds = torch.tensor(preds)
    mean_pred = preds.mean().item()
    std_pred = preds.std().item()
    mean_orig = y_scaler.inverse_transform([[mean_pred]])[0, 0]
    std_orig = y_scaler.inverse_transform([[mean_pred + std_pred]])[0, 0] - mean_orig
    plus_minus_percent = (std_orig / mean_orig * 100) if mean_orig != 0 else 0
    return {
        "predicted_price": mean_orig,
        "plus_minus_percent": plus_minus_percent,
        "confidence_percent": 95.0,
        "lower_bound": mean_orig - std_orig,
        "upper_bound": mean_orig + std_orig,
    }

# ------------------------
# Data Preparation
# ------------------------
def prepare_input_sequence(df, x_scaler, seq_len=90):
    features = df[["Open", "High", "Low", "Close", "Volume"]].values
    X_scaled = x_scaler.transform(features)
    X_seq = X_scaled[-seq_len:]
    return torch.tensor(X_seq, dtype=torch.float32).unsqueeze(0).to(DEVICE)

def prepare_unified_input_sequence(df, x_scaler, seq_len, ticker_idx, num_tickers):
    features = df[["Open", "High", "Low", "Close", "Volume"]].values
    X_scaled = x_scaler.transform(features)
    onehot = np.eye(num_tickers)[ticker_idx]
    onehot_seq = np.repeat(onehot.reshape(1, -1), seq_len, axis=0)
    X_full = np.hstack([X_scaled[-seq_len:], onehot_seq])
    X_tensor = torch.tensor(X_full, dtype=torch.float32).unsqueeze(0).to(DEVICE)
    return X_tensor


# ------------------------
# Load Models
# ------------------------
def load_per_ticker_model(ticker, horizon):
    out_dir = os.path.join(MODELS_DIR, ticker)
    model_path = os.path.join(out_dir, f"{ticker}_{horizon}_model.pth")
    x_scaler_path = os.path.join(out_dir, f"{ticker}_{horizon}_scaler.pkl")
    y_scaler_path = os.path.join(out_dir, f"{ticker}_{horizon}_y_scaler.pkl")
    config_path = os.path.join(out_dir, f"{ticker}_{horizon}_config.json")
    if not all(os.path.exists(p) for p in [model_path, x_scaler_path, y_scaler_path]):
        raise FileNotFoundError(f"Missing model/scalers for {ticker} ({horizon})")
    x_scaler = joblib.load(x_scaler_path)
    y_scaler = joblib.load(y_scaler_path)
    cfg = {}
    if os.path.exists(config_path):
        with open(config_path, "r") as f:
            cfg = json.load(f)
    model = LSTMModel(
        input_size=cfg.get("input_size", len(x_scaler.mean_)),
        hidden_size=cfg.get("hidden_size", 128),
        num_layers=cfg.get("num_layers", 2),
        dropout=cfg.get("dropout", 0.2),
    )
    model.load_state_dict(torch.load(model_path, map_location=DEVICE, weights_only=True))
    model.to(DEVICE)
    seq_len = cfg.get("seq_len", 90)
    return model, x_scaler, y_scaler, seq_len

def load_unified_model_and_scalers(horizon):
    udir = os.path.join(MODELS_DIR, "unified")
    model_path = os.path.join(udir, f"unified_{horizon}_model.pth")
    x_scaler_path = os.path.join(udir, f"unified_{horizon}_scaler.pkl")
    y_scaler_path = os.path.join(udir, f"unified_{horizon}_y_scaler.pkl")
    ticker_map_path = os.path.join(udir, "unified_tickers.pkl")
    if not all(os.path.exists(p) for p in [model_path, x_scaler_path, y_scaler_path, ticker_map_path]):
        raise FileNotFoundError(f"Missing unified model/scalers for {horizon}")
    x_scaler = joblib.load(x_scaler_path)
    y_scaler = joblib.load(y_scaler_path)
    ticker_map = joblib.load(ticker_map_path)
    # Load config if exists
    config_path = os.path.join(udir, f"unified_{horizon}_config.json")
    if os.path.exists(config_path):
        with open(config_path, "r") as f:
            cfg = json.load(f)
        input_size = cfg.get("input_size", len(x_scaler.mean_) + len(ticker_map))
        hidden_size = cfg.get("hidden_size", 128)
        num_layers = cfg.get("num_layers", 2)
        dropout = cfg.get("dropout", 0.2)
        seq_len = cfg.get("seq_len", 90)
    else:
        input_size = len(x_scaler.mean_) + len(ticker_map)
        hidden_size, num_layers, dropout, seq_len = 128, 2, 0.2, 90
    model = LSTMModel(
        input_size=input_size,
        hidden_size=hidden_size,
        num_layers=num_layers,
        dropout=dropout,
    )
    model.load_state_dict(torch.load(model_path, map_location=DEVICE, weights_only=True))
    model.to(DEVICE)
    return model, x_scaler, y_scaler, ticker_map, seq_len

# ------------------------
# Run Prediction
# ------------------------
def run_prediction(ticker, horizon, model_type):
    result = {}
    csv_path = os.path.join(DATA_PATH, f"{ticker}.csv")
    if not os.path.exists(csv_path):
        return {"error": f"No data for {ticker}"}
    df = pd.read_csv(csv_path)
    try:
        if model_type == "Per-Ticker":
            model, x_scaler, y_scaler, seq_len = load_per_ticker_model(ticker, horizon)
            X = prepare_input_sequence(df, x_scaler, seq_len)
            result[ticker] = {horizon: mc_dropout_predict(model, X, y_scaler)}
        else:  # Unified
            model, x_scaler, y_scaler, ticker_map, seq_len = load_unified_model_and_scalers(horizon)
            if ticker not in ticker_map:
                return {"error": f"{ticker} not in unified model"}
            X = prepare_unified_input_sequence(df, x_scaler, seq_len, ticker_map[ticker], len(ticker_map))
            result[ticker] = {horizon: mc_dropout_predict(model, X, y_scaler)}
    except Exception as e:
        result["error"] = str(e)
    return json.dumps(result, indent=2)

# ------------------------
# GRADIO UI
# ------------------------
with gr.Blocks(title="Stock Price Forecast") as demo:
    gr.Markdown("# Stock Price Forecasting with LSTM + MC Dropout")
    gr.Markdown("Select ticker, horizon, and model type, then click Predict.")

    ticker_dropdown = gr.Dropdown(choices=TICKERS, label="Ticker", value=TICKERS[0])
    horizon_dropdown = gr.Dropdown(choices=list(HORIZONS.keys()), label="Forecast Horizon", value="1d")
    model_dropdown = gr.Dropdown(choices=MODEL_TYPES, label="Model Type", value="Per-Ticker")
    predict_btn = gr.Button("Predict")
    output_box = gr.Code(label="JSON Output", language="json")

    predict_btn.click(
        fn=run_prediction,
        inputs=[ticker_dropdown, horizon_dropdown, model_dropdown],
        outputs=[output_box],
    )

demo.launch()