"""Why do the path features rank first by gain yet not improve ranking?

XGBoost puts `path_risk` at the top of the importance table, but adding the
graph block leaves PR-AUC lower than the station-only model. The suspicion is
that a path family's historical fail rate is not stable over time — that the
encoding memorises the training window rather than describing the process.

This script measures that stability directly: for path families with enough
products in both the training and the test window, how well does the training
fail rate predict the test fail rate?

Run: python 03_bosch/diagnose_encoding.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 scipy import stats

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

from common import drift  # noqa: E402

CACHE = HERE / "cache"
FIGS = HERE / "figures"
OUT = HERE / "outputs"
MIN_N = 200  # families smaller than this are too noisy to compare


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

    times = pd.read_parquet(CACHE / "station_times.parquet")
    stations = [c for c in times.columns if c != "Id"]
    T = times[stations].to_numpy(dtype=np.float32)
    visited = ~np.isnan(T)
    y = pd.read_parquet(CACHE / "station_numeric.parquet",
                        columns=["Response"])["Response"].to_numpy()
    pid = np.load(CACHE / "path_ids.npy")

    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_idx = order[:int(0.70 * n)]
    test_idx = order[int(0.70 * n):]

    def rates(idx):
        d = pd.DataFrame({"pid": pid[idx], "y": y[idx]})
        g = d.groupby("pid")["y"].agg(["size", "mean"])
        return g.rename(columns={"size": "n", "mean": "rate"})

    tr, te = rates(train_idx), rates(test_idx)
    both = tr.join(te, how="inner", lsuffix="_train", rsuffix="_test")
    big = both[(both["n_train"] >= MIN_N) & (both["n_test"] >= MIN_N)]

    r_p = stats.pearsonr(big["rate_train"], big["rate_test"])
    r_s = stats.spearmanr(big["rate_train"], big["rate_test"])
    coverage = te.loc[te.index.isin(tr.index), "n"].sum() / te["n"].sum()

    print(f"path families: {len(tr):,} in train, {len(te):,} in test, "
          f"{len(both):,} in both")
    print(f"test products whose path family was seen in training: {coverage:.2%}")
    print(f"families with >= {MIN_N} products in both windows: {len(big):,}")
    print(f"train-vs-test fail-rate correlation: "
          f"Pearson r = {r_p.statistic:.3f} (p={r_p.pvalue:.3g}), "
          f"Spearman rho = {r_s.statistic:.3f} (p={r_s.pvalue:.3g})")

    overall_tr = y[train_idx].mean()
    overall_te = y[test_idx].mean()
    print(f"overall fail rate: train {overall_tr:.3%} -> test {overall_te:.3%} "
          f"({overall_te / overall_tr:.2f}x)")

    fig, ax = plt.subplots(figsize=(6.5, 6))
    ax.scatter(big["rate_train"] * 100, big["rate_test"] * 100,
               s=np.clip(big["n_test"] / 60, 8, 260), alpha=0.6,
               color="steelblue", edgecolors="black", linewidths=0.4)
    lim = max(big["rate_train"].max(), big["rate_test"].max()) * 100 * 1.1
    ax.plot([0, lim], [0, lim], ls="--", color="gray", lw=1,
            label="perfect stability")
    ax.axhline(overall_te * 100, color="crimson", ls=":", lw=1,
               label=f"test overall ({overall_te:.2%})")
    ax.axvline(overall_tr * 100, color="darkorange", ls=":", lw=1,
               label=f"train overall ({overall_tr:.2%})")
    ax.set(xlim=(0, lim), ylim=(0, lim),
           xlabel="path family fail rate, training window (%)",
           ylabel="path family fail rate, test window (%)",
           title="Does a path's historical fail rate carry forward?\n"
                 f"Pearson r = {r_p.statistic:.2f} across {len(big)} families "
                 f"(≥{MIN_N} products each)")
    ax.legend(fontsize=8)
    fig.tight_layout()
    fig.savefig(FIGS / "path_risk_stability.png", dpi=150)
    plt.close(fig)

    # --- feature-level drift, using the same PSI machinery as SECOM
    num = pd.read_parquet(CACHE / "station_numeric.parquet")
    feat_cols = [c for c in num.columns if c.endswith("__mean")]
    psi_df = drift.psi_report(num.iloc[train_idx][feat_cols],
                              num.iloc[test_idx][feat_cols])
    trig = drift.drift_trigger(psi_df)
    print(f"\nfeature drift (PSI, {len(feat_cols)} station means): "
          f"median {psi_df['psi'].median():.4f}, "
          f"{trig['n_action']}/{trig['n_features']} in the action band "
          f"({trig['frac_action']:.1%}) -> retrain trigger: {trig['retrain']}")
    print("  most drifted:", ", ".join(psi_df.head(5)["feature"]))
    psi_df.to_csv(OUT / "feature_drift_psi.csv", index=False)

    big.reset_index().to_csv(OUT / "path_risk_stability.csv", index=False)
    pd.DataFrame([{
        "n_families_train": len(tr), "n_families_test": len(te),
        "n_families_both": len(both), "n_families_compared": len(big),
        "test_coverage_by_seen_paths": coverage,
        "pearson_r": r_p.statistic, "pearson_p": r_p.pvalue,
        "spearman_rho": r_s.statistic, "spearman_p": r_s.pvalue,
        "fail_rate_train": overall_tr, "fail_rate_test": overall_te,
    }]).to_csv(OUT / "path_risk_stability_summary.csv", index=False)


if __name__ == "__main__":
    main()