"""SPC baseline for Bosch — and the structural reason it cannot work here.

Spec rule: establish the control chart a process engineer would reach for
before claiming ML is needed. On Bosch that runs into an assumption failure
worth stating out loud.

Hotelling T² compares a unit's measurement vector against an in-control mean
and covariance. That presumes **every unit is measured the same way**. Bosch
products take 7,927 distinct routes, so most measurement slots are empty for
most products — not missing readings, but stations the product never visited.
Imputing them invents in-control behaviour that never happened.

So the chart is run twice:
  (1) naively, across all products with median imputation — what you get if
      you ignore the routing structure;
  (2) restricted to the single most common route family, where every product
      really does share a measurement vector and T² is legitimate.

Run: python 03_bosch/spc_baseline.py   (needs prepare.py + build_graph.py)
"""
from __future__ import annotations

import sys
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from sklearn.decomposition import PCA

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

from common import metrics, spc  # noqa: E402

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

COST_FN, COST_FP = 50.0, 1.0
ALPHA = 0.01
N_COMPONENTS = 0.80


def temporal_masks(times: pd.DataFrame, stations: list[str]):
    T = times[stations].to_numpy(dtype=np.float32)
    visited = ~np.isnan(T)
    with np.errstate(all="ignore"):
        start = np.where(visited.any(axis=1),
                         np.nanmin(np.nan_to_num(T, nan=np.inf), axis=1), np.nan)
    known = ~np.isnan(start)
    order = np.argsort(np.where(known, start, np.inf), kind="stable")[:known.sum()]
    n = len(order)
    train = np.zeros(len(times), bool); train[order[:int(0.70 * n)]] = True
    test = np.zeros(len(times), bool); test[order[int(0.70 * n):]] = True
    return train, test, visited


