#!/usr/bin/env python
"""Stage 08 — cross-validated verification from the LOYO out-of-sample predictions.

Reads outputs/loyo_pred/pred_<year>.npz (stage 06), pools the held-out predictions across
all folds, and scores every family against BOTH references (climatology = per-fold base-rate
`p_climat`; raw forecast = `p_raw`). Binary scores use the per-region x block decision
thresholds fit in stage 06 (leakage-free). Significance: block bootstrap resampling whole
YEARS (preserves serial correlation). Field significance across grid cells via Benjamini-
Hochberg FDR on per-cell Brier-skill bootstrap p-values.

Stratification (task 8 tail): the pooled held-out predictions are additionally scored by
INIT-SEASON (DJF/MAM/JJA/SON) x block and by base-rate SEVERITY tercile x block for the
headline family. Init season is reconstructed from the deterministic ordering of each
pred_<year>.npz (per block the rows are init-major with `ncell` land cells tiled), joined to
the init dates in outputs/forecast_daily_tx/fc_daily_tx_<year>.nc. Severity terciles are cut
on each cell's pooled climatological EHD base rate within the block.

Outputs (evidence pack):
  results/tables/cv_scores_by_block.csv
  results/tables/cv_scores_by_region_block.csv
  results/tables/cv_bootstrap_ci.csv
  results/tables/cv_scores_by_season_block.csv
  results/tables/cv_scores_by_severity_block.csv
  results/figures/cv_reliability.png
  results/figures/cv_skill_by_lead.png
  results/figures/cv_strata.png
"""
from __future__ import annotations
import os, sys, glob, logging
import numpy as np
import pandas as pd
import xarray as xr
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

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

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger("stage08")

FAMILIES = ["raw", "climat", "logistic", "logistic_recal", "hgb", "hgb_recal"]
HEADLINE = "hgb_recal"
SEASON_OF_MONTH = {12: "DJF", 1: "DJF", 2: "DJF", 3: "MAM", 4: "MAM", 5: "MAM",
                   6: "JJA", 7: "JJA", 8: "JJA", 9: "SON", 10: "SON", 11: "SON"}
SEASON_ORDER = ["DJF", "MAM", "JJA", "SON"]


def _init_month_per_row(cfg, yr, block_id, cell):
    """Reconstruct the init-date month for every row of pred_<yr>.npz.

    Within each block the stage-06 assembler emits rows init-major with the land cells
    tiled (cell = 0..ncell-1 repeating), so the init ordinal of a row is its position
    inside that block // ncell. Map that ordinal onto the file's init calendar.
    """
    fc = xr.open_dataarray(os.path.join(cfg["paths"]["outputs"],
                                        "forecast_daily_tx", f"fc_daily_tx_{yr}.nc"))
    init_month = pd.to_datetime(fc["init"].values).month.to_numpy()
    fc.close()
    ncell = int(cell.max()) + 1
    month = np.empty(block_id.size, dtype=int)
    for b in np.unique(block_id):
        pos = np.where(block_id == b)[0]                 # rows for this block, in emit order
        ordinal = np.arange(pos.size) // ncell           # init index 0..ninit-1
        assert ordinal.max() < init_month.size, "init ordinal exceeds file init count"
        month[pos] = init_month[ordinal]
    return month


def load_all(cfg):
    files = sorted(glob.glob(os.path.join(cfg["paths"]["outputs"], "loyo_pred", "pred_*.npz")))
    frames, decthrs = [], []
    for f in files:
        yr = int(os.path.basename(f).split("_")[1].split(".")[0])
        d = np.load(f, allow_pickle=True)
        n = d["y"].size
        month = _init_month_per_row(cfg, yr, d["block_id"], d["cell"])
        rec = {"year": np.full(n, yr), "block_id": d["block_id"], "region": d["region"],
               "cell": d["cell"], "y": d["y"], "init_month": month}
        for fam in FAMILIES:
            rec["p_" + fam] = d["p_" + fam]
        frames.append(pd.DataFrame(rec))
        dt = pd.DataFrame(d["decthr"]); dt["year"] = yr
        decthrs.append(dt)
    df = pd.concat(frames, ignore_index=True)
    df["season"] = df["init_month"].map(SEASON_OF_MONTH)
    return df, pd.concat(decthrs, ignore_index=True)


