#!/usr/bin/env python
"""Stage 20d — controlled, leakage-safe retrain to test whether the SELECTED extended predictors
(stage 20c) improve out-of-sample skill over the Stage-1 forecast backbone.

For each held-out year (fully nested LOYO, fold-specific threshold), per block, we train two
matched HGB classifiers on identical folds:
  base : the 10 Stage-1 forecast-backbone features only
  sel  : backbone + the block's SELECTED extended features (feature_selection_by_block.csv)
Both are isotonic-recalibrated on TRAIN and get a per-region x block Peirce decision threshold on
TRAIN. Only the feature set differs, so any skill gap is attributable to the added predictors.

Leakage control:
  - base + teleconnection + persistence features and the labels use the FOLD threshold
    (min(Tx90p,35), climatology excluding the held-out year) — recomputed per fold as in Stage 1.
  - gridded predictors are pure forecast fields (threshold-independent), loaded from the stage-20a
    cache; they cannot leak.
  - recalibration and decision thresholds are fit on TRAIN only.

Output: outputs/loyo_pred_v3/pred_v3_<year>.npz with y, region, cell, block_id and
p_base / p_sel (+ flags). Scored by stage 20e.
"""
from __future__ import annotations
import os, sys, logging, argparse, time
import numpy as np
import pandas as pd
import xarray as xr
from sklearn.ensemble import HistGradientBoostingClassifier
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, features_ext as fext, verify as vf

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger("stage20d")
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)
NAME_IDX = {n: i for i, n in enumerate(FEAT_NAMES)}


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 full 28-col candidate matrix + y + cell over `years` (fold thr for base/ext)."""
    ncell = iy.size
    block_names = list(blocks)
    acc = {b: {"X": [], "y": [], "cell": []} 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]["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, rng, metric="Peirce", cap=50000):
    """Per-region Peirce-optimal threshold; the sweep is subsampled (cap) for speed — a scalar
    threshold is stable well below the full ~1e6 rows."""
    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, metric)[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_hgb(Xtr, ytr, Xte):
    clf = HistGradientBoostingClassifier(max_iter=300, learning_rate=0.06, max_leaf_nodes=31,
                                         l2_regularization=1.0, random_state=0)
    clf.fit(Xtr, ytr.astype(int))
    ptr = clf.predict_proba(Xtr)[:, 1]; pte = clf.predict_proba(Xte)[:, 1]
    m = np.isfinite(ptr) & np.isfinite(ytr)
    iso = IsotonicRegression(out_of_bounds="clip").fit(ptr[m], ytr[m])
    return iso.transform(pte), iso.transform(ptr)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--folds", type=int, nargs="*", default=None)
    ap.add_argument("--cap", type=int, default=200000, help="train-row cap per block (fair: same rows base & sel)")
    args = ap.parse_args()
    cfg = load_config(); blocks = lead_blocks(cfg); block_names = list(blocks)
    block_ids = {b: k + 1 for k, b in enumerate(block_names)}
    region_ids = cfg["regions"]["ids"]
    outdir = os.path.join(cfg["paths"]["outputs"], "loyo_pred_v3sel"); os.makedirs(outdir, exist_ok=True)
    folds = args.folds or cfg["cross_validation"]["years"]
    rng = np.random.default_rng(cfg["seed"])
    CAP = args.cap

    sel = pd.read_csv(os.path.join(ROOT, "results", "tables", "feature_selection_by_block.csv"))
    sel_cols = {r["block"]: [NAME_IDX[f] for f in str(r["selected"]).split(";")]
                for _, r in sel.iterrows()}
    base_cols = list(range(10))

    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)

    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": []}
        pc = {"p_base": [], "p_sel": [], "flag_base": [], "flag_sel": []}
        for b in block_names:
            Xtr, ytr, ctr = tr[b]; Xte, yte, cte = te[b]
            reg_tr = region_of_cell[ctr]; reg_te = region_of_cell[cte]
            mtr = np.isfinite(ytr)
            Xtr_m, ytr_m, reg_tr_m = Xtr[mtr], ytr[mtr], reg_tr[mtr]
            # fair train-row cap: identical subsample for base and sel
            if CAP and ytr_m.size > CAP:
                si = rng.choice(ytr_m.size, CAP, replace=False)
                Xtr_m, ytr_m, reg_tr_m = Xtr_m[si], ytr_m[si], reg_tr_m[si]
            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 tag, idx in (("base", base_cols), ("sel", sel_cols[b])):
                pte, ptr = fit_hgb(Xtr_m[:, idx], ytr_m, Xte[:, idx])
                rthr = region_thresholds(ytr_m, ptr, reg_tr_m, region_ids, rng)
                tvec = np.array([rthr[int(r)] for r in reg_te])
                pc[f"p_{tag}"].append(pte.astype(np.float32))
                pc[f"flag_{tag}"].append((pte >= tvec))
        pred = {k: np.concatenate(v) for k, v in cols.items()}
        for k, v in pc.items():
            pred[k] = np.concatenate(v)
        np.savez_compressed(os.path.join(outdir, f"pred_v3_{yr}.npz"), **pred)
        w1 = pred["block_id"] == 1
        bss_b = vf.brier_skill_score(pred["y"][w1], pred["p_base"][w1],
                                     np.full(w1.sum(), np.nanmean(tr["W1"][1])))
        bss_s = vf.brier_skill_score(pred["y"][w1], pred["p_sel"][w1],
                                     np.full(w1.sum(), np.nanmean(tr["W1"][1])))
        log.info("fold %d done in %.0fs | W1 BSS base=%.3f sel=%.3f -> pred_v3_%d.npz",
                 yr, time.time() - t0, bss_b, bss_s, yr)


if __name__ == "__main__":
    main()