"""Track 2: learned trajectory representations (1D-CNN / GRU) vs hand-crafted features.

The spec's question is not "can I train an LSTM" — it is which side of the
bias-variance trade wins at n ≈ 2,000 runs. To make that answerable the two
tracks are held identical everywhere except the representation:

  same group-aware temporal split (by wafer, never straddling)
  same official 424-run competition test set
  same metric (MSE)
  same static inputs (tool wear, stage, chamber, run length)

Only the trajectory encoding differs: 208 hand-crafted time/frequency
descriptors, versus a network reading the resampled raw signals.

Run: python 02_phm_cmp/sequence_models.py   (needs prepare.py + run_model.py)
"""
from __future__ import annotations

import sys
import time
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score

HERE = Path(__file__).parent
sys.path.insert(0, str(HERE.parent))
sys.path.insert(0, str(HERE))

from common import splits  # noqa: E402
from prepare import CONTEXT_COLS, KEY, SIGNAL_COLS, WEAR_COLS  # noqa: E402

CACHE = HERE / "cache"
FIGS = HERE / "figures"
OUT = HERE / "outputs"

SEQ_LEN = 128        # every run resampled onto a common progress axis
MRR_MAX = 200.0      # same corrupt-label bound as the hand-crafted track
EPOCHS = 300
PATIENCE = 30
SEEDS = [0, 1, 2]    # small data -> report spread, not one lucky run

DEVICE = torch.device("cpu")


def resample_runs(name: str) -> tuple[np.ndarray, pd.DataFrame]:
    """-> (X of shape [n_runs, n_channels, SEQ_LEN], run table).

    Runs differ in length, so each channel is linearly interpolated onto a
    normalised 0..1 progress axis. That keeps shape information while making
    every run the same size — the sequence-model analogue of the
    length-invariance the hand-crafted features get for free.
    """
    cached = CACHE / f"{name}_seq{SEQ_LEN}.npz"
    runs_path = CACHE / f"{name}_runs.parquet"
    if cached.exists():
        z = np.load(cached, allow_pickle=True)
        return z["X"], pd.read_parquet(runs_path)

    traj = pd.read_parquet(CACHE / f"{name}_trajectories.parquet")
    runs = pd.read_parquet(runs_path)
    grid = np.linspace(0, 1, SEQ_LEN)
    out = np.zeros((len(runs), len(SIGNAL_COLS), SEQ_LEN), dtype=np.float32)
    index = {tuple(k): i for i, k in enumerate(runs[KEY].to_numpy())}

    for key, g in traj.groupby(KEY, sort=False):
        i = index.get(tuple(key))
        if i is None:
            continue
        g = g.sort_values("TIMESTAMP")
        src = np.linspace(0, 1, len(g))
        for c, col in enumerate(SIGNAL_COLS):
            y = g[col].to_numpy(dtype=np.float64)
            if not np.isfinite(y).any():
                continue
            y = np.nan_to_num(y, nan=np.nanmedian(y[np.isfinite(y)]))
            out[i, c] = np.interp(grid, src, y).astype(np.float32)

    np.savez_compressed(cached, X=out)
    print(f"  cached {cached.name} {out.shape}")
    return out, runs


def static_features(name: str) -> pd.DataFrame:
    """The non-trajectory inputs, identical to what the tree model receives."""
    feats = pd.read_parquet(CACHE / f"{name}_features.parquet")
    cols = ([f"{c}__mean" for c in WEAR_COLS] + CONTEXT_COLS
            + ["STAGE_is_B", "n_samples", "duration", "sample_rate"])
    return feats[KEY + [c for c in cols if c in feats.columns]]


class CNN1D(nn.Module):
    """Three dilated conv blocks, then global average+max pooling."""

    def __init__(self, n_ch: int, n_static: int, width: int = 64):
        super().__init__()
        def block(i, o, d):
            return nn.Sequential(
                nn.Conv1d(i, o, kernel_size=5, padding=2 * d, dilation=d),
                nn.BatchNorm1d(o), nn.ReLU(),
                nn.Conv1d(o, o, kernel_size=5, padding=2 * d, dilation=d),
                nn.BatchNorm1d(o), nn.ReLU(), nn.MaxPool1d(2))
        self.conv = nn.Sequential(block(n_ch, width, 1),
                                  block(width, width, 2),
                                  block(width, width, 4))
        self.head = nn.Sequential(
            nn.Linear(2 * width + n_static, 128), nn.ReLU(), nn.Dropout(0.3),
            nn.Linear(128, 1))

    def forward(self, x, s):
        h = self.conv(x)
        h = torch.cat([h.mean(dim=2), h.amax(dim=2), s], dim=1)
        return self.head(h).squeeze(1)


