#!/usr/bin/env python3
"""Compare the 29/08/2026 EHDE forecast of EC_0.25_SD vs EC_0.25_org (../HW),
both through the SAME Ope1b system, and emit the 21-point contingency tables
(3 points x 7 climate regions) EXP3 asks for.

Sections:
  A. 29/08/2026 forecast: per-week pred count / mean prob / grid agreement,
     plus the EXP1 regional coverage table (% of grid points flagged EHDE).
  B. Reference hindcast skill W1-W6 (416 valid init) from the two final_verify.
  C. 21 contingency tables: per representative point, aggregated over W1-W6 and
     all 416 valid init (the "full sample" of both nEHDE and yEHDE per EXP3),
     SD vs org side by side.

Output: ehd_forecast_sd_vs_org.html
"""
from __future__ import annotations
import json
import os
import sys

import numpy as np
from netCDF4 import Dataset

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
import ehd_common as C
import ehd_metrics as M

ROOT = C.ROOT
SD_FCS = "outputs_sd/rt/forecast_20260829_sd.nc"
ORG_FCS = "/home/tanpv/Progs/claude/HW/outputs/operational/forecast_20260829.nc"
SD_VALID = "outputs_sd/operational_ope1b/forecast_valid.nc"
ORG_VALID = "/home/tanpv/Progs/claude/HW/outputs/operational/forecast_valid.nc"
SD_VERIFY = "outputs_sd/final_verify_sd.json"
ORG_VERIFY = "/home/tanpv/Progs/claude/HW/outputs/AI_Results/final_verify_ope1b.json"
OUT = "ehd_forecast_sd_vs_org.html"

REG_NAME = {1: "S1 Tây Bắc", 2: "S2 Đông Bắc", 3: "S3 ĐB Bắc Bộ", 4: "S4 Bắc Trung Bộ",
            5: "S5 Nam Trung Bộ", 6: "S6 Tây Nguyên", 7: "S7 Nam Bộ"}
WK = ["W1", "W2", "W3", "W4", "W5", "W6"]


def load_pred(path):
    with Dataset(path) as ds:
        return {w: np.asarray(ds["pred_%s" % w][:]).astype(bool) for w in WK}, \
               {w: np.asarray(ds["prob_%s" % w][:]) for w in WK}


def obs_ehde_valid(grid):
    """(416, 6, 541) boolean observed EHDE per valid init / week / point."""
    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, grid.npt)
    iv = np.nonzero(np.isin(yr, C.VALID_YEARS))[0]
    E = np.asarray(T[iv][:, :, 2, :])                        # (416,46,541)
    wk = dict(C.WEEKS)
    out = np.zeros((iv.size, len(WK), grid.npt), bool)
    for j, w in enumerate(WK):
        lts = np.asarray(wk[w])
        out[:, j, :] = np.nansum(E[:, lts, :], axis=1) >= 1.0
    return out


def pick_points(grid, per_region=3):
    """3 spread-out land points per region -> list of (point_index, region)."""
    pts = []
    for r in range(1, 8):
        idx = np.nonzero(grid.reg == r)[0]
        # order by longitude then latitude for a spatial spread, take 3 quantiles
        order = idx[np.lexsort((grid.lat[grid.iy][idx], grid.lon[grid.ix][idx]))]
        sel = [order[int(q * (len(order) - 1))] for q in (0.2, 0.5, 0.8)]
        pts += [(int(p), r) for p in sel]
    return pts


def fmt(x, nd=3):
    return "%.*f" % (nd, x) if np.isfinite(x) else "–"


