#!/usr/bin/env python3
"""Operational EHD bulletin for one CFS initial time, as a self-contained page.

Follows the release rule the evaluation established, and this is where the page
differs from a plain plot: W1 is issued as a BINARY regional forecast, W2-W6 are
issued only as a severity margin, because no threshold at any lead beyond W1 met
the target box. Each window carries a plausibility line putting the forecast area
next to the historical observed area for the same time of year.
"""

from __future__ import annotations

import argparse, base64, datetime as dt, io, json, os, sys
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.colors import BoundaryNorm, ListedColormap, TwoSlopeNorm
import shapefile

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import cfs_common as C
import build_obs as B

SHAPE = "/home/tanpv/GrADS_SCR/VN_SubRegs1/Vietnam7.shp"
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ộ"}
CSS = open(os.path.join(os.path.dirname(os.path.abspath(__file__)),
                        "_report_css.txt")).read()


def borders():
    try:
        r = shapefile.Reader(SHAPE)
    except Exception:
        return []
    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 fig(fn):
    buf = io.BytesIO(); fn()
    plt.savefig(buf, format="png", dpi=130, bbox_inches="tight"); plt.close()
    return "<figure><img src='data:image/png;base64,%s'></figure>" % \
           base64.b64encode(buf.getvalue()).decode()


