#!/usr/bin/env python3
"""Kết xuất BẢN TIN DỰ BÁO NGHIỆP VỤ EHDE thành một file HTML độc lập.

Bám đúng quy ước vẽ của make_maps_html.py: 541 ô lưới đất của
Mask_VN_S2S_EC_0.25.nc, ranh giới bảy vùng từ Vietnam7.shp, hình nhúng base64.

Quy tắc phát hành theo RUNBOOK mục 1 — và đây là chỗ báo cáo này khác một bản vẽ
thuần tuý:
  * W1 phát hành BẢN ĐỒ NHỊ PHÂN (r=0.84, đạt trọn bộ chỉ tiêu);
  * W2-W6 chỉ phát hành TRƯỜNG XÁC SUẤT — không có ngưỡng nào trong dải
    r in [0.02, 1.50] đạt FAR < 0.2, nên trường nhị phân ở các hạn này KHÔNG
    được trình bày như một sản phẩm.
Mỗi cửa sổ đều kèm một dòng kiểm tra độ hợp lý: diện tích dự báo đặt cạnh phân vị
95 của diện tích quan trắc lịch sử.  Vượt p95 thì đánh dấu cảnh báo.
"""
from __future__ import annotations

import argparse
import base64
import datetime as dt
import io
import json
import os
import sys
from html import escape

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import shapefile
from matplotlib.colors import BoundaryNorm, ListedColormap
from netCDF4 import Dataset

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

