#!/usr/bin/env python3
"""EXP2 per grid point: a nEHDE gate and a yEHDE forecast fitted at every point.

The user's design (2026-08-29).  Instead of one pooled national model, every
grid point gets its own pair of models:

  1. the GATE answers "is this sample almost certainly nEHDE?".  Where it says
     yes, no EHD forecast is issued at all.
  2. the yEHDE model is fitted only on the training rows the gate lets through,
     so it sees a cleaner, higher-base-rate subpopulation.

Two labels are offered for the gate, because "nEHDE" can be read at two scales:

  wk   the point-week reading: the gate label is 1 when the (init, week, point)
       contains at least one EHD day.  This is the direct per-grid-point
       analogue of EXP2 and is self-contained at a single point.
  reg  the literal EXP2 rule of CLAUDE.md: the region's P_i = m_i/n_i exceeds
       `--gate-cover` (default 5 %) for that (init, lead).  n_i counts every
       land point of the region, not just the three in the testbed.

Three arms are scored so the two effects can be separated:

  single   one per-point model, no cascade   -> isolates "per-point vs pooled"
  hard     gate + yEHDE model, hard cut      -> the user's design
  soft     P(gate) x P(EHD) with no cut      -> the cascade without information loss
  pooled   the national ensemble restricted to the same 21 points (baseline)

Two rules keep the comparison honest:

  * the gate cut is chosen on OUT-OF-FOLD gate scores, and the yEHDE model is
    fitted on the rows that out-of-fold gate passes.  Selecting stage-2's
    training set with in-sample gate scores would hand it an unrealistically
    clean subpopulation.
  * every table is computed over ALL validation rows.  A row the gate closed on
    that did have an event becomes a miss, exactly as it would operationally.
    Scoring only where the gate opened conditions verification on the forecast
    itself -- the mistake recorded in developing.txt section S.
"""

from __future__ import annotations

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

import numpy as np
from joblib import Parallel, delayed

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")
OUT_JSON = os.path.join(RES_DIR, "pointwise.json")
ENS_SHARDS = os.path.join(RES_DIR, "shards_ens")
NFOLD = 4


# --------------------------------------------------------------- small models
def make_model(kind, pos_weight):
    """Per-point variants: the national settings overfit ~600 events at one point."""
    if kind == "logreg":
        return T.make_model("logreg", pos_weight)
    if kind == "lgbm_s":
        import lightgbm as lgb
        return lgb.LGBMClassifier(
            n_estimators=300, learning_rate=0.05, num_leaves=15,
            min_child_samples=100, subsample=0.8, subsample_freq=1,
            colsample_bytree=0.8, reg_lambda=5.0,
            scale_pos_weight=pos_weight, n_jobs=1, verbose=-1)
    raise ValueError(kind)


def _pos_weight(y):
    r = float(np.mean(y))
    return (1.0 - r) / max(r, 1e-9)


def fit_score(kind, Xtr, ytr, Xte):
    m = make_model(kind, _pos_weight(ytr))
    m.fit(Xtr, ytr)
    return m.predict_proba(Xte)[:, 1]


def oof_score(kind, X, y, year, nfold=NFOLD, seed=0):
    """Out-of-fold probabilities over the training years, split by whole years."""
    years = np.unique(year)
    folds = np.array_split(np.random.default_rng(seed).permutation(years), nfold)
    out = np.zeros(X.shape[0], dtype="f8")
    for f in folds:
        te = np.isin(year, f)
        if te.all() or not te.any() or y[~te].sum() < 5:
            continue
        out[te] = fit_score(kind, X[~te], y[~te], X[te])
    return out


def keep_cut(score, evt, keep):
    """Lowest cut on `score` that still retains `keep` of the events."""
    if evt.sum() == 0:
        return -np.inf
    o = np.argsort(-score, kind="mergesort")
    frac = np.cumsum(evt[o]) / evt.sum()
    i = int(np.searchsorted(frac, keep))
    return float(score[o][min(i, score.size - 1)])


# ------------------------------------------------------------------ gate data
def gate_labels(store, pts, cover):
    """Both gate labels for every (init, lt) and every point of `pts`.

    Returns wk[ninit, 46, npts] and reg[ninit, 46, npts], both boolean, where
    True means "yEHDE" -- the gate should stay open.
    """
    evt_all = np.asarray(store.target[:, :, 2, :])            # (ninit, 46, 541)
    reg_of = store.st["region"]
    P = np.empty(evt_all.shape[:2] + (8,), dtype="f4")
    for r in range(1, 8):
        sel = reg_of == r
        P[:, :, r] = np.nanmean(evt_all[:, :, sel], axis=2) if sel.any() else 0.0
    reg = P[:, :, reg_of[pts]] > cover                        # (ninit, 46, npts)

    e = evt_all[:, :, pts] > 0.5
    wk = np.zeros_like(e)
    for _, lts in C.WEEKS:                                    # any EHD day in the week
        wk[:, lts, :] = np.any(e[:, lts, :], axis=1)[:, None, :]
    return {"wk": wk, "reg": reg}