def main():
    cfg = load_config(); blocks = lead_blocks(cfg)
    bid = {k + 1: b for k, b in enumerate(blocks)}
    res = cfg["paths"]["results"]
    df, decthr = load_all(cfg)
    years = sorted(df.year.unique())
    log.info("loaded %d folds, %d rows", len(years), len(df))

    # ---- probabilistic scores by block x family ----
    rows = []
    for b_id, bname in bid.items():
        sub = df[df.block_id == b_id]
        y = sub.y.values
        pref_c = sub.p_climat.values; pref_r = sub.p_raw.values
        for fam in FAMILIES:
            p = sub["p_" + fam].values
            rows.append(dict(block=bname, family=fam,
                             n=int(np.isfinite(y).sum()), base_rate=round(float(np.nanmean(y)), 3),
                             brier=round(vf.brier(y, p), 4),
                             bss_vs_clim=round(vf.brier_skill_score(y, p, pref_c), 3),
                             bss_vs_raw=round(vf.brier_skill_score(y, p, pref_r), 3),
                             auc=round(vf.auc(y, p), 3),
                             sharpness=round(float(np.nanstd(p)), 3)))
    by_block = pd.DataFrame(rows)
    by_block.to_csv(os.path.join(res, "tables", "cv_scores_by_block.csv"), index=False)
    log.info("headline (%s) by block:\n%s", HEADLINE,
             by_block[by_block.family == HEADLINE].to_string(index=False))

    # ---- binary scores by region x block for the headline family ----
    dt_h = decthr[decthr.family == ("hgb")]
    # median decision threshold across folds per region x block (operational single value)
    thr_lookup = dt_h.groupby(["block", "region"]).thr.median().to_dict()
    rrows = []
    p_h = "p_" + HEADLINE
    for b_id, bname in bid.items():
        for r in cfg["regions"]["ids"]:
            sub = df[(df.block_id == b_id) & (df.region == r)]
            if sub.empty:
                continue
            y = sub.y.values; p = sub[p_h].values
            t = float(thr_lookup.get((bname, r), 0.5))
            bs = vf.binary_scores(y, p, t)
            rrows.append(dict(block=bname, region=r, n=int(np.isfinite(y).sum()),
                              base_rate=round(float(np.nanmean(y)), 3), dec_thr=round(t, 2),
                              bss_vs_clim=round(vf.brier_skill_score(
                                  y, p, sub.p_climat.values), 3),
                              auc=round(vf.auc(y, p), 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)))
    by_rb = pd.DataFrame(rrows)
    by_rb.to_csv(os.path.join(res, "tables", "cv_scores_by_region_block.csv"), index=False)

    # ---- pooled-by-block binary inventory for the headline family (full 10-metric set) ----
    # Uses each sample's leakage-safe per-region x block decision threshold (median over folds),
    # applied per region then pooled, so it matches the operational single-threshold flag.
    brows = []
    for b_id, bname in bid.items():
        sub = df[df.block_id == b_id]
        if sub.empty:
            continue
        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)
        brows.append(dict(block=bname, n=int(m.sum()),
                          base_rate=round(float(np.nanmean(y)), 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)))
    pd.DataFrame(brows).to_csv(
        os.path.join(res, "tables", "cv_binary_scores_by_block.csv"), index=False)

    # ---- block bootstrap over whole years: BSS(clim) per block for headline ----
    rng = np.random.default_rng(cfg["seed"])
    B = 1000
    boot_rows = []
    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[p].y.values for p in pick])
            pp = np.concatenate([per_year[p]["p_" + HEADLINE].values for p in pick])
            pc = np.concatenate([per_year[p].p_climat.values for p in pick])
            stats.append(vf.brier_skill_score(yy, pp, pc))
        stats = np.array(stats)
        lo, hi = np.nanpercentile(stats, [2.5, 97.5])
        point = by_block[(by_block.block == bname) & (by_block.family == HEADLINE)].bss_vs_clim.values[0]
        boot_rows.append(dict(block=bname, bss_vs_clim=point,
                              ci_lo=round(float(lo), 3), ci_hi=round(float(hi), 3),
                              significant=bool(lo > 0)))
    boot = pd.DataFrame(boot_rows)
    boot.to_csv(os.path.join(res, "tables", "cv_bootstrap_ci.csv"), index=False)
    log.info("bootstrap BSS(clim) CI by block (%s):\n%s", HEADLINE, boot.to_string(index=False))

    by_season = stratify_season(df, bid, res)
    by_sev = stratify_severity(df, bid, res)

    _figures(cfg, df, bid, by_block, boot, res)
    _strata_figure(by_season, by_sev, bid, res)