REG_ORDER = [(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ộ")]
P_REGION = 0.20          # EXP1: vùng có EHDE nếu > 20% ô lưới


def load_regions(path=C.SHAPE_7REG):
    r = shapefile.Reader(path)
    out = []
    for shp in r.shapes():
        stops = list(shp.parts) + [len(shp.points)]
        for i in range(len(stops) - 1):
            out.append(np.asarray(shp.points[stops[i]:stops[i + 1]]))
    return out


def _panel(ax, grid, borders, f2d, cmap, norm):
    ax.set_facecolor("#f2f2f2")
    m = ax.pcolormesh(grid.lon, grid.lat, np.ma.masked_invalid(f2d),
                      cmap=cmap, norm=norm, shading="nearest")
    for poly in borders:
        ax.plot(poly[:, 0], poly[:, 1], color="0.25", lw=0.45)
    ax.set_xlim(grid.lon.min() - .125, grid.lon.max() + .125)
    ax.set_ylim(grid.lat.min() - .125, grid.lat.max() + .125)
    ax.set_aspect("equal"); ax.set_xticks([]); ax.set_yticks([])
    for s in ax.spines.values():
        s.set_linewidth(.5); s.set_color("0.6")
    return m


def _png(fig):
    buf = io.BytesIO()
    fig.savefig(buf, format="png", dpi=125, facecolor="white")
    plt.close(fig)
    return base64.b64encode(buf.getvalue()).decode("ascii")


def prob_row(fields, grid, borders, labels, title):
    """Một hàng bản đồ xác suất, chung một thang màu 0-1."""
    lev = np.linspace(0, 1, 11)
    norm = BoundaryNorm(lev, ncolors=256, clip=False)
    n = len(fields)
    fig, axes = plt.subplots(1, n, figsize=(2.2 * n + 0.9, 4.15),
                             gridspec_kw=dict(wspace=.02, left=.012, right=.90,
                                              top=.83, bottom=.03))
    axes = np.atleast_1d(axes)
    for k, ax in enumerate(axes):
        m = _panel(ax, grid, borders, fields[k], "YlOrRd", norm)
        ax.set_title(labels[k], fontsize=10, pad=3)
    cax = fig.add_axes([.915, .05, .014, .74])
    cb = fig.colorbar(m, cax=cax, ticks=lev[::2]); cb.ax.tick_params(labelsize=8)
    cb.outline.set_linewidth(.5)
    fig.suptitle(title, fontsize=12.5, y=.995, x=.012, ha="left", va="top")
    return _png(fig)


def scalar_row(fields, grid, borders, labels, title, cmap, unit, hi=None):
    """Một hàng bản đồ giá trị liên tục, chung một thang màu cho cả sáu tuần."""
    allv = np.concatenate([f[np.isfinite(f)] for f in fields])
    top = hi if hi is not None else max(np.percentile(allv, 98), 1e-6)
    lev = np.linspace(0.0, top, 11)
    norm = BoundaryNorm(lev, ncolors=256, clip=True)
    n = len(fields)
    fig, axes = plt.subplots(1, n, figsize=(2.2 * n + 0.9, 4.15),
                             gridspec_kw=dict(wspace=.02, left=.012, right=.90,
                                              top=.83, bottom=.03))
    axes = np.atleast_1d(axes)
    for k, ax in enumerate(axes):
        m = _panel(ax, grid, borders, fields[k], cmap, norm)
        ax.set_title(labels[k], fontsize=10, pad=3)
    cax = fig.add_axes([.915, .05, .014, .74])
    cb = fig.colorbar(m, cax=cax, ticks=lev[::2]); cb.ax.tick_params(labelsize=8)
    cb.set_label(unit, fontsize=8.5); cb.outline.set_linewidth(.5)
    fig.suptitle(title, fontsize=12.5, y=.995, x=.012, ha="left", va="top")
    return _png(fig)


def binary_map(f2d, grid, borders, title):
    cmap = ListedColormap(["#dfe7ef", "#c1272d"])
    norm = BoundaryNorm([0, .5, 1], ncolors=2)
    fig, ax = plt.subplots(figsize=(4.1, 4.4),
                           gridspec_kw=dict(left=.02, right=.98, top=.90, bottom=.03))
    _panel(ax, grid, borders, f2d, cmap, norm)
    ax.set_title(title, fontsize=11.5, pad=5)
    return _png(fig)


PAGE = """<!doctype html>
<meta charset="utf-8">
<title>%(title)s</title>
<style>
 body{font-family:-apple-system,Segoe UI,Roboto,Helvetica,Arial,sans-serif;
      margin:0 auto;max-width:1240px;padding:28px 22px 60px;color:#1a1a1a}
 h1{font-size:22px;margin:0 0 4px} h2{font-size:15px;font-weight:600;margin:28px 0 6px}
 p.sub{color:#555;font-size:13px;margin:0 0 18px;line-height:1.55}
 figure{margin:0 0 6px} img{max-width:100%%;display:block}
 table{border-collapse:collapse;font-size:12.5px;margin:10px 0 22px}
 th,td{border:1px solid #ddd;padding:3px 9px;text-align:right} th{background:#f6f6f6}
 td:first-child,th:first-child{text-align:left}
 code{background:#f4f4f4;padding:1px 4px;border-radius:3px}
 .yes{background:#fde2e2;font-weight:600} .no{color:#888}
 td.ref,th.ref{background:#f0f0f0;color:#999;font-weight:400}
 td.yes.ref{background:#efe6e6;color:#a08080;font-weight:400}
 .warn{background:#fff4e0} .note{border-left:3px solid #c1272d;background:#fbf6f6;
      padding:9px 13px;font-size:13px;line-height:1.55;margin:14px 0}
</style>
<h1>%(title)s</h1>
<p class="sub">%(subtitle)s</p>
%(body)s
"""


def _degeneracy_txt(cov_mod, cov_obs):
    """Câu mô tả mức quá tán “cả nước cùng bật”, tính từ chính gói đang dùng."""
    if cov_mod is None:
        return "chưa có forecast_valid.nc của gói này để đo"
    ws = [w for w, _ in C.WEEKS][3:]                      # W4-W6
    m = [100.0 * (cov_mod[w] > 0.95).mean() for w in ws]
    o = [100.0 * (cov_obs[w] > 0.95).mean() for w in ws]
    return ("ở %s–%s là %.1f–%.1f%% số lần so với %.1f%% trong quan trắc"
            % (ws[0], ws[-1], min(m), max(m), max(o))).replace(".", ",")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--fcs", default=os.path.join(C.OUT, "operational",
                                                  "forecast_20260829.nc"))
    ap.add_argument("--init", default="2026-08-29")
    ap.add_argument("--valid-fcs", default=os.path.join(C.OUT, "operational",
                                                        "forecast_valid.nc"))
    ap.add_argument("--out", default=os.path.join(C.ROOT, "ehd_forecast.html"))
    ap.add_argument("--verify-json", default=os.path.join(C.OUT, "AI_Results",
                                                          "final_verify.json"),
                    help="bảng kỹ năng tham chiếu của ĐÚNG gói đang dựng bản tin; "
                         "để mặc định thì lấy của PA1/OPE1")
    ap.add_argument("--pack-label", default="outputs/operational/",
                    help="đường dẫn gói tham số, chỉ để ghi ở chân trang")
    args = ap.parse_args()

    init = dt.date.fromisoformat(args.init)
    grid = C.Grid()
    borders = load_regions()
    npt = grid.npt

    with Dataset(args.fcs) as ds:
        prob = {w: np.asarray(ds["prob_%s" % w][:]).ravel() for w, _ in C.WEEKS}
        pred = {w: np.asarray(ds["pred_%s" % w][:]).ravel().astype(bool)
                for w, _ in C.WEEKS}
        pgate = {w: np.asarray(ds["p_gate_%s" % w][:]).ravel() for w, _ in C.WEEKS}
        has_sev = ("sev_acc_W1" in ds.variables)
        sev_acc = ({w: np.asarray(ds["sev_acc_%s" % w][:]).ravel()
                    for w, _ in C.WEEKS} if has_sev else None)
        sev_ai = ({w: np.asarray(ds["sev_ai_%s" % w][:]).ravel()
                   for w, _ in C.WEEKS} if has_sev else None)

    # --- khí hậu học diện tích quan trắc, để kiểm tra độ hợp lý -------------
    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, npt)        # BIEN truoc DIEM
    iv = np.nonzero(np.isin(yr, C.VALID_YEARS))[0]
    E = np.asarray(T[iv][:, :, 2, :])
    cov_obs = {w: (np.nansum(E[:, lts, :], axis=1) >= 1.0).mean(axis=1)
               for w, lts in C.WEEKS}

    # Tần suất trạng thái "cả nước cùng bật" (>95% ô) của CHÍNH gói đang dựng bản
    # tin, đo trên 416 init kiểm định.  Trước đây con số này viết cứng trong văn
    # bản theo PA1; gói khác thì con số khác, nên phải tính lại.
    cov_mod = None
    if os.path.exists(args.valid_fcs):
        with Dataset(args.valid_fcs) as ds:
            cov_mod = {w: np.asarray(ds["pred_%s" % w][:]).astype(bool).mean(axis=1)
                       for w, _ in C.WEEKS}

    body = []

    # ------------------------------------------------ 1. hạn dự báo
    rows = ["<table><tr><th>Cửa sổ</th><th>Ngày hiệu lực</th><th>Số ngày</th>"
            "<th>Sản phẩm phát hành</th></tr>"]
    for w, lts in C.WEEKS:
        d0 = init + dt.timedelta(days=int(lts[0]))
        d1 = init + dt.timedelta(days=int(lts[-1]))
        prod = "bản đồ nhị phân + xác suất" if w == "W1" else "chỉ trường xác suất"
        rows.append("<tr><td>%s</td><td>%s → %s</td><td>%d</td><td>%s</td></tr>"
                    % (w, d0.strftime("%d/%m/%Y"), d1.strftime("%d/%m/%Y"),
                       lts.size, prod))
    rows.append("</table>")
    body.append("<h2>1. Hạn dự báo và sản phẩm</h2>" + "".join(rows))

    labs = ["%s\n%s–%s" % (w,
                            (init + dt.timedelta(days=int(l[0]))).strftime("%d/%m"),
                            (init + dt.timedelta(days=int(l[-1]))).strftime("%d/%m"))
            for w, l in C.WEEKS]

    # ------------------------------------------------ 2. cường độ nắng nóng
    body.append("<h2>2. Cường độ nắng nóng dự báo, W1–W6</h2>")
    if sev_acc is not None:
        png = scalar_row([grid.to_grid(sev_acc[w]) for w, _ in C.WEEKS], grid,
                         borders, labs,
                         "Tổng Sev\u22650 tích luỹ trong tuần — khởi tạo %s"
                         % init.strftime("%d/%m/%Y"), "YlOrRd", "°C·ngày")
        body.append('<figure><img alt="Sev tích luỹ" '
                    'src="data:image/png;base64,%s"></figure>' % png)
        body.append('<div class="note"><b>Đây là bản đồ chính.</b> Mỗi ô lưới là '
                    '&Sigma;<sub>ngày</sub> max(Sev, 0) trong tuần, với '
                    'Sev = Tx<sub>EC</sub> &minus; ngưỡng EHD, lấy thẳng từ lần chạy '
                    'EC. Nó gộp được cả ba trường hợp: một ngày vượt ngưỡng thì giá '
                    'trị nhỏ (trừ khi ngày đó rất nóng), nhiều ngày vượt thì lớn, '
                    'không ngày nào vượt thì <b>đúng bằng 0</b>. Số ô bằng 0 tăng '
                    'dần theo hạn: %s.</div>'
                    % ", ".join("%s %d/541" % (w, int((sev_acc[w] == 0).sum()))
                                for w, _ in C.WEEKS))
        png = scalar_row([grid.to_grid(sev_ai[w]) for w, _ in C.WEEKS], grid,
                         borders, labs,
                         "Sev cực đại tuần, trung bình trên 72 thành viên AI",
                         "YlOrBr", "°C")
        body.append('<figure><img alt="Sev AI" '
                    'src="data:image/png;base64,%s"></figure>' % png)
        body.append('<div class="note">Trường của <b>mô hình AI</b>: trung bình trên '
                    'N = 8 họ &times; 9 phân vị = 72 thành viên của max(Sev, 0). Lưu '
                    'ý tầng 1 được huấn luyện trên Sev <b>cực đại</b> của tuần '
                    '(<code>phase1_arch.py:162</code>), không phải tổng, nên trường '
                    'này đo <i>đỉnh</i> nắng nóng chứ không đo tích luỹ — và vì trung '
                    'bình max(thành viên, 0) trên 72 thành viên có độ tán thì luôn '
                    'dương, nó <b>không thể hiện được</b> trường hợp “không ngày nào '
                    'vượt ngưỡng”. Dùng bản đồ tích luỹ ở trên cho mục đích đó.</div>')
    b = binary_map(grid.to_grid(pred["W1"].astype("f4")), grid, borders,
                   "W1 — EHDE (đỏ = có)")
    body.append("<h3 style='font-size:13.5px;margin:22px 0 6px'>2b. Sản phẩm nhị "
                "phân phát hành (chỉ W1, %s → %s)</h3>"
                % ((init).strftime("%d/%m"),
                   (init + dt.timedelta(days=10)).strftime("%d/%m/%Y")))
    body.append('<figure><img alt="W1 nhị phân" '
                'src="data:image/png;base64,%s"></figure>' % b)

    # ------------------------------------------------ 3. xác suất W1-W6
    png = prob_row([grid.to_grid(prob[w]) for w, _ in C.WEEKS], grid, borders,
                   labs, "Xác suất xảy ra EHDE — khởi tạo %s"
                   % init.strftime("%d/%m/%Y"))
    body.append("<h2>3. Trường xác suất EHDE, W1–W6</h2>"
                '<figure><img alt="xác suất W1-W6" '
                'src="data:image/png;base64,%s"></figure>' % png)

    # ------------------------------------------------ 4. bảng theo vùng
    # Bốn cách xác định yEHDE tại TỪNG Ô LƯỚI; sau đó vùng có EHDE nếu > 20% ô.
    # P(yEHDE) = M/N là tỉ lệ phiếu của N = 72 thành viên tầng 1 (p_gate).
    variants = [("Quy tắc nghiệp vụ (cổng p1 riêng từng ô + tầng 2)",
                 {w: pred[w] for w, _ in C.WEEKS}, True)]
    for thr_mn in (0.5, 0.6, 0.7):
        variants.append(("P(yEHDE) = M/N &gt; %.1f" % thr_mn,
                         {w: pgate[w] > thr_mn for w, _ in C.WEEKS}, False))

    body.append("<h2>4. EHDE theo vùng — bốn phương án xác định yEHDE</h2>")
    body.append('<div class="note">Tiêu chí <b>vùng</b> giữ nguyên theo EXP1: vùng '
                'được coi là có EHDE nếu <b>&gt; 20%</b> số ô lưới của vùng đó có '
                'yEHDE (ô tô đậm). Cái thay đổi giữa bốn bảng là cách xác định '
                'yEHDE tại <b>từng ô lưới</b>. P(yEHDE) = M/N là tỉ lệ phiếu thuận '
                'trong N = 8 họ &times; 9 phân vị = <b>72 thành viên</b> của tầng 1. '
                'Chỉ bảng đầu là quy tắc đang phát hành; ba bảng sau để so sánh độ '
                'nhạy theo ngưỡng.</div>')
    for name, dec, is_op in variants:
        rows = ["<table><tr><th>Vùng</th><th>n ô</th>" +
                "".join("<th%s>%s</th>" % ("" if (w == "W1" or not is_op)
                                           else ' class="ref"', w)
                        for w, _ in C.WEEKS) + "</tr>"]
        for code, nmv in REG_ORDER:
            sel = grid.reg == code
            cells = []
            for w, _ in C.WEEKS:
                P = dec[w][sel].mean()
                cls = "yes" if P > P_REGION else "no"
                if is_op and w != "W1":
                    cls += " ref"
                cells.append('<td class="%s">%.0f%%</td>' % (cls, 100 * P))
            rows.append("<tr><td>%s</td><td>%d</td>%s</tr>"
                        % (escape(nmv), sel.sum(), "".join(cells)))
        nreg = ["%s %d/7" % (w, sum(1 for c, _ in REG_ORDER
                                    if dec[w][grid.reg == c].mean() > P_REGION))
                for w, _ in C.WEEKS]
        rows.append('<tr><td colspan="2"><i>số vùng có EHDE</i></td>' +
                    "".join("<td><i>%s</i></td>" % x.split()[1] for x in nreg) +
                    "</tr></table>")
        tag = " <span style='color:#c1272d'>(đang phát hành)</span>" if is_op else ""
        body.append("<h3 style='font-size:13.5px;margin:20px 0 4px'>%s%s</h3>"
                    % (name, tag) + "".join(rows))

    # ------------------------------------------------ 5. kiểm tra độ hợp lý
    rows = ["<table><tr><th>Cửa sổ</th><th>Diện tích dự báo</th>"
            "<th>Trung vị quan trắc</th><th>Phân vị 95 quan trắc</th>"
            "<th>Đánh giá</th></tr>"]
    for w, _ in C.WEEKS:
        f = pred[w].mean()
        o = cov_obs[w]
        p95 = np.percentile(o, 95)
        over = f > p95
        rows.append('<tr%s><td>%s</td><td>%.1f%%</td><td>%.1f%%</td><td>%.1f%%</td>'
                    '<td>%s</td></tr>'
                    % (' class="warn"' if over else "", w, 100 * f,
                       100 * np.median(o), 100 * p95,
                       "vượt p95 — không phát nhị phân" if over else "trong dải hợp lý"))
    rows.append("</table>")
    body.append("<h2>5. Kiểm tra độ hợp lý của diện tích dự báo</h2>" + "".join(rows))
    body.append('<div class="note">Diện tích dự báo được đặt cạnh phân bố diện '
                'tích <b>quan trắc</b> trên 416 lần khởi tạo của bốn năm kiểm định. '
                'Mô hình quá tán ở mức toàn quốc: nó phát trạng thái “cả nước cùng '
                'bật” nhiều hơn tự nhiên rất nhiều (%s). Đã thử ràng buộc diện tích để chữa và '
                '<b>đã bác</b>: cách đó xoá được suy biến nhưng làm mất kỹ năng, W1 '
                'mất luôn chỉ tiêu FAR&lt;0,2. Vì vậy giữ nguyên quy tắc gốc và chỉ '
                'phát hành nhị phân cho W1.</div>' % _degeneracy_txt(cov_mod, cov_obs))

    # ------------------------------------------------ 6. kỹ năng tham chiếu
    fv = args.verify_json
    if os.path.exists(fv):
        with open(fv) as fh:
            res = json.load(fh)["national"]
        ms = ["POD", "POFD", "FAR", "FBI", "CSI", "SEDI", "AUC"]
        rows = ["<table><tr><th>Chỉ số</th>" +
                "".join("<th>%s</th>" % w for w, _ in C.WEEKS) + "</tr>"]
        for m in ms:
            rows.append("<tr><td>%s</td>%s</tr>"
                        % (m, "".join("<td>%.3f</td>" % res[w][m]
                                      for w, _ in C.WEEKS)))
        rows.append("</table>")
        body.append("<h2>6. Kỹ năng tham chiếu (416 lần khởi tạo, 4 năm kiểm định)"
                    "</h2>" + "".join(rows))
        body.append('<div class="note">Chỉ <b>W1</b> đạt trọn bộ chỉ tiêu '
                    '(POD&gt;0,6; POFD&lt;0,2; FAR&lt;0,2; CSI&gt;0,4). Từ W2 trở đi '
                    'không có ngưỡng nào trong dải r∈[0,02; 1,50] đạt FAR&lt;0,2 — '
                    'đó là trần tín hiệu, không phải lỗi hiệu chỉnh — nên các hạn này '
                    'chỉ phát hành trường xác suất để người dùng tự chọn ngưỡng theo '
                    'chi phí bỏ sót và chi phí báo động sai của mình.</div>')

    sub = ("Khởi tạo <b>%s</b> từ <b>control forecast</b> của ECMWF. "
           "541 ô lưới đất của <code>Mask_VN_S2S_EC_0.25.nc</code>, ranh giới bảy "
           "vùng khí hậu từ <code>Vietnam7.shp</code>. Ngưỡng EHD "
           "<code>min(Tx90p, 35&nbsp;&deg;C)</code>, Tx90p lấy trên cửa sổ 31 ngày. "
           "EHDE của một tuần = có ít nhất một ngày EHD trong tuần đó. "
           "Bộ tham số đóng băng tại <code>%s</code>."
           % (init.strftime("%d/%m/%Y"), escape(args.pack_label)))
    title = "Dự báo nội mùa nắng nóng cực đoan (EHDE) — khởi tạo %s" % \
            init.strftime("%d/%m/%Y")
    with open(args.out, "w") as fh:
        fh.write(PAGE % dict(title=escape(title), subtitle=sub,
                             body="\n".join(body)))
    print("[done] -> %s" % args.out, flush=True)


if __name__ == "__main__":
    main()
