#!/usr/bin/env python3
"""Stage 2 section 7: national model, trained and verified across several nodes.

The national problem is 16 years x 104 initialisations x 46 lead times x 541 land
points = 41.4 million training rows, and assembling those rows out of the memmaps
costs far more than fitting the trees.  The work is therefore split into phases
that each use the cluster in the shape that suits them:

  prep      120 single-core tasks (6 nodes x 20), one shard of initialisations
            each, writing thinned training rows to disk
  fit       5 tasks of 20 cores, run at the same time: four fold models whose
            held-out predictions give the decision thresholds, plus the final
            model fitted on all sixteen training years
  predict   120 single-core tasks again, scoring the held-out years shard by shard
  finalize  one task that pools everything, applies the per-week thresholds and
            writes the gridded metrics

Shards are assigned round-robin over initialisations so every task gets a mix of
seasons and years, and the fold split is derived from a fixed seed so every task
computes the same one without having to communicate.
"""

from __future__ import annotations

import argparse
import glob
import json
import os
import sys
import time

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import ehd_common as C
import ehd_metrics as M
import ai_data as D
import train_ai as T

RES_DIR = os.path.join(C.OUT, "AI_Results")
NFOLD = 4
CHUNK = 52          # initialisations held in memory at once inside one task


# --------------------------------------------------------------------- layout
def shard_dir(tag=""):
    """Scratch directory of one national run; the tag keeps parallel runs apart."""
    return os.path.join(RES_DIR, "shards" + tag)


def shard_of(isel, nshard, k):
    """Round-robin slice of the initialisation list for task k."""
    return np.asarray(isel)[k::nshard]


def year_folds(nfold=NFOLD, seed=0):
    """The same year split in every task, derived from a fixed seed."""
    rng = np.random.default_rng(seed)
    return [np.sort(f) for f in
            np.array_split(rng.permutation(np.array(C.TRAIN_YEARS)), nfold)]


# ------------------------------------------------------------------- strata
# Two ways of cutting the country up, both suggested in CLAUDE.md: by climatic
# sub-region (the 7 zones of the mask) and by season.  A cut can be used two
# ways, and they are independent: a separate MODEL per stratum, and a separate
# decision THRESHOLD per stratum.  The threshold cut matters even when one
# national model is used, because one national threshold is being asked to serve
# climates whose event rates differ by an order of magnitude.
SEASON_NAMES = ["DJF", "MAM", "JJA", "SON"]
_MONTH_END = np.cumsum([31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31])


