#!/usr/bin/env python3
"""EXP3: N mô hình hồi quy Sev_obs tại mỗi ô lưới, quyết định bằng tỷ lệ phiếu p = M/N.

The user's specification (CLAUDE.md, EXP3), implemented literally:

  * every grid point gets its own ensemble of N regression models;
  * the target is Sev_obs = Tx_obs - Tx_thres_obs, kept CONTINUOUS and signed --
    negative days are training data too, not discarded;
  * M of the N members predict Sev >= 0, so p(yEHDE) = M/N and q = 1 - p;
  * stage 1: p < P1 (default 0.3) means nEHDE -- the sample is dropped;
  * stage 2: the survivors are re-fitted with a fresh N-model ensemble (possibly a
    different family), and p2 > P2 (default 0.5) is the final yEHDE call.

Three things are reported alongside the literal design, all free because they come
out of the very same fitted members:

  vote_mean   the ensemble MEAN of predicted Sev instead of the vote count.  The
              vote binarises each member before averaging; the mean does not.
              Section AA measured the same structure at region level and the
              difference was large (AUC 0.739 vs 0.931), so it is worth knowing
              whether p = M/N is costing anything here.
  fbi1        the same continuous score cut at FBI ~ 1 instead of at the fixed
              0.3 / 0.5.  Every earlier table in the project uses the FBI ~ 1
              operating point, so without this the comparison is not like for like.
  soft        p x p2 with no hard cut, since EXP2 measured that hard-cutting the
              gate throws away the gain the product keeps.

Diversity of the N members is deliberate: hyper-parameters, feature subsets and
model family are all varied.  N members of one family fitted on one data set agree
almost perfectly, which would collapse M/N onto 0 and 1 and make the 0.3 threshold
meaningless.

Every table is computed over ALL validation rows.  A row dropped at stage 1 that
did have an event counts as a miss, exactly as it would operationally.
"""

from __future__ import annotations

import argparse
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
import pointwise as PW

RES_DIR = os.path.join(C.OUT, "AI_Results")
OUT_JSON = os.path.join(RES_DIR, "exp3.json")
NFOLD = 4


# ------------------------------------------------------- the N-member ensemble
def make_member(i, seed, members="bagged", n_models=20):
    """One ensemble member.

    `bagged`  the literal EXP3 reading: N regressors of the conditional mean of
              Sev, differing by hyper-parameters, feature subset and family.
    `quantile` the same N-member vote, but member i estimates the (i+0.5)/N
              quantile of Sev.  M/N is then the share of quantile levels lying at
              or above 0, i.e. a direct estimate of P(Sev >= 0).  This is needed
              because a least-squares member predicts the conditional MEAN of Sev,
              which is below 0 almost everywhere (measured: observed Sev >= 0 on
              12.6 % of rows, mean-regression point forecasts on 0.7 %), so the
              bagged vote collapses onto 0 and the p < 0.3 gate drops ~99 % of the
              events.  The vote formula is unchanged; only what a member estimates.
    """
    if members == "quantile":
        import lightgbm as lgb
        return lgb.LGBMRegressor(
            objective="quantile", alpha=(i + 0.5) / n_models,
            n_estimators=200, 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,
            random_state=int(seed) % (2 ** 31), n_jobs=1, verbose=-1)
    rng = np.random.default_rng(seed)
    if i % 5 == 4:                      # one in five is linear: family diversity
        from sklearn.linear_model import Ridge
        from sklearn.pipeline import make_pipeline
        from sklearn.preprocessing import StandardScaler
        return make_pipeline(StandardScaler(),
                             Ridge(alpha=float(rng.choice([1.0, 10.0, 100.0]))))
    import lightgbm as lgb
    return lgb.LGBMRegressor(
        n_estimators=int(rng.choice([200, 300, 400])),
        learning_rate=float(rng.choice([0.03, 0.05, 0.08])),
        num_leaves=int(rng.choice([7, 15, 31])),
        min_child_samples=int(rng.choice([50, 100, 200])),
        subsample=0.7, subsample_freq=1,
        colsample_bytree=float(rng.choice([0.5, 0.7, 0.9])),
        reg_lambda=float(rng.choice([1.0, 5.0, 20.0])),
        random_state=int(seed) % (2 ** 31), n_jobs=1, verbose=-1)


def member_cols(nfeat, seed, frac=0.75, members="bagged"):
    if members == "quantile":
        return np.arange(nfeat)         # quantile members differ by level, not features
    rng = np.random.default_rng(seed + 7919)
    k = max(3, int(round(frac * nfeat)))
    return np.sort(rng.choice(nfeat, size=k, replace=False))


def ens_predict(Xtr, ytr, Xte, n_models, seed0, members="bagged"):
    """Fit the N members on (Xtr, ytr) and return their predicted Sev on Xte."""
    out = np.empty((n_models, Xte.shape[0]), dtype="f4")
    for i in range(n_models):
        cols = member_cols(Xtr.shape[1], seed0 + i, members=members)
        m = make_member(i, seed0 + i, members=members, n_models=n_models)
        m.fit(Xtr[:, cols], ytr)
        out[i] = m.predict(Xte[:, cols])
    return out


