app.py

import os
import warnings
warnings.filterwarnings("ignore")

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

import gradio as gr

from sklearn.preprocessing import MinMaxScaler
from sklearn.metrics import mean_squared_error, mean_absolute_percentage_error

import joblib
from statsmodels.tsa.arima.model import ARIMA


# Prophet can be heavy; we load lazily inside functions to avoid slow import on startup
try:
    from prophet import Prophet
    PROPHET_AVAILABLE = True
except Exception:
    PROPHET_AVAILABLE = False

# TensorFlow (CPU) for LSTM
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
try:
    from tensorflow.keras.models import load_model
    TENSORFLOW_AVAILABLE = True
except Exception:
    TENSORFLOW_AVAILABLE = False


# ----------------------------
# Load data & artifacts
# ----------------------------
DATA_PATH = "data/AAPL_stock_2015_2025.csv"
df = pd.read_csv(DATA_PATH)
df["Date"] = pd.to_datetime(df["Date"])
df = df.sort_values("Date").set_index("Date")

features = ['Open','High','Low','Close','Volume']
for c in features:
    df[c] = pd.to_numeric(df[c], errors="coerce")
df = df.dropna()

series_close = df[["Close"]].copy()

# models
ARIMA_PATH = "models/arima_model.pkl"
PROPHET_PATH = "models/prophet_model.pkl"
LSTM_PATH = "models/lstm_model.h5"

arima_model = None
prophet_model = None
lstm_model = None

if os.path.exists(ARIMA_PATH):
    try:
        arima_model = joblib.load(ARIMA_PATH)
    except Exception:
        arima_model = None

if PROPHET_AVAILABLE and os.path.exists(PROPHET_PATH):
    try:
        prophet_model = joblib.load(PROPHET_PATH)
    except Exception:
        prophet_model = None

if TENSORFLOW_AVAILABLE:
    print("✅ TensorFlow imported successfully")
    print("🔎 Checking for LSTM model at:", LSTM_PATH)
    if os.path.exists(LSTM_PATH):
        try:
            lstm_model = load_model(LSTM_PATH, compile=False)
            print("✅ LSTM model loaded successfully")
        except Exception as e:
            print("❌ LSTM load failed:", e)
            lstm_model = None
    else:
        print("❌ LSTM model file not found!")


# ----------------------------
# Helpers
# ----------------------------
def plot_series_with_forecasts(history_df, forecast_dict, title):
    plt.figure(figsize=(12,5))
    plt.plot(history_df.index, history_df.values, label="Actual", linewidth=1.5)
    for name, (idx, vals) in forecast_dict.items():
        plt.plot(idx, vals, label=name, linewidth=1.5)
    plt.title(title)
    plt.xlabel("Date")
    plt.ylabel("Close (USD)")
    plt.legend()
    plt.tight_layout()
    return plt.gcf()

def backtest_split_last_n(n_days):
    """Return train/test splits for backtest using last n days as test."""
    train = series_close.iloc[:-n_days, 0]
    test = series_close.iloc[-n_days:, 0]
    return train, test

def safe_mape(y_true, y_pred):
    return mean_absolute_percentage_error(y_true, y_pred)


# ----------------------------
# ARIMA
# ----------------------------
def forecast_arima(horizon, retrain_if_needed=True):
    global arima_model
    if arima_model is None and retrain_if_needed:
        # fit a reasonable default quickly
        model = ARIMA(series_close.squeeze(), order=(5,1,0))
        arima_model = model.fit()
    if arima_model is None:
        raise RuntimeError("ARIMA model not available and retrain disabled.")
    fc = arima_model.forecast(steps=horizon)
    idx = pd.date_range(series_close.index[-1] + pd.Timedelta(days=1), periods=horizon, freq="D")
    return idx, fc.values

def backtest_arima(n_days, retrain_if_needed=True):
    train, test = backtest_split_last_n(n_days)
    model = ARIMA(train, order=(5,1,0))
    fit = model.fit()
    pred = fit.forecast(steps=len(test)).values
    rmse = np.sqrt(mean_squared_error(test.values, pred))
    mape = safe_mape(test.values, pred)
    idx = test.index
    return idx, pred, rmse, mape


