#!/usr/bin/env python3
"""Stage 2 of the evaluation: per-instrument and pooled statistics from the wide loss tables.
usage: analyze.py --run main [--period test|validation] [--symbols ...] [--reps 2000]
Outputs (results/metrics/analysis_<run>_<period>/): per_symbol/<sym>.json (pairwise ratios, DM, coverage, MCS, SPA, regimes,
tail risk, failure accounting), pooled.csv/json, mcs_counts.csv, tailrisk.csv, compute.csv, and results/tables/*.csv copies.
Everything is computed on the requested period only (the test period is used only after the protocol freeze)."""
from __future__ import annotations

import argparse
import glob
import json
import sys
from pathlib import Path

import numpy as np
import pandas as pd
import yaml

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))
from wp46.diagnostics import backtest_table, var_es_forecasts  # noqa: E402
from wp46.inference import block_bootstrap_mean_ci, dm_test, mcs, spa  # noqa: E402
from wp46.preprocessing import load_panel  # noqa: E402

PROT = yaml.safe_load(open(ROOT / "configs" / "forecast_protocol.yaml"))
REG = yaml.safe_load(open(ROOT / "configs" / "model_registry.yaml"))["models"]
INST = yaml.safe_load(open(ROOT / "configs" / "instruments.yaml"))
CLASS_OF = {s: c for c, syms in (INST.get("universe") or INST["candidates"]).items() for s in syms}
PRIMARY_PROXY, PRIMARY_LOSS = PROT["primary_proxy"], PROT["primary_loss"]
HORIZONS = tuple(PROT["horizons"])
MIN_COVERAGE_MCS = 0.90


def spec_meta(spec_id: str) -> dict:
    mid, dist, rest = spec_id.split("__")
    m = REG.get(mid, {})
    return {"model_id": mid, "dist": dist, "family": m.get("family", "?"), "tier": m.get("tier", "?"), "variant": rest}


def bench_id(specs) -> str:
    b = [s for s in specs if s.startswith("garch11__normal__")]
    if not b:
        raise RuntimeError("benchmark GARCH(1,1)-normal missing")
    return b[0]


def load_symbol(run: str, sym: str):
    d = ROOT / "results" / "metrics" / f"losses_{run}"
    W = pd.read_parquet(d / f"{sym}.parquet"); W["origin"] = pd.to_datetime(W["origin"])
    T = pd.read_parquet(d / f"targets_{sym}.parquet"); T["origin"] = pd.to_datetime(T["origin"])
    S = pd.read_parquet(d / f"status_{sym}.parquet"); S["origin"] = pd.to_datetime(S["origin"])
    G = pd.read_parquet(d / f"fits_{sym}.parquet") if (d / f"fits_{sym}.parquet").exists() else pd.DataFrame()
    return W, T.set_index("origin"), S.set_index("origin"), G


def block(W: pd.DataFrame, kind: str, proxy: str = "", loss: str = "", h: int = 1) -> pd.DataFrame:
    m = (W["kind"] == kind) & (W["proxy"] == proxy) & (W["loss"] == loss) & (W["h"] == h)
    return W[m].drop(columns=["kind", "proxy", "loss", "h", "symbol"]).set_index("origin").sort_index()


def regimes_for(panel: pd.DataFrame, origins: pd.DatetimeIndex, val_start: str, val_end: str) -> pd.DataFrame:
    p = panel.set_index(pd.to_datetime(panel["session"]))
    rv = p["rv_5m_ss"].where(p["valid_rm"]).fillna(p["r"] ** 2)
    trail22 = rv.rolling(22, min_periods=15).mean()
    trail66 = rv.rolling(66, min_periods=40).mean()
    v = trail22[(trail22.index >= val_start) & (trail22.index <= val_end)]
    q = v.quantile([1 / 3, 2 / 3]).values if len(v) > 50 else trail22.quantile([1 / 3, 2 / 3]).values
    v66 = trail66[(trail66.index >= val_start) & (trail66.index <= val_end)]
    q66 = v66.quantile([1 / 3, 2 / 3]).values if len(v66) > 50 else trail66.quantile([1 / 3, 2 / 3]).values
    R = pd.DataFrame(index=origins)
    t22 = trail22.reindex(origins); t66 = trail66.reindex(origins)
    R["vol_tercile"] = np.where(t22 <= q[0], "low", np.where(t22 <= q[1], "mid", "high"))
    R["vol_tercile66"] = np.where(t66 <= q66[0], "low", np.where(t66 <= q66[1], "mid", "high"))
    R["prev_sign"] = np.where(p["r"].reindex(origins) < 0, "neg", "pos")
    ep = PROT["regimes"]["calendar_episodes"]
    R["episode"] = "none"
    for name, w in ep.items():
        R.loc[(origins >= w["start"]) & (origins <= w["end"]), "episode"] = name
    return R


