#!/usr/bin/env python3
"""Hiệu chỉnh ĐỘ TÁN MỨC TOÀN QUỐC cho hệ EHDE đã đóng băng.

Chẩn đoán 2026-08-31: mô hình vẫn có kỹ năng phân dị không gian (AUC trong-init
0.798 ở W1) nhưng QUÁ TÁN ở thành phần toàn quốc — phát "cả nước cùng bật" gấp
~25 lần tự nhiên ở W4-W6.  Đây là hậu xử lý, KHÔNG huấn luyện lại.

Ý tưởng tách đôi bài toán:
  * "những ô nào"  -> u[i,p] = phân vị của điểm số init i tại ô p, so với chính
    phân bố của ô p.  Giữ nguyên thứ hạng không gian, tức giữ nguyên kỹ năng.
  * "bao nhiêu ô"  -> F[i] = ánh xạ phân vị từ mức hoạt động toàn quốc của init i
    sang PHÂN BỐ KHÍ HẬU của diện tích quan trắc.  Chỗ này chữa quá tán.
Rồi bật đúng round(F[i]*npt) ô có u lớn nhất trong init i.

Hiệu chỉnh và chấm điểm tách năm (2 năm hiệu chỉnh / 2 năm chấm, đổi vai), nên
mọi con số in ra đều NGOÀI MẪU.
"""
from __future__ import annotations
import os, sys
import numpy as np
from netCDF4 import Dataset
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import ehd_common as C
import ehd_metrics as M

FCS = os.path.join(C.OUT, "operational", "forecast_valid.nc")
FOLDS = [([2006, 2016], [2011, 2021]), ([2011, 2021], [2006, 2016])]


def point_percentile(P, ref):
    """u[i,p] = phân vị của P[i,p] trong phân bố tham chiếu ref[:,p] của ô p."""
    u = np.empty_like(P)
    for p in range(P.shape[1]):
        u[:, p] = np.searchsorted(np.sort(ref[:, p]), P[:, p], "left") / ref.shape[0]
    return u


def qmap(x_cal, y_cal, x_new):
    """Ánh xạ phân vị: đưa x_new sang thang phân bố của y_cal, hiệu chuẩn trên x_cal."""
    r = np.searchsorted(np.sort(x_cal), x_new, "left") / max(len(x_cal), 1)
    return np.quantile(y_cal, np.clip(r, 0.0, 1.0))


def main():
    st = np.load(os.path.join(C.OUT, "AI_Inp", "static.npz"), allow_pickle=True)
    yr = np.asarray(st["init_year"])
    T = np.memmap(os.path.join(C.OUT, "AI_Inp", "target.f4"), "f4", "r"
                  ).reshape(2080, C.NLT, 3, 541)          # BIEN truoc DIEM
    iv = np.nonzero(np.isin(yr, C.VALID_YEARS))[0]
    E = np.asarray(T[iv][:, :, 2, :])
    yv = yr[iv]
    ds = Dataset(FCS)
    npt = 541

    print("%-5s %-9s %6s %6s %6s %6s %6s %6s" %
          ("Cuaso", "phuong an", "POD", "POFD", "FAR", "FBI", "CSI", "SEDI"))
    for w, lts in C.WEEKS:
        P = np.asarray(ds["prob_%s" % w][:])
        base = np.asarray(ds["pred_%s" % w][:]).astype(bool)
        Y = (np.nansum(E[:, lts, :], axis=1) >= 1.0)
        new = np.zeros_like(base)
        for cal_y, tst_y in FOLDS:
            c = np.isin(yv, cal_y)
            t = np.isin(yv, tst_y)
            u_all = point_percentile(P, P[c])          # tham chieu = nam hieu chinh
            act_c = u_all[c].mean(axis=1)              # muc hoat dong toan quoc
            act_t = u_all[t].mean(axis=1)
            covc = Y[c].mean(axis=1)                   # dien tich quan trac (khi hau)
            F = qmap(act_c, covc, act_t)
            k = np.rint(F * npt).astype(int).clip(0, npt)
            idx = np.nonzero(t)[0]
            for j, i in enumerate(idx):
                if k[j] > 0:
                    sel = np.argsort(-u_all[i])[:k[j]]
                    new[i, sel] = True
        for nm, Q in (("goc", base), ("hieuchinh", new)):
            g = M.scores(*M.contingency(Q.ravel(), Y.ravel()))
            print("%-5s %-9s %6.3f %6.3f %6.3f %6.3f %6.3f %6.3f" %
                  (w, nm, g["POD"], g["POFD"], g["FAR"], g["FBI"], g["CSI"], g["SEDI"]))
        fo, fb, fn = Y.mean(1), base.mean(1), new.mean(1)
        print("      do tan: quan trac >95%%=%.1f%% | goc %.1f%% | hieuchinh %.1f%%"
              % (100*(fo > .95).mean(), 100*(fb > .95).mean(), 100*(fn > .95).mean()))
    ds.close()


if __name__ == "__main__":
    main()