#!/usr/bin/env python
"""Stage 28 — LOYO runner for the TRUE-SPATIAL full-field ResU-Net (v4).

Consumes the stage-27 full-field cache (outputs/fullfield_cache/ff_<year>.npz) and trains the
two-branch ffnet.FullFieldResUNet per fold x block, leakage-safe (climatology/threshold,
network weights, isotonic map, decision threshold all fit without the held-out year). Writes
outputs/loyo_pred_v4ff/pred_v4ff_<year>.npz with block_id/region/cell/y matching the tabular
zoo's (init-major, cell-minor) row order + a p_ffresunet column + per-region x block decthr, so
it merges straight into the diversity stack:

  scripts/23_v3_diversity_stack.py --dl-dir outputs/loyo_pred_v4ff --dl-pattern pred_v4ff_*.npz

VN-branch channels = the 51 stage-27 `vn` fields STACKED with the 10 Stage-1 per-cell forecast
features (feat.FEATURE_NAMES); ext-branch channels = the 4 extended-domain fields. Preserves the
v1 hard constraints (six fixed blocks, EHD=min(Tx90p,35), LOYO leakage rule).
"""
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 ffnet

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger("stage28")
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
NFEAT = len(feat.FEATURE_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)
    ny, nx = reg.values.shape
    return da, obs_vals, obs_lut, fc, iy, ix, region_of_cell, ny, nx


def _load_ff(cfg, yr):
    return np.load(os.path.join(cfg["paths"]["outputs"], "fullfield_cache", f"ff_{yr}.npz"),
                   allow_pickle=True)


def year_block_images(cfg, yr, bi, thr, blocks, obs_vals, obs_lut, fc, iy, ix, ny, nx):
    """Return (vn_img, ext_img, lab, mask, y_flat, cell_flat) for one year & block index bi.

    vn_img : (ninit, 51+NFEAT, ny, nx)   ext_img: (ninit, 4, ne, nw)
    lab/mask: (ninit, ny, nx)            y_flat/cell_flat: (ninit*ncell,) init-major cell-minor
    """
    ncell = iy.size
    thr_i = thr.reindex(lat=fc[yr].lat, lon=fc[yr].lon, method="nearest")
    out = feat.assemble_year(fc[yr], obs_vals, obs_lut, thr_i, blocks, iy, ix)
    bname = list(blocks)[bi]
    X = out[bname]["X"]; y = out[bname]["y_primary"]; cell = out[bname]["cell_id"]
    ninit = y.size // ncell
    Xr = X.reshape(ninit, ncell, X.shape[1])                 # init-major, cell-minor
    yr_arr = y.reshape(ninit, ncell)

    ff = _load_ff(cfg, yr)
    vn_ff = ff["vn"][:, bi].transpose(0, 3, 1, 2).astype("float32")   # (ninit,51,ny,nx)
    ext_img = ff["ext"][:, bi].transpose(0, 3, 1, 2).astype("float32")  # (ninit,4,ne,nw)
    assert vn_ff.shape[0] == ninit, f"{yr}: ff ninit {vn_ff.shape[0]} != feat {ninit}"

    # rasterise the 10 forecast features onto the VN grid as extra channels
    feat_img = np.full((ninit, NFEAT, ny, nx), np.nan, "float32")
    for c in range(NFEAT):
        feat_img[:, c, iy, ix] = Xr[:, :, c]
    vn_img = np.concatenate([vn_ff, feat_img], axis=1)        # (ninit,51+NFEAT,ny,nx)

    lab = np.zeros((ninit, ny, nx), "float32"); mask = np.zeros((ninit, ny, nx), "float32")
    valid = np.isfinite(yr_arr)
    lab[:, iy, ix] = np.where(valid, yr_arr, 0.0)
    mask[:, iy, ix] = valid.astype("float32")
    return vn_img, ext_img, lab, mask, y, cell


def stack_years(cfg, years, bi, thr, blocks, obs_vals, obs_lut, fc, iy, ix, ny, nx):
    V, E, L, M, Y, C = [], [], [], [], [], []
    for y in years:
        vn_img, ext_img, lab, mask, yf, cf = year_block_images(
            cfg, y, bi, thr, blocks, obs_vals, obs_lut, fc, iy, ix, ny, nx)
        V.append(vn_img); E.append(ext_img); L.append(lab); M.append(mask)
        Y.append(yf); C.append(cf)
    return (np.concatenate(V), np.concatenate(E), np.concatenate(L), np.concatenate(M),
            np.concatenate(Y), np.concatenate(C))


def region_thresholds(y, p, regions, region_ids, rng, cap=50000):
    def _bt(yy, pp):
        ok = np.isfinite(yy) & np.isfinite(pp); yy, pp = yy[ok], pp[ok]
        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 main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--folds", type=int, nargs="*", default=None)
    ap.add_argument("--blocks", nargs="*", default=None)
    ap.add_argument("--epochs", type=int, default=40)
    ap.add_argument("--ch", type=int, default=32)
    ap.add_argument("--batch", type=int, default=32)
    ap.add_argument("--train-years", type=int, default=0, help="cap #train years (0=all, for smoke)")
    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_v4ff"); os.makedirs(outdir, exist_ok=True)
    folds = args.folds or cfg["cross_validation"]["years"]
    rng = np.random.default_rng(cfg["seed"])

    da, obs_vals, obs_lut, fc, iy, ix, region_of_cell, ny, nx = preload(cfg)
    log.info("grid %dx%d land=%d | blocks=%s epochs=%d ch=%d", ny, nx, iy.size,
             block_names, args.epochs, args.ch)

    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]
        if args.train_years:
            tr_years = tr_years[:args.train_years]

        cols = {"block_id": [], "region": [], "cell": [], "y": []}
        prob_cols = {"p_ffresunet": []}; dec_records = []
        for b in block_names:
            bi = list(blocks).index(b)
            Vtr, Etr, Ltr, Mtr, ytr, ctr = stack_years(
                cfg, tr_years, bi, thr, blocks, obs_vals, obs_lut, fc, iy, ix, ny, nx)
            Vte, Ete, Lte, Mte, yte, cte = year_block_images(
                cfg, yr, bi, thr, blocks, obs_vals, obs_lut, fc, iy, ix, ny, nx)
            reg_tr = region_of_cell[ctr]; 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)

            model = ffnet.FullFieldResUNet(iy, ix, ny, nx, ch=args.ch, max_epoch=args.epochs,
                                           batch=args.batch, seed=cfg["seed"] % 1000)
            model.fit(Vtr, Etr, Ltr, Mtr)
            p_te = model.predict_proba(Vte, Ete)
            p_tr = model.predict_proba(Vtr, Etr)
            prob_cols["p_ffresunet"].append(p_te.astype("float32"))
            rthr = region_thresholds(ytr, p_tr, reg_tr, region_ids, rng)
            for r in region_ids:
                dec_records.append(("ffresunet", b, r, rthr[r]))
            cl = np.full(yte.size, float(np.nanmean(ytr)), "float32")
            log.info("  fold %d %s: BSS=%.3f (n_te=%d, ep<=%d) %.0fs", yr, b,
                     vf.brier_skill_score(yte, p_te, cl), n, args.epochs, time.time() - t0)

        pred = {k: np.concatenate(v).astype("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_v4ff_{yr}.npz"), decthr=dt, **pred)
        log.info("fold %d done %.0fs -> pred_v4ff_%d.npz", yr, time.time() - t0, yr)


if __name__ == "__main__":
    main()