#!/usr/bin/env python3
"""Đánh giá MÔ HÌNH CUỐI CÙNG (gói vận hành đã đóng băng) và vẽ bản đồ chỉ số.

Đầu vào là chính đầu ra của src/operational_forecast.py trên 416 lần khởi tạo của
bốn năm kiểm định -- tức là chấm đúng cái sẽ đem đi chạy nghiệp vụ, không phải chấm
một biến thể nghiên cứu nào khác.

Xuất ra:
  * chỉ số cho TỪNG Ô LƯỚI trong 541 ô, từng tuần W1-W6, lưu ra NetCDF;
  * bảng tổng hợp toàn quốc và theo bảy vùng khí hậu;
  * bản đồ POD, POFD, FAR, FBI, CSI, ETS, EDI, SEDI và base rate.

Quan trắc chỉ dùng để CHẤM ĐIỂM, không bao giờ đi vào đặc trưng.
"""

from __future__ import annotations

import json
import os
import 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
import ai_data as D

FCS = os.path.join(C.OUT, "operational", "forecast_valid.nc")
OUT_NC = os.path.join(C.OUT, "final_metrics.nc")
OUT_JSON = os.path.join(C.OUT, "AI_Results", "final_verify.json")
METRICS = ["POD", "POFD", "FAR", "FBI", "CSI", "ETS", "EDI", "SEDI", "base_rate"]


def observations(store, isel, points, chunk=64):
    """Số ngày EHD quan trắc trong từng tuần, (n_init, 6, n_point).

    Chỉ đọc mảng target chứ không dựng cả ma trận đặc trưng: rẻ hơn nhiều lần.
    """
    out = np.empty((isel.size, len(C.WEEKS), points.size), dtype="f4")
    for a in range(0, isel.size, chunk):
        sl = isel[a:a + chunk]
        tgt = store._gather(store.target, sl, points, 3)      # (n,46,3,npt)
        evt = tgt[:, :, 2, :]
        for k, (_, lts) in enumerate(C.WEEKS):
            out[a:a + sl.size, k] = np.nansum(evt[:, lts, :], axis=1)
    return out


def main():
    import argparse
    ap = argparse.ArgumentParser()
    ap.add_argument("--fcs", default=FCS)
    ap.add_argument("--out-nc", default=OUT_NC)
    ap.add_argument("--out-json", default=OUT_JSON)
    args = ap.parse_args()

    store = D.Store()
    grid = store.grid
    points = np.arange(grid.reg.size)
    isel = np.nonzero(np.isin(store.init_year, C.VALID_YEARS))[0]
    print("[fv] %d lần khởi tạo x %d ô lưới" % (isel.size, points.size), flush=True)

    cnt = observations(store, isel, points)
    obs = cnt >= 1.0                                    # yEHDE tuần = >= 1 ngày EHD

    with Dataset(args.fcs) as ds:
        pred = np.stack([np.array(ds["pred_%s" % w][:]) > 0
                         for w in C.WEEK_NAMES], axis=1)   # (n_init, 6, npt)
        prob = np.stack([np.array(ds["prob_%s" % w][:])
                         for w in C.WEEK_NAMES], axis=1)
    assert pred.shape == obs.shape, (pred.shape, obs.shape)

    # ---------------------------------------------------------- từng ô lưới
    fld = {m: np.full((len(C.WEEKS), points.size), np.nan, "f4") for m in METRICS}
    fld["AUC"] = np.full((len(C.WEEKS), points.size), np.nan, "f4")
    for k in range(len(C.WEEKS)):
        for p in range(points.size):
            o, q = obs[:, k, p], pred[:, k, p]
            ok = np.isfinite(o.astype("f4"))
            g = M.scores(*M.contingency(q[ok], o[ok]))
            for m in METRICS:
                fld[m][k, p] = g[m]
            if 0 < o[ok].mean() < 1:
                fld["AUC"][k, p] = M.auc_score(prob[ok, k, p], o[ok])
        print("  %s xong" % C.WEEK_NAMES[k], flush=True)

    with Dataset(args.out_nc, "w") as ds:
        ds.createDimension("week", len(C.WEEKS))
        ds.createDimension("lat", grid.nlat)
        ds.createDimension("lon", grid.nlon)
        ds.createVariable("lat", "f8", ("lat",))[:] = grid.lat
        ds.createVariable("lon", "f8", ("lon",))[:] = grid.lon
        ds.weeks = " ".join(C.WEEK_NAMES)
        ds.source = ("Đánh giá mô hình CUỐI CÙNG (gói vận hành đã đóng băng) trên "
                     "416 lần khởi tạo của 4 năm kiểm định, 541 ô lưới đất.")
        for m, arr in fld.items():
            v = ds.createVariable(m, "f4", ("week", "lat", "lon"), zlib=True,
                                  fill_value=np.nan)
            v[:] = grid.to_grid(arr)
    print("[fv] -> %s" % args.out_nc, flush=True)

    # ------------------------------------------------- tổng hợp và theo vùng
    res = {"national": {}, "by_region": {}, "n_init": int(isel.size)}
    for k, w in enumerate(C.WEEK_NAMES):
        g = M.scores(*M.contingency(pred[:, k, :].ravel(), obs[:, k, :].ravel()))
        g["AUC"] = float(np.nanmean(fld["AUC"][k]))
        res["national"][w] = g
        res["by_region"][w] = {}
        for rg in range(1, 8):
            sel = grid.reg == rg
            gg = M.scores(*M.contingency(pred[:, k, sel].ravel(),
                                         obs[:, k, sel].ravel()))
            gg["AUC"] = float(np.nanmean(fld["AUC"][k, sel]))
            res["by_region"][w]["S%d" % rg] = gg
        print("[%s] POD %.3f POFD %.3f FAR %.3f FBI %.3f CSI %.3f ETS %.3f "
              "EDI %.3f SEDI %.3f AUC %.3f"
              % (w, g["POD"], g["POFD"], g["FAR"], g["FBI"], g["CSI"], g["ETS"],
                 g["EDI"], g["SEDI"], g["AUC"]), flush=True)
    with open(args.out_json, "w") as fh:
        json.dump(res, fh, indent=1, default=float)
    print("[fv] -> %s" % args.out_json, flush=True)


if __name__ == "__main__":
    main()