#!/usr/bin/env python
"""Stage 12b — LOYO for the torch-CPU families (CNN + deep ANN), v2 Phase B.

Mirrors stage 12's leakage-safe assembly but trains only the two new families and writes their
out-of-sample P(yEHDE) so stage 14 can fold them into EHD_ens and stages 15/16 can map/report
them. Single-level (no cascade; cutoff 1.0) to match the ensemble members.

  ann_torch  deeper torch MLP on the block-aggregated tabular features (train-capped).
  cnn        fully-convolutional net on the (init, channel, lat, lon) grid (full field, no cap;
             convolutions let each cell use its neighbours).

Writes outputs/loyo_pred_v2_cnn/pred_v2cnn_<year>.npz: block_id, region, cell, y,
p_cnn_c1, p_anntorch_c1, and per-region decision thresholds (fit on TRAIN, leakage-safe).
"""
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 gridmodels as gm

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


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)
    ny, nx = ref.sizes["lat"], ref.sizes["lon"]
    return da, obs_vals, obs_lut, fc, iy, ix, region_of_cell, ny, nx


def assemble(years, thr, blocks, obs_vals, obs_lut, fc, iy, ix):
    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_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, metric="Peirce"):
    g_thr, _ = 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_thr
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--folds", type=int, nargs="*", default=None)
    ap.add_argument("--families", nargs="*", default=["cnn", "ann_torch"])
    ap.add_argument("--train-cap", type=int, default=200000, help="cap for ann_torch (cnn uncapped)")
    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_cnn")
    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, ny, nx = preload(cfg)
    ncell = iy.size
    geom = dict(iy=iy, ix=ix, ny=ny, nx=nx, ncell=ncell)
    log.info("families=%s grid=%dx%d ncell=%d", args.families, ny, nx, ncell)

    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 = {f"p_{f}_c1": [] for f in args.families}
        dec_records = []
        for b in blocks:
            Xtr, ytr, ctr = tr[b]; Xte, yte, cte = te[b]
            reg_te = region_of_cell[cte]
            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:
                if fam == "cnn":
                    model = gm.make_torch_model("cnn", geom=geom)   # full field, no subsample
                    model.fit(Xtr, ytr)
                    p_te = model.predict_proba(Xte); p_tr = model.predict_proba(Xtr)
                else:  # ann_torch (train-capped, seeded)
                    rng = np.random.default_rng(seed + yr * 10 + block_ids[b])
                    sel = (np.arange(ytr.size) if ytr.size <= args.train_cap
                           else np.sort(rng.choice(ytr.size, args.train_cap, replace=False)))
                    model = gm.make_torch_model("ann_torch")
                    model.fit(Xtr[sel], ytr[sel])
                    p_te = model.predict_proba(Xte); p_tr = model.predict_proba(Xtr)
                prob_cols[f"p_{fam}_c1"].append(p_te.astype(np.float32))
                mtr = np.isfinite(ytr) & np.isfinite(p_tr)
                rthr = region_thresholds(ytr[mtr], p_tr[mtr], region_of_cell[ctr][mtr], region_ids)
                for r in region_ids:
                    dec_records.append((f"{fam}_c1", b, r, rthr[r]))
            log.info("  fold %d %s done (%.0fs elapsed)", yr, b, time.time() - t0)

        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_v2cnn_{yr}.npz"), decthr=dt, **pred)
        w1 = pred["block_id"] == 1
        for fam in args.families:
            p = pred[f"p_{fam}_c1"]
            log.info("fold %d %s done in %.0fs | W1 AUC=%.3f -> pred_v2cnn_%d.npz",
                     yr, fam, time.time() - t0, vf.auc(pred["y"][w1], p[w1]), yr)


if __name__ == "__main__":
    main()