"""P7 inference: pairwise contrasts between saved models (DM-HAC + day-cluster bootstrap, Holm/BH), subgroup
analysis of a contrast, and stress-day calibration (H3). Works from per_origin_*.npz + origins_<tag>.csv.
Usage: pairwise.py --exp main --tag test --pairs mrpt_small:mrpt_small_1m_only,mrpt_small:har_rv,... [--subgroups mrpt_small:har_rv]"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import numpy as np
import pandas as pd

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))
from chronos_x import pipeline as P  # noqa: E402
from chronos_x.eval.statistics import holm, benjamini_hochberg, cluster_bootstrap_mean  # noqa: E402
from chronos_x.eval.subgroups import add_subgroups  # noqa: E402
from chronos_x.eval.calibration import pit_histogram  # noqa: E402


def load(exp: str, name: str, tag: str) -> dict:
    return np.load(ROOT / "results" / "metrics" / exp / f"per_origin_{name}_{tag}.npz")


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--exp", default="main"); ap.add_argument("--tag", default="test"); ap.add_argument("--pairs", required=True)
    ap.add_argument("--subgroups", default=None); ap.add_argument("--stride", type=int, default=15)
    ap.add_argument("--metrics", default="crps,wis,qlike"); ap.add_argument("--horizons", default="5,15,30,60")
    ap.add_argument("--universe", default="data/manifests/universe_main.json"); ap.add_argument("--out-suffix", default="")
    a = ap.parse_args()
    d = ROOT / "results" / "metrics" / a.exp
    tab = pd.read_csv(d / f"origins_{a.tag}.csv")
    man = json.loads((ROOT / a.universe).read_text())
    meta = add_subgroups(tab, man)
    H = [int(h) for h in a.horizons.split(",")]
    rows = []
    cache = {}
    def per(name):
        if name not in cache:
            cache[name] = load(a.exp, name, a.tag)
        return cache[name]
    for pair in a.pairs.split(","):
        m, ref = pair.split(":")
        for metric in a.metrics.split(","):
            for h in H:
                r = P.compare(per(m)[f"{metric}_{h}"], per(ref)[f"{metric}_{h}"], tab, h, a.stride)
                rows.append({"model": m, "reference": ref, "metric": metric, **r})
    df = pd.DataFrame(rows)
    for metric in a.metrics.split(","):
        sel = df["metric"] == metric
        df.loc[sel, "p_holm"] = holm(df.loc[sel, "dm_p"].tolist()); df.loc[sel, "p_bh"] = benjamini_hochberg(df.loc[sel, "dm_p"].tolist())
    out = d / f"pairwise_{a.tag}{a.out_suffix}.csv"; df.to_csv(out, index=False)
    print(df[df.metric == "crps"][["model", "reference", "H", "rel_diff", "lo", "hi", "dm_t", "dm_p", "p_holm"]].to_string(index=False, float_format=lambda x: f"{x:.4g}"))
    # subgroups for one contrast (relative CRPS difference with day-cluster CI inside each group)
    if a.subgroups:
        m, ref = a.subgroups.split(":")
        srows = []
        for col in ("asset_class", "adv_tercile", "vol_tercile", "group", "tod", "year", "stress_day"):
            for g, idx in meta.groupby(col, observed=True).indices.items():
                for h in H:
                    pa_, pb_ = per(m)[f"crps_{h}"][idx], per(ref)[f"crps_{h}"][idx]
                    cl = pd.factorize(meta["session_date"].values[idx])[0]
                    cb = cluster_bootstrap_mean(pa_ - pb_, cl, n_boot=500)
                    base = pb_.mean()
                    srows.append({"col": col, "group": str(g), "H": h, "n": int(len(idx)), "rel": cb["mean"] / base, "lo": cb["lo"] / base, "hi": cb["hi"] / base,
                                  "crps_model": pa_.mean(), "crps_ref": base})
        sdf = pd.DataFrame(srows); sdf.to_csv(d / f"subgroups_{m}_vs_{ref}_{a.tag}.csv", index=False)
        print(sdf[sdf.H == 60].to_string(index=False, float_format=lambda x: f"{x:.4g}"))
    # stress-day calibration (H3): coverage and PIT dispersion on stress vs normal days for every model in the pairs
    names = sorted({x for pair in a.pairs.split(",") for x in pair.split(":")})
    crows = []
    for name in names:
        p = per(name)
        for h in H:
            for lab, idx in (("stress", np.flatnonzero(meta["stress_day"].values)), ("normal", np.flatnonzero(~meta["stress_day"].values))):
                u = p[f"pit_{h}"][idx]
                ph = pit_histogram(u)
                crows.append({"model": name, "H": h, "days": lab, "n": int(len(idx)), "cov_50": float(((u >= .25) & (u <= .75)).mean()),
                              "cov_80": float(((u >= .1) & (u <= .9)).mean()), "cov_95": float(((u >= .025) & (u <= .975)).mean()),
                              "pit_var": ph["var"], "pit_ks": ph["ks_stat"], "crps": float(p[f"crps_{h}"][idx].mean()),
                              "crps_scaled": float(p[f"crps_scaled_{h}"][idx].mean())})
    cdf = pd.DataFrame(crows); cdf.to_csv(d / f"stress_calibration_{a.tag}{a.out_suffix}.csv", index=False)
    print(cdf[cdf.H == 60].to_string(index=False, float_format=lambda x: f"{x:.4g}"))


if __name__ == "__main__":
    main()