def pooled_baseline(pts):
    """The national ensemble score at `pts`: out-of-fold on train, direct on valid."""
    out = {}
    for phase, pat in (("oof", "oofpred_*.npz"), ("val", "pred_*.npz")):
        acc = {k: [] for k in ("score", "evt", "week", "lt", "point", "init",
                               "capable", "doy", "region")}
        for f in sorted(glob.glob(os.path.join(ENS_SHARDS, pat))):
            z = np.load(f)
            m = np.isin(z["point"], pts)
            if not m.any():
                continue
            for k in acc:
                acc[k].append(z[k][m])
        out[phase] = {k: np.concatenate(v) for k, v in acc.items()} if acc["score"] else None
    return out


# ------------------------------------------------------------- one grid point
def run_point(ip, Xtr, ytr, wtr, Xte, yte, wte, gtr, gte, kinds, keeps, gates):
    """Every arm at one grid point.  Returns {arm: (oof score, valid score)}."""
    sd = Xtr.std(axis=0)
    col = sd > 1e-12                       # lat/lon/region are constant at one point
    Xtr, Xte = np.ascontiguousarray(Xtr[:, col]), np.ascontiguousarray(Xte[:, col])
    res, t0 = {}, time.time()

    for kind in kinds:
        s_oof = oof_score(kind, Xtr, ytr, wtr)
        s_val = fit_score(kind, Xtr, ytr, Xte)
        res["%s/single" % kind] = (s_oof, s_val)

        for gname in gates:
            g_tr, g_te = gtr[gname], gte[gname]
            if g_tr.sum() < 20 or g_tr.all():
                continue
            g_oof = oof_score(kind, Xtr, g_tr.astype("f8"), wtr)
            g_val = fit_score(kind, Xtr, g_tr.astype("f8"), Xte)
            # the gate is a forecast in its own right: model 1 of the pair
            res["%s/GATE_%s" % (kind, gname)] = (g_oof, g_val)
            res["%s/soft_%s" % (kind, gname)] = (g_oof * s_oof, g_val * s_val)

            for keep in keeps:
                tau = keep_cut(g_oof, ytr, keep)
                p_tr, p_te = g_oof >= tau, g_val >= tau
                if p_tr.sum() < 500 or ytr[p_tr].sum() < 30:
                    continue
                o2 = np.full(Xtr.shape[0], -1.0)
                o2[p_tr] = oof_score(kind, Xtr[p_tr], ytr[p_tr], wtr[p_tr])
                v2 = np.full(Xte.shape[0], -1.0)
                s2 = fit_score(kind, Xtr[p_tr], ytr[p_tr], Xte[p_te])
                v2[p_te] = s2
                res["%s/hard_%s_k%03d" % (kind, gname, round(1000 * keep))] = (o2, v2)

    print("[point %3d] %d arms, %d train rows, base %.4f, %.0fs"
          % (ip, len(res), Xtr.shape[0], ytr.mean(), time.time() - t0), flush=True)
    return ip, res


# ------------------------------------------------------------------ reporting
def tables(score_oof, evt_oof, wk_oof, pt_oof, score, evt, wk, pt, cap):
    """FBI~1 thresholds from the out-of-fold field, then the usual score tables."""
    w1 = np.ones(evt_oof.size)
    th = T.bias_one_thresholds(score_oof, evt_oof, wk_oof, w1)
    pred = T.apply_thresholds(np.nan_to_num(score, nan=-1e18), wk, th)

    # option 2 folded in for free: one FBI~1 cut per (point, week)
    thp = T.stratified_thresholds(score_oof, evt_oof, wk_oof, pt_oof, w1,
                                  min_events=30)
    pred_p = T.apply_stratified(np.nan_to_num(score, nan=-1e18), wk, pt, thp)

    obs = evt > 0.5
    out = {}
    for lbl, sel in (("daily_all", np.ones(obs.size, bool)),
                     ("daily_capable", cap),
                     ("W4plus", wk >= 3)):
        rows = {}
        for tag, pr in (("thr_week", pred), ("thr_point_week", pred_p)):
            s = M.score_binary(pr[sel], obs[sel])
            s["AUC"] = M.auc_score(score[sel], obs[sel])
            rows[tag] = s
        out[lbl] = rows
    per_week = {}
    for w in range(len(C.WEEKS)):
        m = wk == w
        if m.sum():
            s = M.score_binary(pred[m], obs[m])
            s["AUC"] = M.auc_score(score[m], obs[m])
            per_week[C.WEEK_NAMES[w]] = s
    out["per_week"] = per_week
    return out


