#!/usr/bin/env python
"""Stage 12 — v2 leave-one-year-out engine: severity predictand + expanded zoo + cascade.

For each held-out reforecast year EVERYTHING is fit without that year (leakage rule), exactly
as stage 06, but v2 adds:
  * PREDICTAND = severity `sev = block-max(Tx_obs - Thr)` for the regressor families
    (ridge_sev, hgb_sev, mlp_sev); classifiers (logistic, hgb, rf, mlp) still learn y directly.
    sev>=0 <=> yEHDE (>=1 EHD in block), so v2 scores the SAME event as v1 (comparable).
  * TWO-LEVEL CASCADE (src/ehd/cascade.py): an LV1 hgb classifier (fit once per fold x block,
    isotonic-recalibrated) screens "nearly certain nEHDE" samples (P(yEHDE) <= 1-cutoff) out of
    the LV2 training/scoring pool. `--cutoffs 1.0 0.90 0.95 0.99` (1.0 = no screening baseline).
  * TRAIN SUBSAMPLING (`--train-cap`, seeded) so the expanded zoo x cutoff sweep is feasible;
    the TEST set is never subsampled, so verification (stage 13) is on the whole dataset.

Writes one NPZ of out-of-sample predictions per held-out year to outputs/loyo_pred_v2/
pred_v2_<year>.npz: block_id, region, cell, y, one prob column `p_<family>_c<cutoff>` per
(family, cutoff), and the per-region decision thresholds (fit on TRAIN combined probs). Scoring
and the whole-dataset binary skill table (POD/POFD/FAR/EDI/SEDI...) are stage 13.
"""
from __future__ import annotations
import os, sys, logging, argparse, time
import numpy as np
import xarray as xr

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))
from ehd.config import load_config, lead_blocks
from ehd import climatology as clim, features as feat, verify as vf
from ehd import cascade as csc

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

ZOO = ["logistic", "hgb", "rf", "mlp", "xgb", "lgbm",
       "ridge_sev", "hgb_sev", "mlp_sev", "xgb_sev", "lgbm_sev"]
CUTOFFS = [1.0, 0.90, 0.95, 0.99]
LV1_FAMILY = "hgb"


def preload(cfg):
    da = clim.load_gridded_obs(cfg["paths"]["gridded_obs"])
    obs_vals, obs_lut = feat.build_obs_index(da)
    fc, ref = {}, None
    for y in cfg["forecast"]["init_years"]:
        fc[y] = xr.open_dataarray(os.path.join(cfg["paths"]["outputs"],
                                               "forecast_daily_tx", f"fc_daily_tx_{y}.nc")).load()
        ref = fc[y]
    reg = xr.open_dataset(cfg["paths"]["mask"])["reg"].reindex(
        lat=ref.lat, lon=ref.lon, method="nearest")
    iy, ix = np.where(np.isfinite(reg.values))
    region_of_cell = reg.values[iy, ix].astype(int)
    return da, obs_vals, obs_lut, fc, iy, ix, region_of_cell


def assemble(years, thr, blocks, obs_vals, obs_lut, fc, iy, ix):
    """Return per-block (X, y, sev, cell) concatenated over `years`."""
    acc = {b: {"X": [], "y": [], "sev": [], "cell": []} for b in blocks}
    for y in years:
        thr_i = thr.reindex(lat=fc[y].lat, lon=fc[y].lon, method="nearest")
        out = feat.assemble_year(fc[y], obs_vals, obs_lut, thr_i, blocks, iy, ix)
        for b in blocks:
            acc[b]["X"].append(out[b]["X"])
            acc[b]["y"].append(out[b]["y_primary"])
            acc[b]["sev"].append(out[b]["sev"])
            acc[b]["cell"].append(out[b]["cell_id"])
    return {b: (np.concatenate(v["X"]), np.concatenate(v["y"]),
                np.concatenate(v["sev"]), np.concatenate(v["cell"]))
            for b, v in acc.items()}


def region_thresholds(y, p, regions, region_ids, metric="Peirce"):
    g_thr, _ = vf.best_threshold(y, p, metric)
    out = {}
    for r in region_ids:
        m = regions == r
        if m.sum() > 200 and np.nansum(y[m]) > 10:
            out[r], _ = vf.best_threshold(y[m], p[m], metric)
        else:
            out[r] = g_thr
    return out


