#!/usr/bin/env python
"""Stage 21 — unified v3 LOYO engine: expanded zoo + LEAD-DEPENDENT predictors + OOF calibration.

This folds together three Stage-2 threads into one leakage-safe member generator:

  * EXPANDED ZOO (task 1 / doc §2): logistic, hgb, rf, mlp, xgb, lgbm classifiers plus the
    severity regressors ridge_sev/hgb_sev/mlp_sev/xgb_sev/lgbm_sev (predictand sev = block-max
    observed Tx - min(Tx90p,35); sev>=0 <=> yEHDE). Boosting diversity is the cheapest high-value
    recommendation of the hybrid-design note.
  * LEAD-DEPENDENT PREDICTORS (task 2 / doc §3): each block trains on the config `predictors`
    slice of the full base+ext+grid candidate matrix. Short lead = local weather state
    (forecast Tx, T850, CAPE, antecedent persistence); long lead = slow large-scale controls
    (ENSO, MJO, upper-level Z500). Sets are fixed a priori per block -> no selection leakage.
  * OOF CALIBRATION (the Stage-20f lesson): every family's probabilities are isotonic-calibrated
    on inner K-fold-by-year OUT-OF-FOLD predictions, never in-sample, so the calibrated BSS is
    honest. Per-region x block Peirce decision thresholds are fit on those OOF train probs.

Each held-out year is scored by a model that never saw it (climatology/threshold, features,
calibration, decision thresholds all fit on the other 19 years). Writes one NPZ per fold to
outputs/loyo_pred_v3zoo/pred_v3zoo_<year>.npz: block_id, region, cell, y, one calibrated prob
column p_<family> per family, and the decision thresholds (decthr). The multi-model stack and
whole-dataset binary verification are downstream stages (23/25).
"""
from __future__ import annotations
import os, sys, logging, argparse, time
import numpy as np
import xarray as xr
from sklearn.isotonic import IsotonicRegression
from sklearn.model_selection import KFold

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))
from ehd.config import load_config, lead_blocks, block_predictor_index
from ehd import climatology as clim, features as feat, features_ext as fext, verify as vf
from ehd import calibration as cal

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger("stage21")
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
FEAT_NAMES = list(feat.FEATURE_NAMES) + list(fext.EXT_FEATURE_NAMES) + list(fext.GRID_FEATURE_NAMES)

ZOO = ["logistic", "hgb", "rf", "mlp", "xgb", "lgbm",
       "ridge_sev", "hgb_sev", "mlp_sev", "xgb_sev", "lgbm_sev"]


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)
    extgrid = {y: fext.load_extgrid(cfg["paths"]["outputs"], y) for y in cfg["forecast"]["init_years"]}
    return da, obs_vals, obs_lut, fc, iy, ix, region_of_cell, extgrid


def assemble(years, thr, blocks, obs_vals, obs_lut, fc, iy, ix, nino34, omi, extgrid):
    """Per-block (X[full], y, sev, cell, yr) over `years` (per-row year label for inner CV)."""
    ncell = iy.size; block_names = list(blocks)
    acc = {b: {"X": [], "y": [], "sev": [], "cell": [], "yr": []} for b in block_names}
    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)
        ext = fext.ext_features_year(fc[y], obs_vals, obs_lut, thr_i, iy, ix, nino34, omi)["X_ext"]
        eg = extgrid[y]
        for bi, b in enumerate(block_names):
            Xg = fext.ext_grid_block(eg, bi, ncell)
            X = np.column_stack([out[b]["X"], ext, Xg]).astype(np.float32)
            acc[b]["X"].append(X); acc[b]["y"].append(out[b]["y_primary"])
            acc[b]["sev"].append(out[b]["sev"]); acc[b]["cell"].append(out[b]["cell_id"])
            acc[b]["yr"].append(np.full(X.shape[0], y))
    return {b: (np.concatenate(v["X"]), np.concatenate(v["y"]), np.concatenate(v["sev"]),
                np.concatenate(v["cell"]), np.concatenate(v["yr"])) for b, v in acc.items()}


def region_thresholds(y, p, regions, region_ids, rng, cap=50000):
    def _bt(yy, pp):
        if cap and yy.size > cap:
            s = rng.choice(yy.size, cap, replace=False); yy, pp = yy[s], pp[s]
        return vf.best_threshold(yy, pp, "Peirce")[0]
    g = _bt(y, p); out = {}
    for r in region_ids:
        m = regions == r
        out[r] = _bt(y[m], p[m]) if (m.sum() > 200 and np.nansum(y[m]) > 10) else g
    return out


def _fit_predict_raw(name, Xtr, ytr, sevtr, Xq):
    """Fit one family on train, return RAW (uncalibrated) score on Xq.

    Classifiers -> P(y=1); severity regressors -> predicted severity (monotone in P via the
    downstream isotonic map). A fresh model instance is created every call (no state reuse).
    """
    if cal.is_sev_family(name):
        base = cal.make_model(name).base          # underlying regressor, bypass built-in isotonic
        m = np.isfinite(sevtr)
        base.fit(Xtr[m], sevtr[m])
        return base.predict(Xq)
    m = cal.make_model(name)
    m.fit(Xtr, ytr)
    return m.predict_proba(Xq)