def ens_oof(X, y, year, n_models, seed0, nfold=NFOLD, members="bagged"):
    """Leave-years-out predictions of every member, for the whole training set."""
    years = np.unique(year)
    folds = np.array_split(np.random.default_rng(0).permutation(years), nfold)
    out = np.zeros((n_models, X.shape[0]), dtype="f4")
    for f in folds:
        te = np.isin(year, f)
        if te.all() or not te.any():
            continue
        out[:, te] = ens_predict(X[~te], y[~te], X[te], n_models, seed0, members)
    return out


def vote(pred):
    """p(yEHDE) = M/N, the share of members predicting Sev >= 0."""
    return (pred >= 0.0).mean(axis=0)


# ------------------------------------------------------------- one grid point
def run_point(ip, Xtr, sev_tr, wtr, Xte, wte, n_models, p1, seed0,
              members="bagged"):
    """Both EXP3 stages at one grid point.  Returns the scores each arm produces."""
    sd = Xtr.std(axis=0)
    col = sd > 1e-12                      # lat/lon/region are constant at one point
    Xtr = np.ascontiguousarray(Xtr[:, col]); Xte = np.ascontiguousarray(Xte[:, col])
    t0 = time.time()

    # ---- stage 1: the whole sample, Sev as target
    o1 = ens_oof(Xtr, sev_tr, wtr, n_models, seed0, members=members)
    v1 = ens_predict(Xtr, sev_tr, Xte, n_models, seed0, members=members)
    p_oof, p_val = vote(o1), vote(v1)
    m_oof, m_val = o1.mean(axis=0), v1.mean(axis=0)

    # ---- stage 2: only the survivors of the stage-1 gate, chosen out-of-fold
    keep_tr = p_oof >= p1
    res = {}
    if keep_tr.sum() >= 500:
        o2 = np.full((n_models, Xtr.shape[0]), -9.0, dtype="f4")
        o2[:, keep_tr] = ens_oof(Xtr[keep_tr], sev_tr[keep_tr], wtr[keep_tr],
                                 n_models, seed0 + 500, members=members)
        v2 = ens_predict(Xtr[keep_tr], sev_tr[keep_tr], Xte, n_models, seed0 + 500,
                         members=members)
        p2_oof, p2_val = vote(o2), vote(v2)
        p2_oof[~keep_tr] = 0.0
        m2_oof, m2_val = o2.mean(axis=0), v2.mean(axis=0)
        m2_oof[~keep_tr] = -9.0

        res["exp3_vote"] = (p_oof, p_val, p2_oof, p2_val)
        res["exp3_mean"] = (p_oof, p_val, m2_oof, m2_val)
        res["exp3_soft"] = (p_oof * p2_oof, p_val * p2_val, None, None)
    res["stage1_vote"] = (p_oof, p_val, None, None)
    res["stage1_mean"] = (m_oof, m_val, None, None)
    print("[point %3d] %d hang train, cong giu %.3f, %.0fs"
          % (ip, Xtr.shape[0], keep_tr.mean(), time.time() - t0), flush=True)
    return ip, res


# ------------------------------------------------------------------ reporting
def scores_at(pred, obs):
    s = M.score_binary(pred, obs)
    return s


def evaluate(score_oof, evt_oof, wk_oof, score, evt, wk, pred_fixed=None):
    """FBI~1 operating point, plus the fixed-threshold one EXP3 specifies."""
    out = {}
    w1 = np.ones(evt_oof.size)
    th = T.bias_one_thresholds(np.nan_to_num(score_oof, nan=-1e18),
                               evt_oof, wk_oof, w1)
    pr = T.apply_thresholds(np.nan_to_num(score, nan=-1e18), wk, th)
    obs = evt > 0.5
    for lbl, sel in (("W4plus", wk >= 3), ("daily_all", np.ones(obs.size, bool))):
        s = scores_at(pr[sel], obs[sel]); s["AUC"] = M.auc_score(score[sel], obs[sel])
        out[lbl + "_fbi1"] = s
        if pred_fixed is not None:
            s2 = scores_at(pred_fixed[sel], obs[sel])
            s2["AUC"] = M.auc_score(score[sel], obs[sel])
            out[lbl + "_exp3thr"] = s2
    return out


