#!/usr/bin/env python3
"""EXP3 hai tầng đầy đủ, quét theo ĐỘ DÀI CỬA SỔ gộp.

Hai câu hỏi được trả lời trong cùng một lần chạy.

(1) TẦNG 2 sau khi cổng loại các mẫu gần như chắc chắn nEHDE.  Cổng chốt bảo thủ
    (rò rỉ <= 5 % số sự kiện) đúng theo yêu cầu "loại ít".  Tập còn lại được khớp
    bằng LỚP MÔ HÌNH KHÁC tầng 1 -- phân loại trực tiếp yEHDE thay vì hồi quy Sev --
    vì đặc tả EXP3 cho phép tầng 2 dùng loại mô hình khác.  Đánh giá trên TOÀN BỘ
    tập kiểm định: mẫu bị cổng đóng xuất ra nEHDE, mẫu còn lại lấy quyết định tầng 2.

(2) CỬA SỔ DÀI CÓ NHIỀU TÍN HIỆU HƠN KHÔNG.  Van Straaten et al. (2023) bỏ cửa sổ
    hai tuần để chuyển sang trung bình 31 ngày vì "the predictable process is better
    detectable on the longer time scale".  Ở đây quét 13 cửa sổ từ 1 tuần tới 6 tuần.

HAI ĐỊNH NGHĨA TARGET, và lý do phải có cả hai.  Tiêu chí "có ít nhất MỘT ngày EHD"
BÃO HOÀ khi cửa sổ dài ra: base rate đi từ 0.329 (1 tuần) lên 0.685 (W3-6) và 0.814
(W1-6).  So AUC giữa các cửa sổ có base rate khác nhau như vậy là so hai bài toán
khác nhau.  Vì thế:
  T1 "so1"  -- >= 1 ngày EHD trong cửa sổ.  Tiêu chí người dùng đã chọn, giữ nguyên.
  T2 "eqbase" -- >= k ngày EHD, với k chọn TRÊN CÁC NĂM HUẤN LUYỆN sao cho tần suất
     gần 1/3 nhất.  Với cửa sổ một tuần thì k = 1, tức T2 trùng T1; với cửa sổ dài
     hơn k lớn dần.  Nhờ giữ độ hiếm không đổi, AUC giữa các cửa sổ so được với nhau.
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time

import numpy as np
from joblib import Parallel, delayed
from scipy.stats import norm

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 exp3_zoo as Z
import exp3_offset as OFF
from exp3_multi import conditional_spread, apply_spread
from exp3_week import MAX_VARS
from exp3_week_s1 import gate_scores, pick_p1

RES_DIR = os.path.join(C.OUT, "AI_Results")
OUT_JSON = os.path.join(RES_DIR, "exp3_window2.json")
OUT_NPZ = os.path.join(RES_DIR, "exp3_window2_scores.npz")
NFOLD = 4
NQ = 9
QLEV = (np.arange(NQ) + 0.5) / NQ
ZQ = norm.ppf(QLEV)
TARGET_RATE = 1.0 / 3.0

WINDOWS = [("W1", [0]), ("W2", [1]), ("W3", [2]), ("W4", [3]), ("W5", [4]),
           ("W6", [5]), ("W1-2", [0, 1]), ("W3-4", [2, 3]), ("W5-6", [4, 5]),
           ("W1-3", [0, 1, 2]), ("W4-6", [3, 4, 5]), ("W3-6", [2, 3, 4, 5]),
           ("W1-6", [0, 1, 2, 3, 4, 5])]


def stage2_zoo():
    """Phân loại trực tiếp yEHDE -- lớp mô hình khác hẳn tầng 1 (hồi quy Sev)."""
    from sklearn.linear_model import LogisticRegression
    from sklearn.ensemble import ExtraTreesClassifier, RandomForestClassifier
    from sklearn.pipeline import make_pipeline
    from sklearn.preprocessing import StandardScaler
    import lightgbm as lgb
    return [
        ("LogReg_L2", "hồi quy logistic, phạt L2",
         lambda: make_pipeline(StandardScaler(),
                               LogisticRegression(C=0.3, max_iter=800))),
        ("LogReg_EN", "hồi quy logistic, phạt L1+L2 (saga)",
         lambda: make_pipeline(StandardScaler(),
                               LogisticRegression(penalty="elasticnet", C=0.5,
                                                  l1_ratio=0.5, solver="saga",
                                                  max_iter=800))),
        ("LGBM_clf", "cây tăng cường phân loại",
         lambda: lgb.LGBMClassifier(n_estimators=250, num_leaves=15,
                                    learning_rate=0.05, min_child_samples=60,
                                    colsample_bytree=0.8, reg_lambda=5.0,
                                    n_jobs=1, verbose=-1)),
        ("LGBM_clf_slow", "cây tăng cường học chậm, lá nông",
         lambda: lgb.LGBMClassifier(n_estimators=400, num_leaves=7,
                                    learning_rate=0.02, min_child_samples=80,
                                    colsample_bytree=0.7, reg_lambda=10.0,
                                    n_jobs=1, verbose=-1)),
        ("ExtraTrees_clf", "cây cực kỳ ngẫu nhiên, phân loại",
         lambda: ExtraTreesClassifier(n_estimators=200, min_samples_leaf=15,
                                      max_features=0.5, n_jobs=1, random_state=0)),
        ("RandomForest_clf", "rừng ngẫu nhiên, phân loại",
         lambda: RandomForestClassifier(n_estimators=200, min_samples_leaf=10,
                                        max_features=0.4, n_jobs=1, random_state=0)),
        ("ANN_16_clf", "ANN 1 lớp ẩn 16 nút (offset = 0, tức không neo)",
         lambda: None),          # xử lý riêng bằng OffsetNet với offset 0
    ]


S2 = stage2_zoo()


def fit_stage2(Xtr, ytr, Xte):
    """Xác suất yEHDE trung bình của các họ tầng 2, và số phiếu M/N."""
    ps = []
    for nm, _, make in S2:
        try:
            if nm == "ANN_16_clf":
                m = OFF.OffsetNet(hidden=(16,), act="relu", epochs=150, seed=0)
                m.fit(Xtr, ytr, np.zeros(Xtr.shape[0]))
                p = m.predict_proba(Xte, np.zeros(Xte.shape[0]))
            else:
                m = make()
                m.fit(Xtr, np.asarray(ytr).astype(int))
                p = m.predict_proba(Xte)[:, 1]
            ps.append(np.clip(np.asarray(p, dtype="f8"), 0.0, 1.0))
        except Exception as e:
            print("    [warn] s2 %s: %s" % (nm, e), flush=True)
    if not ps:
        return np.full(Xte.shape[0], np.nan), np.full(Xte.shape[0], np.nan)
    P = np.vstack(ps)
    vote = (P[:, :, None] >= QLEV[None, None, :]).sum(axis=(0, 2)) / (P.shape[0] * NQ)
    return P.mean(axis=0), vote


def stage1_votes(zoo, Xtr, ztr, Xte, thr):
    """Phiếu tầng 1: thành viên (họ, phân vị) bỏ phiếu khi giá trị dự báo >= thr."""
    out = np.zeros((len(zoo), Xte.shape[0]), dtype="i1")
    for i, fam in enumerate(zoo):
        try:
            m = fam[-1]()
            m.fit(Xtr, ztr)
            mu_tr = np.asarray(m.predict(Xtr)).ravel()
            edges, sd = conditional_spread(mu_tr, ztr - mu_tr)
            mu_te = np.asarray(m.predict(Xte)).ravel()
            sg = apply_spread(mu_te, edges, sd)
            out[i] = ((mu_te[:, None] + sg[:, None] * ZQ[None, :]) >= thr).sum(axis=1)
        except Exception as e:
            print("    [warn] s1 %s: %s" % (fam[0], e), flush=True)
            out[i] = -1
    return out


def p_of(v):
    good = v >= 0
    return np.where(good.any(axis=0),
                    np.where(good, v, 0).sum(axis=0)
                    / np.maximum(good.sum(axis=0) * NQ, 1), np.nan)


def fbi1_threshold(score, obs):
    r = float(obs.mean())
    if r <= 0.0 or r >= 1.0:
        return np.inf
    return float(np.quantile(score, 1.0 - r))


def run_cell(ip, wname, tname, Xtr, ztr, ytr, yr, Xte, zthr, names):
    """Cả hai tầng cho một (ô lưới, cửa sổ, định nghĩa target)."""
    t0 = time.time()
    sd = Xtr.std(axis=0)
    col = sd > 1e-12
    nm = [n for n, k in zip(names, col) if k]
    Xtr = np.ascontiguousarray(Xtr[:, col]); Xte = np.ascontiguousarray(Xte[:, col])
    zoo = OFF.resid_families(Z.zoo_full(), nm)
    ys = np.unique(yr)
    folds = np.array_split(np.random.default_rng(0).permutation(ys), NFOLD)

    # --- tầng 1
    v_oof = np.zeros((len(zoo), Xtr.shape[0]), dtype="i1")
    for f in folds:
        te = np.isin(yr, f)
        if te.all() or not te.any():
            continue
        v_oof[:, te] = stage1_votes(zoo, Xtr[~te], ztr[~te], Xtr[te], zthr)
    q_oof = p_of(v_oof)
    q_val = p_of(stage1_votes(zoo, Xtr, ztr, Xte, zthr))
    p1 = pick_p1(q_oof, ytr)

    # --- tầng 2 CHỈ trên tập cổng giữ lại
    keep_tr = q_oof >= p1
    s2_oof = np.zeros(Xtr.shape[0]); s2_val = np.zeros(Xte.shape[0])
    n_keep = int(keep_tr.sum())
    if n_keep >= 100 and 0 < ytr[keep_tr].sum() < n_keep:
        for f in folds:
            te = np.isin(yr, f) & keep_tr
            trn = (~np.isin(yr, f)) & keep_tr
            if not te.any() or trn.sum() < 50 or np.unique(ytr[trn]).size < 2:
                continue
            s2_oof[te], _ = fit_stage2(Xtr[trn], ytr[trn], Xtr[te])
        s2_val, _ = fit_stage2(Xtr[keep_tr], ytr[keep_tr], Xte)
    print("[%s %s p%3d] cổng giữ %d/%d mẫu (base %.3f -> %.3f), %.0fs"
          % (tname, wname, ip, n_keep, Xtr.shape[0], ytr.mean(),
             ytr[keep_tr].mean() if n_keep else np.nan, time.time() - t0), flush=True)
    return (ip, wname, tname, q_oof, q_val, s2_oof, s2_val, p1, n_keep)


def build(store, points, years, fset, lts, names_full):
    """Đặc trưng gộp trên cửa sổ tuỳ ý, cùng Sev cực đại và số ngày EHD."""
    isel = np.nonzero(np.isin(store.init_year, years))[0]
    X, y, info = store.build_isel(points, isel, fset)
    ni, nk, npt = isel.size, C.NLT, points.size
    nf = X.shape[1]
    names = store.feature_names(fset)
    X = X.reshape(ni, nk, npt, nf)[:, lts, :, :]
    sev = y["sev"].reshape(ni, nk, npt)[:, lts, :]
    evt = y["evt"].reshape(ni, nk, npt)[:, lts, :]
    imax = [names.index(v) for v in MAX_VARS if v in names]
    isv = names.index("sv_mod")
    agg = np.concatenate(
        [np.nanmean(X, axis=1), np.nanmax(X[:, :, :, imax], axis=1),
         (X[:, :, :, isv] >= 0.0).sum(axis=1, dtype="f4")[:, :, None]],
        axis=2).reshape(ni * npt, -1)
    return dict(agg=agg,
                sev=np.nanmax(sev, axis=1).reshape(ni * npt),
                cnt=np.nansum(evt, axis=1).reshape(ni * npt),
                year=np.broadcast_to(store.init_year[isel][:, None],
                                     (ni, npt)).reshape(ni * npt),
                point=np.broadcast_to(points[None, :], (ni, npt)).reshape(ni * npt))


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("--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)
    names = store.feature_names(args.fset)
    imax = [names.index(v) for v in MAX_VARS if v in names]
    agg_names = (["m_" + s for s in names] + ["x_" + names[i] for i in imax]
                 + ["n_hotday_mod"])

    jobs, meta = [], []
    for wname, ws in WINDOWS:
        lts = np.concatenate([C.WEEKS[i][1] for i in ws])
        tr = build(store, pts, C.TRAIN_YEARS, args.fset, lts, agg_names)
        te = build(store, pts, C.VALID_YEARS, args.fset, lts, agg_names)
        for p in pts:
            mt, mv = tr["point"] == p, te["point"] == p
            ok_t = (np.isfinite(tr["agg"][mt]).all(axis=1) & np.isfinite(tr["sev"][mt])
                    & np.isfinite(tr["cnt"][mt]))
            ok_v = (np.isfinite(te["agg"][mv]).all(axis=1)
                    & np.isfinite(te["cnt"][mv]) & np.isfinite(te["sev"][mv]))
            A, Av = tr["agg"][mt][ok_t], te["agg"][mv][ok_v]
            cnt, cntv = tr["cnt"][mt][ok_t], te["cnt"][mv][ok_v]
            sev = tr["sev"][mt][ok_t]
            yr = tr["year"][mt][ok_t]
            for tname in ("so1", "eqbase"):
                if tname == "so1":
                    k = 1
                    ztr, zthr = sev, 0.0          # hồi quy Sev, ngưỡng phiếu 0
                else:
                    # k chọn TRÊN NĂM HUẤN LUYỆN cho tần suất gần 1/3 nhất
                    ks = np.arange(1, int(max(cnt.max(), 1)) + 1)
                    if ks.size == 0:
                        continue
                    rates = np.array([(cnt >= kk).mean() for kk in ks])
                    k = int(ks[np.argmin(np.abs(rates - TARGET_RATE))])
                    ztr, zthr = cnt.astype("f8"), float(k)   # hồi quy SỐ NGÀY
                    if k == 1 and len(ws) == 1:
                        continue                  # trùng hệt so1, khỏi chạy lại
                ytr = (cnt >= k).astype("f8")
                yva = (cntv >= k)
                jobs.append((int(p), wname, tname, A, ztr, ytr, yr, Av, zthr,
                             agg_names))
                meta.append((int(p), wname, tname, ytr > 0.5, yva, k))
    print("[win2] %d công việc = %d cửa sổ x %d điểm x định nghĩa target"
          % (len(jobs), len(WINDOWS), pts.size), flush=True)

    got = Parallel(n_jobs=args.jobs, verbose=5)(delayed(run_cell)(*j) for j in jobs)

    obs_tr = {(p, w, t): o for p, w, t, o, _, _ in meta}
    obs_va = {(p, w, t): o for p, w, t, _, o, _ in meta}
    kval = {(p, w, t): k for p, w, t, _, _, k in meta}

    res = {"config": {"points": pts.tolist(),
                      "windows": [w for w, _ in WINDOWS],
                      "fset": args.fset, "n_quantile": NQ,
                      "n_family_s1": len(OFF.resid_families(Z.zoo_full(), agg_names)),
                      "n_family_s2": len(S2)},
           "models_s2": [{"name": n, "desc": d} for n, d, _ in S2],
           "by_window": {}, "by_cell": {}}

    byk = {}
    dump = {}
    for ip, w, t, qo, qv, so, sv, p1, nk in got:
        byk.setdefault((w, t), []).append((ip, qo, qv, so, sv, p1, nk))
        # điểm số cuối cùng, để quét lại ngưỡng mà không phải khớp lại 441 ô
        kk = "%d|%s|%s" % (ip, w, t)
        dump[kk + "|s_oof"] = np.where(qo >= p1, so, -1.0)
        dump[kk + "|s_val"] = np.where(qv >= p1, sv, -1.0)
        dump[kk + "|obs_oof"] = obs_tr[(ip, w, t)]
        dump[kk + "|obs_val"] = obs_va[(ip, w, t)]
    np.savez_compressed(OUT_NPZ, **dump)
    print("[dump] điểm số -> %s" % OUT_NPZ, flush=True)

    for (w, t), cells in sorted(byk.items()):
        pv_all, ov_all, po_all, oo_all = [], [], [], []
        for ip, qo, qv, so, sv, p1, nk in cells:
            oo, ov = obs_tr[(ip, w, t)], obs_va[(ip, w, t)]
            # điểm số cuối: cổng đóng thì loại hẳn, cổng mở thì lấy tầng 2
            s_o = np.where(qo >= p1, so, -1.0)
            s_v = np.where(qv >= p1, sv, -1.0)
            th = fbi1_threshold(s_o, oo)
            pred_v, pred_o = s_v > th, s_o > th
            g = M.scores(*M.contingency(pred_v, ov))
            g.update({"region": int(store.grid.reg[ip]),
                      "p1": p1, "n_keep_train": nk, "k_days": kval[(ip, w, t)],
                      "closure_val": float((qv < p1).mean())})
            res["by_cell"]["%d|%s|%s" % (ip, w, t)] = g
            pv_all.append(pred_v); ov_all.append(ov)
            po_all.append(pred_o); oo_all.append(oo)
        pv = np.concatenate(pv_all); ov = np.concatenate(ov_all)
        g = M.scores(*M.contingency(pv, ov))
        g["PC"] = float((pv == ov).mean())
        g["k_days"] = kval[(cells[0][0], w, t)]
        g["closure_val"] = float(np.mean([res["by_cell"]["%d|%s|%s" % (c[0], w, t)]
                                          ["closure_val"] for c in cells]))
        g["n_keep_train"] = int(np.mean([c[6] for c in cells]))
        # AUC của điểm số liên tục, để so cửa sổ với nhau không phụ thuộc ngưỡng
        sv_all = np.concatenate([np.where(c[2] >= c[5], c[4], -1.0) for c in cells])
        g["AUC"] = float(M.auc_score(sv_all, ov))
        res["by_window"]["%s|%s" % (w, t)] = g
        print("[%s %s] k=%d base %.3f | POD %.3f POFD %.3f FAR %.3f FBI %.3f "
              "CSI %.3f SEDI %.3f AUC %.3f (cổng đóng %.0f%%)"
              % (w, t, g["k_days"], g["base_rate"], g["POD"], g["POFD"], g["FAR"],
                 g["FBI"], g["CSI"], g["SEDI"], g["AUC"], 100 * g["closure_val"]),
              flush=True)

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


if __name__ == "__main__":
    main()