def oof_calibrated(name, Xtr, ytr, sevtr, yrtr, Xte, inner, rng):
    """Inner K-fold-by-year OOF calibration for one family.

    Returns (p_te, p_tr_oof, ok) where p_te is the calibrated test probability, p_tr_oof the
    calibrated OOF train probability (for leakage-safe decision thresholds), and ok the finite
    mask over the train rows used. Isotonic is fit ONLY on OOF scores (never in-sample).
    """
    yi = ytr.astype(int)
    uyears = np.array(sorted(set(yrtr.tolist())))
    kf = KFold(n_splits=min(inner, len(uyears)), shuffle=True, random_state=0)
    oof = np.full(ytr.size, np.nan)
    for tr_idx, va_idx in kf.split(uyears):
        va = np.isin(yrtr, uyears[va_idx]); trn = ~va
        oof[va] = _fit_predict_raw(name, Xtr[trn], yi[trn], sevtr[trn], Xtr[va])
    ok = np.isfinite(oof)
    iso = IsotonicRegression(out_of_bounds="clip").fit(oof[ok], yi[ok])
    p_tr_oof = iso.transform(oof[ok])
    raw_te = _fit_predict_raw(name, Xtr, yi, sevtr, Xte)   # refit on FULL train
    return iso.transform(raw_te), p_tr_oof, ok


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--folds", type=int, nargs="*", default=None)
    ap.add_argument("--families", nargs="*", default=ZOO)
    ap.add_argument("--blocks", nargs="*", default=None, help="subset of block names (smoke test)")
    ap.add_argument("--cap", type=int, default=200000, help="train-row cap per block (0 = no cap)")
    ap.add_argument("--inner", type=int, default=4, help="inner K-fold (by year) for OOF calibration")
    args = ap.parse_args()
    cfg = load_config(); blocks = lead_blocks(cfg); block_names = args.blocks or list(blocks)
    block_ids = {b: k + 1 for k, b in enumerate(list(blocks))}
    region_ids = cfg["regions"]["ids"]
    outdir = os.path.join(cfg["paths"]["outputs"], "loyo_pred_v3zoo"); os.makedirs(outdir, exist_ok=True)
    folds = args.folds or cfg["cross_validation"]["years"]
    rng = np.random.default_rng(cfg["seed"]); CAP = args.cap

    # lead-dependent predictor column indices (config-driven, fixed a priori -> no selection leakage)
    blk_idx = {b: block_predictor_index(cfg, b, FEAT_NAMES) for b in block_names}
    for b in block_names:
        from ehd.config import block_predictors
        log.info("block %s predictors (%d): %s", b, len(blk_idx[b]), block_predictors(cfg, b))

    nino34 = fext.load_nino34(os.environ.get("NINO34_TXT",
                              os.path.join(ROOT, "..", "HW", "data", "external", "nino34.txt")))
    omi = fext.load_omi(os.environ.get("OMI_TXT",
                        os.path.join(ROOT, "..", "HW", "data", "external", "omi.txt")))
    da, obs_vals, obs_lut, fc, iy, ix, region_of_cell, extgrid = preload(cfg)
    log.info("families=%s cap=%s inner=%s blocks=%s", args.families, CAP, args.inner, block_names)

    for yr in folds:
        t0 = time.time()
        base_years = [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_years, 90, 15), 35.0, "min")
        tr_years = [y for y in cfg["cross_validation"]["years"] if y != yr]
        tr = assemble(tr_years, thr, blocks, obs_vals, obs_lut, fc, iy, ix, nino34, omi, extgrid)
        te = assemble([yr], thr, blocks, obs_vals, obs_lut, fc, iy, ix, nino34, omi, extgrid)

        cols = {"block_id": [], "region": [], "cell": [], "y": []}
        prob_cols = {}; dec_records = []
        for b in block_names:
            Xtr, ytr, sevtr, ctr, yrtr = tr[b]
            Xte, yte, cte = te[b][0], te[b][1], te[b][3]
            reg_tr = region_of_cell[ctr]; reg_te = region_of_cell[cte]
            m = np.isfinite(ytr)
            Xtr, ytr, sevtr, yrtr, reg_tr = Xtr[m], ytr[m], sevtr[m], yrtr[m], reg_tr[m]
            if CAP and ytr.size > CAP:
                si = rng.choice(ytr.size, CAP, replace=False)
                Xtr, ytr, sevtr, yrtr, reg_tr = Xtr[si], ytr[si], sevtr[si], yrtr[si], reg_tr[si]
            idx = blk_idx[b]
            Xtr_b, Xte_b = Xtr[:, idx], Xte[:, idx]

            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 args.families:
                p_te, p_tr_oof, ok = oof_calibrated(fam, Xtr_b, ytr, sevtr, yrtr, Xte_b,
                                                    args.inner, rng)
                prob_cols.setdefault("p_" + fam, []).append(p_te.astype(np.float32))
                rthr = region_thresholds(ytr[ok].astype(int), p_tr_oof, reg_tr[ok], region_ids, rng)
                for r in region_ids:
                    dec_records.append((fam, 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_v3zoo_{yr}.npz"), decthr=dt, **pred)

        # quick log: a long-lead block where large-scale predictors should matter
        wsel = "W4" if "W4" in block_names else block_names[0]
        w = pred["block_id"] == block_ids[wsel]
        cl = np.full(int(w.sum()), float(np.nanmean(pred["y"][w])))
        best = max(args.families,
                   key=lambda f: vf.brier_skill_score(pred["y"][w], pred["p_" + f][w], cl))
        log.info("fold %d %.0fs | %s best=%s BSS=%.3f AUC=%.3f -> pred_v3zoo_%d.npz",
                 yr, time.time() - t0, wsel, best,
                 vf.brier_skill_score(pred["y"][w], pred["p_" + best][w], cl),
                 vf.auc(pred["y"][w], pred["p_" + best][w]), yr)


if __name__ == "__main__":
    main()