src/forecasting/tft_model.py

"""
TFT training and inference wrapper.

Uses Nixtla's neuralforecast library which follows the (unique_id, ds, y) convention.
The model is trained once and cached; the Streamlit app calls forecast() at runtime.
"""

import os
import pickle
from pathlib import Path
from typing import List, Optional

import numpy as np
import pandas as pd
from neuralforecast import NeuralForecast
from neuralforecast.models import TFT
from neuralforecast.losses.pytorch import MQLoss

from src.forecasting.config import TFTConfig, DEFAULT_CONFIG


# ── Feature engineering ───────────────────────────────────────────────────────

FASHION_WEEK_MONTHS = {(1, 15), (1, 20), (1, 25), (5, 10), (5, 18),
                       (9, 25), (10, 2), (10, 8), (10, 15)}

EVENT_MAP = {
    "FashionWeek": 1, "ProductDrop": 2, "Christmas": 3,
    "ValentinesDay": 4, "ChineseNewYear": 5, "MothersDaySurge": 6,
}


def prepare_features(df: pd.DataFrame) -> pd.DataFrame:
    """
    Encode categorical and temporal features required by the TFT config.
    Operates on the full long-format DataFrame.
    """
    df = df.copy()
    df["ds"] = pd.to_datetime(df["ds"])

    # Temporal features
    df["day_of_week"] = df["ds"].dt.dayofweek / 6.0          # [0, 1]
    df["month"] = (df["ds"].dt.month - 1) / 11.0             # [0, 1]
    df["event_flag"] = df["event_name"].map(EVENT_MAP).fillna(0).astype(float)
    df["is_fashion_week"] = (
        df[["ds"]].assign(
            key=list(zip(df["ds"].dt.month, df["ds"].dt.day))
        )["key"].isin(FASHION_WEEK_MONTHS)
    ).astype(float)

    # Static categorical encodings (label-encode per series)
    category_enc = {c: i for i, c in enumerate(df["category"].unique())}
    store_enc = {s: i for i, s in enumerate(df["store"].unique())}
    df["category_enc"] = df["category"].map(category_enc).astype(float)
    df["store_enc"] = df["store"].map(store_enc).astype(float)

    # Price normalisation (min-max per item)
    df["price"] = df.groupby("unique_id")["price"].transform(
        lambda x: (x - x.min()) / (x.max() - x.min() + 1e-8)
    )

    return df


def train(
    df: pd.DataFrame,
    cfg: TFTConfig = DEFAULT_CONFIG,
    val_size: int = 56,     # last 8 weeks as validation
) -> NeuralForecast:
    """
    Train TFT on the full dataset and persist the model.
    Returns the fitted NeuralForecast object.
    """
    df = prepare_features(df)

    model = TFT(
        h=cfg.h,
        input_size=cfg.input_size,
        hidden_size=cfg.hidden_size,
        n_head=cfg.n_head,
        attn_dropout=cfg.attn_dropout,
        dropout=cfg.dropout,
        loss=MQLoss(level=[80, 95]),
        learning_rate=cfg.learning_rate,
        max_steps=cfg.max_steps,
        batch_size=cfg.batch_size,
        val_check_steps=cfg.val_check_steps,
        stat_exog_list=cfg.stat_exog_list,
        hist_exog_list=cfg.hist_exog_list,
        futr_exog_list=cfg.futr_exog_list,
        scaler_type="standard",
    )

    nf = NeuralForecast(models=[model], freq="D")
    nf.fit(df=df, val_size=val_size)

    Path(cfg.model_dir).mkdir(exist_ok=True)
    nf.save(path=f"{cfg.model_dir}/{cfg.model_name}", overwrite=True)

    return nf


def load_model(cfg: TFTConfig = DEFAULT_CONFIG) -> NeuralForecast:
    """Load a previously trained model from disk."""
    return NeuralForecast.load(path=f"{cfg.model_dir}/{cfg.model_name}")


def forecast(
    nf: NeuralForecast,
    df: pd.DataFrame,
    unique_ids: Optional[List[str]] = None,
    cfg: TFTConfig = DEFAULT_CONFIG,
) -> pd.DataFrame:
    """
    Run inference for the given series.

    Returns a DataFrame with columns:
        unique_id, ds, TFT (point), TFT-lo-80, TFT-hi-80, TFT-lo-95, TFT-hi-95
    """
    df = prepare_features(df)

    if unique_ids is not None:
        df = df[df["unique_id"].isin(unique_ids)]

    # Build future exogenous DataFrame for the forecast horizon
    futr_df = _build_future_exog(df, cfg.h)

    preds = nf.predict(df=df, futr_df=futr_df)
    return preds.reset_index()


def _build_future_exog(df: pd.DataFrame, h: int) -> pd.DataFrame:
    """
    Extend the known covariates (event_flag, day_of_week, month, is_fashion_week)
    into the forecast horizon for each series.
    """
    last_dates = df.groupby("unique_id")["ds"].max()
    records = []

    for uid, last_date in last_dates.items():
        future_dates = pd.date_range(
            start=last_date + pd.Timedelta(days=1), periods=h, freq="D"
        )
        fw_flags = pd.Series(future_dates).apply(
            lambda d: float((d.month, d.day) in FASHION_WEEK_MONTHS)
        )
        records.append(pd.DataFrame({
            "unique_id": uid,
            "ds": future_dates,
            "day_of_week": future_dates.dayofweek / 6.0,
            "month": (future_dates.month - 1) / 11.0,
            "event_flag": 0.0,          # conservative: no known future events
            "is_fashion_week": fw_flags.values,
        }))

    return pd.concat(records, ignore_index=True)


# ── Benchmark utilities ───────────────────────────────────────────────────────

def compute_metrics(actuals: pd.Series, predictions: pd.Series) -> dict:
    """MAE, RMSE, MASE (vs. seasonal naïve baseline with period=7)."""
    mae = np.abs(actuals - predictions).mean()
    rmse = np.sqrt(((actuals - predictions) ** 2).mean())
    naive = np.abs(actuals.values[7:] - actuals.values[:-7]).mean() + 1e-8
    mase = mae / naive
    return {"MAE": round(mae, 2), "RMSE": round(rmse, 2), "MASE": round(mase, 3)}


# ── Entrypoint ────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    from src.data.loader import generate_dataset

    print("Generating dataset...")
    df = generate_dataset()
    print(f"  {len(df):,} rows | {df['unique_id'].nunique()} series")

    print("Training TFT model...")
    nf = train(df)
    print("  Model saved to models/tft_luxury")