#!/usr/bin/env python3
"""GĐ2: chấm điểm hệ thống trên TOÀN MIỀN 541 ô lưới.

Ngưỡng được chốt theo tỷ lệ phát báo r, ước lượng trên trường NGOÀI MẪU của 16 năm
huấn luyện, riêng cho từng ô lưới, và KÈM BIÊN AN TOÀN.  GĐ1 (AI27) cho thấy r tối
ưu theo CSI luôn nằm sát biên chỉ tiêu nên không chuyển được sang ô lưới chưa thấy;
chốt kèm biên đổi khoảng 0.02-0.03 CSI lấy độ bền.

Xuất ra: chỉ số toàn quốc, chỉ số theo từng vùng S1-S7, và trường chỉ số theo từng
ô lưới để vẽ bản đồ.
"""

from __future__ import annotations

import json
import os
import sys

import numpy as np

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

RES = os.path.join(C.OUT, "AI_Results", "phase1")
OUT = os.path.join(C.OUT, "AI_Results", "phase2_summary.json")
WINDOWS = ["W1", "W1-2", "W1-3", "W1-6"]
RGRID = np.round(np.arange(0.50, 1.31, 0.01), 3)
MARGIN = 0.02


def preds(z, pts, r):
    """Dự báo nhị phân cho từng ô lưới ở mức phát báo r (bội của base rate)."""
    out = {}
    for p in pts:
        k = str(int(p))
        if k + "|s_val" not in z:
            continue
        so, oo = z[k + "|s_oof"], z[k + "|obs_oof"]
        sv, ov = z[k + "|s_val"], z[k + "|obs_val"]
        rate = float(np.clip(r * oo.mean(), 1e-4, 0.999))
        th = float(np.quantile(so, 1.0 - rate))
        out[int(p)] = (sv > th, ov, so > th, oo)
    return out


def pooled(d, sel=None):
    pv = np.concatenate([v[0] for p, v in d.items() if sel is None or p in sel])
    ov = np.concatenate([v[1] for p, v in d.items() if sel is None or p in sel])
    return M.scores(*M.contingency(pv, ov))


def pooled_oof(d):
    pv = np.concatenate([v[2] for v in d.values()])
    ov = np.concatenate([v[3] for v in d.values()])
    return M.scores(*M.contingency(pv, ov))


def ok(g, m=0.0):
    return (g["POD"] > 0.60 + m and g["POFD"] < 0.20 - m
            and g["FAR"] < 0.20 - m and g["CSI"] > 0.40)


def main():
    grid = D.Store().grid
    pts = np.arange(grid.reg.size)
    res = {"config": {"windows": WINDOWS, "margin": MARGIN, "n_point": int(pts.size)},
           "chosen": {}, "national": {}, "by_region": {}, "per_point": {}}

    combos = [("national", "so1", ""), ("region", "so1", ""),
              ("national", "geq_k", "_geq_k")]
    for variant, target, suf in combos:
        for w in WINDOWS:
            f = os.path.join(RES, "full_%s_%s%s.npz" % (variant, w, suf))
            if not os.path.exists(f):
                print("[bỏ qua] chưa có %s" % os.path.basename(f), flush=True)
                continue
            z = np.load(f, allow_pickle=True)
            # --- chốt r TRÊN TRƯỜNG NGOÀI MẪU, kèm biên; không bao giờ trên kiểm định
            best = None
            for r in RGRID:
                d = preds(z, pts, r)
                g = pooled_oof(d)
                if ok(g, MARGIN) and (best is None or g["CSI"] > best[1]["CSI"]):
                    best = (r, g)
            if best is None:                       # không có r nào đạt kèm biên
                for r in RGRID:
                    d = preds(z, pts, r)
                    g = pooled_oof(d)
                    if ok(g) and (best is None or g["CSI"] > best[1]["CSI"]):
                        best = (r, g)
            if best is None:
                cand = [(r, pooled_oof(preds(z, pts, r))) for r in RGRID]
                cand = [c for c in cand if c[1]["POD"] > 0.60] or cand
                best = min(cand, key=lambda c: c[1]["FAR"])
            r = float(best[0])

            d = preds(z, pts, r)
            key = "%s|%s" % (variant if target == "so1"
                                 else variant + "-" + target, w)
            gnat = pooled(d)
            res["chosen"][key] = {"r": r, "oof": best[1]}
            res["national"][key] = gnat
            res["by_region"][key] = {}
            for rg in range(1, 8):
                sel = set(int(p) for p in pts[grid.reg == rg])
                res["by_region"][key]["S%d" % rg] = pooled(d, sel)
            pp = {m: [] for m in ("POD", "FAR", "CSI", "SEDI", "FBI", "base_rate")}
            for p in pts:
                v = d.get(int(p))
                g = (M.scores(*M.contingency(v[0], v[1])) if v is not None
                     else {m: np.nan for m in pp})
                for m in pp:
                    pp[m].append(float(g[m]) if g[m] == g[m] else None)
            res["per_point"][key] = pp
            print("[%s %s %s] r=%.2f | toàn quốc POD %.3f POFD %.3f FAR %.3f "
                  "FBI %.3f CSI %.3f SEDI %.3f %s"
                  % (variant, target, w, r, gnat["POD"], gnat["POFD"], gnat["FAR"],
                     gnat["FBI"], gnat["CSI"], gnat["SEDI"],
                     "ĐẠT" if ok(gnat) else "hụt"), flush=True)
            for rg in range(1, 8):
                g = res["by_region"][key]["S%d" % rg]
                print("      S%d %-22s POD %.3f POFD %.3f FAR %.3f CSI %.3f "
                      "SEDI %.3f %s"
                      % (rg, C.REGION_NAMES.get(rg, "")[3:], g["POD"], g["POFD"],
                         g["FAR"], g["CSI"], g["SEDI"], "ĐẠT" if ok(g) else "hụt"),
                      flush=True)

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


if __name__ == "__main__":
    main()