"""Bước 1-2 của kế hoạch p_raw error-dressing (đã duyệt 2026-09-11).

Ý tưởng: EC ở đây chỉ có CONTROL (không tổ hợp nhiều thành viên), nên không lấy
được p_raw kiểu 'tỷ lệ thành viên vượt ngưỡng'. Thay vào đó KHOÁC SAI SỐ lên dự báo
severity của control:

    EHDE  <=>  max_tuan( Sev_obs )  >= 0
    dự báo control cho đại lượng đó là  mu = max_tuan( sv_mod )   (cột 'x_sv_mod')
    học phân bố sai số  resid = Sev_obs_max_tuan - mu  CHỈ TRÊN TRAIN, theo thập
    phân vị của mu (bắt heteroscedasticity + bias có điều kiện), rồi

        p_raw = P( mu + resid >= 0 ) = Phi( (mu + bias(mu)) / sd(mu) ).

Khoác lên MAX-TUẦN (một biến), KHÔNG khoác từng ngày rồi 1-Prod(1-p) (các ngày sai
số tương quan mạnh -> thổi phồng xác suất).

So sánh trên 21 điểm (3 điểm x 7 vùng), train = TRAIN_YEARS, valid = VALID_YEARS:
  * p_raw thô (dressing, chưa hiệu chuẩn lại),
  * p_raw + isotonic (hiệu chuẩn lại, fit trên train),
  * baseline ML tầng-2 (LogReg toàn đặc trưng) -- để đối chiếu AUC hiện có.
Chỉ tiêu đọc: base rate, AUC, Brier skill score, độ tin cậy (reliability), max-SEDI.
Kỳ vọng (nói thẳng): p_raw KHÔNG phá trần AUC; giá trị của nó là XÁC SUẤT HIỆU CHUẨN
TỐT + neo vật lý + diễn giải được, đúng thứ phát hành được cho W2-W6.
"""
import json
import os
import sys

import numpy as np
from scipy.stats import norm
from sklearn.isotonic import IsotonicRegression
from sklearn.linear_model import LogisticRegression

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, best_threshold, reliability, brier_skill_score

NBIN = 10            # số thập phân vị của mu để ước lượng bias + sd
MIN_BIN = 50         # số mẫu tối thiểu trong một bin mới tin được bias/sd riêng
OUT = os.path.join(C.OUT, "AI_Results", "praw")


def pick_points(store, per_region=3):
    """3 điểm mỗi vùng, rải đều theo chỉ số ô lưới trong vùng (giống ablation)."""
    reg = np.asarray(store.st["region"]).astype(int)
    pts = []
    for r in range(1, 8):
        idx = np.nonzero(reg == r)[0]
        if idx.size == 0:
            continue
        sel = idx[np.linspace(0, idx.size - 1, per_region).round().astype(int)]
        pts.extend(sel.tolist())
    return np.array(sorted(set(pts)), dtype=int)


def dress_stats(mu, truth):
    """Học bias + sd của resid = truth - mu theo thập phân vị của mu (CHỈ train).

    Trả về (edges, bias, sd): dùng để dựng phân bố dự báo dạng N(mu+bias, sd^2).
    """
    edges = np.quantile(mu, np.linspace(0, 1, NBIN + 1))
    edges[0] -= 1e-6
    edges[-1] += 1e-6
    b = np.clip(np.digitize(mu, edges) - 1, 0, NBIN - 1)
    resid = truth - mu
    g_bias = float(np.mean(resid))
    g_sd = float(np.std(resid)) if resid.size else 1.0
    bias = np.full(NBIN, g_bias)
    sd = np.full(NBIN, g_sd)
    for k in range(NBIN):
        m = b == k
        if m.sum() >= MIN_BIN:
            bias[k] = float(np.mean(resid[m]))
            sd[k] = float(np.std(resid[m]))
    return edges, bias, np.maximum(sd, 1e-3)


def dress_prob(mu, edges, bias, sd):
    """p_raw = P(mu + resid >= 0) = Phi((mu + bias(mu)) / sd(mu))."""
    b = np.clip(np.digitize(mu, edges) - 1, 0, NBIN - 1)
    return norm.cdf((mu + bias[b]) / sd[b])


def rel_gap(prob, obs, nbin=10):
    """Sai số hiệu chuẩn có trọng số (reliability component of the Brier score)."""
    pts = reliability(prob, obs, nbin)
    num = sum(n * (f - o) ** 2 for f, o, n in pts if n)
    den = sum(n for _, _, n in pts if n)
    return num / den if den else np.nan


def prep_ml(Atr, Ate):
    """Impute (train median) + standardize (train stats) cho baseline ML."""
    med = np.nanmedian(Atr, axis=0)
    Xtr = np.where(np.isnan(Atr), med, Atr)
    Xte = np.where(np.isnan(Ate), med, Ate)
    mu = Xtr.mean(0)
    sd = Xtr.std(0)
    sd[sd == 0] = 1.0
    return (Xtr - mu) / sd, (Xte - mu) / sd


