"""FORECAST-OF-OPPORTUNITY (trụ 3 của đổi-trục W3-6).

Câu hỏi: kỹ năng W3-6 KHÔNG đều theo init. Có tồn tại một TẬP CON init "cơ hội"
— nhận biết được TẠI init, KHÔNG dùng Tx — mà trên đó kỹ năng cao hẳn (đạt chỉ tiêu),
đổi lấy việc chỉ phát cảnh báo chắc trên một phần các tuần?

KHÁC với PA1 (AH7, ÂM): PA1 hỏi "MJO/ENSO có nâng kỹ năng TRUNG BÌNH không" -> không.
Ở đây hỏi "MJO có PHÂN TẦNG kỹ năng không" -> đó là câu hỏi forecast-of-opportunity,
chưa từng đo. Biên độ MJO tại init có thể vô dụng làm feature nhưng vẫn tách được
tuần dễ dự báo khỏi tuần khó.

Tín hiệu cơ hội (đều biết tại init, KHÔNG từ Tx):
  - MJO active: x_omi_amp >= 1.0  (pha MJO mạnh -> khả dự báo nội mùa cao hơn).
  - Sharpness mô hình: |p - base| lớn (mô hình "tự tin"); đây là subset theo p, đánh
    giá reliability có điều kiện — bổ trợ, vì confidence-from-Tx bị nghi ngờ.

Đánh giá CÓ ĐIỀU KIỆN: trong mỗi subset, ngưỡng FBI=1 riêng của subset (quantile theo
base rate của chính subset), rồi POD/POFD/FAR/CSI/SEDI. 21 điểm gộp, LogReg + LGBM.
Chạy cục bộ, không cần SLURM.
"""
import json
import os
import sys

import numpy as np
from sklearn.linear_model import LogisticRegression
import lightgbm as lgb

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import ehd_common as C
from ai_data import Store
from phase1_arch import build_window
from ehd_metrics import auc_score, scores as sc_scores, contingency

OUT = os.path.join(C.OUT, "AI_Results", "foo")


def pick_points(store, per_region=3):
    reg = np.asarray(store.st["region"]).astype(int)
    pts = []
    for r in range(1, 8):
        idx = np.nonzero(reg == r)[0]
        if idx.size:
            pts.extend(idx[np.linspace(0, idx.size - 1, per_region)
                           .round().astype(int)].tolist())
    return np.array(sorted(set(pts)), dtype=int)


def prep(Xtr, Xte):
    med = np.nanmedian(Xtr, axis=0)
    Xtr = np.where(np.isnan(Xtr), med, Xtr)
    Xte = np.where(np.isnan(Xte), med, Xte)
    mu, sd = Xtr.mean(0), Xtr.std(0)
    sd[sd == 0] = 1.0
    return (Xtr - mu) / sd, (Xte - mu) / sd


def op(score, y):
    """AUC + bảng contingency tại điểm vận hành FBI=1 (ngưỡng = quantile base rate)."""
    if y.size < 30 or y.sum() < 5 or y.sum() == y.size:
        return dict(n=int(y.size), base=float(y.mean()), AUC=np.nan,
                    POD=np.nan, POFD=np.nan, FAR=np.nan, FBI=np.nan,
                    CSI=np.nan, SEDI=np.nan)
    auc = auc_score(score, y)
    r = float(y.mean())
    t = np.quantile(score, 1.0 - r)
    a, b, c, d = contingency(score >= t, y)
    s = sc_scores(a, b, c, d)
    return dict(n=int(y.size), base=r, AUC=auc, POD=s["POD"], POFD=s["POFD"],
                FAR=s["FAR"], FBI=s["FBI"], CSI=s["CSI"], SEDI=s["SEDI"])


def fit_models(Xtr, ytr, Xte):
    m1 = LogisticRegression(C=0.3, max_iter=2000).fit(Xtr, ytr)
    p1 = m1.predict_proba(Xte)[:, 1]
    m2 = lgb.LGBMClassifier(n_estimators=300, learning_rate=0.05, num_leaves=31,
                            min_child_samples=200, colsample_bytree=0.8,
                            reg_lambda=5.0, n_jobs=8, verbose=-1).fit(Xtr, ytr)
    p2 = m2.predict_proba(Xte)[:, 1]
    return p1, p2


