#!/usr/bin/env python
"""Stage 13 — v2 whole-dataset verification (2003-2022) with binary contingency skill.

Consumes the out-of-sample predictions from stage 12 (outputs/loyo_pred_v2/pred_v2_*.npz) and
scores every (family, cutoff) x lead block over the WHOLE pooled dataset of 20 held-out years.
Because each test sample used the decision threshold fit on ITS OWN fold's training data (per
region x block, leakage-safe), the binary flag is rebuilt per sample from the stored per-fold
per-region thresholds, then hits/false-alarms/misses/correct-negs are POOLED and scored once.

Metrics per (family, cutoff, block):
  probabilistic : Brier, BSS vs climatology (pooled base rate), AUC
  binary        : POD, POFD, FAR, CSI, ETS, HSS, Peirce, EDI, SEDI, frequency bias
Outputs:
  results/tables/v2_cv_scores_by_model_block.csv     (every family x cutoff x block)
  results/tables/v2_best_model_by_block.csv          (top model per block by SEDI)
  results/figures/v2_binary_skill_by_lead.png        (SEDI/Peirce vs lead, best cutoff)
"""
from __future__ import annotations
import os, sys, glob
import numpy as np
import pandas as pd

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))
from ehd.config import load_config, lead_blocks
from ehd import verify as vf

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
TABLES = os.path.join(ROOT, "results", "tables")
FIGS = os.path.join(ROOT, "results", "figures")


def parse_tag(tag):
    fam, cut = tag.rsplit("_c", 1)
    return fam, float(cut)