class GRUNet(nn.Module):
    """Bidirectional GRU over the resampled trajectory."""

    def __init__(self, n_ch: int, n_static: int, hidden: int = 96):
        super().__init__()
        self.rnn = nn.GRU(n_ch, hidden, num_layers=2, batch_first=True,
                          bidirectional=True, dropout=0.2)
        self.head = nn.Sequential(
            nn.Linear(2 * hidden + n_static, 128), nn.ReLU(), nn.Dropout(0.3),
            nn.Linear(128, 1))

    def forward(self, x, s):
        h, _ = self.rnn(x.transpose(1, 2))
        h = torch.cat([h.mean(dim=1), s], dim=1)
        return self.head(h).squeeze(1)


def train_one(model, Xf, Sf, yf, Xv, Sv, yv, seed: int) -> tuple[nn.Module, int]:
    torch.manual_seed(seed)
    opt = torch.optim.AdamW(model.parameters(), lr=2e-3, weight_decay=1e-4)
    sched = torch.optim.lr_scheduler.ReduceLROnPlateau(opt, factor=0.5, patience=10)
    loss_fn = nn.MSELoss()
    n = len(yf)
    best, best_state, bad, best_epoch = np.inf, None, 0, 0
    g = torch.Generator().manual_seed(seed)

    for epoch in range(EPOCHS):
        model.train()
        perm = torch.randperm(n, generator=g)
        for i in range(0, n, 64):
            idx = perm[i:i + 64]
            opt.zero_grad()
            loss = loss_fn(model(Xf[idx], Sf[idx]), yf[idx])
            loss.backward()
            nn.utils.clip_grad_norm_(model.parameters(), 5.0)
            opt.step()
        model.eval()
        with torch.no_grad():
            v = float(loss_fn(model(Xv, Sv), yv))
        sched.step(v)
        if v < best - 1e-4:
            best, bad, best_epoch = v, 0, epoch
            best_state = {k: t.clone() for k, t in model.state_dict().items()}
        else:
            bad += 1
            if bad >= PATIENCE:
                break
    model.load_state_dict(best_state)
    return model, best_epoch


def scores(y_true, y_pred) -> dict:
    return {"MSE": mean_squared_error(y_true, y_pred),
            "RMSE": float(np.sqrt(mean_squared_error(y_true, y_pred))),
            "MAE": mean_absolute_error(y_true, y_pred),
            "R2": r2_score(y_true, y_pred)}


