#!/usr/bin/env python
"""Stage 20 — item 2: EXPAND the member set (registry v3), leakage-safe LOYO.

Adds new CONFIGURATIONS to families that earned their place in Stage 1, targeting W2-W6 where
operational value lives and Stage-1 skill is thinnest. Two kinds of new member, both single-level
(cutoff 1.0; the cascade screen was shown ~neutral):

  <fam>_tele   : the family trained on the Stage-1 forecast features AUGMENTED with the
                 subseasonal expansion predictors (ehd.features_ext): ENSO NINO3.4, MJO OMI
                 PC1/PC2/amplitude, and antecedent-persistence (obs EHD frequency + mean severity
                 over the 15 days before init). Families: logistic, hgb, rf, mlp, hgb_sev.
  <fam>_bal    : class-imbalance-aware configurations (class_weight balanced) on the SAME
                 augmented features, for logistic and rf (ehd.expand_models).

Everything is fit without the held-out year (leakage rule), exactly as stage 12: fold-safe Tx90p
threshold, per-region x block decision thresholds fit on TRAIN, isotonic recalibration on TRAIN.
Writes outputs/loyo_pred_v3/pred_v3_<year>.npz with one `p_<tag>` column per new member plus
block_id/region/cell/y and the per-region thresholds. Downstream stages merge v2 + cnn + v3.

Run per fold on SLURM (submit_stage2_loyo.slurm); heavy compute never on the login node.
"""
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 features_ext as fext, calibration as cal, expand_models as xm
from sklearn.isotonic import IsotonicRegression

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

# new members: (kind, family) -> tag <family>_<kind>_c1
TELE_FAMILIES = ["logistic", "hgb", "rf", "mlp", "hgb_sev"]
BAL_FAMILIES = ["logistic", "rf"]
NINO34 = os.environ.get("NINO34_TXT", "/home/tanpv/Progs/claude/HW/data/external/nino34.txt")
OMI = os.environ.get("OMI_TXT", "/home/tanpv/Progs/claude/HW/data/external/omi.txt")


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))
    return da, obs_vals, obs_lut, fc, iy, ix, reg.values[iy, ix].astype(int)


def assemble_aug(years, thr, blocks, obs_vals, obs_lut, fc, iy, ix, nino34, omi):
    """Per-block (X_aug, y, sev, cell) over `years`; X_aug = Stage-1 X + ext features."""
    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)
        ext = fext.ext_features_year(fc[y], obs_vals, obs_lut, thr_i, iy, ix, nino34, omi)["X_ext"]
        for b in blocks:
            acc[b]["X"].append(np.column_stack([out[b]["X"], ext]))   # ext aligns row-for-row
            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, _ = vf.best_threshold(y, p, metric)
    out = {}
    for r in region_ids:
        m = regions == r
        out[r] = vf.best_threshold(y[m], p[m], metric)[0] if (m.sum() > 200 and np.nansum(y[m]) > 10) else g
    return out


def _recal(p_tr, y, p_te):
    m = np.isfinite(p_tr) & np.isfinite(y)
    iso = IsotonicRegression(out_of_bounds="clip").fit(p_tr[m], y[m].astype(int))
    return iso.transform(p_tr), iso.transform(p_te)


def fit_config(kind, fam, Xtr, ytr, sevtr, Xte):
    """Return (p_te, p_tr) for one expansion config, isotonic-recalibrated on TRAIN."""
    if kind == "tele":
        model = cal.make_model(fam)
        if cal.is_sev_family(fam):
            model.fit(Xtr, sevtr, ytr)                 # sev regressors self-calibrate (isotonic)
            return model.predict_proba(Xte), model.predict_proba(Xtr)
        model.fit(Xtr, ytr)
    else:                                              # class-balanced
        model = xm.make_balanced(fam).fit(Xtr, ytr)
    p_tr, p_te = model.predict_proba(Xtr), model.predict_proba(Xte)
    p_tr, p_te = _recal(p_tr, ytr, p_te)
    return p_te, p_tr


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


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

    log.info("expansion predictors: nino34=%s omi=%s", NINO34, OMI)
    nino34 = fext.load_nino34(NINO34); omi = fext.load_omi(OMI)
    log.info("preloading obs + forecasts ...")
    da, obs_vals, obs_lut, fc, iy, ix, region_of_cell = preload(cfg)
    configs = [("tele", f) for f in TELE_FAMILIES] + [("bal", f) for f in BAL_FAMILIES]
    log.info("new members: %s", [f"{f}_{'tele' if k=='tele' else 'bal'}_c1" for k, f in configs])

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

        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]
            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)
            for kind, fam in configs:
                tag = f"{fam}_{'tele' if kind=='tele' else 'bal'}_c1"
                p_te, p_tr = fit_config(kind, fam, Xtr, ytr, sevtr, Xte)
                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", "U24"), ("block", "U4"),
                                          ("region", "i4"), ("thr", "f4")])
        np.savez_compressed(os.path.join(outdir, f"pred_v3_{yr}.npz"), decthr=dt, **pred)
        log.info("fold %d done in %.0fs -> pred_v3_%d.npz (%d new members)",
                 yr, time.time() - t0, yr, len(configs))


if __name__ == "__main__":
    main()