def main():
    cfg = load_config()
    blocks = lead_blocks(cfg)
    block_names = list(blocks.keys())
    bid_of = {k + 1: b for k, b in enumerate(block_names)}
    files = sorted(glob.glob(os.path.join(ROOT, "outputs", "loyo_pred_v2", "pred_v2_*.npz")))
    assert files, "no v2 LOYO predictions found (run stage 12)"
    # optional torch families (Phase B) live in a sibling dir with identical sample ordering
    cnn_files = {int(os.path.basename(p).split("_")[-1][:-4]): p
                 for p in glob.glob(os.path.join(ROOT, "outputs", "loyo_pred_v2_cnn",
                                                 "pred_v2cnn_*.npz"))}
    print(f"[verify_v2] {len(files)} folds; cnn/ann folds: {len(cnn_files)}")

    tags = None
    Y = {b: [] for b in block_names}                 # pooled y per block, in file order
    P = {}                                            # (tag, block) -> list of p arrays
    ABCD = {}                                         # (tag, block) -> [a,b,c,d]

    for f in files:
        d = np.load(f, allow_pickle=True)
        yr = int(os.path.basename(f).split("_")[-1][:-4])
        dce = np.load(cnn_files[yr], allow_pickle=True) if yr in cnn_files else None
        probs = {k[2:]: d[k] for k in d.files if k.startswith("p_")}
        dt = d["decthr"]
        if dce is not None:
            probs.update({k[2:]: dce[k] for k in dce.files if k.startswith("p_")})
            dt = np.concatenate([dt, dce["decthr"]])
        if tags is None:
            tags = sorted(probs.keys())
        bid = d["block_id"].astype(int); reg = d["region"].astype(int); y = d["y"]
        for bcode, bname in bid_of.items():
            Y[bname].append(y[bid == bcode])
        for tag in tags:
            fam, cut = parse_tag(tag)
            p = probs[tag]
            # per (tag, block) threshold-by-region array
            for bcode, bname in bid_of.items():
                bm = bid == bcode
                yb, pb, rb = y[bm], p[bm], reg[bm]
                P.setdefault((tag, bname), []).append(pb)
                sel = (dt["tag"] == tag) & (dt["block"] == bname)
                thr_by_reg = np.full(int(reg.max()) + 1, 0.5, dtype=float)
                for row in dt[sel]:
                    thr_by_reg[int(row["region"])] = float(row["thr"])
                thr_vec = thr_by_reg[rb]
                flag = pb >= thr_vec
                valid = np.isfinite(yb) & np.isfinite(pb)
                yy = yb[valid]; ff = flag[valid]
                a = int(np.sum(ff & (yy == 1))); b = int(np.sum(ff & (yy == 0)))
                c = int(np.sum(~ff & (yy == 1))); dd = int(np.sum(~ff & (yy == 0)))
                acc = ABCD.setdefault((tag, bname), [0, 0, 0, 0])
                acc[0] += a; acc[1] += b; acc[2] += c; acc[3] += dd

    yb_all = {b: np.concatenate(Y[b]) for b in block_names}
    rows = []
    for tag in tags:
        fam, cut = parse_tag(tag)
        for bname in block_names:
            y = yb_all[bname]
            p = np.concatenate(P[(tag, bname)])
            m = np.isfinite(y) & np.isfinite(p)
            yv, pv = y[m], p[m]
            base = float(np.mean(yv))
            brier = float(np.mean((pv - yv) ** 2))
            bss = 1.0 - brier / (base * (1 - base)) if 0 < base < 1 else np.nan
            auc = vf.auc(yv, pv)
            a, b, c, dd = ABCD[(tag, bname)]
            bs = vf.binary_scores_from_counts(a, b, c, dd)
            rows.append(dict(family=fam, cutoff=cut, block=bname, n=int(m.sum()),
                             base_rate=round(base, 3), brier=round(brier, 4),
                             bss_vs_clim=round(bss, 3), auc=round(auc, 3),
                             POD=round(bs["POD"], 3), POFD=round(bs["POFD"], 3),
                             FAR=round(bs["FAR"], 3), CSI=round(bs["CSI"], 3),
                             ETS=round(bs["ETS"], 3), HSS=round(bs["HSS"], 3),
                             Peirce=round(bs["Peirce"], 3), EDI=round(bs["EDI"], 3),
                             SEDI=round(bs["SEDI"], 3), freq_bias=round(bs["freq_bias"], 3)))
    out = pd.DataFrame(rows).sort_values(["block", "family", "cutoff"]).reset_index(drop=True)
    os.makedirs(TABLES, exist_ok=True)
    out.to_csv(os.path.join(TABLES, "v2_cv_scores_by_model_block.csv"), index=False)
    print(f"[verify_v2] wrote {len(out)} rows -> v2_cv_scores_by_model_block.csv")

    # best model per block by SEDI (a proper score for rare/extreme events)
    best = out.loc[out.groupby("block")["SEDI"].idxmax()].reset_index(drop=True)
    best.to_csv(os.path.join(TABLES, "v2_best_model_by_block.csv"), index=False)
    print("\n[verify_v2] best model per block (by SEDI):")
    print(best[["block", "family", "cutoff", "SEDI", "EDI", "Peirce", "POD", "FAR",
                "bss_vs_clim", "auc"]].to_string(index=False))

    _figure(out, block_names)


def _figure(out, block_names):
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    # choose, per family, its best cutoff by mean SEDI across blocks
    fams = sorted(out["family"].unique())
    x = np.arange(len(block_names))
    fig, ax = plt.subplots(1, 2, figsize=(12, 4.5), sharex=True)
    for fam in fams:
        sub = out[out.family == fam]
        cut = sub.groupby("cutoff")["SEDI"].mean().idxmax()
        s = sub[sub.cutoff == cut].set_index("block").reindex(block_names)
        ax[0].plot(x, s["SEDI"], marker="o", label=f"{fam} (c={cut:g})")
        ax[1].plot(x, s["Peirce"], marker="o", label=f"{fam} (c={cut:g})")
    for a, t in zip(ax, ["SEDI", "Peirce (POD-POFD)"]):
        a.set_xticks(x); a.set_xticklabels(block_names); a.axhline(0, color="k", lw=.6)
        a.set_title(f"v2 binary skill by lead — {t}"); a.set_xlabel("lead block"); a.grid(alpha=.3)
    ax[0].legend(fontsize=8, ncol=2)
    os.makedirs(FIGS, exist_ok=True)
    fig.tight_layout(); fig.savefig(os.path.join(FIGS, "v2_binary_skill_by_lead.png"), dpi=130)
    print("[verify_v2] wrote figure -> results/figures/v2_binary_skill_by_lead.png")


if __name__ == "__main__":
    main()