# ----------------------------
# Prophet
# ----------------------------
def forecast_prophet(horizon, retrain_if_needed=True):
    global prophet_model
    if not PROPHET_AVAILABLE:
        raise RuntimeError("Prophet not installed in this Space.")

    if prophet_model is None and retrain_if_needed:
        p_df = df.reset_index()[["Date","Close"]]
        p_df.columns = ["ds","y"]
        model = Prophet(daily_seasonality=True)
        model.fit(p_df)
        prophet_model = model

    if prophet_model is None:
        raise RuntimeError("Prophet model not available and retrain disabled.")

    future = pd.date_range(series_close.index[-1] + pd.Timedelta(days=1), periods=horizon, freq="D")
    future_df = pd.DataFrame({"ds": future})
    forecast = prophet_model.predict(future_df)
    return future, forecast["yhat"].values

def backtest_prophet(n_days, retrain_if_needed=True):
    if not PROPHET_AVAILABLE:
        raise RuntimeError("Prophet not installed in this Space.")

    p_df = df.reset_index()[["Date","Close"]]
    p_df.columns = ["ds","y"]
    train = p_df.iloc[:-n_days]
    test = p_df.iloc[-n_days:]

    model = Prophet(daily_seasonality=True)
    model.fit(train)
    future = model.make_future_dataframe(periods=len(test))
    fc = model.predict(future).tail(len(test))["yhat"].values

    rmse = np.sqrt(mean_squared_error(test["y"].values, fc))
    mape = safe_mape(test["y"].values, fc)
    return test["ds"].values, fc, rmse, mape


# ----------------------------
# LSTM (multivariate) – recursive forecast with naive covariates
# ----------------------------
WINDOW = 60

def prepare_scaled_features():
    scaler = MinMaxScaler().fit(df[features].values)
    scaled = scaler.transform(df[features].values)
    return scaler, scaled

def forecast_lstm(horizon):
    if not (TENSORFLOW_AVAILABLE and lstm_model is not None):
        raise RuntimeError("LSTM model not available in this Space.")

    scaler, scaled = prepare_scaled_features()
    window = scaled[-WINDOW:].copy()
    preds_scaled_close = []

    # naive strategy
    for _ in range(horizon):
        x = np.expand_dims(window, axis=0)  # (1, 60, 5)
        pred_scaled_close = lstm_model.predict(x, verbose=0)[0,0]
        preds_scaled_close.append(pred_scaled_close)

        next_vec = window[-1].copy()
        next_vec[3] = pred_scaled_close  # Close
        next_vec[0] = pred_scaled_close  # Open ~ Close
        next_vec[1] = pred_scaled_close  # High ~ Close
        next_vec[2] = pred_scaled_close  # Low  ~ Close
        # Volume unchanged
        window = np.vstack([window[1:], next_vec])

    # inverse-transform Close
    dummy = np.zeros((len(preds_scaled_close), len(features)))
    dummy[:,3] = np.array(preds_scaled_close)
    preds_close = scaler.inverse_transform(dummy)[:,3]

    idx = pd.date_range(series_close.index[-1] + pd.Timedelta(days=1), periods=horizon, freq="D")
    return idx, preds_close

def backtest_lstm(n_days):
    if not (TENSORFLOW_AVAILABLE and lstm_model is not None):
        raise RuntimeError("LSTM model not available in this Space.")

    scaler, scaled = prepare_scaled_features()
    test_scaled = scaled[-(n_days + WINDOW):]
    preds = []
    gts = []

    for i in range(WINDOW, len(test_scaled)):
        window = test_scaled[i-WINDOW:i].copy()
        x = np.expand_dims(window, axis=0)
        pred_scaled_close = lstm_model.predict(x, verbose=0)[0,0]
        preds.append(pred_scaled_close)
        gts.append(test_scaled[i,3])  # true close (scaled)

    preds = np.array(preds)
    gts = np.array(gts)

    # inverse-transform Close
    dummy_p = np.zeros((len(preds), len(features))); dummy_p[:,3] = preds
    dummy_t = np.zeros((len(gts), len(features)));   dummy_t[:,3] = gts
    pred_close = scaler.inverse_transform(dummy_p)[:,3]
    true_close = scaler.inverse_transform(dummy_t)[:,3]

    rmse = np.sqrt(mean_squared_error(true_close, pred_close))
    mape = safe_mape(true_close, pred_close)
    idx = df.index[-n_days:]
    return idx, pred_close[-n_days:], rmse, mape


