#!/usr/bin/env python
"""Stage 08b — verify the SECONDARY target (P(heat-wave, >=3 consecutive EHDs)) from the
stage-06b out-of-sample LOYO predictions. Mirrors the primary headline (stage 08):

  results/tables/cv_scores_secondary_by_block.csv   base rate, Brier, BSS vs clim & vs raw, AUC
  results/tables/cv_bootstrap_secondary_ci.csv      whole-year block bootstrap 95% CI on BSS_clim
  results/tables/cv_binary_secondary_by_block.csv   full 10-metric binary inventory (hgb_recal)

Leakage-safe: every prediction is out-of-sample (its year was held out in stage 06b); the binary
flag uses the per-region x block decision threshold that stage 06b fit on training years only
(median over folds -> the operational single value, as in stage 08).
"""
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")
HEADLINE = "hgb_recal"


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


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, decthr = load_pred(cfg)
    p_h = "p_" + HEADLINE

    # ---- pooled continuous/prob scores by block ----
    rows = []
    for b_id, bname in bid.items():
        sub = df[df.block_id == b_id]
        y = sub.y.values; p = sub[p_h].values
        base = float(np.nanmean(y))
        brier = float(np.nanmean((p - y) ** 2))
        bss_c = vf.brier_skill_score(y, p, sub.p_climat.values)
        bss_r = vf.brier_skill_score(y, p, sub.p_raw.values)
        rows.append(dict(block=bname, family=HEADLINE, n=int(np.isfinite(y).sum()),
                         base_rate=round(base, 4), brier=round(brier, 4),
                         bss_vs_clim=round(bss_c, 3), bss_vs_raw=round(bss_r, 3),
                         auc=round(vf.auc(y, p), 3)))
    by_block = pd.DataFrame(rows)
    by_block.to_csv(os.path.join(TABLES, "cv_scores_secondary_by_block.csv"), index=False)
    print("[08b] secondary by-block scores:\n", by_block.to_string(index=False))

    # ---- whole-year block bootstrap on BSS vs clim ----
    rng = np.random.default_rng(cfg["seed"]); B = 1000
    brows = []
    for b_id, bname in bid.items():
        sub = df[df.block_id == b_id]
        per_year = {yv: sub[sub.year == yv] for yv in years}
        stats = []
        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])
            pp = np.concatenate([per_year[q][p_h].values for q in pick])
            pc = np.concatenate([per_year[q].p_climat.values for q in pick])
            stats.append(vf.brier_skill_score(yy, pp, pc))
        lo, hi = np.nanpercentile(stats, [2.5, 97.5])
        point = by_block[by_block.block == bname].bss_vs_clim.values[0]
        brows.append(dict(block=bname, bss_vs_clim=point, ci_lo=round(float(lo), 3),
                          ci_hi=round(float(hi), 3), significant=bool(lo > 0)))
    pd.DataFrame(brows).to_csv(os.path.join(TABLES, "cv_bootstrap_secondary_ci.csv"), index=False)

    # ---- full binary inventory by block, per-region x block median threshold ----
    dt_h = decthr[decthr.family == "hgb"]
    thr_lookup = dt_h.groupby(["block", "region"]).thr.median().to_dict()
    irows = []
    for b_id, bname in bid.items():
        sub = df[df.block_id == b_id]
        y = sub.y.values; p = sub[p_h].values
        tvec = np.array([float(thr_lookup.get((bname, int(r)), 0.5)) for r in sub.region.values])
        flag = p >= tvec
        m = np.isfinite(y) & np.isfinite(p)
        yv, fv = y[m], flag[m].astype(bool)
        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)
        irows.append(dict(block=bname, n=int(m.sum()), base_rate=round(float(np.nanmean(y)), 4),
                          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)))
    pd.DataFrame(irows).to_csv(os.path.join(TABLES, "cv_binary_secondary_by_block.csv"), index=False)
    print("[08b] wrote cv_scores_secondary_by_block.csv, cv_bootstrap_secondary_ci.csv, "
          "cv_binary_secondary_by_block.csv")


if __name__ == "__main__":
    main()