def main():
    FIGS.mkdir(exist_ok=True)
    OUT.mkdir(exist_ok=True)

    print("resampling trajectories ...", flush=True)
    Xtr_all, runs_tr = resample_runs("training")
    Xte_all, runs_te = resample_runs("test")
    Str_all, Ste_all = static_features("training"), static_features("test")

    tr = runs_tr.merge(Str_all, on=KEY, how="left")
    te = runs_te.merge(Ste_all, on=KEY, how="left")
    keep = (tr["AVG_REMOVAL_RATE"] <= MRR_MAX).to_numpy()
    Xtr_all, tr = Xtr_all[keep], tr[keep].reset_index(drop=True)
    print(f"training {Xtr_all.shape}, test {Xte_all.shape}")

    stat_cols = [c for c in tr.columns
                 if c not in set(KEY) | {"AVG_REMOVAL_RATE", "t_start", "t_end"}
                 and tr[c].dtype != object]
    stat_cols = [c for c in stat_cols if c in te.columns]

    tr_idx, va_idx = splits.group_temporal_split(tr, "WAFER_ID", "t_start",
                                                 train_frac=0.8)
    fit_m = tr.index.isin(tr_idx)
    val_m = tr.index.isin(va_idx)
    assert set(tr.loc[fit_m, "WAFER_ID"]).isdisjoint(set(tr.loc[val_m, "WAFER_ID"]))
    print(f"internal split: fit {fit_m.sum():,} / val {val_m.sum():,} runs")

    # standardise using the fit window only, per channel and per static column
    mu = Xtr_all[fit_m].mean(axis=(0, 2), keepdims=True)
    sd = Xtr_all[fit_m].std(axis=(0, 2), keepdims=True)
    sd[sd < 1e-6] = 1.0
    norm = lambda A: ((A - mu) / sd).astype(np.float32)

    S_all = tr[stat_cols].to_numpy(dtype=np.float64)
    S_te_all = te[stat_cols].to_numpy(dtype=np.float64)
    smu, ssd = np.nanmean(S_all[fit_m], 0), np.nanstd(S_all[fit_m], 0)
    ssd[ssd < 1e-6] = 1.0
    snorm = lambda A: np.nan_to_num((A - smu) / ssd).astype(np.float32)

    T = lambda a: torch.tensor(a, device=DEVICE)
    Xf, Xv, Xt = T(norm(Xtr_all[fit_m])), T(norm(Xtr_all[val_m])), T(norm(Xte_all))
    Sf, Sv, St = T(snorm(S_all[fit_m])), T(snorm(S_all[val_m])), T(snorm(S_te_all))
    y = tr["AVG_REMOVAL_RATE"].to_numpy(dtype=np.float32)
    ymu, ysd = float(y[fit_m].mean()), float(y[fit_m].std())
    yf = T((y[fit_m] - ymu) / ysd)
    yv = T((y[val_m] - ymu) / ysd)
    y_val_raw, y_te_raw = y[val_m], te["AVG_REMOVAL_RATE"].to_numpy(dtype=np.float32)

    rows, preds = [], {}
    for name, ctor in [("1D-CNN", CNN1D), ("BiGRU", GRUNet)]:
        val_p, te_p, epochs = [], [], []
        for seed in SEEDS:
            torch.manual_seed(seed)
            t0 = time.time()
            model = ctor(len(SIGNAL_COLS), len(stat_cols)).to(DEVICE)
            model, ep = train_one(model, Xf, Sf, yf, Xv, Sv, yv, seed)
            model.eval()
            with torch.no_grad():
                val_p.append(model(Xv, Sv).cpu().numpy() * ysd + ymu)
                te_p.append(model(Xt, St).cpu().numpy() * ysd + ymu)
            epochs.append(ep)
            print(f"  {name} seed {seed}: best epoch {ep}, "
                  f"{time.time() - t0:.0f}s, "
                  f"test MSE {mean_squared_error(y_te_raw, te_p[-1]):.2f}", flush=True)
        # seed ensemble is the honest single number to compare against XGBoost
        te_mean = np.mean(te_p, axis=0)
        per_seed = [mean_squared_error(y_te_raw, p) for p in te_p]
        rec = {"model": f"{name} (ensemble of {len(SEEDS)} seeds)"}
        rec.update({f"val_{k}": v for k, v in
                    scores(y_val_raw, np.mean(val_p, axis=0)).items()})
        rec.update({f"test_{k}": v for k, v in scores(y_te_raw, te_mean).items()})
        rec["test_MSE_seed_spread"] = f"{min(per_seed):.2f}-{max(per_seed):.2f}"
        rec["mean_best_epoch"] = float(np.mean(epochs))
        rows.append(rec)
        preds[name] = te_mean

    seq = pd.DataFrame(rows).set_index("model")
    print("\n=== sequence models on the official test set ===")
    print(seq.round(3).to_string())

    # --- fair comparison against the hand-crafted track
    hand = pd.read_csv(OUT / "model_comparison.csv").set_index("model")
    combined = pd.concat([
        hand[["val_MSE", "val_MAE", "val_R2", "test_MSE", "test_MAE", "test_R2"]],
        seq[["val_MSE", "val_MAE", "val_R2", "test_MSE", "test_MAE", "test_R2"]],
    ])
    combined["track"] = ["baseline"] * 2 + ["hand-crafted features"] * 3 + \
                        ["learned representation"] * 2
    combined = combined.sort_values("test_MSE")
    combined.to_csv(OUT / "track_comparison.csv")
    print("\n=== both tracks, one split, one metric ===")
    print(combined.round(3).to_string())

    fig, axes = plt.subplots(1, 2, figsize=(14, 5))
    ax = axes[0]
    colors = {"baseline": "#bbbbbb", "hand-crafted features": "steelblue",
              "learned representation": "darkorange"}
    d = combined.iloc[::-1]
    ax.barh(d.index, d["test_MSE"], color=[colors[t] for t in d["track"]],
            edgecolor="black")
    ax.set_xscale("log")
    for i, v in enumerate(d["test_MSE"]):
        ax.text(v * 1.12, i, f"{v:.1f}", va="center", fontsize=8)
    ax.set_xlim(right=d["test_MSE"].max() * 3)
    ax.set(xlabel="test MSE, log scale (lower is better)",
           title="Hand-crafted vs learned, same split and metric")
    ax.tick_params(labelsize=8)
    handles = [plt.Rectangle((0, 0), 1, 1, color=c) for c in colors.values()]
    ax.legend(handles, colors.keys(), fontsize=8, loc="lower right")

    ax = axes[1]
    best_seq = min(preds, key=lambda k: mean_squared_error(y_te_raw, preds[k]))
    ax.scatter(y_te_raw, preds[best_seq], s=16, alpha=0.6, color="darkorange",
               edgecolors="black", linewidths=0.3, label=best_seq)
    lo, hi = y_te_raw.min() - 5, y_te_raw.max() + 5
    ax.plot([lo, hi], [lo, hi], ls="--", color="gray", lw=1)
    ax.set(xlim=(lo, hi), ylim=(lo, hi), xlabel="actual MRR",
           ylabel="predicted MRR", title=f"{best_seq} on the official test set")
    ax.legend(fontsize=8)
    fig.tight_layout()
    fig.savefig(FIGS / "track_comparison.png", dpi=150)
    plt.close(fig)


if __name__ == "__main__":
    main()