def season_of(doy):
    """0=DJF 1=MAM 2=JJA 3=SON from a 1-based day of year."""
    month = np.searchsorted(_MONTH_END, np.clip(doy, 1, 365), side="left") + 1
    return ((month % 12) // 3).astype("i4")


def stratum_of(region, doy, mode):
    """Integer stratum label per row for the requested cut."""
    if mode in (None, "", "none"):
        return np.zeros(np.asarray(region).size, "i4")
    if mode == "region":
        return np.asarray(region).astype("i4")
    if mode == "season":
        return season_of(np.asarray(doy))
    if mode == "region_season":
        return np.asarray(region).astype("i4") * 4 + season_of(np.asarray(doy))
    raise ValueError(mode)


def train_inits(store):
    return np.nonzero(np.isin(store.init_year, C.TRAIN_YEARS))[0]


def valid_inits(store):
    return np.nonzero(np.isin(store.init_year, C.VALID_YEARS))[0]


# ----------------------------------------------------------------- phase prep
def cmd_prep(args):
    store = D.Store()
    points = np.arange(store.grid.npt)
    mine = shard_of(train_inits(store), args.nshard, args.shard)
    rng = np.random.default_rng(1000 + args.shard)

    Xs, evt, sev, week, year, cap = [], [], [], [], [], []
    reg, doy, ini, lt = [], [], [], []
    for a in range(0, mine.size, CHUNK):
        X, y, info = store.build_isel(points, mine[a:a + CHUNK], args.fset)
        ok = np.isfinite(X).all(axis=1) & np.isfinite(y["evt"])
        pos = ok & (y["evt"] > 0.5)
        neg = ok & ~pos
        if args.neg_keep < 1.0:
            neg &= rng.random(neg.size) < args.neg_keep
        take = pos | neg
        Xs.append(X[take]); evt.append(y["evt"][take]); sev.append(y["sev"][take])
        week.append(info["week"][take]); year.append(info["year"][take])
        cap.append(info["capable"][take])
        reg.append(info["region"][take]); doy.append(info["doy"][take])
        # (init, lead, region) is what joins a row to its region-day, which EXP2
        # needs to tell the yEHD rows from the nEHD ones
        ini.append(info["init"][take]); lt.append(info["lt"][take])
        del X, y, info

    os.makedirs(shard_dir(args.tag), exist_ok=True)
    np.savez(os.path.join(shard_dir(args.tag), "train_%04d.npz" % args.shard),
             X=np.concatenate(Xs).astype("f4"),
             evt=np.concatenate(evt).astype("f4"),
             sev=np.concatenate(sev).astype("f4"),
             week=np.concatenate(week).astype("i2"),
             year=np.concatenate(year).astype("i4"),
             capable=np.concatenate(cap).astype("i1"),
             region=np.concatenate(reg).astype("i1"),
             doy=np.concatenate(doy).astype("i2"),
             init=np.concatenate(ini).astype("i4"),
             lt=np.concatenate(lt).astype("i1"))
    print("[prep %03d] %d initialisations -> %d rows"
          % (args.shard, mine.size, sum(x.shape[0] for x in Xs)), flush=True)


def load_train(tag=""):
    """Concatenate every training shard."""
    files = sorted(glob.glob(os.path.join(shard_dir(tag), "train_*.npz")))
    if not files:
        raise SystemExit("no training shards in %s" % shard_dir(tag))
    parts = [np.load(f) for f in files]
    X = np.concatenate([p["X"] for p in parts])
    y = {"evt": np.concatenate([p["evt"] for p in parts]),
         "sev": np.concatenate([p["sev"] for p in parts])}
    g = {"week": np.concatenate([p["week"] for p in parts]).astype("i4"),
         "year": np.concatenate([p["year"] for p in parts])}
    if "capable" in parts[0].files:
        g["capable"] = np.concatenate([p["capable"] for p in parts]).astype(bool)
    for k in ("region", "doy", "init", "lt"):
        if k in parts[0].files:
            g[k] = np.concatenate([p[k] for p in parts]).astype("i4")
    return X, y, g


def group_mask(g, gate_file, group):
    """Rows of the yEHD or the nEHD half of EXP2.

    The split is observational, not modelled: a region-day is yEHD when the
    fraction of its cells that actually had EHD=1 exceeds the coverage of the
    EXP1 rule.  The labels come from the same npz the gate was written to, keyed
    by (init, lead, region), so the two halves line up exactly with what the gate
    is trained to separate.
    """
    if group == "all":
        return np.ones(g["week"].size, bool)
    if not all(k in g for k in ("init", "lt")):
        raise SystemExit("training shards have no init/lt: re-run prep for this tag")
    G = np.load(gate_file)
    key = (G["tr_init"].astype("i8") * C.NLT + G["tr_lead"].astype("i8")) * 8 \
        + G["tr_region"]
    o = np.argsort(key, kind="stable")
    ks, ys = key[o], (G["tr_Pobs"][o] > float(G["cover"]))
    mine = (g["init"].astype("i8") * C.NLT + g["lt"].astype("i8")) * 8 + g["region"]
    pos = np.clip(np.searchsorted(ks, mine), 0, ks.size - 1)
    hit = ks[pos] == mine
    keep = np.zeros(mine.size, bool)
    keep[hit] = ys[pos[hit]] if group == "yEHD" else ~ys[pos[hit]]
    print("[group %s] giữ %d / %d hàng (%.1f%%), %d hàng không ghép được cổng"
          % (group, int(keep.sum()), keep.size, 100 * keep.mean(),
             int((~hit).sum())), flush=True)
    return keep


# ------------------------------------------------------------------ phase fit
def cmd_fit(args):
    import joblib
    # EXP2 fits two models from one set of training shards, so the shards it reads
    # and the directory the model lands in are allowed to differ.
    os.makedirs(shard_dir(args.tag), exist_ok=True)
    X, y, g = load_train(args.train_tag or args.tag)
    if args.group != "all":
        keep = group_mask(g, args.gate_file, args.group)
        X = X[keep]
        y = {k: v[keep] for k, v in y.items()}
        g = {k: v[keep] for k, v in g.items()}
    strat = stratum_of(g.get("region"), g.get("doy"), args.stratify)
    svals = np.unique(strat)
    print("[fit %s] %s rows, %d stratum(s) by '%s'"
          % (args.fold, X.shape, svals.size, args.stratify or "none"), flush=True)
    t0 = time.time()

    def weight(mask):
        """Positive-class weight from the event rate inside this stratum."""
        rate = float(y["evt"][mask].mean())
        return (1 - rate) / max(rate, 1e-9) * args.neg_keep

    if args.fold == "full":
        models = {}
        for sv in svals:
            m = strat == sv
            if y["evt"][m].sum() == 0:
                continue
            _, mod = T.fit_predict(args.model, X[m], {k: v[m] for k, v in y.items()},
                                   X[m][:1], weight(m))
            models[int(sv)] = mod
        joblib.dump({"models": models, "stratify": args.stratify,
                     "kind": args.model, "fset": args.fset},
                    os.path.join(shard_dir(args.tag), "model.joblib"))
        print("[fit full] %d model(s), %.0fs" % (len(models), time.time() - t0), flush=True)
        return

    f = int(args.fold)
    hold = year_folds()[f]
    te = np.isin(g["year"], hold)
    s_all = np.full(X.shape[0], np.nan)
    fold_models = {}
    for sv in svals:
        ms = strat == sv
        tr, ts = ms & ~te, ms & te
        if tr.sum() == 0 or ts.sum() == 0 or y["evt"][tr].sum() == 0:
            continue
        sc, mod = T.fit_predict(args.model, X[tr], {k: v[tr] for k, v in y.items()},
                                X[ts], weight(tr))
        s_all[ts] = sc
        fold_models[sv] = mod

    # The fold models are what makes an unthinned out-of-fold field possible:
    # a later phase re-scores every training row with the model that held its
    # year out, giving a complete spatial field with no leakage.
    joblib.dump({"models": {int(sv): m for sv, m in fold_models.items()},
                 "stratify": args.stratify, "kind": args.model, "fset": args.fset},
                os.path.join(shard_dir(args.tag), "model_fold%d.joblib" % f))

    extra = {k: g[k][te] for k in ("region", "doy") if k in g}
    if "capable" in g:
        extra["capable"] = g["capable"][te].astype("i1")
    np.savez(os.path.join(shard_dir(args.tag), "oof_%d.npz" % f),
             score=s_all[te].astype("f8"), evt=y["evt"][te], week=g["week"][te], **extra)
    print("[fit fold %d] held out %s, %d rows, %.0fs"
          % (f, hold.tolist(), int(te.sum()), time.time() - t0), flush=True)


# -------------------------------------------------------------- phase predict
def apply_bundle(bundle, X, info, bad):
    """Score rows with the right stratum model, leaving unusable rows as NaN."""
    kind = bundle["kind"]
    models = bundle.get("models") or {0: bundle["model"]}
    st = stratum_of(info["region"], info["doy"], bundle.get("stratify", "none"))
    s = np.full(X.shape[0], np.nan)
    for sv, mod in models.items():
        m = st == sv
        if not m.any():
            continue
        s[m] = (mod.predict(X[m]) if kind == "lgbm_reg"
                else mod.predict_proba(X[m])[:, 1])
    s[bad] = np.nan
    return s


def cmd_predict_oof(args):
    """Score every training-year row with the fold model that held its year out.

    The `prep` phase throws away 85% of the non-events, so the out-of-fold rows
    it leaves behind cannot be aggregated over space -- a regional coverage
    fraction needs every cell present.  This pass rebuilds the full field: no row
    is ever scored by a model that saw its year.
    """
    import joblib
    store = D.Store()
    points = np.arange(store.grid.npt)
    folds = year_folds()
    bundles = {f: joblib.load(os.path.join(shard_dir(args.tag),
                                           "model_fold%d.joblib" % f))
               for f in range(NFOLD)}
    mine = shard_of(train_inits(store), args.nshard, args.shard)

    sc, evt, keep = [], [], {k: [] for k in
                             ("week", "lt", "point", "capable", "init", "region", "doy")}
    for a in range(0, mine.size, CHUNK):
        X, y, info = store.build_isel(points, mine[a:a + CHUNK], args.fset)
        bad = ~np.isfinite(X).all(axis=1)
        X[bad] = 0.0
        s = np.full(X.shape[0], np.nan)
        for f, hold in enumerate(folds):
            m = np.isin(info["year"], hold)
            if m.any():
                s[m] = apply_bundle(bundles[f], X[m], {k: v[m] for k, v in info.items()},
                                    bad[m])
        sc.append(s.astype("f4")); evt.append(y["evt"].astype("f4"))
        for k in keep:
            keep[k].append(info[k])
        del X, y, info

    np.savez(os.path.join(shard_dir(args.tag), "oofpred_%04d.npz" % args.shard),
             score=np.concatenate(sc), evt=np.concatenate(evt),
             **{k: np.concatenate(v) for k, v in keep.items()})
    print("[oofpred %03d] %d initialisations -> %d rows"
          % (args.shard, mine.size, sum(x.size for x in sc)), flush=True)


def cmd_predict(args):
    import joblib
    bundle = joblib.load(os.path.join(shard_dir(args.tag), "model.joblib"))
    store = D.Store()
    points = np.arange(store.grid.npt)
    mine = shard_of(valid_inits(store), args.nshard, args.shard)

    sc, evt, keep = [], [], {k: [] for k in
                             ("week", "lt", "point", "capable", "init", "region", "doy")}
    for a in range(0, mine.size, CHUNK):
        X, y, info = store.build_isel(points, mine[a:a + CHUNK], args.fset)
        bad = ~np.isfinite(X).all(axis=1)
        X[bad] = 0.0
        s = apply_bundle(bundle, X, info, bad)
        sc.append(s.astype("f4")); evt.append(y["evt"].astype("f4"))
        for k in keep:
            keep[k].append(info[k])
        del X, y, info

    np.savez(os.path.join(shard_dir(args.tag), "pred_%04d.npz" % args.shard),
             score=np.concatenate(sc), evt=np.concatenate(evt),
             **{k: np.concatenate(v) for k, v in keep.items()})
    print("[pred %03d] %d initialisations -> %d rows"
          % (args.shard, mine.size, sum(x.size for x in sc)), flush=True)


# ------------------------------------------------------------- phase finalize
def cmd_finalize(args):
    import joblib
    store = D.Store()
    grid = store.grid
    try:
        names = store.feature_names(args.fset)
    except Exception:
        names = []          # the CNN has no tabular feature list to report

    oof = [np.load(f) for f in sorted(glob.glob(os.path.join(shard_dir(args.tag), "oof_*.npz")))]
    s_oof = np.concatenate([o["score"] for o in oof])
    e_oof = np.concatenate([o["evt"] for o in oof])
    w_oof = np.concatenate([o["week"] for o in oof]).astype("i4")
    c_oof = (np.concatenate([o["capable"] for o in oof]).astype(bool)
             if "capable" in oof[0].files else None)
    r_oof = (np.concatenate([o["region"] for o in oof]).astype("i4")
             if "region" in oof[0].files else None)
    d_oof = (np.concatenate([o["doy"] for o in oof]).astype("i4")
             if "doy" in oof[0].files else None)
    keep_oof = np.isfinite(s_oof)
    if not keep_oof.all():
        s_oof, e_oof, w_oof = s_oof[keep_oof], e_oof[keep_oof], w_oof[keep_oof]
        c_oof = c_oof[keep_oof] if c_oof is not None else None
        r_oof = r_oof[keep_oof] if r_oof is not None else None
        d_oof = d_oof[keep_oof] if d_oof is not None else None
    # The training rows had their non-events thinned, so the out-of-fold score
    # distribution carries a base rate several times the real one.  Reweighting
    # restores the original balance before any threshold is read off it; without
    # this the thresholds come out far too low and the run over-forecasts ~3x.
    wt_oof = T.population_weights(e_oof, args.neg_keep)
    print("[weights] out-of-fold base rate %.4f raw -> %.4f after undoing the thinning"
          % (float(e_oof.mean()),
             float((wt_oof * e_oof).sum() / wt_oof.sum())), flush=True)
    th_csi = T.csi_thresholds(s_oof, e_oof, w_oof, wt_oof)
    th_fbi = T.bias_one_thresholds(s_oof, e_oof, w_oof, wt_oof)
    print("[thr] rule A, max weighted CSI : %s"
          % {C.WEEK_NAMES[k]: round(v, 4) for k, v in th_csi.items()}, flush=True)
    print("[thr] rule B, FBI ~ 1          : %s"
          % {C.WEEK_NAMES[k]: round(v, 4) for k, v in th_fbi.items()}, flush=True)
    # Rule B is the one the gridded fields use: the brief asks for FBI near 1, and
    # on the first national run it also scored the higher CSI of the two.
    th = th_fbi

    # A threshold read off all rows gives FBI ~ 1 over all rows, but the
    # heat-capable cells carry a much higher base rate, so the same threshold
    # over-forecasts badly inside that subset.  Whichever
    # frame a score is reported in, its threshold has to be tuned in that frame, so
    # the capable frame gets its own pair of thresholds read off capable rows only.
    th_cap_fbi = th_cap_csi = None
    if c_oof is not None and c_oof.any():
        th_cap_fbi = T.bias_one_thresholds(s_oof[c_oof], e_oof[c_oof],
                                           w_oof[c_oof], wt_oof[c_oof])
        th_cap_csi = T.csi_thresholds(s_oof[c_oof], e_oof[c_oof],
                                      w_oof[c_oof], wt_oof[c_oof])
        print("[thr] capable frame, FBI ~ 1  : %s"
              % {C.WEEK_NAMES[k]: round(v, 4) for k, v in th_cap_fbi.items()}, flush=True)

    preds = [np.load(f) for f in sorted(glob.glob(os.path.join(shard_dir(args.tag), "pred_*.npz")))]
    score = np.concatenate([p["score"] for p in preds]).astype("f8")
    y_te = np.concatenate([p["evt"] for p in preds])
    keys = [k for k in ("week", "lt", "point", "capable", "init", "region", "doy")
            if k in preds[0].files]
    info = {k: np.concatenate([p[k] for p in preds]) for k in keys}
    info["week"] = info["week"].astype("i4")
    info["init"] = info["init"].astype("i8")
    del preds

    ok = np.isfinite(score) & np.isfinite(y_te)
    pred = T.apply_thresholds(np.nan_to_num(score, nan=-1e18), info["week"], th)
    sub = {k: v[ok] for k, v in info.items()}
    res = {"daily_all": T.evaluate(score[ok], pred[ok], y_te[ok], sub),
           "daily_capable": T.evaluate(score[ok], pred[ok], y_te[ok], sub,
                                       info["capable"][ok].astype(bool))}
    ws, wp, wo, wi = T.weekly_event(score[ok], pred[ok], y_te[ok], sub)
    res["weekly_all"] = T.evaluate(ws, wp, wo, wi)
    for name, tab in res.items():
        print("\n" + T.fmt_table(name, tab), flush=True)

    # The other decision rule, reported alongside so the trade-off is visible.
    pred_a = T.apply_thresholds(np.nan_to_num(score, nan=-1e18), info["week"], th_csi)
    res_a = {"daily_all": T.evaluate(score[ok], pred_a[ok], y_te[ok], sub),
             "daily_capable": T.evaluate(score[ok], pred_a[ok], y_te[ok], sub,
                                         info["capable"][ok].astype(bool))}
    ws_a, wp_a, wo_a, wi_a = T.weekly_event(score[ok], pred_a[ok], y_te[ok], sub)
    res_a["weekly_all"] = T.evaluate(ws_a, wp_a, wo_a, wi_a)
    # The capable frame re-scored with its own thresholds.
    cap_mask = info["capable"][ok].astype(bool)
    if th_cap_fbi is not None:
        for lbl, tc in (("daily_capable_own_fbi1", th_cap_fbi),
                        ("daily_capable_own_csi", th_cap_csi)):
            pc = T.apply_thresholds(np.nan_to_num(score, nan=-1e18), info["week"], tc)
            res[lbl] = T.evaluate(score[ok], pc[ok], y_te[ok], sub, cap_mask)
            print("\n" + T.fmt_table(lbl, res[lbl]), flush=True)

    # One threshold per (week, stratum): the national threshold has to serve all
    # seven climatic zones and all four seasons at once, and their event rates are
    # nothing like each other.
    if r_oof is not None and d_oof is not None and "region" in info:
        for cut in ("region", "season"):
            st_oof = stratum_of(r_oof, d_oof, cut)
            st_prd = stratum_of(info["region"], info["doy"], cut)
            tsx = T.stratified_thresholds(s_oof, e_oof, w_oof, st_oof, wt_oof)
            px = T.apply_stratified(np.nan_to_num(score, nan=-1e18),
                                    info["week"], st_prd, tsx)
            lbl = "daily_all_by_%s" % cut
            res[lbl] = T.evaluate(score[ok], px[ok], y_te[ok], sub)
            print("\n" + T.fmt_table(lbl, res[lbl]), flush=True)

            if c_oof is not None and c_oof.any():
                tsc = T.stratified_thresholds(s_oof[c_oof], e_oof[c_oof], w_oof[c_oof],
                                              st_oof[c_oof], wt_oof[c_oof])
                pc = T.apply_stratified(np.nan_to_num(score, nan=-1e18),
                                        info["week"], st_prd, tsc)
                lbl = "daily_capable_by_%s" % cut
                if cut == "region":
                    pred_reg = pc
                res[lbl] = T.evaluate(score[ok], pc[ok], y_te[ok], sub, cap_mask)
                print("\n" + T.fmt_table(lbl, res[lbl]), flush=True)

    # Per-region picture: the national scores hide that the seven zones behave
    # very differently, so each is scored separately over the target weeks.
    if "region" in info:
        rr = info["region"][ok].astype("i4")
        tgt = np.isin(sub["week"], [3, 4, 5])          # W4-W6, where the targets apply
        res_reg = {}
        print("\n  by region, mean over W4-W6")
        print("    %-8s %8s %7s %7s %7s %7s %7s %7s"
              % ("region", "n_evt", "POD", "POFD", "FAR", "FBI", "CSI", "AUC"))
        # A national threshold serves the whole country at once, and the zones do
        # not behave alike, so each is also scored under its own threshold.
        pr = locals().get("pred_reg")
        cases = [("daily_all", np.ones(rr.size, bool), pred),
                 ("daily_capable", cap_mask, pred)]
        if pr is not None:
            cases.append(("daily_capable_own_region_thr", cap_mask, pr))
        for frame, mask, pp in cases:
            res_reg[frame] = {}
            if frame != "daily_all":
                print("    -- %s --" % frame, flush=True)
            for r in range(1, 8):
                m = (rr == r) & tgt & mask
                if m.sum() == 0:
                    continue
                sc = M.score_binary(pp[ok][m], y_te[ok][m] > 0.5)
                sc["AUC"] = M.auc_score(score[ok][m], y_te[ok][m] > 0.5)
                sc["n_events"] = float((y_te[ok][m] > 0.5).sum())
                res_reg[frame]["S%d" % r] = sc
                if frame != "daily_all":
                    print("    S%-7d %8d %7.3f %7.3f %7.3f %7.3f %7.3f %7.3f"
                          % (r, int(sc["n_events"]), sc["POD"], sc["POFD"], sc["FAR"],
                             sc["FBI"], sc["CSI"], sc["AUC"]), flush=True)

    print("\n########## rule A - threshold maximises weighted CSI ##########")
    for name, tab in res_a.items():
        print("\n" + T.fmt_table(name, tab), flush=True)

    # Probabilistic verification, which stage 1 could not provide because the raw
    # EC control run is a single deterministic member.
    prob = T.platt_calibrate(s_oof, e_oof, wt_oof, score[ok])
    print("\n[calibration] mean forecast probability %.4f vs observed rate %.4f"
          % (float(prob.mean()), float((y_te[ok] > 0.5).mean())), flush=True)
    prob_res = {}
    for w in range(len(C.WEEKS)):
        m = sub["week"] == w
        prob_res[C.WEEK_NAMES[w]] = {
            "reliability": M.reliability(prob[m], y_te[ok][m] > 0.5),
            "bss": M.brier_skill_score(prob[m], y_te[ok][m] > 0.5),
            "auc": M.auc_score(score[ok][m], y_te[ok][m] > 0.5)}
    print("[probabilistic] Brier skill score against sample climatology")
    for w, v in prob_res.items():
        print("    %-4s BSS = %7.4f   AUC = %.3f" % (w, v["bss"], v["auc"]))

    nw, npt = len(C.WEEKS), grid.npt
    acc = {}
    for nm, cond in (("hit", pred & (y_te > 0.5)), ("fal", pred & (y_te <= 0.5)),
                     ("mis", ~pred & (y_te > 0.5)), ("cor", ~pred & (y_te <= 0.5))):
        arr = np.zeros((nw, npt))
        for w in range(nw):
            m = ok & (info["week"] == w) & cond
            arr[w] = np.bincount(info["point"][m], minlength=npt)
        acc[nm] = arr
    a, b, c, d = acc["hit"], acc["fal"], acc["mis"], acc["cor"]
    tot = a + b + c + d

    def rate(n, dd):
        o = np.full(n.shape, np.nan, "f4"); g = dd > 0; o[g] = n[g] / dd[g]; return o

    ar = rate((a + b) * (a + c), tot)
    fields = {"pod": rate(a, a + c), "pofd": rate(b, b + d), "far": rate(b, a + b),
              "fbi": rate(a + b, a + c), "csi": rate(a, a + b + c),
              "ets": rate(a - ar, a + b + c - ar), "base_rate": rate(a + c, tot)}

    auc = np.full((nw, npt), np.nan, "f4")
    rows = np.nonzero(ok)[0]
    key = info["week"][rows].astype("i8") * npt + info["point"][rows]
    order = np.argsort(key, kind="stable")
    key = key[order]; sc = score[rows][order]; ob = y_te[rows][order] > 0.5
    edges = np.nonzero(np.diff(key))[0] + 1
    for lo, hi in zip(np.r_[0, edges], np.r_[edges, key.size]):
        gkey = int(key[lo])
        auc[gkey // npt, gkey % npt] = M.auc_score(sc[lo:hi], ob[lo:hi])
    fields["auc"] = auc

    out_nc = os.path.join(C.OUT, "ai_verification%s.nc" % args.tag)
    C.write_nc(out_nc,
               [("week", np.arange(1, nw + 1, dtype="i4")),
                ("lat", grid.lat), ("lon", grid.lon)],
               [(k, ("week", "lat", "lon"), grid.to_grid(v), {}) for k, v in fields.items()],
               {"weeks": " ".join(C.WEEK_NAMES), "model": args.model,
                "features": args.fset, "valid_years": str(C.VALID_YEARS),
                "note": "AI model verification on held-out years, land points only"})

    payload = {"summary": res, "summary_rule_csi": res_a,
               "by_region": res_reg if "region" in info else None,
               "probabilistic": prob_res,
               "thresholds": {str(k): float(v) for k, v in th.items()},
               "thresholds_csi": {str(k): float(v) for k, v in th_csi.items()},
               "thresholds_fbi1": {str(k): float(v) for k, v in th_fbi.items()},
               "thresholds_capable_fbi1": (
                   {str(k): float(v) for k, v in th_cap_fbi.items()}
                   if th_cap_fbi else None),
               "decision_rule": "FBI~1 (rule B) on reweighted out-of-fold scores",
               "neg_keep": args.neg_keep, "features": names, "model": args.model}
    mpath = os.path.join(shard_dir(args.tag), "model.joblib")
    bundle = joblib.load(mpath) if os.path.exists(mpath) else {"models": {}}
    mods = list((bundle.get("models") or ({0: bundle["model"]} if "model" in bundle
                                          else {})).values())
    if mods and names and hasattr(mods[0], "feature_importances_"):
        # averaged over the stratum models when the run was stratified
        imp = np.mean([np.asarray(m.feature_importances_, "f8") for m in mods], axis=0)
        imp /= max(imp.sum(), 1e-9)
        order = np.argsort(imp)[::-1][:25]
        payload["top_features"] = [(names[i], float(imp[i])) for i in order]
        print("\n[importance] normalised tree gain, top 25")
        for i in order:
            print("    %-16s %.4f" % (names[i], imp[i]))
    with open(os.path.join(RES_DIR, "national_%s_%s%s.json"
                           % (args.model, args.fset, args.tag)), "w") as fh:
        json.dump(payload, fh, indent=1, default=float)
    print("\n[done] -> %s" % out_nc, flush=True)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("phase", choices=["prep", "fit", "predict", "predict_oof",
                                     "finalize"])
    ap.add_argument("--shard", type=int, default=0)
    ap.add_argument("--nshard", type=int, default=120)
    ap.add_argument("--fold", default="full")
    ap.add_argument("--model", default="lgbm_clf")
    ap.add_argument("--fset", default="F3")
    ap.add_argument("--neg-keep", type=float, default=0.15)
    ap.add_argument("--tag", default="")
    ap.add_argument("--train-tag", default="",
                    help="read the training shards of this tag instead of --tag; the "
                         "model still lands in --tag")
    ap.add_argument("--group", default="all", choices=["all", "yEHD", "nEHD"],
                    help="EXP2: fit only on the region-days that did (yEHD) or did "
                         "not (nEHD) have a regional EHDE")
    ap.add_argument("--gate-file",
                    default=os.path.join(RES_DIR, "gate_ens.npz"),
                    help="npz written by regional_model.py --save-gate; supplies the "
                         "observational yEHD / nEHD label of each region-day")
    ap.add_argument("--stratify", default="none",
                    choices=["none", "region", "season", "region_season"])
    args = ap.parse_args()
    {"prep": cmd_prep, "fit": cmd_fit, "predict_oof": cmd_predict_oof,
     "predict": cmd_predict, "finalize": cmd_finalize}[args.phase](args)


if __name__ == "__main__":
    main()