def fmt(title, rows):
    keys = ["base_rate", "POD", "POFD", "FAR", "FBI", "CSI", "ETS", "AUC"]
    lines = ["  " + title, "    %-18s" % "" + "".join("%9s" % k for k in keys)]
    for name, s in rows.items():
        lines.append("    %-18s" % 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("--fset", default="F3")
    ap.add_argument("--n-models", type=int, default=20)
    ap.add_argument("--members", default="bagged", choices=["bagged", "quantile"],
                    help="ban chat cua N thanh vien: N mo hinh bagging (dung nguyen "
                         "van EXP3) hay N muc phan vi cua Sev")
    ap.add_argument("--p1", type=float, default=0.3, help="nguong cong tang 1")
    ap.add_argument("--p2", type=float, default=0.5, help="nguong quyet dinh tang 2")
    ap.add_argument("--p1-sweep", type=float, nargs="*",
                    default=[0.1, 0.2, 0.3, 0.4, 0.5])
    ap.add_argument("--p2-sweep", type=float, nargs="*",
                    default=[0.3, 0.4, 0.5, 0.6, 0.7])
    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 diem: %s" % (pts.size, pts.tolist()), flush=True)
    print("[exp3] N = %d thanh vien/diem (%s), target = Sev_obs, p1 = %.2f, p2 = %.2f"
          % (args.n_models, args.members, args.p1, args.p2), flush=True)
    os.makedirs(RES_DIR, exist_ok=True)

    Xtr, ytr_, itr = store.build(pts, C.TRAIN_YEARS, args.fset)
    Xte, yte_, ite = store.build(pts, C.VALID_YEARS, args.fset)
    ok_tr = np.isfinite(Xtr).all(axis=1) & np.isfinite(ytr_["sev"]) & np.isfinite(ytr_["evt"])
    ok_te = np.isfinite(Xte).all(axis=1) & np.isfinite(yte_["evt"])
    print("[data] train %d hang, valid %d hang, base rate %.4f"
          % (ok_tr.sum(), ok_te.sum(), ytr_["evt"][ok_tr].mean()), flush=True)

    jobs = []
    for p in pts:
        mtr = ok_tr & (itr["point"] == p)
        mte = ok_te & (ite["point"] == p)
        jobs.append(delayed(run_point)(
            int(p), Xtr[mtr], ytr_["sev"][mtr], itr["year"][mtr],
            Xte[mte], ite["year"][mte], args.n_models, args.p1,
            args.seed + 13 * int(p), args.members))
    got = Parallel(n_jobs=args.jobs, verbose=5)(jobs)

    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}
    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)}
    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])
    eo, ev = ytr_["evt"][sel_o], yte_["evt"][sel_v]
    wo, wv = itr["week"][sel_o], ite["week"][sel_v]

    def stitch(arm, k):
        a = np.full(offo[-1], np.nan); b = np.full(offv[-1], np.nan)
        n = 0
        for ip, r in got:
            if arm not in r or r[arm][k] is None:
                continue
            j = order[ip]
            a[offo[j]:offo[j + 1]] = r[arm][k]
            b[offv[j]:offv[j + 1]] = r[arm][k + 1]
            n += 1
        return (a, b) if n == pts.size else (None, None)

    results = {"config": {"n_models": args.n_models, "members": args.members,
                          "p1": args.p1, "p2": args.p2,
                          "fset": args.fset, "points": pts.tolist()}}

    # ---- the literal EXP3 decision, and the sweeps the spec asks for
    p1o, p1v = stitch("stage1_vote", 0)
    p2o, p2v = stitch("exp3_vote", 2)
    if p1v is not None and p2v is not None:
        pred = (p1v >= args.p1) & (p2v > args.p2)
        results["exp3_literal"] = evaluate(
            np.where(p1o >= args.p1, p2o, -1.0), eo, wo,
            np.where(p1v >= args.p1, p2v, -1.0), ev, wv, pred_fixed=pred)
        sweep = {}
        for a in args.p1_sweep:
            for b in args.p2_sweep:
                pr = (p1v >= a) & (p2v > b)
                s = scores_at(pr[wv >= 3], ev[wv >= 3] > 0.5)
                sweep["p1=%.2f,p2=%.2f" % (a, b)] = s
        results["exp3_sweep_W4plus"] = sweep
        gate_keep = float((p1v >= args.p1).mean())
        gate_evt = float((ev[p1v >= args.p1] > 0.5).sum() / max((ev > 0.5).sum(), 1))
        results["gate_stats"] = {"rows_kept": gate_keep, "events_kept": gate_evt}
        print("\n[cong tang 1] giu %.1f%% so hang, giu %.1f%% so su kien"
              % (100 * gate_keep, 100 * gate_evt), flush=True)

    # ---- the free comparison arms
    for arm, key in (("exp3_soft", 0), ("stage1_vote", 0), ("stage1_mean", 0)):
        a, b = stitch(arm, key)
        if a is not None:
            results[arm] = evaluate(a, eo, wo, b, ev, wv)
    a, b = stitch("exp3_mean", 2)
    if a is not None:
        results["exp3_mean"] = evaluate(
            np.where(p1o >= args.p1, a, -99.0), eo, wo,
            np.where(p1v >= args.p1, b, -99.0), ev, wv)

    base = PW.pooled_baseline(pts)
    if base["oof"] is not None:
        o, v = base["oof"], base["val"]
        results["pooled_ensemble"] = evaluate(o["score"], o["evt"], o["week"],
                                              v["score"], v["evt"], v["week"])

    for k in [x for x in results if x not in ("config", "exp3_sweep_W4plus",
                                              "gate_stats")]:
        print("\n=== %s ===" % k, flush=True)
        print(fmt("W4 tro di", {a: b for a, b in results[k].items()
                                if a.startswith("W4plus")}), flush=True)

    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()