#!/usr/bin/env python3
"""Chấm GĐ1: so ba kiến trúc, và kiểm tra tính khái quát không gian.

Hai câu hỏi tách bạch:
  (1) KIẾN TRÚC -- per-point, gộp vùng hay gộp toàn quốc cho kết quả tốt nhất trên
      cùng tập điểm đánh giá và cùng danh mục họ mô hình?
  (2) KHÁI QUÁT KHÔNG GIAN -- điểm vận hành tìm được trên 21 điểm seed 0 có chuyển
      sang 21 điểm seed 1 hay không?  Mọi kết luận từ AI18 tới AI25 đều dựa trên
      seed 0, nên nếu không chuyển được thì không triển khai toàn miền được.
"""

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

RES = os.path.join(C.OUT, "AI_Results", "phase1")
OUT = os.path.join(C.OUT, "AI_Results", "phase1_summary.json")
RGRID = np.round(np.arange(0.50, 1.31, 0.01), 3)
VARIANTS = ["point", "region", "national"]
WINDOWS = ["W1", "W1-2", "W1-3", "W1-6"]


def meets(g):
    return (g["POD"] > 0.60 and g["POFD"] < 0.20 and g["FAR"] < 0.20
            and g["CSI"] > 0.40)


def evaluate(z, pts, r):
    pv, ov = [], []
    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, o = z[k + "|s_val"], z[k + "|obs_val"]
        rate = float(np.clip(r * oo.mean(), 1e-4, 0.999))
        pv.append(sv > float(np.quantile(so, 1.0 - rate)))
        ov.append(o)
    if not pv:
        return None
    return M.scores(*M.contingency(np.concatenate(pv), np.concatenate(ov)))


def best_point(z, pts):
    """Điểm vận hành đạt chỉ tiêu có CSI cao nhất; nếu không có thì lấy FAR nhỏ nhất
    trong các cấu hình còn POD > 0.6."""
    rows = [(r, evaluate(z, pts, r)) for r in RGRID]
    rows = [(r, g) for r, g in rows if g is not None]
    ok = [(r, g) for r, g in rows if meets(g)]
    if ok:
        r, g = max(ok, key=lambda x: x[1]["CSI"])
        return r, g, True
    cand = [(r, g) for r, g in rows if g["POD"] > 0.60]
    if cand:
        r, g = min(cand, key=lambda x: x[1]["FAR"])
        return r, g, False
    r, g = max(rows, key=lambda x: x[1]["CSI"])
    return r, g, False


def main():
    res = {"by_arch": {}, "generalisation": {}}
    print("=== (1) SO KIẾN TRÚC (đánh giá trên cả 42 điểm) ===")
    print("Cửa sổ  kiến trúc    r     POD    POFD   FAR    FBI    CSI    SEDI   ĐẠT")
    for w in WINDOWS:
        for v in VARIANTS:
            f = os.path.join(RES, "%s_%s.npz" % (v, w))
            if not os.path.exists(f):
                print("%-7s %-11s (chưa có kết quả)" % (w, v)); continue
            z = np.load(f, allow_pickle=True)
            pts = np.unique(np.concatenate([z["eval_seed0"], z["eval_seed1"]]))
            r, g, ok = best_point(z, pts)
            g["r"] = float(r); g["feasible"] = bool(ok)
            res["by_arch"]["%s|%s" % (w, v)] = g
            print("%-7s %-11s %.2f  %.3f  %.3f  %.3f  %.3f  %.3f  %.3f  %s"
                  % (w, v, r, g["POD"], g["POFD"], g["FAR"], g["FBI"], g["CSI"],
                     g["SEDI"], "ĐẠT" if ok else "hụt"))

    print()
    print("=== (2) KHÁI QUÁT KHÔNG GIAN: chốt r trên seed 0, áp sang seed 1 ===")
    print("Cửa sổ  kiến trúc    r(seed0)  CSI seed0  CSI seed1   SEDI s0  SEDI s1  ĐẠT s1")
    for w in WINDOWS:
        for v in VARIANTS:
            f = os.path.join(RES, "%s_%s.npz" % (v, w))
            if not os.path.exists(f):
                continue
            z = np.load(f, allow_pickle=True)
            p0, p1 = z["eval_seed0"], z["eval_seed1"]
            p1 = np.setdiff1d(p1, p0)          # chỉ các điểm THỰC SỰ mới
            r, g0, _ = best_point(z, p0)
            g1 = evaluate(z, p1, r)
            if g1 is None:
                continue
            res["generalisation"]["%s|%s" % (w, v)] = {
                "r": float(r), "seed0": g0, "seed1": g1,
                "seed1_feasible": bool(meets(g1)), "n_new_points": int(p1.size)}
            print("%-7s %-11s %.2f      %.3f      %.3f       %.3f    %.3f    %s"
                  % (w, v, r, g0["CSI"], g1["CSI"], g0["SEDI"], g1["SEDI"],
                     "ĐẠT" if meets(g1) else "hụt"))

    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()