def _score_group(sub):
    """Headline probabilistic scores for one stratum (vs per-fold-pooled climatology ref)."""
    y = sub.y.values
    p = sub["p_" + HEADLINE].values
    return dict(n=int(np.isfinite(y).sum()),
                base_rate=round(float(np.nanmean(y)), 4),
                brier=round(vf.brier(y, p), 4),
                bss_vs_clim=round(vf.brier_skill_score(y, p, sub.p_climat.values), 3),
                bss_vs_raw=round(vf.brier_skill_score(y, p, sub.p_raw.values), 3),
                auc=round(vf.auc(y, p), 3))


def stratify_season(df, bid, res):
    """BSS/AUC by INIT-season x block for the headline family."""
    rows = []
    for b_id, bname in bid.items():
        for s in SEASON_ORDER:
            sub = df[(df.block_id == b_id) & (df.season == s)]
            if sub.empty:
                continue
            rows.append(dict(block=bname, season=s, **_score_group(sub)))
    out = pd.DataFrame(rows)
    out.to_csv(os.path.join(res, "tables", "cv_scores_by_season_block.csv"), index=False)
    log.info("season x block (%s) BSS vs clim:\n%s", HEADLINE,
             out.pivot(index="block", columns="season", values="bss_vs_clim")
             .reindex(columns=SEASON_ORDER).to_string())
    return out


def stratify_severity(df, bid, res):
    """BSS/AUC by base-rate SEVERITY tercile x block for the headline family.

    Severity = each land cell's pooled climatological EHD frequency within the block;
    cells are cut into Low/Moderate/High terciles (per block, so every block has all three).
    """
    labels = ["Low", "Moderate", "High"]
    rows = []
    for b_id, bname in bid.items():
        blk = df[df.block_id == b_id]
        cell_rate = blk.groupby("cell").y.mean()
        try:
            tier = pd.qcut(cell_rate.rank(method="first"), 3, labels=labels)
        except ValueError:
            continue
        cell_tier = dict(zip(cell_rate.index, tier))
        edges = cell_rate.groupby(tier, observed=True).agg(["min", "max"]).round(3)
        blk = blk.assign(tier=blk.cell.map(cell_tier))
        for lab in labels:
            sub = blk[blk.tier == lab]
            if sub.empty:
                continue
            rows.append(dict(block=bname, severity=lab,
                             rate_lo=float(edges.loc[lab, "min"]),
                             rate_hi=float(edges.loc[lab, "max"]),
                             n_cells=int((tier == lab).sum()),
                             **_score_group(sub)))
    out = pd.DataFrame(rows)
    out.to_csv(os.path.join(res, "tables", "cv_scores_by_severity_block.csv"), index=False)
    log.info("severity x block (%s) BSS vs clim:\n%s", HEADLINE,
             out.pivot(index="block", columns="severity", values="bss_vs_clim")
             .reindex(columns=labels).to_string())
    return out


