"""Bước 3 của kế hoạch p_raw error-dressing (đã duyệt): neo tầng-2 lên p_raw.

Van Straaten (2023): ML KHÔNG học lại từ đầu mà chỉ học phần SỬA cái tiên nghiệm:

    logit(p_cuoi) = logit(p_raw) + f_ML(dac_trung)

Ở đây p_raw là tiên nghiệm VẬT LÝ dựng bằng error-dressing của control (bước 1-2,
praw_dress.py): p_raw = Phi((mu + bias(mu)) / sd(mu)), mu = max-tuan(sv_mod). Đây là
tiên nghiệm dung lượng cực thấp (10 bin bias/sd), đúng tinh thần "prior chứ không
phải mô hình thứ hai".

So sánh trên 21 điểm (3 điểm x 7 vùng), train = TRAIN_YEARS, valid = VALID_YEARS:
  * prior   -- chỉ p_raw + isotonic (không có ML), lấy từ bước 1-2;
  * ML0     -- ML KHÔNG neo (offset = 0): độ phân biệt thuần của ML;
  * MLanchor-- ML neo lên logit(p_raw): kỳ vọng = độ phân biệt của ML + hiệu chuẩn
               + diễn giải của p_raw, không tụt AUC.
Ba họ neo (đều cài sẵn trong exp3_offset): Off_Logit (logistic có offset, '0 hidden'
của bài báo), Off_ANN_4 (ANN 1 lớp 4 nút -- kiến trúc thắng cuộc), Off_LGBM_gbdt.
Chỉ số: AUC, Brier skill, reliability gap, max-SEDI.
"""
import json
import os
import sys

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

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, brier_skill_score
import exp3_offset as OFF
from praw_dress import pick_points, dress_stats, dress_prob, rel_gap, prep_ml

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

# Chế độ big-data cho ~35k mẫu/tuần: ít epoch, batch lớn (giống phase1_arch).
FAMILIES = [
    ("Off_Logit", lambda: OFF.OffsetNet(hidden=(), epochs=60, batch=2048, seed=0)),
    ("Off_ANN_4", lambda: OFF.OffsetNet(hidden=(4,), act="elu", epochs=60,
                                        batch=2048, seed=1)),
    ("Off_LGBM", lambda: OFF.OffsetLGBM(n_estimators=300, num_leaves=31,
                                        learning_rate=0.05, min_child_samples=200,
                                        colsample_bytree=0.8, reg_lambda=5.0)),
]


def metrics(p, y):
    _, sedi = best_threshold(p, y, metric="SEDI")
    return dict(auc=auc_score(p, y), bss=brier_skill_score(p, y),
                relgap=rel_gap(p, y), sedi=sedi)


def main():
    store = Store()
    pts = pick_points(store)
    os.makedirs(OUT, exist_ok=True)
    print("21-point testbed:", pts.size, "points", flush=True)

    results = {}
    for wname, lts in C.WEEKS:
        Atr, Str, Ntr, _, _, names = build_window(
            store, pts, C.TRAIN_YEARS, "F3", lts)
        Ate, _, Nte, _, _, _ = build_window(
            store, pts, C.VALID_YEARS, "F3", lts)
        ix = names.index("x_sv_mod")
        ytr = (Ntr >= 1).astype(int)
        yte = (Nte >= 1).astype(int)

        # --- p_raw error-dressing (học trên train) ---
        mu_tr, mu_te = Atr[:, ix], Ate[:, ix]
        ok_tr = np.isfinite(mu_tr) & np.isfinite(Str)
        edges, bias, sd = dress_stats(mu_tr[ok_tr], Str[ok_tr])
        praw_tr = dress_prob(np.nan_to_num(mu_tr, nan=np.nanmedian(mu_tr[ok_tr])),
                             edges, bias, sd)
        praw_te = dress_prob(np.nan_to_num(mu_te, nan=np.nanmedian(mu_tr[ok_tr])),
                             edges, bias, sd)
        off_tr = OFF.logit(praw_tr)
        off_te = OFF.logit(praw_te)

        # --- prior thuần: p_raw + isotonic ---
        iso = IsotonicRegression(out_of_bounds="clip")
        iso.fit(praw_tr, ytr)
        piso = iso.predict(praw_te)

        # --- đặc trưng ML (chuẩn hoá theo train, impute median train) ---
        Xtr, Xte = prep_ml(Atr, Ate)
        zeros_tr = np.zeros(Xtr.shape[0])
        zeros_te = np.zeros(Xte.shape[0])

        wres = {"base": float(yte.mean()), "prior_iso": metrics(piso, yte)}
        for fname, mk in FAMILIES:
            # ML không neo
            m0 = mk(); m0.fit(Xtr, ytr, zeros_tr)
            p0 = m0.predict_proba(Xte, zeros_te)
            # ML neo lên logit(p_raw)
            ma = mk(); ma.fit(Xtr, ytr, off_tr)
            pa = ma.predict_proba(Xte, off_te)
            wres[fname] = {"ML0": metrics(p0, yte), "anchor": metrics(pa, yte)}
        results[wname] = wres

        line = "%-4s base=%.3f | prior AUC=%.3f BSS=%+.3f rel=%.4f" % (
            wname, wres["base"], wres["prior_iso"]["auc"],
            wres["prior_iso"]["bss"], wres["prior_iso"]["relgap"])
        for fname, _ in FAMILIES:
            a = wres[fname]["ML0"]; b = wres[fname]["anchor"]
            line += ("\n     %-10s ML0  AUC=%.3f BSS=%+.3f rel=%.4f SEDI=%.3f | "
                     "anchor AUC=%.3f BSS=%+.3f rel=%.4f SEDI=%.3f") % (
                fname, a["auc"], a["bss"], a["relgap"], a["sedi"],
                b["auc"], b["bss"], b["relgap"], b["sedi"])
        print(line, flush=True)

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


if __name__ == "__main__":
    main()