def analyze_symbol(run: str, sym: str, period: str, reps: int) -> dict:
    W, T, S, G = load_symbol(run, sym)
    panel = load_panel(sym)
    per = PROT["periods"][period]
    lo, hi = pd.Timestamp(per["start"]), pd.Timestamp(per["end"])
    out = {"symbol": sym, "asset_class": CLASS_OF.get(sym), "period": period, "horizons": {}, "regimes": {}, "secondary": {}}
    specs = None
    for h in HORIZONS:
        L = block(W, "loss", PRIMARY_PROXY, PRIMARY_LOSS, h)
        L = L[(L.index >= lo) & (L.index <= hi)]
        specs = list(L.columns)
        b = bench_id(specs)
        y = T[f"__target__{PRIMARY_PROXY}_{h}"].reindex(L.index)
        tv = np.isfinite(y.values)
        n_target = int(tv.sum())
        res = {"n_origins": int(len(L)), "n_valid_targets": n_target, "benchmark": b, "specs": {}}
        # explosive-forecast diagnostic (A-06): cumulative forecast above 100x the trailing 22-day mean of the proxy at the origin
        Fh = block(W, "forecast", "", "", h).reindex(L.index)
        y1 = T[f"__target__{PRIMARY_PROXY}_1"].reindex(L.index)
        trail = pd.Series(y1.values, index=L.index).shift(1).rolling(22, min_periods=10).mean().values * h
        for s in specs:
            pair = np.isfinite(L[s].values) & np.isfinite(L[b].values)
            cov = float(np.isfinite(L[s].values[tv]).mean()) if n_target else np.nan
            fx = Fh[s].values if s in Fh else np.full(len(L), np.nan)
            okf = np.isfinite(fx) & np.isfinite(trail)
            expl = float(np.mean(fx[okf] > 100.0 * trail[okf])) if okf.sum() else np.nan
            if pair.sum() >= 30:
                ls, lb = L[s].values[pair], L[b].values[pair]
                dm = dm_test(ls, lb, h)
                k = max(1, int(0.01 * len(ls)))     # 1 % trimmed means (both series trimmed on their own order statistics)
                lt, bt = np.sort(ls)[k:-k], np.sort(lb)[k:-k]
                res["specs"][s] = {"n_pair": int(pair.sum()), "coverage": cov, "mean": float(ls.mean()), "mean_bench": float(lb.mean()),
                                   "ratio": float(ls.mean() / lb.mean()), "dm_stat": dm["stat"], "dm_p": dm["p"],
                                   "ratio_trim1": float(lt.mean() / bt.mean()), "ratio_median": float(np.median(ls) / np.median(lb)),
                                   "explosive_share": expl,
                                   "own_mean": float(np.nanmean(L[s].values)), "own_n": int(np.isfinite(L[s].values).sum())}
            else:
                res["specs"][s] = {"n_pair": int(pair.sum()), "coverage": cov, "ratio": np.nan, "dm_stat": np.nan, "dm_p": np.nan}
        # MCS / SPA per tier on the common-success sample of well-covered specs
        for tier in ("A", "B"):
            cand = [s for s in specs if spec_meta(s)["tier"] == tier and res["specs"][s]["coverage"] >= MIN_COVERAGE_MCS
                    and REG.get(s.split("__")[0], {}).get("status") == "implemented"]   # pre-registered set only
            if tier == "B":
                cand = cand + ([b] if b not in cand else [])   # GARCH(1,1) reference inside the realized tier (labelled)
            if len(cand) < 2:
                continue
            M = L[cand].dropna()
            res[f"common_sample_{tier}"] = {"n": int(len(M)), "specs": cand, "mean": {s: float(M[s].mean()) for s in cand},
                                            "ratio": {s: float(M[s].mean() / M[b].mean()) for s in cand} if b in M else {}}
            if len(M) >= 100:
                try:
                    mc = mcs(M, size=PROT["inference"]["mcs"]["alpha"], reps=reps, block=PROT["inference"]["mcs"]["block_length"])
                    res[f"mcs_{tier}"] = {r.model: {"p": float(r.pvalue), "in": bool(r.included)} for r in mc.itertuples()}
                except Exception as e:  # noqa: BLE001
                    res[f"mcs_{tier}"] = {"error": repr(e)[:200]}
                if b in M and len(cand) > 1:
                    try:
                        res[f"spa_{tier}"] = spa(M[b].values, M[[c for c in cand if c != b]], reps=reps, block=PROT["inference"]["spa"]["block_length"])
                    except Exception as e:  # noqa: BLE001
                        res[f"spa_{tier}"] = {"error": repr(e)[:200]}
        # exploratory: MCS on one representative specification per architecture (Student-t where available, else the
        # single available law) to reduce the model count and sharpen the set (not pre-registered; labelled as such)
        rep = {}
        for sp in specs:
            mid, dist, _ = sp.split("__")
            if res["specs"][sp]["coverage"] < MIN_COVERAGE_MCS or REG.get(mid, {}).get("status") != "implemented":
                continue
            pref = {"t": 0, "normal": 1, "empirical": 2, "ged": 3, "skewt": 4}[dist]
            if mid not in rep or pref < rep[mid][0]:
                rep[mid] = (pref, sp)
        for tier in ("A", "B"):
            cand = [v[1] for v in rep.values() if spec_meta(v[1])["tier"] == tier]
            if tier == "B" and b not in cand:
                cand.append(b)
            M = L[cand].dropna() if len(cand) >= 2 else None
            if M is not None and len(M) >= 100:
                try:
                    mc = mcs(M, size=PROT["inference"]["mcs"]["alpha"], reps=reps, block=PROT["inference"]["mcs"]["block_length"])
                    res[f"mcs_rep_{tier}"] = {r.model: {"p": float(r.pvalue), "in": bool(r.included)} for r in mc.itertuples()}
                    res[f"mcs_rep_{tier}_n"] = int(len(M))
                except Exception as e:  # noqa: BLE001
                    res[f"mcs_rep_{tier}"] = {"error": repr(e)[:200]}
        out["horizons"][h] = res
        # regimes (pairwise ratios + DM within regime)
        R = regimes_for(panel, L.index, PROT["periods"]["validation"]["start"], PROT["periods"]["validation"]["end"])
        reg = {}
        for col in ("vol_tercile", "vol_tercile66", "prev_sign", "episode"):
            for lab in sorted(R[col].unique()):
                idx = R.index[R[col] == lab]
                Lr = L.loc[idx]
                if len(Lr) < (30 if col == "episode" else 60):
                    continue
                d = {}
                for s in specs:
                    pair = np.isfinite(Lr[s].values) & np.isfinite(Lr[b].values)
                    if pair.sum() >= 30:
                        dm = dm_test(Lr[s].values[pair], Lr[b].values[pair], h)
                        d[s] = {"n": int(pair.sum()), "ratio": float(Lr[s].values[pair].mean() / Lr[b].values[pair].mean()), "dm_p": dm["p"], "dm_stat": dm["stat"]}
                reg[f"{col}={lab}"] = d
        out["regimes"][h] = reg
    # secondary losses / proxies at every horizon (pairwise ratios only)
    b = bench_id(specs)
    for (kind, proxy, loss) in [("loss", PRIMARY_PROXY, ln) for ln in PROT["secondary_losses"]] + [("loss", px, PRIMARY_LOSS) for px in PROT["secondary_proxies"]]:
        for h in HORIZONS:
            L = block(W, kind, proxy, loss, h); L = L[(L.index >= lo) & (L.index <= hi)]
            if L.empty or b not in L:
                continue
            d = {}
            for s in L.columns:
                pair = np.isfinite(L[s].values) & np.isfinite(L[b].values)
                if pair.sum() >= 30:
                    d[s] = float(L[s].values[pair].mean() / L[b].values[pair].mean()) if loss != "mse" or True else np.nan
            out["secondary"][f"{proxy}|{loss}|{h}"] = d
    # tail risk, h = 1
    F1 = block(W, "forecast", "", "", 1); F1 = F1[(F1.index >= lo) & (F1.index <= hi)]
    MU = block(W, "mu", "", "", 1).reindex(F1.index); D1 = block(W, "dp1", "", "", 1).reindex(F1.index); D2 = block(W, "dp2", "", "", 1).reindex(F1.index)
    p = panel.set_index(pd.to_datetime(panel["session"]))["r"]
    r_next = p.shift(-1).reindex(F1.index).values    # realized return of the day after the origin
    tail = {}
    for s in F1.columns:
        meta = spec_meta(s)
        dist = meta["dist"]
        h1 = F1[s].values; mu = MU[s].values if s in MU else np.zeros(len(F1))
        ok = np.isfinite(h1) & np.isfinite(r_next)
        if ok.sum() < 100:
            continue
        emp = None
        if dist == "empirical" and len(G):
            g = G[G["spec_id"] == s].sort_values("origin")
            if "emp_q01" in g:
                gi = g.set_index(pd.to_datetime(g["origin"]))
                qq = {}
                for lvl, k in ((0.01, "01"), (0.05, "05")):
                    qq[lvl] = {"q": gi[f"emp_q{k}"].reindex(F1.index, method="ffill").values, "es": gi[f"emp_es{k}"].reindex(F1.index, method="ffill").values}
                emp = qq
        try:
            fc = var_es_forecasts(np.nan_to_num(mu), np.where(ok, h1, np.nan), dist, np.nan_to_num(D1[s].values if s in D1 else np.full(len(F1), np.nan), nan=5.0),
                                  np.nan_to_num(D2[s].values if s in D2 else np.full(len(F1), np.nan)), emp_q=emp)
            tail[s] = {str(lvl): backtest_table(np.where(ok, r_next, np.nan), fc, lvl) for lvl in (0.01, 0.05)}
        except Exception as e:  # noqa: BLE001
            tail[s] = {"error": repr(e)[:200]}
    out["tailrisk"] = tail
    # failure and compute accounting (refits in period)
    acc = {}
    if len(G):
        G["origin"] = pd.to_datetime(G["origin"])
        Gp = G[(G["origin"] >= lo) & (G["origin"] <= hi)]
        for s, g in Gp.groupby("spec_id"):
            acc[s] = {"n_refits": int(len(g)), "failed": float((g["status"] == "failed").mean()), "inadmissible": float((g["status"] == "inadmissible").mean()),
                      "fit_time_mean": float(g["time_s"].mean()), "fit_time_median": float(g["time_s"].median()), "iter_mean": float(g["n_iter"].mean()),
                      "restarts_mean": float(g["n_restarts"].mean())}
    Sp = S[(S.index >= lo) & (S.index <= hi)]
    for s in Sp.columns:
        acc.setdefault(s, {})
        vc = Sp[s].value_counts(normalize=True).to_dict()
        acc[s].update({"status_" + k: float(v) for k, v in vc.items()})
    out["accounting"] = acc
    return out