def fmt(d):
    if not np.isfinite(d["AUC"]):
        return "n=%5d base=%.3f   (too few events)" % (d["n"], d["base"])
    return ("n=%5d base=%.3f AUC=%.3f POD=%.3f POFD=%.3f FAR=%.3f FBI=%.3f "
            "CSI=%.3f SEDI=%.3f" % (d["n"], d["base"], d["AUC"], d["POD"],
            d["POFD"], d["FAR"], d["FBI"], d["CSI"], d["SEDI"]))


def main():
    store = Store()
    pts = pick_points(store)
    npt = pts.size
    os.makedirs(OUT, exist_ok=True)

    # MJO amplitude tại init (đã có sẵn trong external.npz, cột x_omi_amp)
    idx_names = [str(s) for s in store.ext["idx_names"]]
    j_amp = idx_names.index("x_omi_amp")
    amp_by_init = store.ext["idx"][:, j_amp]              # (ninit,)

    # init toàn cục cho hàng VALID: build_window duyệt init theo thứ tự tăng dần của
    # isel_all, mỗi init phát npt hàng -> repeat. (initstate_probe KHÔNG bỏ hàng NaN.)
    isel_val = np.nonzero(np.isin(store.init_year, C.VALID_YEARS))[0]
    amp_val_row = np.repeat(amp_by_init[isel_val], npt)  # (nrow_val,)

    print("FoO probe — 21 điểm; opportunity = MJO active (x_omi_amp>=1) tại init",
          flush=True)
    print("MJO amp: mean %.2f, active-init share %.2f\n"
          % (np.nanmean(amp_by_init), np.mean(amp_by_init >= 1.0)), flush=True)

    res = {}
    for wi in range(1, 6):                                # W2..W6
        wname, lt_w = C.WEEKS[wi]
        Ac_tr, _, N_tr, _, _, _ = build_window(store, pts, C.TRAIN_YEARS, "F3", lt_w)
        Ac_te, _, N_te, _, _, _ = build_window(store, pts, C.VALID_YEARS, "F3", lt_w)
        assert Ac_te.shape[0] == amp_val_row.shape[0], (
            Ac_te.shape[0], amp_val_row.shape[0])
        ytr = (N_tr >= 1).astype(int)
        yte = (N_te >= 1).astype(int)
        Xtr, Xte = prep(Ac_tr, Ac_te)
        p_lr, p_gb = fit_models(Xtr, ytr, Xte)

        act = amp_val_row >= 1.0                          # MJO active tại init
        res[wname] = {}
        print("=== %s (VALID base %.3f) ===" % (wname, yte.mean()), flush=True)
        for mname, p in (("LogReg", p_lr), ("LGBM", p_gb)):
            d_all = op(p, yte)
            d_act = op(p[act], yte[act])
            d_ina = op(p[~act], yte[~act])
            # sharpness FoO: 1/3 hàng "tự tin nhất" theo |p - base|
            conf = np.abs(p - p.mean())
            hi = conf >= np.quantile(conf, 2.0 / 3.0)
            d_hi = op(p[hi], yte[hi])
            res[wname][mname] = dict(all=d_all, mjo_active=d_act,
                                     mjo_inactive=d_ina, sharp_top3=d_hi)
            print("  %-6s ALL        %s" % (mname, fmt(d_all)), flush=True)
            print("  %-6s MJO-active %s" % (mname, fmt(d_act)), flush=True)
            print("  %-6s MJO-inact. %s" % (mname, fmt(d_ina)), flush=True)
            print("  %-6s sharp-top3 %s" % (mname, fmt(d_hi)), flush=True)
        print("", flush=True)

    with open(os.path.join(OUT, "foo_probe.json"), "w") as f:
        json.dump(res, f, indent=1, default=float)
    print("[done] ->", os.path.join(OUT, "foo_probe.json"), flush=True)


if __name__ == "__main__":
    main()