def run_chart(label: str, X: np.ndarray, y: np.ndarray,
              train_m: np.ndarray, test_m: np.ndarray,
              max_fit: int = 120_000, seed: int = 0) -> dict:
    """Fit T² on in-control (passing) training units, score the test period.

    The in-control set runs to ~800k rows against 150 columns; the control
    limits are estimated from a random 120k of them, which is far more than
    enough for a 150-dimensional mean and covariance and keeps the SVD cheap.
    """
    in_control = X[train_m & (y == 0)]
    if len(in_control) > max_fit:
        rng = np.random.default_rng(seed)
        in_control = in_control[rng.choice(len(in_control), max_fit, replace=False)]
    pca = PCA(n_components=N_COMPONENTS, svd_solver="full").fit(in_control)
    chart = spc.HotellingT2().fit(pca.transform(in_control))
    t2 = chart.score(pca.transform(X[test_m]))
    ucl = chart.ucl(alpha=ALPHA)
    y_te = y[test_m]
    pred = (t2 > ucl).astype(int)
    tp = int(((pred == 1) & (y_te == 1)).sum())
    fp = int(((pred == 1) & (y_te == 0)).sum())
    fn = int(((pred == 0) & (y_te == 1)).sum())
    return {
        "baseline": label,
        "n_test": int(test_m.sum()),
        "n_components": int(pca.n_components_),
        "recall@UCL": tp / max(1, tp + fn),
        "false_alarm_rate": fp / max(1, (y_te == 0).sum()),
        "pr_auc": metrics.pr_auc(y_te, t2),
        # a random ranker scores PR-AUC = prevalence; without it the number
        # above cannot be read
        "pr_auc_random": float(y_te.mean()),
        "capture@20%": metrics.decile_capture(y_te, t2, 0.2),
        "expected_cost": metrics.expected_cost(y_te, pred, COST_FN, COST_FP),
        "_t2": t2, "_ucl": float(ucl), "_y": y_te,
    }


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

    times = pd.read_parquet(CACHE / "station_times.parquet")
    num = pd.read_parquet(CACHE / "station_numeric.parquet")
    stations = [c for c in times.columns if c != "Id"]
    y = num["Response"].to_numpy().astype(np.int64)
    train_m, test_m, visited = temporal_masks(times, stations)
    print(f"{len(y):,} products; train {train_m.sum():,} / test {test_m.sum():,}")

    feat_cols = [c for c in num.columns if c.endswith(("__mean", "__max", "__min"))]
    X = num[feat_cols].to_numpy(dtype=np.float32)

    results = []

    # (1) naive: impute everything with the training median, ignore routing
    med = np.nanmedian(X[train_m], axis=0)
    Xi = np.where(np.isnan(X), med, X)
    sd = Xi[train_m].std(axis=0); sd[sd < 1e-9] = 1.0
    Xi = (Xi - Xi[train_m].mean(axis=0)) / sd
    results.append(run_chart("T² over all products (median-imputed)",
                             Xi, y, train_m, test_m))

    # (2) restricted: the largest route family, where a common measurement
    # vector genuinely exists and nothing has to be invented
    pid = np.load(CACHE / "path_ids.npy")
    biggest = pd.Series(pid[train_m]).value_counts().idxmax()
    fam = pid == biggest
    fam_stations = np.flatnonzero(visited[fam].mean(axis=0) > 0.99)
    fam_cols = [i for i, c in enumerate(feat_cols)
                if c.split("__")[0] in {stations[j] for j in fam_stations}]
    Xf = X[:, fam_cols]
    complete = ~np.isnan(Xf).any(axis=1)
    sel = fam & complete
    print(f"largest route family: {fam.sum():,} products "
          f"({sel.sum():,} with complete readings), "
          f"{len(fam_cols)} measurement columns, "
          f"fail rate {y[sel].mean():.3%}")
    sd = Xf[sel & train_m].std(axis=0); sd[sd < 1e-9] = 1.0
    Xf = (Xf - Xf[sel & train_m].mean(axis=0)) / sd
    results.append(run_chart("T² within the largest route family",
                             Xf, y, sel & train_m, sel & test_m))

    table = pd.DataFrame([{k: v for k, v in r.items() if not k.startswith("_")}
                          for r in results]).set_index("baseline")
    table.to_csv(OUT / "spc_baseline.csv")
    print("\n=== SPC baselines (test period) ===")
    print(table.round(4).to_string())

    ml = pd.read_csv(OUT / "model_comparison.csv").set_index("model")
    print(f"\nfor comparison, XGBoost model A: PR-AUC {ml.loc[ml.index[0], 'pr_auc']:.4f}, "
          f"capture@20% {ml.loc[ml.index[0], 'capture@20%']:.4f}")

    fig, axes = plt.subplots(1, 2, figsize=(14, 4.6))
    for ax, r in zip(axes, results):
        t2, ucl, y_te = r["_t2"], r["_ucl"], r["_y"]
        step = max(1, len(t2) // 6000)      # thin for legibility, keep all fails
        idx = np.arange(0, len(t2), step)
        ax.plot(idx, np.clip(t2[idx], 1e-3, None), lw=0.4, color="steelblue",
                alpha=0.8, label="T² (test period)")
        f = np.flatnonzero(y_te == 1)
        ax.scatter(f, np.clip(t2[f], 1e-3, None), s=7, color="crimson",
                   zorder=3, label=f"true failures (n={len(f)})")
        ax.axhline(ucl, color="darkorange", ls="--", lw=1.2,
                   label=f"UCL (α={ALPHA})")
        ax.set(yscale="log", xlabel="product (time order)", ylabel="Hotelling T²",
               title=f"{r['baseline']}\nrecall {r['recall@UCL']:.1%} at "
                     f"{r['false_alarm_rate']:.1%} false alarms")
        ax.legend(fontsize=7.5)
    fig.tight_layout()
    fig.savefig(FIGS / "spc_baseline.png", dpi=150)
    plt.close(fig)


if __name__ == "__main__":
    main()