#!/usr/bin/env python3
"""Stage-2 model that predicts the regional EHDE event directly.

`regional_ehde.py` collapses the point forecast to one number per region-day --
the mean predicted probability -- and thresholds it.  That framing raised CSI
(0.259 -> 0.349 at W4) but, as section Z7 of developing.txt says plainly, the
AUC did NOT improve: the gain came from redefining the event to one with twice
the base rate, not from forecasting it better.

This module attacks the part Z left on the table.  The mean is only one summary
of the probability field inside a region; the shape of that field carries more.
A day where 5 % of the cells sit at 0.8 is a different day from one where every
cell sits at 0.04, and both have the same mean.  So each region-day becomes one
row with ~25 features (quantiles of the probability field, soft coverage
fractions at several probability levels, the capable fraction, region, lead,
season) and a second-stage classifier is fitted on the training years to predict

    EHDE_i = 1[ m_i / n_i > cover ]

Two thresholding rules are compared on top of it: one threshold per forecast
week, as before, and one per (region, week).  Section Y7 showed the national
threshold over-forecasts by 2.2x in S7, the region where the model is best, so
the pooled number was hiding a real effect.

Leakage control, in the same style as the rest of the pipeline:
  * the point-level probabilities of the training years come from `oofpred_*`,
    where each row was scored by the fold model that held its year out;
  * the stage-2 model is fitted on training years only;
  * its decision thresholds are read off a grouped (by year) cross-validated
    prediction inside the training years, never off the validation years.
"""

from __future__ import annotations

import argparse
import json
import os
import sys

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
from train_national import RES_DIR
from regional_ehde import NREG

KEYS = ("score", "evt", "week", "lt", "init", "region", "capable", "doy")


def load(tag, pattern):
    """Every shard of one phase, concatenated; `doy` is needed here on top of
    what regional_ehde.load reads, so this keeps its own key list."""
    import glob
    from train_national import shard_dir
    files = sorted(glob.glob(os.path.join(shard_dir(tag), pattern)))
    if not files:
        raise SystemExit("no %s for tag '%s'" % (pattern, tag))
    parts = [np.load(f) for f in files]
    miss = [k for k in KEYS if k not in parts[0].files]
    if miss:
        raise SystemExit("%s lacks %s" % (pattern, ", ".join(miss)))
    d = {k: np.concatenate([p[k] for p in parts]) for k in KEYS}
    ok = np.isfinite(d["score"]) & np.isfinite(d["evt"])
    return {k: v[ok] for k, v in d.items()}

QUANTS = (0.50, 0.75, 0.90, 0.95, 0.99)
LEVELS = (0.05, 0.10, 0.20, 0.30, 0.50)
MIN_EVT_REGION = 40      # below this a region-week falls back to the pooled cut


def rank_within_week(X, week, ref_week=None, ref=None):
    """Replace every feature by its rank fraction inside its forecast week.

    The training rows are scored by the fold models (12 years each) while the
    validation rows are scored by the full model (16 years), so the two sets of
    probabilities do not share a scale.  A tree fitted on absolute values in one
    scale and applied to the other splits in the wrong place.  Ranking each
    feature inside its own dataset and week removes any monotone part of that
    shift; it uses no labels, only the score distribution, which an operational
    system would have from its own hindcast climatology.
    """
    out = np.empty_like(X, dtype="f4")
    for w in np.unique(week):
        m = week == w
        for j in range(X.shape[1]):
            base = X[m, j] if ref is None else ref[ref_week == w, j]
            order = np.argsort(base, kind="mergesort")
            srt = base[order]
            out[m, j] = np.searchsorted(srt, X[m, j], side="right") / float(srt.size)
    return out


def _q(vals, st, en, p):
    """Quantile p inside each [st, en) run of a sorted-within-group array."""
    idx = st + np.minimum((en - st - 1), np.floor(p * (en - st)).astype("i8"))
    return vals[idx]