def plausibility(r):
    z = np.load(B.TX_NPZ); oord = z["ordinal"]
    ehd, _ = B.ehd_of("train")
    out = {}
    for w, (wn, idx) in enumerate(C.WEEKS):
        q, nd = r["weeks"][wn]["doy"], len(idx)
        fr = []
        for y in range(1961, 2024):
            c = dt.date(y, 1, 1).toordinal() + q - 1
            rows = [np.where(oord == c + k - nd // 2)[0] for k in range(nd)]
            rows = [x[0] for x in rows if len(x)]
            if len(rows) < nd:
                continue
            fr.append(ehd[rows].any(0).mean())
        fr = np.array(fr) * 100
        f = r["weeks"][wn]["frac_point"] * 100
        p50, p90, p95 = np.percentile(fr, [50, 90, 95])
        out[wn] = (f, p50, p90, p95,
                   "bình thường" if f <= p90 else
                   ("cao" if f <= p95 else "VƯỢT p95 — cảnh báo"))
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--init", default="2026082900")
    a = ap.parse_args()
    r = json.load(open(os.path.join(C.OUT, "RT", "fcst_%s.json" % a.init)))
    z = np.load(B.TX_NPZ)
    reg = z["reg"]; lat = z["lat"][z["jj"]]; lon = z["lon"][z["ii"]]
    bd = borders()
    pl = plausibility(r)
    init_txt = "%s-%s-%s %sZ" % (a.init[:4], a.init[4:6], a.init[6:8], a.init[8:10])

    H = ["<!doctype html><html lang='vi'><meta charset='utf-8'>",
         "<title>Bản tin EHD %s</title><style>%s</style>" % (init_txt, CSS),
         "<h1>Bản tin dự báo ngày nóng cực đoan — hạn nội mùa</h1>",
         "<p class='sub'>Nguồn: CFS nghiệp vụ NOMADS, initial <b>%s</b>, thành viên "
         "01. Hạ quy mô BCSD về lưới VnGAD 0.25°, 541 ô lưới đất, 7 vùng.</p>"
         % init_txt,
         "<div class='stamp'>Chuỗi xử lý: TMAX 6 giờ → Tx ngày (căn A) → cực đại "
         "tuần tại điểm → BCSD bằng quantile map trên <b>năm huấn luyện</b> trong "
         "cửa sổ ±30 ngày quanh doy của tuần → phát báo khi vượt ngưỡng EHD địa "
         "phương + <b>δ = %.2f°C</b> → vùng có EHDE khi hơn <b>%.0f%%</b> số ô lưới "
         "phát báo (EXP1). Không có tham số nào được khớp lại trên chính bản dự báo "
         "này. Dựng lúc %s.</div>"
         % (r["delta"], 100 * r["reg_frac"],
            dt.datetime.now().strftime("%Y-%m-%d %H:%M")),
         "<div class='warnbox'><b>Quy tắc phát hành.</b> Chỉ <b>W1</b> được phát "
         "dưới dạng bản đồ và kết luận nhị phân: đó là cửa sổ duy nhất đạt trọn bộ "
         "chỉ tiêu khi kiểm định ngoài mẫu (POD 0.65, POFD 0.19, FAR 0.13, CSI 0.59). "
         "<b>W2–W6 không được chứng nhận</b> — không tồn tại ngưỡng nào ở các hạn đó "
         "cho FAR &lt; 0.2, nên chúng chỉ được trình bày như <b>biên độ vượt ngưỡng</b> "
         "(°C) để tham khảo xu thế, tuyệt đối không dùng như dự báo có/không.</div>"]

    # ---- regional table
    H.append("<h2>1. Kết luận theo vùng</h2>")
    H.append("<table><caption>P = tỷ lệ ô lưới trong vùng vượt ngưỡng. "
             "Vùng có EHDE khi P &gt; 20%.</caption><tr><th>Vùng</th>" +
             "".join("<th>%s</th>" % w for w, _ in C.WEEKS) + "</tr>")
    for rg in sorted(REG_NAME):
        cells = []
        for wn, _ in C.WEEKS:
            d = r["weeks"][wn]; P = d["P"][str(rg)]
            hit = P > r["reg_frac"]
            cls = ("good" if wn == "W1" else "warn") if hit else ""
            cells.append("<td%s>%.0f%%%s</td>"
                         % (" class='%s'" % cls if cls else "", 100 * P,
                            " ✓" if hit else ""))
        H.append("<tr><td>%s</td>%s</tr>" % (REG_NAME[rg], "".join(cells)))
    H.append("</table>")
    w1 = [REG_NAME[k] for k in sorted(REG_NAME)
          if r["weeks"]["W1"]["P"][str(k)] > r["reg_frac"]]
    H.append("<div class='note'><b>W1 (%s … %s): có EHDE ở %d/7 vùng — %s.</b></div>"
             % (r["weeks"]["W1"]["valid_from"], r["weeks"]["W1"]["valid_to"],
                len(w1), ", ".join(w1) if w1 else "không vùng nào"))

    # ---- maps
    H.append("<h2>2. Bản đồ</h2>")

    def draw():
        fg, axes = plt.subplots(2, 3, figsize=(13, 9))
        for k, (wn, _) in enumerate(C.WEEKS):
            ax = axes[k // 3][k % 3]
            d = r["weeks"][wn]
            m = np.array(d["margin"])
            if wn == "W1":
                v = (m >= r["delta"]).astype(float)
                sc = ax.scatter(lon, lat, c=v, s=13, marker="s", vmin=0, vmax=1,
                                cmap=ListedColormap(["#dfe7ef", "#c0392b"]))
            else:
                sc = ax.scatter(lon, lat, c=m, s=13, marker="s", cmap="RdBu_r",
                                norm=TwoSlopeNorm(vcenter=r["delta"],
                                                  vmin=min(m.min(), -3),
                                                  vmax=max(m.max(), 3)))
                plt.colorbar(sc, ax=ax, shrink=.8, label="biên độ (°C)")
            for b in bd:
                ax.plot(b[:, 0], b[:, 1], color="#444", lw=.4)
            ttl = "%s  %s…%s" % (wn, d["valid_from"][5:], d["valid_to"][5:])
            ax.set_title(ttl + ("" if wn == "W1" else "  (tham khảo)"),
                         fontsize=10, color="k" if wn == "W1" else "#777")
            ax.set_xlim(101.5, 110.5); ax.set_ylim(8, 24)
            ax.set_aspect(1.0); ax.tick_params(labelsize=7)
    H.append(fig(draw))
    H.append("<p class='sub'>W1: đỏ = phát báo EHD. W2–W6: biên độ vượt ngưỡng, "
             "trắng tại đúng ngưỡng phát báo δ = +%.2f°C.</p>" % r["delta"])

    # ---- plausibility
    H.append("<h2>3. Kiểm tra độ hợp lý</h2>")
    H.append("<table><caption>Diện tích phát báo đặt cạnh diện tích quan trắc lịch "
             "sử 1961–2023 cùng thời điểm trong năm</caption>"
             "<tr><th>Tuần</th><th>Dự báo</th><th>QT trung vị</th><th>QT p90</th>"
             "<th>QT p95</th><th>Đánh giá</th></tr>")
    for wn, _ in C.WEEKS:
        f, p50, p90, p95, verdict = pl[wn]
        H.append("<tr><td>%s</td><td>%.1f%%</td><td>%.1f%%</td><td>%.1f%%</td>"
                 "<td>%.1f%%</td><td class='%s'>%s</td></tr>"
                 % (wn, f, p50, p90, p95,
                    "good" if verdict == "bình thường" else "bad", verdict))
    H.append("</table>")
    H.append("<div class='note'>Không cửa sổ nào vượt p95. Đáng lưu ý vì dự án "
             "ECMWF tiền nhiệm gặp hiện tượng suy biến “cả nước cùng bật” (&gt;95% "
             "số ô) ở W4–W6; chuỗi BCSD ở đây không có hiện tượng đó.</div>")
    H.append("<div class='warnbox'><b>Chưa kiểm chứng được.</b> Bộ quan trắc VnGAD "
             "chỉ có đến 2023 nên bản tin này không thể đối chiếu với thực tế đã "
             "xảy ra. Các con số chất lượng trích ở trên là kết quả kiểm định trên "
             "2001/2006/2009, không phải cho đợt dự báo này.</div>")
    H.append("</html>")
    out = "/home/tanpv/Progs/claude/HW_CFS_raw/ehd_cfs_forecast_%s.html" % a.init
    open(out, "w").write("\n".join(H))
    print("viet %s (%.2f MB)" % (out, os.path.getsize(out) / 1e6))


if __name__ == "__main__":
    main()