def point_table(score_oof, evt_oof, wk_oof, pt_oof,
                score, evt, wk, pt, week_min=3):
    """One row per grid point, W4 onward, under its own (point, week) FBI~1 cut."""
    w1 = np.ones(evt_oof.size)
    thp = T.stratified_thresholds(score_oof, evt_oof, wk_oof, pt_oof, w1,
                                  min_events=30)
    pred = T.apply_stratified(np.nan_to_num(score, nan=-1e18), wk, pt, thp)
    sel = wk >= week_min
    obs = evt > 0.5
    out = {}
    for p in np.unique(pt[sel]):
        m = sel & (pt == p)
        if m.sum() < 50 or obs[m].sum() < 5:
            continue
        s = M.score_binary(pred[m], obs[m])
        s["AUC"] = M.auc_score(score[m], obs[m])
        s["n_events"] = int(obs[m].sum())
        out[int(p)] = s
    return out


def fmt(title, rows):
    keys = ["base_rate", "POD", "POFD", "FAR", "FBI", "CSI", "ETS", "AUC"]
    lines = ["  " + title, "    %-16s" % "" + "".join("%9s" % k for k in keys)]
    for name, s in rows.items():
        lines.append("    %-16s" % name
                     + "".join("%9.3f" % s.get(k, np.nan) for k in keys))
    return "\n".join(lines)