def build_table(d, npts, init_year):
    """One row per (init, lead, region) with the features of the stage-2 model."""
    key = (d["init"].astype("i8") * C.NLT + d["lt"].astype("i8")) * 8 + d["region"]
    # lexsort: groups ordered by key, and sorted by probability inside each group,
    # so the quantiles below are a single gather instead of 670k little sorts
    o = np.lexsort((d["score"], key))
    k = key[o]
    cut = np.flatnonzero(np.diff(k))
    st = np.r_[0, cut + 1]
    en = np.r_[cut + 1, k.size]
    cnt = (en - st).astype("f8")

    reg = (k[st] % 8).astype(int)
    lead = ((k[st] // 8) % C.NLT).astype(int)
    init = (k[st] // (8 * C.NLT)).astype("i8")
    n = npts[reg - 1].astype("f8")

    s = d["score"][o].astype("f8")
    evt = (d["evt"][o] > 0.5).astype("f8")
    cap = (d["capable"][o] > 0.5).astype("f8")
    doy = d["doy"][o].astype("f8")

    ss = s                                   # already sorted inside each group
    mean = np.add.reduceat(s, st) / cnt
    feats = {
        "Pbar": mean,
        "Pstd": np.sqrt(np.maximum(np.add.reduceat(s * s, st) / cnt - mean ** 2, 0.0)),
        "Pmax": ss[en - 1],
        "Pmin": ss[st],
        "Pcap": np.add.reduceat(cap, st) / n,
        "Pcapmean": np.add.reduceat(s * cap, st) / np.maximum(np.add.reduceat(cap, st), 1.0),
        "lead": lead.astype("f8"),
        "region": reg.astype("f8"),
        "sin_doy": np.sin(2 * np.pi * doy[st] / 365.0),
        "cos_doy": np.cos(2 * np.pi * doy[st] / 365.0),
    }
    for p in QUANTS:
        feats["q%02d" % int(100 * p)] = _q(ss, st, en, p)
    for lv in LEVELS:
        feats["cov%02d" % int(100 * lv)] = np.add.reduceat((s >= lv).astype("f8"), st) / n
    # the same soft-coverage idea, but relative to the field's own spread
    feats["cov_mean2x"] = np.add.reduceat(
        (s >= np.repeat(2.0 * mean, (en - st))).astype("f8"), st) / n

    names = sorted(feats)
    X = np.column_stack([feats[k_] for k_ in names]).astype("f4")
    return {
        "X": X, "names": names,
        "P_obs": np.add.reduceat(evt, st) / n,
        "P_cap": feats["Pcap"],
        "Pbar": feats["Pbar"],
        "region": reg, "lead": lead, "init": init,
        "week": np.searchsorted(np.array([11, 18, 25, 32, 39]), lead, side="right"),
        "year": init_year[init],
    }


def fit_logreg(X, y, seed=0):
    """Low-variance alternative to the trees.

    16 years x 7 regions x 46 leads of heavily autocorrelated rows give a much
    smaller effective sample than the 535k row count suggests, so a 400-tree
    model can chase noise even when one of its own inputs (Pbar) already scores
    better on its own.  A linear model on the same features is the control.
    """
    from sklearn.linear_model import LogisticRegression
    from sklearn.pipeline import make_pipeline
    from sklearn.preprocessing import StandardScaler
    m = make_pipeline(StandardScaler(),
                      LogisticRegression(C=0.1, max_iter=2000, random_state=seed))
    m.fit(X, y)
    return m


def fit_model(X, y, kind, seed=0):
    return fit_logreg(X, y, seed) if kind == "logreg" else fit_lgbm(X, y, seed)


def fit_lgbm(X, y, seed=0):
    import lightgbm as lgb
    m = lgb.LGBMClassifier(
        n_estimators=400, learning_rate=0.05, num_leaves=63, max_depth=-1,
        min_child_samples=60, subsample=0.8, subsample_freq=1,
        colsample_bytree=0.8, reg_lambda=1.0, n_jobs=int(os.environ.get("OMP_NUM_THREADS", 20)),
        random_state=seed, verbose=-1)
    m.fit(X, y)
    return m


def cv_scores(X, y, years, nfold=4, kind="lgbm"):
    """Grouped-by-year cross-validated probabilities inside the training set."""
    uy = np.unique(years)
    out = np.full(y.size, np.nan)
    for f in range(nfold):
        held = uy[f::nfold]
        te = np.isin(years, held)
        if te.sum() == 0 or (~te).sum() == 0:
            continue
        m = fit_model(X[~te], y[~te], kind, seed=f)
        out[te] = m.predict_proba(X[te])[:, 1]
        print("      cv fold %d: %d năm giữ lại, %d hàng" % (f, held.size, int(te.sum())),
              flush=True)
    return out


def cuts_on_target(score, week, region, base, per_region=False):
    """Threshold so the forecast fires at the climatological frequency.

    The FBI~1 cut read off the cross-validated training field arrives ~25% too
    loose here: the four fold models are calibrated slightly differently, so
    pooling their scores widens the distribution and its (1 - base) quantile
    sits too low.  Taking the same quantile of the field actually being
    classified fixes that.  It uses no labels -- only the base rate estimated on
    the training years and the shape of the forecast distribution, both of which
    an operational system has in hand.
    """
    th = {}
    for w in range(len(C.WEEKS)):
        if not per_region:
            m = week == w
            th[w] = float(np.quantile(score[m], 1.0 - base[w])) if m.sum() else np.inf
    if not per_region:
        return th
    for r in range(1, NREG + 1):
        for w in range(len(C.WEEKS)):
            m = (region == r) & (week == w)
            th[(r, w)] = float(np.quantile(score[m], 1.0 - base[(r, w)])) \
                if m.sum() and (r, w) in base else np.inf
    return th


def base_rates(obs, week, region, per_region=False):
    if not per_region:
        return {w: float(obs[week == w].mean()) if (week == w).sum() else 0.0
                for w in range(len(C.WEEKS))}
    out = {}
    for r in range(1, NREG + 1):
        for w in range(len(C.WEEKS)):
            m = (region == r) & (week == w)
            if m.sum() and obs[m].sum() >= MIN_EVT_REGION:
                out[(r, w)] = float(obs[m].mean())
    return out


def cuts_by_week(score, obs, week):
    th = {}
    for w in range(len(C.WEEKS)):
        m = week == w
        th[w] = float(np.quantile(score[m], 1.0 - obs[m].mean())) \
            if m.sum() and obs[m].sum() else np.inf
    return th


def cuts_by_region_week(score, obs, week, region, pooled):
    """FBI~1 cut per (region, week), falling back to the pooled week cut."""
    th = {}
    for r in range(1, NREG + 1):
        for w in range(len(C.WEEKS)):
            m = (region == r) & (week == w)
            th[(r, w)] = float(np.quantile(score[m], 1.0 - obs[m].mean())) \
                if m.sum() and obs[m].sum() >= MIN_EVT_REGION else pooled[w]
    return th


def apply_cuts(score, week, region, th):
    p = np.zeros(score.size, bool)
    if isinstance(next(iter(th)), tuple):
        for (r, w), t in th.items():
            m = (region == r) & (week == w)
            p[m] = score[m] >= t
    else:
        for w, t in th.items():
            m = week == w
            p[m] = score[m] >= t
    return p


def table(pred, obs, week, score, title):
    rows = {}
    print("\n  %s" % title)
    print("        %9s %7s %7s %7s %7s %7s %7s %7s"
          % ("base_rate", "n_evt", "POD", "POFD", "FAR", "FBI", "CSI", "AUC"))
    for w in list(range(len(C.WEEKS))) + ["ALL"]:
        m = (week == w) if w != "ALL" else np.ones(obs.size, bool)
        if m.sum() == 0 or obs[m].sum() == 0:
            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())
        name = C.WEEK_NAMES[w] if w != "ALL" else "ALL"
        rows[name] = s
        print("    %-4s %9.4f %7d %7.3f %7.3f %7.3f %7.3f %7.3f %7.3f"
              % (name, s["base_rate"], s["n_events"], s["POD"], s["POFD"],
                 s["FAR"], s["FBI"], s["CSI"], s["AUC"]))
    return rows


def table_by_region(pred, obs, region, score, title):
    rows = {}
    print("\n  %s" % title)
    print("    %-18s %9s %7s %7s %7s %7s %7s %7s %7s"
          % ("vùng", "base_rate", "n_evt", "POD", "POFD", "FAR", "FBI", "CSI", "AUC"))
    for r in range(1, NREG + 1):
        m = region == r
        if m.sum() == 0 or obs[m].sum() == 0:
            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())
        rows[C.REGION_NAMES[r]] = s
        print("    %-18s %9.4f %7d %7.3f %7.3f %7.3f %7.3f %7.3f %7.3f"
              % (C.REGION_NAMES[r], s["base_rate"], s["n_events"], s["POD"],
                 s["POFD"], s["FAR"], s["FBI"], s["CSI"], s["AUC"]))
    return rows


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--tag", default="_ens")
    ap.add_argument("--name", default="ensemble_F3_ens")
    ap.add_argument("--covers", type=float, nargs="+", default=[0.05, 0.10, 0.20])
    ap.add_argument("--cap-min", type=float, default=0.20)
    ap.add_argument("--nfold", type=int, default=4)
    ap.add_argument("--rank", action="store_true",
                    help="rank-transform the features inside each week, separately in "
                         "the training and the validation set")
    ap.add_argument("--per-week", action="store_true",
                    help="one stage-2 model per forecast week instead of one pooled")
    ap.add_argument("--drop", nargs="*", default=[],
                    help="feature names to leave out of the stage-2 model")
    ap.add_argument("--model", default="lgbm", choices=["lgbm", "logreg"])
    ap.add_argument("--suffix", default="")
    ap.add_argument("--save-gate", default="",
                    help="fit the gate of EXP2 on every day (not just the warm season) "
                         "and write its probability for each region-day here")
    ap.add_argument("--gate-cover", type=float, default=0.05)
    ap.add_argument("--save-scores", default=None,
                    help="thu muc de luu truong diem so tang 2 cua tung cau hinh phu")
    args = ap.parse_args()

    import ai_data
    init_year = ai_data.Store().init_year

    grid = C.Grid()
    npts = np.array([(grid.reg == r).sum() for r in range(1, NREG + 1)])

    va = build_table(load(args.tag, "pred_*.npz"), npts, init_year)
    tr = build_table(load(args.tag, "oofpred_*.npz"), npts, init_year)
    print("[regmod] %d region-day huấn luyện, %d kiểm định, %d đặc trưng: %s"
          % (tr["X"].shape[0], va["X"].shape[0], len(tr["names"]),
             ", ".join(tr["names"])), flush=True)

    # how far apart are the two probability scales?  (fold models vs full model)
    print("\n[chẩn đoán] lệch phân bố giữa oofpred (mô hình fold) và pred (mô hình đủ)")
    print("    %-10s %8s %8s %8s %8s %8s %8s"
          % ("đặc trưng", "tr_mean", "va_mean", "tr_q90", "va_q90", "tr_std", "va_std"))
    for nm in ("Pbar", "q90", "Pmax", "cov20"):
        j = tr["names"].index(nm)
        a, b = tr["X"][:, j], va["X"][:, j]
        print("    %-10s %8.4f %8.4f %8.4f %8.4f %8.4f %8.4f"
              % (nm, a.mean(), b.mean(), np.quantile(a, .9), np.quantile(b, .9),
                 a.std(), b.std()))

    if args.drop:
        keep = [i for i, n in enumerate(tr["names"]) if n not in args.drop]
        print("\n[regmod] bỏ đặc trưng: %s" % ", ".join(args.drop), flush=True)
        for t in (tr, va):
            t["X"] = t["X"][:, keep]
        tr["names"] = [tr["names"][i] for i in keep]
        va["names"] = tr["names"]
    if args.rank:
        print("[regmod] xếp hạng đặc trưng trong từng tuần (khử lệch thang đo)", flush=True)
        tr["X"] = rank_within_week(tr["X"], tr["week"])
        va["X"] = rank_within_week(va["X"], va["week"])

    if args.save_gate:
        # EXP2 step 1: is an EHDE possible in this region-day at all?  The gate must
        # answer for every day of the year, winter included, so unlike the tables
        # below it is fitted without the warm-season restriction.
        yt = tr["P_obs"] > args.gate_cover
        yv = va["P_obs"] > args.gate_cover
        print("\n[gate] phủ > %.0f%%, mọi ngày: base rate huấn luyện %.4f (%d), "
              "kiểm định %.4f (%d)"
              % (100 * args.gate_cover, yt.mean(), yt.sum(), yv.mean(), yv.sum()),
              flush=True)
        scv = cv_scores(tr["X"], yt, tr["year"], args.nfold, args.model)
        gmdl = fit_model(tr["X"], yt, args.model)
        gv = gmdl.predict_proba(va["X"])[:, 1]
        print("[gate] AUC theo tuần — kiểm định: %s"
              % ", ".join("%s %.3f" % (C.WEEK_NAMES[w],
                                       M.auc_score(gv[va["week"] == w],
                                                   yv[va["week"] == w]))
                          for w in range(len(C.WEEKS))), flush=True)
        np.savez(args.save_gate,
                 tr_init=tr["init"], tr_lead=tr["lead"], tr_region=tr["region"],
                 tr_gate=scv, tr_obs=yt, tr_Pobs=tr["P_obs"], tr_Pcap=tr["P_cap"],
                 va_init=va["init"], va_lead=va["lead"], va_region=va["region"],
                 va_gate=gv, va_obs=yv, va_Pobs=va["P_obs"], va_Pcap=va["P_cap"],
                 cover=args.gate_cover)
        print("[gate] -> %s" % args.save_gate, flush=True)

    out = {"cap_min": args.cap_min, "features": tr["names"], "results": {}}
    for cov in args.covers:
        for sname, mt, mv in (
                ("mùa nóng", tr["P_cap"] > args.cap_min, va["P_cap"] > args.cap_min),
                ("mọi ngày", np.ones(tr["P_cap"].size, bool),
                 np.ones(va["P_cap"].size, bool))):
            head = "phủ > %.0f%% · %s" % (100 * cov, sname)
            print("\n" + "=" * 78 + "\n  %s" % head, flush=True)
            key = "cover%02d_%s" % (int(100 * cov),
                                    "season" if sname == "mùa nóng" else "all")

            yt, yv = tr["P_obs"][mt] > cov, va["P_obs"][mv] > cov
            Xt, Xv = tr["X"][mt], va["X"][mv]
            wt, wv = tr["week"][mt], va["week"][mv]
            rt, rv = tr["region"][mt], va["region"][mv]
            if yt.sum() < 200 or yv.sum() < 50:
                print("    bỏ qua: quá ít sự kiện (%d / %d)" % (yt.sum(), yv.sum()))
                continue
            print("    base rate huấn luyện %.4f (%d sự kiện), kiểm định %.4f (%d)"
                  % (yt.mean(), yt.sum(), yv.mean(), yv.sum()), flush=True)

            # (0) baseline of section Z: mean regional probability, weekly cut
            th0 = cuts_by_week(tr["Pbar"][mt], yt, wt)
            out["results"][key + "_Pbar_week"] = table(
                apply_cuts(va["Pbar"][mv], wv, rv, th0), yv, wv, va["Pbar"][mv],
                "MỐC SO SÁNH — xác suất trung bình vùng, ngưỡng theo tuần")
            th0r = cuts_by_region_week(tr["Pbar"][mt], yt, wt, rt, th0)
            out["results"][key + "_Pbar_regweek"] = table(
                apply_cuts(va["Pbar"][mv], wv, rv, th0r), yv, wv, va["Pbar"][mv],
                "xác suất trung bình vùng, ngưỡng riêng theo (vùng, tuần)")

            # (1) stage-2 model; thresholds off the grouped-CV field of the train years
            print("    huấn luyện mô hình tầng 2 ...", flush=True)
            if args.per_week:
                scv = np.full(yt.size, np.nan)
                sv = np.full(yv.size, np.nan)
                mdl = None
                for w in range(len(C.WEEKS)):
                    a, b = wt == w, wv == w
                    if a.sum() < 500 or yt[a].sum() < 50 or b.sum() == 0:
                        continue
                    scv[a] = cv_scores(Xt[a], yt[a], tr["year"][mt][a], args.nfold,
                                       args.model)
                    mw = fit_model(Xt[a], yt[a], args.model)
                    sv[b] = mw.predict_proba(Xv[b])[:, 1]
                    if w == 3:
                        mdl = mw          # keep W4's model for the gain table
                keep_v = np.isfinite(sv)
            else:
                scv = cv_scores(Xt, yt, tr["year"][mt], args.nfold, args.model)
                mdl = fit_model(Xt, yt, args.model)
                sv = mdl.predict_proba(Xv)[:, 1]
                keep_v = np.ones(yv.size, bool)
            ok = np.isfinite(scv)
            pbar_v = va["Pbar"][mv]
            if not keep_v.all():
                yv, wv, rv, sv = yv[keep_v], wv[keep_v], rv[keep_v], sv[keep_v]
                pbar_v = pbar_v[keep_v]

            print("    AUC trên chính các năm huấn luyện (chéo theo năm): %s"
                  % ", ".join("%s %.3f" % (C.WEEK_NAMES[w],
                                           M.auc_score(scv[ok & (wt == w)],
                                                       yt[ok & (wt == w)]))
                              for w in range(len(C.WEEKS))
                              if (ok & (wt == w)).sum() and yt[ok & (wt == w)].sum()),
                  flush=True)
            if args.save_scores:
                # the stage-2 score field itself, so the operating point can be
                # chosen afterwards without refitting anything
                np.savez(os.path.join(args.save_scores, "score_%s.npz" % key),
                         va_score=sv, va_obs=yv, va_week=wv, va_region=rv,
                         tr_score=scv[ok], tr_obs=yt[ok], tr_week=wt[ok],
                         tr_region=rt[ok], cover=np.array([cov]))

            th1 = cuts_by_week(scv[ok], yt[ok], wt[ok])
            out["results"][key + "_stack_week"] = table(
                apply_cuts(sv, wv, rv, th1), yv, wv, sv,
                "MÔ HÌNH TẦNG 2 cấp vùng, ngưỡng theo tuần (dò trên trường CV)")
            th1r = cuts_by_region_week(scv[ok], yt[ok], wt[ok], rt[ok], th1)
            out["results"][key + "_stack_regweek"] = table(
                apply_cuts(sv, wv, rv, th1r), yv, wv, sv,
                "MÔ HÌNH TẦNG 2, ngưỡng riêng theo (vùng, tuần), dò trên trường CV")

            b_w = base_rates(yt, wt, rt)
            th2 = cuts_on_target(sv, wv, rv, b_w)
            out["results"][key + "_stack_week_freq"] = table(
                apply_cuts(sv, wv, rv, th2), yv, wv, sv,
                "MÔ HÌNH TẦNG 2, ngưỡng đặt theo tần suất khí hậu (phân vị của chính "
                "trường dự báo)")
            b_rw = base_rates(yt, wt, rt, per_region=True)
            th2r = cuts_on_target(sv, wv, rv, b_rw, per_region=True)
            th2r = {k: (v if np.isfinite(v) else th2[k[1]]) for k, v in th2r.items()}
            out["results"][key + "_stack_regweek_freq"] = table(
                apply_cuts(sv, wv, rv, th2r), yv, wv, sv,
                "MÔ HÌNH TẦNG 2, tần suất khí hậu riêng theo (vùng, tuần)")
            thp = cuts_on_target(pbar_v, wv, rv, b_rw, per_region=True)
            thp = {k: (v if np.isfinite(v) else cuts_on_target(pbar_v, wv, rv, b_w)[k[1]])
                   for k, v in thp.items()}
            out["results"][key + "_Pbar_regweek_freq"] = table(
                apply_cuts(pbar_v, wv, rv, thp), yv, wv, pbar_v,
                "MỐC SO SÁNH Pbar, tần suất khí hậu riêng theo (vùng, tuần)")
            th1r = th2r

            late = wv >= 3
            if late.sum() and yv[late].sum():
                out["results"][key + "_stack_regweek_byregion"] = table_by_region(
                    apply_cuts(sv, wv, rv, th1r)[late], yv[late], rv[late], sv[late],
                    "MÔ HÌNH TẦNG 2 + ngưỡng (vùng, tuần) — theo vùng, trung bình W4–W6")
                out["results"][key + "_Pbar_week_byregion"] = table_by_region(
                    apply_cuts(va["Pbar"][mv], wv, rv, th0)[late], yv[late], rv[late],
                    va["Pbar"][mv][late],
                    "MỐC SO SÁNH (Pbar, ngưỡng tuần) — theo vùng, trung bình W4–W6")

            if mdl is None or args.model != "lgbm":
                continue
            imp = mdl.booster_.feature_importance(importance_type="gain")
            imp = imp / max(imp.sum(), 1e-9)
            order = np.argsort(imp)[::-1][:10]
            out["results"][key + "_gain"] = {tr["names"][i]: float(imp[i]) for i in order}
            print("\n    độ quan trọng (gain) 10 đặc trưng đầu: %s"
                  % ", ".join("%s %.3f" % (tr["names"][i], imp[i]) for i in order),
                  flush=True)

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


if __name__ == "__main__":
    main()