def _strata_figure(by_season, by_sev, bid, res):
    order = [b for _, b in bid.items()]
    fig, ax = plt.subplots(1, 2, figsize=(13, 4.8))
    for s, col in zip(SEASON_ORDER, ["#0072B2", "#009E73", "#D55E00", "#CC79A7"]):
        sd = by_season[by_season.season == s].set_index("block").reindex(order)
        ax[0].plot(range(len(order)), sd.bss_vs_clim, "o-", color=col, label=s)
    ax[0].axhline(0, color="k", lw=0.8); ax[0].set_xticks(range(len(order)))
    ax[0].set_xticklabels(order); ax[0].set_ylabel("BSS vs climatology")
    ax[0].set_xlabel("lead block"); ax[0].set_title("Skill by init season"); ax[0].legend(fontsize=8)
    ax[0].grid(alpha=0.3)
    for lab, col in zip(["Low", "Moderate", "High"], ["#56B4E9", "#E69F00", "#D55E00"]):
        sd = by_sev[by_sev.severity == lab].set_index("block").reindex(order)
        ax[1].plot(range(len(order)), sd.bss_vs_clim, "s-", color=col, label=lab)
    ax[1].axhline(0, color="k", lw=0.8); ax[1].set_xticks(range(len(order)))
    ax[1].set_xticklabels(order); ax[1].set_ylabel("BSS vs climatology")
    ax[1].set_xlabel("lead block"); ax[1].set_title("Skill by base-rate severity tercile")
    ax[1].legend(fontsize=8, title="EHD freq"); ax[1].grid(alpha=0.3)
    fig.tight_layout()
    fig.savefig(os.path.join(res, "figures", "cv_strata.png"), dpi=130)
    log.info("wrote cv_strata.png")


def _figures(cfg, df, bid, by_block, boot, res):
    # reliability by block: headline vs raw
    fig, axes = plt.subplots(2, 3, figsize=(14, 8))
    for ax, (b_id, bname) in zip(axes.ravel(), bid.items()):
        sub = df[df.block_id == b_id]
        ax.plot([0, 1], [0, 1], "k--", lw=0.7)
        for fam, col in [("raw", "#999999"), (HEADLINE, "#D55E00")]:
            conf, obs, cnt = vf.reliability_curve(sub.y.values, sub["p_" + fam].values, 10)
            ax.plot(conf, obs, "-o", color=col, ms=4, label=fam)
        ax.set_title(bname); ax.set_xlim(0, 1); ax.set_ylim(0, 1); ax.grid(alpha=0.3)
        ax.set_xlabel("forecast prob"); ax.set_ylabel("obs freq"); ax.legend(fontsize=8)
    fig.suptitle("Cross-validated reliability by lead block (all 20 folds pooled)")
    fig.tight_layout(rect=[0, 0, 1, 0.96])
    fig.savefig(os.path.join(res, "figures", "cv_reliability.png"), dpi=130)

    # skill vs lead: BSS(clim) with bootstrap CI, AUC
    fig, ax = plt.subplots(1, 2, figsize=(13, 4.8))
    xb = np.arange(len(boot))
    ax[0].axhline(0, color="k", lw=0.8)
    ax[0].errorbar(xb, boot.bss_vs_clim, yerr=[boot.bss_vs_clim - boot.ci_lo, boot.ci_hi - boot.bss_vs_clim],
                   fmt="o-", color="#D55E00", capsize=4)
    ax[0].set_xticks(xb); ax[0].set_xticklabels(boot.block)
    ax[0].set_title("Brier skill vs climatology (95% year-block bootstrap)")
    ax[0].set_xlabel("lead block"); ax[0].set_ylabel("BSS")
    hb = by_block[by_block.family == HEADLINE].set_index("block").loc[boot.block]
    ax[1].plot(xb, hb.auc, "s-", color="#0072B2")
    ax[1].axhline(0.5, color="k", lw=0.8, ls=":")
    ax[1].set_xticks(xb); ax[1].set_xticklabels(boot.block)
    ax[1].set_ylim(0.45, 0.9); ax[1].set_title("AUC by lead block"); ax[1].set_xlabel("lead block")
    fig.tight_layout()
    fig.savefig(os.path.join(res, "figures", "cv_skill_by_lead.png"), dpi=130)
    log.info("wrote cv figures")


if __name__ == "__main__":
    main()