# --------------------------------------------------------------------- driver
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--per-region", type=int, default=3)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--fsets", nargs="*", default=["F3"])
    ap.add_argument("--kinds", nargs="*", default=["logreg", "lgbm_s"])
    ap.add_argument("--gates", nargs="*", default=["wk", "reg"])
    ap.add_argument("--keeps", type=float, nargs="*", default=[0.995, 0.99, 0.95])
    ap.add_argument("--gate-cover", type=float, default=0.05)
    ap.add_argument("--jobs", type=int, default=20)
    ap.add_argument("--out", default=OUT_JSON)
    args = ap.parse_args()

    store = D.Store()
    pts = D.sample_points(store.grid, args.per_region, args.seed)
    print("[points] %d testbed points: %s" % (pts.size, pts.tolist()), flush=True)
    print("[regions] %s" % np.bincount(store.st["region"][pts], minlength=8)[1:].tolist(),
          flush=True)

    gl = gate_labels(store, pts, args.gate_cover)
    for g in args.gates:
        print("[gate %s] yEHDE fraction = %.4f" % (g, gl[g].mean()), flush=True)

    base = pooled_baseline(pts)
    os.makedirs(RES_DIR, exist_ok=True)
    scores_dir = os.path.join(RES_DIR, "pointwise_scores")
    os.makedirs(scores_dir, exist_ok=True)
    results = {}

    # ---- the pooled national ensemble at the same points, same threshold rule
    if base["oof"] is not None and base["val"] is not None:
        o, v = base["oof"], base["val"]
        results["pooled_ensemble"] = tables(
            o["score"], o["evt"], o["week"], o["point"],
            v["score"], v["evt"], v["week"], v["point"], v["capable"])
        results["pooled_ensemble"]["by_point"] = point_table(
            o["score"], o["evt"], o["week"], o["point"],
            v["score"], v["evt"], v["week"], v["point"])
        results["pooled_ensemble"]["target"] = "EHD"
        print("\n=== pooled_ensemble (baseline, national model at these points) ===",
              flush=True)
        print(fmt("daily, all rows", results["pooled_ensemble"]["daily_all"]), flush=True)
        print(fmt("W4 onward", results["pooled_ensemble"]["W4plus"]), flush=True)

    for fset in args.fsets:
        t0 = time.time()
        Xtr, ytr_, itr = store.build(pts, C.TRAIN_YEARS, fset)
        Xte, yte_, ite = store.build(pts, C.VALID_YEARS, fset)
        ok_tr = np.isfinite(Xtr).all(axis=1) & np.isfinite(ytr_["evt"])
        ok_te = np.isfinite(Xte).all(axis=1) & np.isfinite(yte_["evt"])
        print("\n=== %s | train %d rows | valid %d rows | dropped %d/%d non-finite "
              "| %.0fs" % (fset, ok_tr.sum(), ok_te.sum(), (~ok_tr).sum(),
                           (~ok_te).sum(), time.time() - t0), flush=True)

        # build() lays its rows out as (init, lt, point) and gl is already
        # (ninit, 46, npts), so a plain reshape lines the labels up row for row.
        gi_tr, gi_te = {}, {}
        m_tr = np.isin(store.init_year, C.TRAIN_YEARS)
        m_te = np.isin(store.init_year, C.VALID_YEARS)
        for g in args.gates:
            gi_tr[g] = gl[g][m_tr].reshape(-1)
            gi_te[g] = gl[g][m_te].reshape(-1)

        jobs = []
        for j, p in enumerate(pts):
            mtr = ok_tr & (itr["point"] == p)
            mte = ok_te & (ite["point"] == p)
            jobs.append(delayed(run_point)(
                int(p), Xtr[mtr], ytr_["evt"][mtr], itr["year"][mtr],
                Xte[mte], yte_["evt"][mte], ite["year"][mte],
                {g: gi_tr[g][mtr] for g in args.gates},
                {g: gi_te[g][mte] for g in args.gates},
                args.kinds, args.keeps, args.gates))
        got = Parallel(n_jobs=args.jobs, verbose=5)(jobs)

        # ---- stitch the per-point arms back into national-style columns
        idx_tr = {int(p): (ok_tr & (itr["point"] == p)) for p in pts}
        idx_te = {int(p): (ok_te & (ite["point"] == p)) for p in pts}
        arms = sorted({a for _, r in got for a in r})
        for arm in arms:
            so = np.full(ok_tr.sum(), np.nan); sv = np.full(ok_te.sum(), np.nan)
            offo = np.cumsum([0] + [idx_tr[int(p)].sum() for p in pts])
            offv = np.cumsum([0] + [idx_te[int(p)].sum() for p in pts])
            order = {int(p): j for j, p in enumerate(pts)}
            filled = 0
            for ip, r in got:
                if arm not in r:
                    continue
                j = order[ip]
                so[offo[j]:offo[j + 1]] = r[arm][0]
                sv[offv[j]:offv[j + 1]] = r[arm][1]
                filled += 1
            if filled < pts.size:
                print("[skip] %s: only %d/%d points produced it"
                      % (arm, filled, pts.size), flush=True)
                continue
            key = "%s/%s" % (fset, arm)
            sel_o = np.concatenate([np.nonzero(idx_tr[int(p)])[0] for p in pts])
            sel_v = np.concatenate([np.nonzero(idx_te[int(p)])[0] for p in pts])
            # a GATE arm forecasts the nEHD/yEHD label, not EHD itself, so it has
            # to be scored against that label or the numbers are meaningless
            gn = arm.split("/")[-1][len("GATE_"):] if "/GATE_" in arm else None
            if gn is not None:
                yo = gi_tr[gn][sel_o].astype("f4")
                yv = gi_te[gn][sel_v].astype("f4")
            else:
                yo, yv = ytr_["evt"][sel_o], yte_["evt"][sel_v]
            results[key] = tables(
                so, yo, itr["week"][sel_o], itr["point"][sel_o],
                sv, yv, ite["week"][sel_v], ite["point"][sel_v],
                ite["capable"][sel_v])
            results[key]["by_point"] = point_table(
                so, yo, itr["week"][sel_o], itr["point"][sel_o],
                sv, yv, ite["week"][sel_v], ite["point"][sel_v])
            results[key]["target"] = ("nEHD/yEHD (%s)" % gn) if gn else "EHD"
            np.savez(os.path.join(scores_dir, key.replace("/", "__") + ".npz"),
                     oof=so.astype("f4"), val=sv.astype("f4"),
                     point_val=ite["point"][sel_v], week_val=ite["week"][sel_v],
                     evt_val=yv)
            print("\n=== %s  (muc tieu: %s) ===" % (key, results[key]["target"]),
                  flush=True)
            print(fmt("daily, all rows", results[key]["daily_all"]), flush=True)
            print(fmt("W4 onward", results[key]["W4plus"]), flush=True)

        del Xtr, Xte, ytr_, yte_, itr, ite

    with open(args.out, "w") as fh:
        json.dump(results, fh, indent=1, default=float)
    print("\n[done] -> %s" % args.out, flush=True)


if __name__ == "__main__":
    main()