def pooled(per_sym: dict, run: str, period: str, reps: int, outdir: Path) -> None:
    """Pooled ratios (equal weight per instrument) with a joint block bootstrap over dates; counts of significant DM wins/losses
    (Holm across instruments per (spec, h)); MCS inclusion counts."""
    rows = []
    syms = sorted(per_sym)
    specs = sorted({s for sym in syms for s in per_sym[sym]["horizons"][HORIZONS[0]]["specs"]})
    # per-origin loss matrices are needed for the bootstrap: reload the primary blocks
    cache = {}
    for sym in syms:
        W, T, S, G = load_symbol(run, sym)
        per = PROT["periods"][period]
        for h in HORIZONS:
            L = block(W, "loss", PRIMARY_PROXY, PRIMARY_LOSS, h)
            cache[(sym, h)] = L[(L.index >= per["start"]) & (L.index <= per["end"])]
    for h in HORIZONS:
        b = bench_id(specs)
        for s in specs:
            for cls in ["all"] + sorted(set(CLASS_OF.values())):
                members = [sym for sym in syms if cls == "all" or CLASS_OF.get(sym) == cls]
                ratios, dms, n = [], [], 0
                D = {}
                for sym in members:
                    r = per_sym[sym]["horizons"][h]["specs"].get(s, {})
                    if np.isfinite(r.get("ratio", np.nan)):
                        ratios.append(r["ratio"]); dms.append(r["dm_p"]); n += 1
                        L = cache[(sym, h)]
                        if s in L:
                            pair = np.isfinite(L[s].values) & np.isfinite(L[b].values)
                            D[sym] = pd.Series((L[s].values - L[b].values) / np.nanmean(L[b].values[pair]), index=L.index).where(pair)
                if n == 0:
                    continue
                dmp = np.array([p for p in dms if np.isfinite(p)])
                # Holm across instruments
                m = len(dmp); order = np.argsort(dmp); adj = np.minimum.accumulate((dmp[order] * (m - np.arange(m)))[::-1])[::-1] if m else np.array([])
                holm = np.empty(m); holm[order] = np.minimum(adj, 1.0) if m else holm
                wins = sum(1 for sym, p_ in zip([x for x in members if np.isfinite(per_sym[x]["horizons"][h]["specs"].get(s, {}).get("ratio", np.nan))], holm)
                           if p_ < 0.10 and per_sym[sym]["horizons"][h]["specs"][s]["ratio"] < 1)
                losses_ = sum(1 for sym, p_ in zip([x for x in members if np.isfinite(per_sym[x]["horizons"][h]["specs"].get(s, {}).get("ratio", np.nan))], holm)
                              if p_ < 0.10 and per_sym[sym]["horizons"][h]["specs"][s]["ratio"] > 1)
                boot = {"mean": np.nan, "ci_low": np.nan, "ci_high": np.nan, "p": np.nan}
                if D and s != b:
                    Dm = pd.DataFrame(D)
                    Dm = Dm.dropna(how="all")
                    if len(Dm) > 100:
                        boot = block_bootstrap_mean_ci(Dm, reps=min(reps, 1000), block=22)
                mcs_in = sum(1 for sym in members for tier in ("A", "B") if per_sym[sym]["horizons"][h].get(f"mcs_{tier}", {}).get(s, {}).get("in"))
                trims = [per_sym[sym]["horizons"][h]["specs"][s].get("ratio_trim1", np.nan) for sym in members if np.isfinite(per_sym[sym]["horizons"][h]["specs"].get(s, {}).get("ratio", np.nan))]
                expls = [per_sym[sym]["horizons"][h]["specs"][s].get("explosive_share", np.nan) for sym in members if np.isfinite(per_sym[sym]["horizons"][h]["specs"].get(s, {}).get("ratio", np.nan))]
                rows.append({"h": h, "spec_id": s, **spec_meta(s), "asset_class": cls, "n_instruments": n, "mean_ratio": float(np.mean(ratios)),
                             "median_ratio": float(np.median(ratios)), "share_ratio_below_1": float(np.mean(np.array(ratios) < 1)),
                             "mean_ratio_trim1": float(np.nanmean(trims)) if trims else np.nan, "explosive_share": float(np.nanmean(expls)) if expls else np.nan,
                             "dm_wins_holm10": wins, "dm_losses_holm10": losses_, "boot_mean_rel_diff": boot["mean"], "boot_ci_low": boot["ci_low"],
                             "boot_ci_high": boot["ci_high"], "boot_p": boot["p"], "mcs_inclusions": mcs_in})
    P = pd.DataFrame(rows)
    P.to_csv(outdir / "pooled.csv", index=False)
    return P


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--run", default="main")
    ap.add_argument("--period", default="test", choices=["test", "validation"])
    ap.add_argument("--symbols", default="")
    ap.add_argument("--reps", type=int, default=2000)
    a = ap.parse_args()
    d = ROOT / "results" / "metrics" / f"losses_{a.run}"
    syms = a.symbols.split(",") if a.symbols else sorted(Path(f).stem for f in glob.glob(str(d / "*.parquet")) if not Path(f).stem.startswith(("targets_", "status_", "fits_")))
    outdir = ROOT / "results" / "metrics" / f"analysis_{a.run}_{a.period}"; (outdir / "per_symbol").mkdir(parents=True, exist_ok=True)
    per_sym = {}
    for sym in syms:
        res = analyze_symbol(a.run, sym, a.period, a.reps)
        (outdir / "per_symbol" / f"{sym}.json").write_text(json.dumps(res, indent=0, default=float))
        per_sym[sym] = res
        h1 = res["horizons"][1]
        cands = [(s, v["ratio"]) for s, v in h1["specs"].items() if np.isfinite(v.get("ratio", np.nan))]
        best = min(cands, key=lambda x: x[1]) if cands else ("--", np.nan)
        print(f"{sym:7s} n={h1['n_origins']} valid={h1['n_valid_targets']} best@h1={best[0]} ratio={best[1]:.3f} mcsA={len([k for k,v in h1.get('mcs_A',{}).items() if isinstance(v,dict) and v.get('in')])}", flush=True)
    P = pooled(per_sym, a.run, a.period, a.reps, outdir)
    tabs = ROOT / "results" / "tables"; tabs.mkdir(exist_ok=True)
    P.to_csv(tabs / f"pooled_{a.run}_{a.period}.csv", index=False)
    print("pooled table", outdir / "pooled.csv", len(P))


if __name__ == "__main__":
    main()