def subsample(n, cap, rng):
    if cap is None or n <= cap:
        return np.arange(n)
    return np.sort(rng.choice(n, size=cap, replace=False))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--folds", type=int, nargs="*", default=None)
    ap.add_argument("--families", nargs="*", default=ZOO)
    ap.add_argument("--cutoffs", type=float, nargs="*", default=CUTOFFS)
    ap.add_argument("--train-cap", type=int, default=200000,
                    help="max TRAIN rows per block (seeded subsample); test never capped")
    args = ap.parse_args()
    cfg = load_config(); blocks = lead_blocks(cfg)
    block_ids = {b: k + 1 for k, b in enumerate(blocks)}
    region_ids = cfg["regions"]["ids"]
    outdir = os.path.join(cfg["paths"]["outputs"], "loyo_pred_v2")
    os.makedirs(outdir, exist_ok=True)
    folds = args.folds or cfg["cross_validation"]["years"]
    seed = int(cfg.get("seed", 0))

    log.info("preloading obs + %d forecast years ...", len(cfg["forecast"]["init_years"]))
    da, obs_vals, obs_lut, fc, iy, ix, region_of_cell = preload(cfg)
    log.info("families=%s cutoffs=%s train_cap=%s", args.families, args.cutoffs, args.train_cap)

    for yr in folds:
        t0 = time.time()
        base = [y for y in range(cfg["obs"]["gridded_years"][0], cfg["obs"]["gridded_years"][1] + 1)
                if y != yr]
        thr = clim.ehd_threshold(clim.tx90p_climatology(da, base, 90, 15), 35.0, "min")
        train_years = [y for y in cfg["cross_validation"]["years"] if y != yr]
        tr = assemble(train_years, thr, blocks, obs_vals, obs_lut, fc, iy, ix)
        te = assemble([yr], thr, blocks, obs_vals, obs_lut, fc, iy, ix)

        cols = {"block_id": [], "region": [], "cell": [], "y": []}
        prob_cols = {}
        dec_records = []
        for b in blocks:
            Xtr, ytr, sevtr, ctr = tr[b]
            Xte, yte, sevte, cte = te[b]
            reg_te = region_of_cell[cte]
            # seeded train subsample (per fold x block) for feasibility
            rng = np.random.default_rng(seed + yr * 10 + block_ids[b])
            sel = subsample(ytr.size, args.train_cap, rng)
            Xtr, ytr, sevtr, ctr = Xtr[sel], ytr[sel], sevtr[sel], ctr[sel]
            reg_tr = region_of_cell[ctr]

            n = yte.size
            cols["block_id"].append(np.full(n, block_ids[b]))
            cols["region"].append(reg_te)
            cols["cell"].append(cte)
            cols["y"].append(yte)

            # LV1 once per fold x block (independent of family & cutoff)
            p1_tr, p1_te = csc.screen_lv1(Xtr, ytr, Xte, LV1_FAMILY)
            for cutoff in args.cutoffs:
                for fam in args.families:
                    p_te, p_tr, info = csc.run_cascade(
                        Xtr, ytr, sevtr, Xte, fam, cutoff, p1_tr, p1_te)
                    tag = f"{fam}_c{cutoff:g}"
                    prob_cols.setdefault("p_" + tag, []).append(p_te.astype(np.float32))
                    mtr = np.isfinite(ytr) & np.isfinite(p_tr)
                    rthr = region_thresholds(ytr[mtr], p_tr[mtr], reg_tr[mtr], region_ids)
                    for r in region_ids:
                        dec_records.append((tag, b, r, rthr[r]))

        pred = {k: np.concatenate(v).astype(np.float32) for k, v in cols.items()}
        for k, lst in prob_cols.items():
            pred[k] = np.concatenate(lst)
        dt = np.array(dec_records,
                      dtype=[("tag", "U20"), ("block", "U4"), ("region", "i4"), ("thr", "f4")])
        np.savez_compressed(os.path.join(outdir, f"pred_v2_{yr}.npz"), decthr=dt, **pred)

        w1 = pred["block_id"] == 1
        best = f"p_hgb_c{args.cutoffs[-1]:g}" if f"p_hgb_c{args.cutoffs[-1]:g}" in pred else \
               ("p_" + args.families[0] + f"_c{args.cutoffs[0]:g}")
        clr = np.full(w1.sum(), np.nanmean(tr["W1"][1]))
        bss = vf.brier_skill_score(pred["y"][w1], pred[best][w1], clr)
        log.info("fold %d done in %.0fs | W1 %s BSS_clim=%.3f AUC=%.3f | -> pred_v2_%d.npz",
                 yr, time.time() - t0, best, bss, vf.auc(pred["y"][w1], pred[best][w1]), yr)


if __name__ == "__main__":
    main()