# ----------------------------
# Gradio UI
# ----------------------------
def do_forecast(horizon, arima, prophet, lstm):
    forecasts = {}
    messages = []

    if arima:
        try:
            idx, vals = forecast_arima(horizon)
            forecasts["ARIMA"] = (idx, vals)
        except Exception as e:
            messages.append(f"ARIMA error: {e}")

    if prophet:
        try:
            idx, vals = forecast_prophet(horizon)
            forecasts["Prophet"] = (idx, vals)
        except Exception as e:
            messages.append(f"Prophet error: {e}")

    if lstm:
        try:
            idx, vals = forecast_lstm(horizon)
            forecasts["LSTM (Multivariate)"] = (idx, vals)
        except Exception as e:
            messages.append(f"LSTM error: {e}")

    if not forecasts:
        return None, "No model produced a forecast. Check build logs or dependencies."

    fig = plot_series_with_forecasts(series_close, forecasts, f"Forecast {horizon} days ahead")
    return fig, "\n".join(messages) if messages else "OK"

def do_backtest(test_days, arima, prophet, lstm):
    rows = []
    overlays = {}

    if arima:
        try:
            idx, pred, rmse, mape = backtest_arima(test_days)
            overlays["ARIMA"] = (idx, pred)
            rows.append(["ARIMA", rmse, mape])
        except:
            rows.append(["ARIMA", None, None])

    if prophet:
        try:
            idx, pred, rmse, mape = backtest_prophet(test_days)
            overlays["Prophet"] = (pd.to_datetime(idx), pred)
            rows.append(["Prophet", rmse, mape])
        except:
            rows.append(["Prophet", None, None])

    if lstm:
        try:
            idx, pred, rmse, mape = backtest_lstm(test_days)
            overlays["LSTM (Multivariate)"] = (idx, pred)
            rows.append(["LSTM (Multivariate)", rmse, mape])
        except:
            rows.append(["LSTM (Multivariate)", None, None])

    hist = series_close.iloc[-test_days:]
    fig = plot_series_with_forecasts(hist, overlays, f"Backtest on last {test_days} days")

    table = pd.DataFrame(rows, columns=["Model", "RMSE", "MAPE"])
    return fig, table


with gr.Blocks(title="DataSynthis_ML_JobTask") as demo:
    gr.Markdown("# Multivariate Stock Forecasting (AAPL)\nCompare ARIMA, Prophet, and LSTM.")
    with gr.Tab("Forecast"):
        with gr.Row():
            horizon = gr.Slider(7, 90, value=30, step=1, label="Forecast horizon (days)")
        with gr.Row():
            arima_c = gr.Checkbox(True, label="ARIMA")
            prophet_c = gr.Checkbox(True, label="Prophet")
            lstm_c = gr.Checkbox(True, label="LSTM (Multivariate)")
        run_btn = gr.Button("Run Forecast")
        out_plot = gr.Plot()
        out_msg = gr.Textbox(label="Status / Notes")
        run_btn.click(fn=do_forecast, inputs=[horizon, arima_c, prophet_c, lstm_c], outputs=[out_plot, out_msg])

    with gr.Tab("Backtest"):
        with gr.Row():
            test_days = gr.Slider(30, 180, value=60, step=5, label="Backtest period (last N days)")
        with gr.Row():
            arima_b = gr.Checkbox(True, label="ARIMA")
            prophet_b = gr.Checkbox(True, label="Prophet")
            lstm_b = gr.Checkbox(True, label="LSTM (Multivariate)")
        back_btn = gr.Button("Run Backtest")
        back_plot = gr.Plot()
        back_table = gr.Dataframe(headers=["Model", "RMSE", "MAPE"])
        back_btn.click(fn=do_backtest, inputs=[test_days, arima_b, prophet_b, lstm_b], outputs=[back_plot, back_table])

demo.launch(share=True)