#!/usr/bin/env python3
"""So sanh HAI cach ha quy mo tren CUNG init 29/8, tren toan 541 diem:

    BCSD  quantile-map coarse Tx -> fine (chuoi da chung nhan, d1)
    AI    LightGBM hoc anh xa (coarse Tx, clim, mua) -> fine Tx (d2)

Ca hai:
  * nhan CUNG mot dau vao: Tx tho nearest-cell cua init 29/8 (d0_op);
  * dung CUNG nguong thr;
  * hieu chuan delta DOI XUNG bang frequency-matching tren train-years
    (delta = phan vi (1-base) cua margin train), nen exc_ratio~1 theo cau truc
    va khac biet con lai CHI do cach ha quy mo, khong do calibration.

D2 o day la Tx-only (init 29/8 chi tai ve TMAX tu NOMADS, khong co pgb/16 bien
flux), nen KHAC voi D2 da bien trong bcsd.json. Muc dich la ca nghien cuu tren
mot ca that, khong phai tai chung nhan D2.

Xuat: outputs/RT/cmp_ai_bcsd_<init>.json  va  ehd_cfs_29aug_bcsd_vs_ai.html
"""

from __future__ import annotations

import argparse, base64, datetime as dt, io, json, os, sys
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import cfs_common as C
import build_obs as B
from rt_forecast import read_rt, REG_NAME, REG_FRAC

FIELDS = os.path.join(C.OUT, "AI_Results", "bcsd_fluxonly.json.w%d.fields.npz")
CSS = open(os.path.join(os.path.dirname(__file__), "_report_css.txt")).read()
TARGET = {"POD": (">", .6), "POFD": ("<", .2), "FAR": ("<", .2), "CSI": (">", .4)}


def contingency(fc, ev):
    a = float((fc & ev).sum()); b = float((fc & ~ev).sum())
    c = float((~fc & ev).sum()); d = float((~fc & ~ev).sum())
    pod = a / max(a + c, 1); pofd = b / max(b + d, 1)
    far = b / max(a + b, 1)
    return dict(a=a, b=b, c=c, d=d, base=float(ev.mean()),
                POD=pod, POFD=pofd, FAR=far,
                FBI=(a + b) / max(a + c, 1), CSI=a / max(a + b + c, 1),
                exc_ratio=float(fc.mean()) / max(float(ev.mean()), 1e-9))


def fit_delta(margin, ev):
    """delta sao cho tan suat vuot = tan suat quan trac tren train (frequency match)."""
    m = margin[np.isfinite(margin)]
    base = float(ev.mean())
    if m.size == 0 or base <= 0:
        return 0.0
    return float(np.quantile(m, 1.0 - base))