def main():
    grid = C.Grid()
    lat = grid.lat[grid.iy]; lon = grid.lon[grid.ix]
    sd_pred, sd_prob = load_pred(SD_FCS)
    org_pred, org_prob = load_pred(ORG_FCS)
    obs = obs_ehde_valid(grid)                               # (416,6,541)
    sd_vpred, _ = load_pred(SD_VALID)
    org_vpred, _ = load_pred(ORG_VALID)
    sdv = json.load(open(SD_VERIFY)); orgv = json.load(open(ORG_VERIFY))

    css = ("<style>body{font-family:system-ui,Arial;margin:24px;color:#1a1a1a}"
           "table{border-collapse:collapse;margin:10px 0;font-size:13px}"
           "th,td{border:1px solid #ccc;padding:3px 7px;text-align:center}"
           "th{background:#f0efe9}.sd{background:#eef5ff}.wk{font-weight:700;background:#faf8f2}"
           "h1{font-size:21px}h2{font-size:16px;margin-top:24px}h3{font-size:14px}"
           ".grid21{display:grid;grid-template-columns:repeat(3,1fr);gap:10px}"
           ".pos{color:#127a2e}.neg{color:#b21f2d}.small{color:#555;font-size:12px}</style>")
    h = ["<!doctype html><meta charset=utf-8><title>EHDE 29/08/2026 — EC_0.25_SD vs EC_0.25_org</title>",
         css,
         "<h1>Dự báo EHDE khởi tạo 29/08/2026 — EC_0.25_SD so với EC_0.25_org</h1>",
         "<p class=small>Cùng hệ <b>Ope1b</b> (hệ này là hiện thực nghiệp vụ của EXP3: "
         "tầng 1 cổng p=M/N trên 72 thành viên, tầng 2 phân loại + ngưỡng riêng từng ô). "
         "Khác biệt DUY NHẤT giữa hai cột là tập số liệu EC đầu vào: 0.25° tải thẳng "
         "(EC_0.25_org) hay downscale BCSD từ 1.5° (EC_0.25_SD).</p>"]

    # ---- A. 29/08/2026 forecast -------------------------------------------- #
    h.append("<h2>A. Bản tin 29/08/2026: số ô báo EHDE và độ đồng thuận</h2>")
    h.append("<table><tr><th>Tuần</th><th>ô EHDE org</th><th>ô EHDE SD</th>"
             "<th>P̄ org</th><th>P̄ SD</th><th>đồng thuận ô lưới</th></tr>")
    for w in WK:
        po, ps = org_pred[w].ravel(), sd_pred[w].ravel()
        agr = 100.0 * (po == ps).mean()
        h.append("<tr><td class=wk>%s</td><td>%d</td><td class=sd>%d</td>"
                 "<td>%s</td><td class=sd>%s</td><td>%.1f%%</td></tr>"
                 % (w, po.sum(), ps.sum(), fmt(np.nanmean(org_prob[w])),
                    fmt(np.nanmean(sd_prob[w])), agr))
    h.append("</table>")

    # EXP1 regional coverage (% points flagged) for the 29/08 forecast
    h.append("<h3>Phủ vùng (EXP1): %% ô lưới báo EHDE — org / SD</h3>")
    h.append("<table><tr><th>Vùng</th><th>n ô</th>"
             + "".join("<th>%s</th>" % w for w in WK) + "</tr>")
    for r in range(1, 8):
        m = grid.reg == r
        cells = ["<td class=wk>%s</td><td>%d</td>" % (REG_NAME[r], m.sum())]
        for w in WK:
            co = 100.0 * org_pred[w].ravel()[m].mean()
            cs = 100.0 * sd_pred[w].ravel()[m].mean()
            cells.append("<td>%.0f / <span class=sd>%.0f</span></td>" % (co, cs))
        h.append("<tr>" + "".join(cells) + "</tr>")
    h.append("</table><p class=small>Ngưỡng vùng: P>20%% ⇒ vùng có EHDE.</p>")

    # ---- B. reference hindcast skill --------------------------------------- #
    h.append("<h2>B. Kỹ năng tham chiếu (416 init kiểm định, hindcast)</h2>")
    h.append("<table><tr><th>Tuần</th>"
             "<th>POD org/SD</th><th>FAR org/SD</th><th>FBI org/SD</th>"
             "<th>CSI org/SD</th><th>SEDI org/SD</th></tr>")
    for w in WK:
        o, s = orgv["national"][w], sdv["national"][w]
        h.append("<tr><td class=wk>%s</td>" % w
                 + "".join("<td>%s / <span class=sd>%s</span></td>"
                           % (fmt(o[k]), fmt(s[k])) for k in
                           ("POD", "FAR", "FBI", "CSI", "SEDI")) + "</tr>")
    h.append("</table>")

    # ---- C. 21 contingency tables ------------------------------------------ #
    h.append("<h2>C. 21 bảng contingency (3 điểm × 7 vùng), gộp W1–W6 × 416 init</h2>")
    h.append("<p class=small>Mỗi bảng: dự báo nhị phân của điểm đó trên toàn bộ mẫu "
             "kiểm định (cả nEHDE và yEHDE). H=hit, F=báo động sai, M=bỏ sót, "
             "C=đúng-không. So EC_0.25_SD (nền xanh) với EC_0.25_org.</p>")
    pts = pick_points(grid)
    h.append("<div class=grid21>")
    for p, r in pts:
        po = org_vpred; ps = sd_vpred
        # aggregate over all weeks
        ao = np.concatenate([po[w][:, p] for w in WK])
        as_ = np.concatenate([ps[w][:, p] for w in WK])
        yo = np.concatenate([obs[:, j, p] for j in range(len(WK))])
        go = M.scores(*M.contingency(ao, yo))
        gs = M.scores(*M.contingency(as_, yo))
        ho, fo, mo, co = M.contingency(ao, yo)
        hs, fs, ms, cs = M.contingency(as_, yo)
        h.append(
            "<div><h3>%s — điểm #%d (%.2f°N, %.2f°E)</h3>" % (REG_NAME[r], p, lat[p], lon[p])
            + "<table><tr><th></th><th>obs=1</th><th>obs=0</th></tr>"
            + "<tr><td class=wk>báo=1</td><td>%d/<span class=sd>%d</span></td>"
              "<td>%d/<span class=sd>%d</span></td></tr>" % (ho, hs, fo, fs)
            + "<tr><td class=wk>báo=0</td><td>%d/<span class=sd>%d</span></td>"
              "<td>%d/<span class=sd>%d</span></td></tr>" % (mo, ms, co, cs)
            + "</table>"
            + "<table><tr><th>POD</th><th>FAR</th><th>FBI</th><th>CSI</th></tr>"
            + "<tr><td>%s/<span class=sd>%s</span></td><td>%s/<span class=sd>%s</span></td>"
              "<td>%s/<span class=sd>%s</span></td><td>%s/<span class=sd>%s</span></td></tr>"
              % (fmt(go["POD"], 2), fmt(gs["POD"], 2), fmt(go["FAR"], 2), fmt(gs["FAR"], 2),
                 fmt(go["FBI"], 2), fmt(gs["FBI"], 2), fmt(go["CSI"], 2), fmt(gs["CSI"], 2))
            + "</table></div>")
    h.append("</div>")

    h.append("<p class=small>Bản tin đầy đủ EC_0.25_SD (bản đồ W1–W6, trường xác suất, "
             "bảng vùng): xem <a href='ehd_forecast_ope1b_sd.html'>ehd_forecast_ope1b_sd.html</a>. "
             "Bản tin gốc EC_0.25_org: hpc.meteo.edu.vn/~tanpv/For_Me/claude/HW/ehd_forecast_ope1b.html</p>")

    open(OUT, "w").write("\n".join(h))
    print("[report] ->", OUT)


if __name__ == "__main__":
    main()