#!/usr/bin/env python
"""Stage 20e — verify the Stage-2 selected-predictor retrain (stage 20d), out-of-sample.

Pools the LOYO v3 predictions and scores, per block, the matched pair:
  base : Stage-1 forecast backbone (10 features)
  sel  : backbone + selected extended predictors (ENSO/MJO/Z500/T850/persistence per block)
against a per-fold base-rate climatology, with a whole-year block bootstrap on the BSS gain
(sel - base) to test whether the added predictors help significantly. Also the full binary
inventory for both.

Outputs:
  results/tables/v3_scores_by_block.csv          base vs sel: base_rate, Brier, BSS_clim, AUC, SEDI
  results/tables/v3_bss_gain_bootstrap.csv        per block: mean BSS gain (sel-base) + 95% CI
  results/tables/v3_binary_by_block.csv           full 10-metric inventory, base and sel
  results/figures/v3_base_vs_selected.png         BSS & AUC by lead, base vs sel
"""
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

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


def load(cfg):
    paths = sorted(glob.glob(os.path.join(cfg["paths"]["outputs"], "loyo_pred_v3sel", "pred_v3_*.npz")))
    if not paths:
        raise SystemExit("no outputs/loyo_pred_v3sel/pred_v3_*.npz — run stage 20d first")
    frames = []
    for p in paths:
        yr = int(os.path.basename(p).split("_")[-1][:-4])
        d = np.load(p, allow_pickle=True)
        df = pd.DataFrame({k: d[k] for k in d.files}); df["year"] = yr
        frames.append(df)
    return pd.concat(frames, ignore_index=True)


def _binom_clim(df, b_id):
    sub = df[df.block_id == b_id]
    return np.full(len(sub), float(np.nanmean(sub.y.values)))


def main():
    cfg = load_config(); blocks = lead_blocks(cfg)
    bid = {k + 1: b for k, b in enumerate(blocks)}
    years = list(cfg["cross_validation"]["years"])
    df = load(cfg)

    rows, brows, irows = [], [], []
    rng = np.random.default_rng(cfg["seed"]); B = 1000
    for b_id, bname in bid.items():
        sub = df[df.block_id == b_id]
        y = sub.y.values; pc = np.full(len(sub), float(np.nanmean(y)))
        for tag in ("base", "sel"):
            p = sub[f"p_{tag}"].values; fl = sub[f"flag_{tag}"].values.astype(bool)
            m = np.isfinite(y) & np.isfinite(p); yv, pv, fv = y[m], p[m], fl[m]
            a = int(np.sum(fv & (yv == 1))); bb = int(np.sum(fv & (yv == 0)))
            c = int(np.sum(~fv & (yv == 1))); d = int(np.sum(~fv & (yv == 0)))
            bs = vf.binary_scores_from_counts(a, bb, c, d)
            rows.append(dict(block=bname, model=tag, n=int(m.sum()),
                             base_rate=round(float(np.nanmean(yv)), 3),
                             brier=round(float(np.nanmean((pv - yv) ** 2)), 4),
                             bss_vs_clim=round(vf.brier_skill_score(y, p, pc), 3),
                             auc=round(vf.auc(yv, pv), 3), SEDI=round(bs["SEDI"], 3)))
            irows.append(dict(block=bname, model=tag, base_rate=round(float(np.nanmean(yv)), 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"], 2)))
        # bootstrap the BSS gain (sel - base) over whole years
        per_year = {yv: sub[sub.year == yv] for yv in years}
        gains = []
        for _ in range(B):
            pick = rng.choice(years, size=len(years), replace=True)
            yy = np.concatenate([per_year[q].y.values for q in pick])
            pb = np.concatenate([per_year[q].p_base.values for q in pick])
            ps = np.concatenate([per_year[q].p_sel.values for q in pick])
            cc = np.full(len(yy), float(np.nanmean(yy)))
            gains.append(vf.brier_skill_score(yy, ps, cc) - vf.brier_skill_score(yy, pb, cc))
        lo, hi = np.nanpercentile(gains, [2.5, 97.5])
        base_bss = [r for r in rows if r["block"] == bname and r["model"] == "base"][0]["bss_vs_clim"]
        sel_bss = [r for r in rows if r["block"] == bname and r["model"] == "sel"][0]["bss_vs_clim"]
        brows.append(dict(block=bname, bss_base=base_bss, bss_sel=sel_bss,
                          bss_gain=round(sel_bss - base_bss, 3), gain_ci_lo=round(float(lo), 3),
                          gain_ci_hi=round(float(hi), 3), significant=bool(lo > 0)))

    pd.DataFrame(rows).to_csv(os.path.join(TABLES, "v3_scores_by_block.csv"), index=False)
    pd.DataFrame(brows).to_csv(os.path.join(TABLES, "v3_bss_gain_bootstrap.csv"), index=False)
    pd.DataFrame(irows).to_csv(os.path.join(TABLES, "v3_binary_by_block.csv"), index=False)
    print("[20e] base vs selected:\n", pd.DataFrame(brows).to_string(index=False))
    _fig(pd.DataFrame(rows), list(blocks))


def _fig(df, block_names):
    import matplotlib; matplotlib.use("Agg"); import matplotlib.pyplot as plt
    x = np.arange(len(block_names))
    fig, ax = plt.subplots(1, 2, figsize=(11, 4.3))
    for tag, col in (("base", "0.5"), ("sel", "tab:red")):
        s = df[df.model == tag].set_index("block").reindex(block_names)
        ax[0].plot(x, s.bss_vs_clim, marker="o", color=col, label=tag)
        ax[1].plot(x, s.auc, marker="o", color=col, label=tag)
    for a, t in zip(ax, ["BSS vs climatology", "AUC"]):
        a.set_xticks(x); a.set_xticklabels(block_names); a.axhline(0 if t.startswith("BSS") else 0.5,
                                                                   color="k", lw=.6)
        a.set_title(f"Stage-2 base vs selected — {t}"); a.grid(alpha=.3); a.legend()
    os.makedirs(FIGS, exist_ok=True); fig.tight_layout()
    fig.savefig(os.path.join(FIGS, "v3_base_vs_selected.png"), dpi=130)
    print("[20e] wrote results/figures/v3_base_vs_selected.png")


if __name__ == "__main__":
    main()