#!/usr/bin/env python3
"""Quét ngưỡng cho TỪNG TUẦN W1..W6 và tìm điểm vận hành tốt nhất có thể.

Không có một điểm vận hành "tốt nhất" duy nhất khi chưa nói tốt theo nghĩa gì, nên
script báo cáo bốn chế độ, tất cả đều chốt ngưỡng trên trường NGOÀI MẪU của 16 năm
huấn luyện rồi áp sang tập kiểm định:

  ĐỦ    -- điểm đạt trọn bộ chỉ tiêu (POD>0.6, POFD<0.2, FAR<0.2, CSI>0.4), lấy
           cái có CSI cao nhất.  Nếu không có thì bỏ trống.
  CSI   -- CSI cao nhất, bất kể chỉ tiêu khác.
  FAR   -- CSI cao nhất trong số các điểm giữ được FAR < 0.2 (ưu tiên độ chính xác:
           mỗi lần phát báo phải đáng tin).
  POD   -- CSI cao nhất trong số các điểm giữ được POD > 0.6 (ưu tiên độ bắt: không
           bỏ sót đợt nóng).
  SEDI  -- SEDI cao nhất, tức điểm vận hành có kỹ năng thực cao nhất khi đã gạt bỏ
           ảnh hưởng của base rate.

Nhắc lại ràng buộc đã chứng minh ở AI24: ở FBI = 1 thì FAR = 1 - POD, nên ĐỦ và FAR
luôn đòi FBI < 1, tức phải dự báo thiếu.  Lưới r ở đây kéo xuống tận 0.30 chính là
để cho phép điều đó.
"""

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", "week_tune.json")
WEEKS = ["W1", "W2", "W3", "W4", "W5", "W6"]
TARGETS = [("so1", ""), ("geq_k", "_geq_k")]
RGRID = np.round(np.arange(0.30, 1.51, 0.01), 3)


def curve(z, pts):
    """Toàn bộ đường cong chỉ số theo r, ngưỡng chốt ngoài mẫu, chấm trên kiểm định."""
    keys = [str(int(p)) for p in pts if str(int(p)) + "|s_val" in z]
    so = [z[k + "|s_oof"] for k in keys]
    oo = [z[k + "|obs_oof"] for k in keys]
    sv = [z[k + "|s_val"] for k in keys]
    ov = np.concatenate([z[k + "|obs_val"] for k in keys])
    br = [float(o.mean()) for o in oo]
    rows = []
    for r in RGRID:
        pv = []
        for s_o, s_v, b in zip(so, sv, br):
            rate = float(np.clip(r * b, 1e-4, 0.999))
            pv.append(s_v > float(np.quantile(s_o, 1.0 - rate)))
        g = M.scores(*M.contingency(np.concatenate(pv), ov))
        g["r"] = float(r)
        rows.append(g)
    return rows


def pick(rows):
    full = [g for g in rows if g["POD"] > 0.60 and g["POFD"] < 0.20
            and g["FAR"] < 0.20 and g["CSI"] > 0.40]
    far = [g for g in rows if g["FAR"] < 0.20]
    pod = [g for g in rows if g["POD"] > 0.60]
    best = lambda L, k="CSI": max(L, key=lambda g: g[k]) if L else None
    return {"ĐỦ": best(full), "CSI": best(rows), "FAR": best(far),
            "POD": best(pod), "SEDI": best(rows, "SEDI")}


def fmt(g):
    if g is None:
        return "%-46s" % "        (không có điểm nào)"
    return ("r=%.2f POD %.3f POFD %.3f FAR %.3f FBI %.3f CSI %.3f SEDI %.3f"
            % (g["r"], g["POD"], g["POFD"], g["FAR"], g["FBI"], g["CSI"],
               g["SEDI"]))


def main():
    grid = D.Store().grid
    pts = np.arange(grid.reg.size)
    res = {}
    for tname, suf in TARGETS:
        head = ("≥ 1 ngày EHD" if tname == "so1"
                else "≥ k ngày EHD (k riêng từng ô, ~1/3 tần suất)")
        print("\n=== TARGET: %s ===" % head)
        for w in WEEKS:
            f = os.path.join(RES, "full_national_%s%s.npz" % (w, suf))
            if not os.path.exists(f):
                print("%-4s (chưa có kết quả)" % w); continue
            z = np.load(f, allow_pickle=True)
            rows = curve(z, pts)
            got = pick(rows)
            res["%s|%s" % (tname, w)] = {"curve": rows,
                                         "picks": {k: v for k, v in got.items()}}
            base = rows[0]["base_rate"]
            print("%-4s base %.3f" % (w, base))
            for mode in ("ĐỦ", "CSI", "FAR", "POD", "SEDI"):
                print("     %-5s %s" % (mode, fmt(got[mode])))
    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()