def main():
    store = Store()
    pts = pick_points(store)
    reg = np.asarray(store.st["region"]).astype(int)
    os.makedirs(OUT, exist_ok=True)
    print("21-point testbed (%d):" % pts.size,
          {r: int((reg[pts] == r).sum()) for r in range(1, 8)}, flush=True)

    results = {}
    print("\n%-4s %6s | %-34s | %-22s | %-8s"
          % ("W", "base", "AUC  ctrl / p_raw / p_iso / ML",
             "BSS  p_raw / p_iso / ML", "SEDI iso"), flush=True)
    for wname, lts in C.WEEKS:
        Atr, Str, Ntr, _, Ptr, names = build_window(
            store, pts, C.TRAIN_YEARS, "F3", lts)
        Ate, Ste, Nte, _, Pte, _ = build_window(
            store, pts, C.VALID_YEARS, "F3", lts)
        ix = names.index("x_sv_mod")

        mu_tr = Atr[:, ix]
        mu_te = Ate[:, ix]
        ytr = (Ntr >= 1).astype(int)
        yte = (Nte >= 1).astype(int)

        # loại các hàng dự báo control không xác định (toàn NaN trong tuần)
        ok_tr = np.isfinite(mu_tr) & np.isfinite(Str)
        ok_te = np.isfinite(mu_te)
        mu_tr, Str_v, ytr_v = mu_tr[ok_tr], Str[ok_tr], ytr[ok_tr]
        mu_te_v, yte_v = mu_te[ok_te], yte[ok_te]

        # --- error dressing: học trên train, áp lên valid ---
        edges, bias, sd = dress_stats(mu_tr, Str_v)
        praw_tr = dress_prob(mu_tr, edges, bias, sd)
        praw_te = dress_prob(mu_te_v, edges, bias, sd)

        # --- hiệu chuẩn lại (isotonic, fit trên train) ---
        iso = IsotonicRegression(out_of_bounds="clip")
        iso.fit(praw_tr, ytr_v)
        piso_te = iso.predict(praw_te)

        # --- baseline ML tầng-2: LogReg toàn đặc trưng, để đối chiếu AUC ---
        Xtr, Xte = prep_ml(Atr, Ate)
        mlc = LogisticRegression(C=0.3, max_iter=2000)
        mlc.fit(Xtr, ytr)
        pml_te = mlc.predict_proba(Xte)[:, 1][ok_te]

        base = float(yte_v.mean())
        auc_ctrl = auc_score(mu_te_v, yte_v)      # AUC của chính sv_mod thô (xếp hạng)
        auc_raw = auc_score(praw_te, yte_v)
        auc_iso = auc_score(piso_te, yte_v)
        auc_ml = auc_score(pml_te, yte_v)
        bss_raw = brier_skill_score(praw_te, yte_v)
        bss_iso = brier_skill_score(piso_te, yte_v)
        bss_ml = brier_skill_score(pml_te, yte_v)
        _, sedi_iso = best_threshold(piso_te, yte_v, metric="SEDI")
        _, sedi_ml = best_threshold(pml_te, yte_v, metric="SEDI")
        relgap_raw = rel_gap(praw_te, yte_v)
        relgap_iso = rel_gap(piso_te, yte_v)

        results[wname] = dict(
            base=base, n_valid=int(yte_v.size),
            auc_ctrl=auc_ctrl, auc_raw=auc_raw, auc_iso=auc_iso, auc_ml=auc_ml,
            bss_raw=bss_raw, bss_iso=bss_iso, bss_ml=bss_ml,
            sedi_iso=sedi_iso, sedi_ml=sedi_ml,
            relgap_raw=relgap_raw, relgap_iso=relgap_iso,
            bias=bias.tolist(), sd=sd.tolist(),
            reliab_raw=reliability(praw_te, yte_v),
            reliab_iso=reliability(piso_te, yte_v),
        )
        print("%-4s %6.3f | AUC %.3f/%.3f/%.3f/%.3f | BSS %+.3f/%+.3f/%+.3f | %.3f"
              % (wname, base, auc_ctrl, auc_raw, auc_iso, auc_ml,
                 bss_raw, bss_iso, bss_ml, sedi_iso), flush=True)

    # bảng độ tin cậy chi tiết
    print("\n=== Reliability (p_raw sau isotonic): fcst -> obs (n) ===", flush=True)
    for wname, _ in C.WEEKS:
        row = " ".join("%.2f>%.2f" % (f, o) if n else "  -  "
                       for f, o, n in results[wname]["reliab_iso"])
        print("%-5s relgap raw=%.4f iso=%.4f | %s"
              % (wname, results[wname]["relgap_raw"],
                 results[wname]["relgap_iso"], row), flush=True)

    with open(os.path.join(OUT, "praw_dress.json"), "w") as f:
        json.dump(results, f, indent=2)
    print("\nsaved ->", os.path.join(OUT, "praw_dress.json"), flush=True)


if __name__ == "__main__":
    main()