def train_ai(w, ncore):
    """Huan luyen LGBM Tx-downscaler cho tuan w; tra ve (model, feat_op_builder)."""
    import lightgbm as lgb
    z = np.load(FIELDS % w)
    d0, thr, O, E = z["d0"], z["thr"], z["O"], z["E"]      # (ninit, npt)
    yrs, doy, reg = z["yrs"], z["doy"], z["reg"]
    d1 = z["d1"]
    train = np.isin(yrs, C.TRAIN_YEARS)
    npt = d0.shape[1]
    fine_clim = np.nanmean(np.where(train[:, None], O, np.nan), 0)   # (npt,)
    ptidx = np.arange(npt, dtype="float32")

    def feats(d0_2d, doy_1d):
        # d0_2d: (n, npt) Tx tho; doy_1d: (n,) doy CUA TUNG INIT (broadcast theo diem)
        d0_2d = np.atleast_2d(d0_2d); doy_1d = np.atleast_1d(doy_1d).astype(float)
        n = d0_2d.shape[0]
        cos = np.broadcast_to(np.cos(2 * np.pi * doy_1d / 365.0)[:, None], (n, npt))
        sin = np.broadcast_to(np.sin(2 * np.pi * doy_1d / 365.0)[:, None], (n, npt))
        fc = np.broadcast_to(fine_clim, (n, npt))
        pi = np.broadcast_to(ptidx, (n, npt))
        return np.stack([d0_2d, fc, cos, sin, pi], -1).reshape(-1, 5)

    Ftr = feats(d0[train], doy[train])
    ytr = O[train].ravel()
    ok = np.isfinite(Ftr).all(1) & np.isfinite(ytr)
    gb = lgb.LGBMRegressor(n_estimators=600, learning_rate=0.05, num_leaves=127,
                           min_child_samples=100, subsample=0.8,
                           colsample_bytree=0.8, verbose=-1, n_jobs=ncore)
    gb.fit(Ftr[ok], ytr[ok])

    # d2 tren toan bo init (de fit delta tren train)
    d2 = gb.predict(feats(d0, doy)).reshape(d0.shape)
    return dict(gb=gb, feats=feats, d0=d0, d1=d1, d2=d2, thr=thr, E=E,
                train=train, reg=reg, fine_clim=fine_clim, ptidx=ptidx)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--init", default="2026082900")
    ap.add_argument("--ncore", type=int,
                    default=int(os.environ.get("SLURM_CPUS_PER_TASK", 8)))
    a = ap.parse_args()

    # --- Tx tho cua init du bao (d0_op), nearest cell, weekly max ---
    indir = os.path.join(C.OUT, "RT", "cfs_%s" % a.init)
    step = read_rt(indir, a.init)
    daily = np.stack([np.nanmax(step[4 * k:4 * k + 4], 0)
                      for k in range(C.NLT)]) - 273.15       # (46, ny, nx)
    z = np.load(B.TX_NPZ)
    reg = z["reg"]; olat, olon = z["lat"][z["jj"]], z["lon"][z["ii"]]
    g = C.flux_grid()
    pj = np.abs(olat[:, None] - g["wlat"][None, :]).argmin(1)
    pi = np.abs(olon[:, None] - g["wlon"][None, :]).argmin(1)
    txp = daily[:, pj, pi]                                    # (46, npt)

    # d1_op (BCSD) tu ban tin da chung nhan
    fc0 = json.load(open(os.path.join(C.OUT, "RT", "fcst_%s.json" % a.init)))

    d0 = dt.date(int(a.init[:4]), int(a.init[4:6]), int(a.init[6:8]))
    out = {"init": a.init, "reg_frac": REG_FRAC, "weeks": {}}
    per_week_fields = {}
    print("So sanh BCSD vs AI (Tx-only) — init %s\n" % a.init)
    for w, (wname, idx) in enumerate(C.WEEKS):
        days = [d0 + dt.timedelta(days=int(k)) for k in idx]
        q = days[len(days) // 2].timetuple().tm_yday
        d0_op = np.nanmax(txp[idx], 0)                        # (npt) Tx tho nearest cell

        T = train_ai(w, a.ncore)
        thr_op = np.array(fc0["weeks"][wname]["thr"])
        d1_op = np.array(fc0["weeks"][wname]["tx"])           # BCSD forecast field
        d2_op = T["gb"].predict(T["feats"](d0_op[None, :], np.array([q]))).ravel()

        row = {"doy": int(q), "valid_from": days[0].isoformat(),
               "valid_to": days[-1].isoformat()}
        fields = {"thr": thr_op, "d1": d1_op, "d2": d2_op, "d0": d0_op}
        for tag, fcst_all, op in (("bcsd", T["d1"], d1_op), ("ai", T["d2"], d2_op)):
            m_all = fcst_all[T["train"]] - T["thr"][T["train"]]
            delta = fit_delta(m_all, T["E"][T["train"]].astype(bool))
            # ky nang tren train voi delta nay
            fc_tr = (fcst_all[T["train"]] - T["thr"][T["train"]]) >= delta
            sk = contingency(fc_tr.ravel(),
                             T["E"][T["train"]].astype(bool).ravel())
            # su kien tren init du bao
            ev_op = (op - thr_op) >= delta
            P = {int(r): float(ev_op[reg == r].mean()) for r in np.unique(reg)}
            row[tag] = {"delta": delta, "train_skill": sk,
                        "n_point_event": int(ev_op.sum()),
                        "frac_point": float(ev_op.mean()),
                        "P": P, "EHDE": {str(r): bool(P[r] > REG_FRAC) for r in P},
                        "event": ev_op.tolist()}
        out["weeks"][wname] = row
        per_week_fields[wname] = fields
        eb = np.array(row["bcsd"]["event"]); ea = np.array(row["ai"]["event"])
        print("%s  BCSD %3d o (delta %+.2f)  AI %3d o (delta %+.2f)  trung %3d  "
              "chi BCSD %3d  chi AI %3d"
              % (wname, eb.sum(), row["bcsd"]["delta"], ea.sum(), row["ai"]["delta"],
                 (eb & ea).sum(), (eb & ~ea).sum(), (~eb & ea).sum()))

    outjson = os.path.join(C.OUT, "RT", "cmp_ai_bcsd_%s.json" % a.init)
    json.dump(out, open(outjson, "w"))
    print("\nsaved", outjson)
    build_html(a.init, out, per_week_fields, olat, olon, reg)


def cell(k, v):
    if k in TARGET and np.isfinite(v):
        op, r = TARGET[k]
        ok = v > r if op == ">" else v < r
        return "<td class='%s'>%.3f</td>" % ("good" if ok else "bad", v)
    return "<td>%s</td>" % ("%.3f" % v if np.isfinite(v) else "—")


def fig(fn):
    buf = io.BytesIO(); fn()
    plt.savefig(buf, format="png", dpi=115, bbox_inches="tight"); plt.close()
    return "<figure><img src='data:image/png;base64,%s'></figure>" % \
           base64.b64encode(buf.getvalue()).decode()


def build_html(init, out, F, olat, olon, reg):
    WK = C.WEEK_NAMES
    H = ["<!doctype html><html lang='vi'><meta charset='utf-8'>",
         "<title>29/8 — BCSD vs AI (Tx-only)</title><style>%s</style>" % CSS,
         "<h1>Dự báo EHD 29/8/2026 — BCSD so với AI-regression trên cùng đầu vào Tx</h1>",
         "<p class='sub'>Cùng Tx thô nearest-cell của init %s, cùng ngưỡng, cùng "
         "cách hiệu chuẩn delta (frequency-matching trên năm huấn luyện), chỉ khác "
         "cách hạ quy mô: BCSD (quantile-map) hay AI (LightGBM học ánh xạ). AI ở đây "
         "là Tx-only vì init 29/8 chỉ tải TMAX từ NOMADS — khác D2 đa biến trong "
         "bcsd.json.</p>" % init]

    # 1. train-year skill
    H.append("<h2>1. Kỹ năng trên năm huấn luyện (đã khớp tần suất) — hạ quy mô nào "
             "đặt sự kiện đúng chỗ hơn</h2>")
    H.append("<p class='sub'>Cả hai đã bằng tần suất nền theo cấu trúc (exc_ratio≈1), "
             "nên POD/FAR/CSI phản ánh riêng chất lượng KHÔNG GIAN của trường hạ quy "
             "mô. Xanh/đỏ theo mục tiêu.</p>")
    H.append("<table><tr><th>Tuần</th><th>PP</th><th>delta</th><th>nền</th>"
             "<th>POD</th><th>POFD</th><th>FAR</th><th>FBI</th><th>CSI</th>"
             "<th>exc_ratio</th></tr>")
    for w in WK:
        for tag, nm in (("bcsd", "BCSD"), ("ai", "AI")):
            s = out["weeks"][w][tag]["train_skill"]; dl = out["weeks"][w][tag]["delta"]
            H.append("<tr><td>%s</td><td>%s</td><td>%+.2f</td><td>%.3f</td>"
                     "%s%s%s<td>%.3f</td>%s<td>%.3f</td></tr>"
                     % (w if tag == "bcsd" else "", nm, dl, s["base"],
                        cell("POD", s["POD"]), cell("POFD", s["POFD"]),
                        cell("FAR", s["FAR"]), s["FBI"], cell("CSI", s["CSI"]),
                        s["exc_ratio"]))
    H.append("</table>")

    # 2. forecast 29/8 side by side
    H.append("<h2>2. Bản tin 29/8 — số ô EHD và độ trùng khớp không gian</h2>")
    H.append("<table><tr><th>Tuần</th><th>BCSD số ô</th><th>AI số ô</th>"
             "<th>trùng</th><th>chỉ BCSD</th><th>chỉ AI</th>"
             "<th>vùng EHDE BCSD</th><th>vùng EHDE AI</th></tr>")
    for w in WK:
        r = out["weeks"][w]
        eb = np.array(r["bcsd"]["event"]); ea = np.array(r["ai"]["event"])
        nb = sum(r["bcsd"]["EHDE"].values()); na = sum(r["ai"]["EHDE"].values())
        H.append("<tr><td>%s</td><td>%d</td><td>%d</td><td>%d</td><td>%d</td>"
                 "<td>%d</td><td>%d/7</td><td>%d/7</td></tr>"
                 % (w, eb.sum(), ea.sum(), (eb & ea).sum(), (eb & ~ea).sum(),
                    (~eb & ea).sum(), nb, na))
    H.append("</table>")

    # 3. region P table (W1 only, certified week)
    H.append("<h2>3. Xác suất vùng P — tuần W1 (tuần được chứng nhận)</h2>")
    H.append("<table><tr><th>Vùng</th><th>P BCSD</th><th>EHDE</th>"
             "<th>P AI</th><th>EHDE</th></tr>")
    r1 = out["weeks"]["W1"]
    for rr in sorted(int(k) for k in r1["bcsd"]["P"]):
        pb = r1["bcsd"]["P"][rr] if rr in r1["bcsd"]["P"] else r1["bcsd"]["P"][str(rr)]
        pa = r1["ai"]["P"][str(rr)] if str(rr) in r1["ai"]["P"] else r1["ai"]["P"][rr]
        eb = r1["bcsd"]["EHDE"][str(rr)]; ea = r1["ai"]["EHDE"][str(rr)]
        H.append("<tr><td>%s</td><td>%.1f%%</td><td class='%s'>%s</td>"
                 "<td>%.1f%%</td><td class='%s'>%s</td></tr>"
                 % (REG_NAME[rr], 100 * pb, "good" if eb else "", "EHDE" if eb else "-",
                    100 * pa, "good" if ea else "", "EHDE" if ea else "-"))
    H.append("</table>")

    # 4. maps: event points per week, both methods
    H.append("<h2>4. Bản đồ điểm EHD 29/8 — BCSD (xanh) vs AI (đỏ)</h2>")

    def mapw():
        fig_, ax = plt.subplots(2, 3, figsize=(11, 7.5))
        for i, w in enumerate(WK):
            axi = ax[i // 3][i % 3]
            eb = np.array(out["weeks"][w]["bcsd"]["event"])
            ea = np.array(out["weeks"][w]["ai"]["event"])
            axi.scatter(olon[~(eb | ea)], olat[~(eb | ea)], s=3, c="#ddd")
            axi.scatter(olon[eb & ~ea], olat[eb & ~ea], s=9, c="#1a6", label="chỉ BCSD")
            axi.scatter(olon[ea & ~eb], olat[ea & ~eb], s=9, c="#d33", label="chỉ AI")
            axi.scatter(olon[eb & ea], olat[eb & ea], s=9, c="#333", label="cả hai")
            axi.set_title("%s (BCSD %d / AI %d)" % (w, eb.sum(), ea.sum()), fontsize=9)
            axi.set_xticks([]); axi.set_yticks([])
        ax[0][0].legend(fontsize=7, loc="lower left")
        plt.tight_layout()
    H.append(fig(mapw))

    H.append("<div class='note'><b>Kết luận (đã hiệu đính theo số liệu):</b> khi hiệu "
             "chuẩn delta ĐỐI XỨNG (cả hai exc_ratio=1, FBI=1), kỹ năng không gian "
             "cho thấy điều NGƯỢC với trực giác ở Mục 7: <b>W1 hoà</b> (POD/CSI gần "
             "hệt), còn <b>W2–W6 nhánh AI phân biệt EHD tốt hơn đều</b> (POD +0.05…"
             "+0.08, CSI +0.04…+0.07, POFD thấp hơn). Vẻ 'BCSD trội hẳn ở tầng giá "
             "trị' trong Mục 7 (exc_ratio D2 chỉ 0.56) phần lớn là <b>lỗi HIỆU CHUẨN</b> "
             "do hồi quy co về trung bình — chỉnh lại ngưỡng (delta≈−0.32) thì trường "
             "AI hồi phục và vượt lên ở hạn dài. Điều này nhất quán với việc tầng sự "
             "kiện EXP3 HOÀ (ensemble tự tái hiệu chuẩn).</div>")
    H.append("<div class='note'><b>Hàm ý nghiệp vụ (cân bằng, không cường điệu):</b> "
             "(1) BCSD vẫn ưu ở chỗ trường 0.25° dùng được NGAY, không cần bước tái "
             "hiệu chuẩn ngưỡng riêng (delta≈0), bảo toàn phân bố vật lý, tất định, "
             "dễ giải thích — an toàn cho vận hành. (2) AI-reg cần một offset hiệu "
             "chuẩn (fit trên năm huấn luyện) nhưng sau đó phân biệt nhỉnh hơn ở "
             "W2–W6. (3) Vì tầng ensemble EXP3 hoà, chọn nhánh nào cho khối ensemble "
             "không quyết định; nên giữ BCSD làm trường nền nghiệp vụ và có thể thêm "
             "AI-reg (đã hiệu chuẩn) như một đầu vào bổ trợ ở hạn dài. LƯU Ý: AI ở "
             "đây là Tx-only; D2 đa biến + pgb có thể còn sắc hơn nếu tải đủ dữ liệu "
             "cho init thực.</div>")
    H.append("</html>")
    outhtml = "/home/tanpv/Progs/claude/HW_CFS_raw/ehd_cfs_29aug_bcsd_vs_ai.html"
    open(outhtml, "w").write("\n".join(H))
    print("viet", outhtml)


if __name__ == "__main__":
    main()