#!/usr/bin/env python
"""Stage 06b — LOYO engine for the SECONDARY target: P(heat-wave = >=3 consecutive EHDs).

Identical machinery and leakage rule to stage 06, but the predictand is the block heat-wave
flag `y_secondary` (from src/ehd/targets.py, straddle rule) instead of the primary >=1-EHD flag.
Everything is still fit without the held-out year: Tx90p threshold, calibration model, isotonic
recalibration, and the per-region x block Peirce decision threshold.

Writes one NPZ of out-of-sample predictions per held-out year to
outputs/loyo_pred_secondary/pred_hw_<year>.npz. Scoring is stage 08b (kept separate).

Note: the heat-wave event is rarer than the primary event, so some region x block cells have
too few positives for a stable per-region threshold and fall back to the global one (as in
stage 06) — this is expected and reported honestly.
"""
from __future__ import annotations
import os, sys, logging, argparse, time
import numpy as np
import xarray as xr
from sklearn.isotonic import IsotonicRegression

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, calibration as cal, verify as vf

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

CAL_FAMILIES = ["logistic", "hgb"]     # trainable calibrators (each also recalibrated)
REF_FAMILIES = ["raw", "climat"]       # non-trainable references


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(cfg, years, thr, blocks, obs_vals, obs_lut, fc, iy, ix):
    """Assemble features + the SECONDARY (heat-wave) label per block."""
    acc = {b: {"X": [], "y": [], "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_secondary"])
            acc[b]["cell"].append(out[b]["cell_id"])
    return {b: (np.concatenate(v["X"]), np.concatenate(v["y"]), 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, g_thr


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--folds", type=int, nargs="*", default=None)
    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_secondary")
    os.makedirs(outdir, exist_ok=True)
    folds = args.folds or cfg["cross_validation"]["years"]

    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)

    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(cfg, train_years, thr, blocks, obs_vals, obs_lut, fc, iy, ix)
        te = assemble(cfg, [yr], thr, blocks, obs_vals, obs_lut, fc, iy, ix)

        cols = {"block_id": [], "region": [], "cell": [], "y": []}
        prob_cols = {}
        dec_thr_records = []
        for b in blocks:
            Xtr, ytr, _ = tr[b]
            Xte, yte, cte = te[b]
            reg_te = region_of_cell[cte]
            reg_tr = region_of_cell[tr[b][2]]
            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)

            for fam in REF_FAMILIES:
                mdl = cal.make_model(fam).fit(Xtr, ytr)
                prob_cols.setdefault(fam, []).append(mdl.predict_proba(Xte))
            for fam in CAL_FAMILIES:
                mdl = cal.make_model(fam).fit(Xtr, ytr)
                ptr = mdl.predict_proba(Xtr); pte = mdl.predict_proba(Xte)
                prob_cols.setdefault(fam, []).append(pte)
                mtr = np.isfinite(ytr) & np.isfinite(ptr)
                iso = IsotonicRegression(out_of_bounds="clip").fit(ptr[mtr], ytr[mtr])
                pte_rc = iso.transform(pte); ptr_rc = iso.transform(ptr)
                prob_cols.setdefault(fam + "_recal", []).append(pte_rc)
                rthr, gthr = region_thresholds(ytr[mtr], ptr_rc[mtr], reg_tr[mtr], region_ids)
                for r in region_ids:
                    dec_thr_records.append((fam, b, r, rthr[r]))

        pred = {k: np.concatenate(v).astype(np.float32) for k, v in cols.items()}
        for fam, lst in prob_cols.items():
            pred["p_" + fam] = np.concatenate(lst).astype(np.float32)
        dt = np.array(dec_thr_records,
                      dtype=[("family", "U12"), ("block", "U4"), ("region", "i4"), ("thr", "f4")])
        np.savez_compressed(os.path.join(outdir, f"pred_hw_{yr}.npz"), decthr=dt, **pred)
        w1 = pred["block_id"] == 1
        base_rate = float(np.nanmean(pred["y"]))
        log.info("fold %d done in %.0fs | overall HW base rate=%.4f | -> pred_hw_%d.npz",
                 yr, time.time() - t0, base_rate, yr)


if __name__ == "__main__":
    main()