#!/usr/bin/env python3
"""Assemble the session's results into a single self-contained HTML report.

Reads whatever result JSONs exist under outputs/downscale/ and renders:
  - downscaling skill by forecast week (RMSE / MAE / ME / CORR)
  - EHD event scores at grid-point and regional (EXP1) level, against the targets
  - feature importance of the full-predictor model
  - maps of per-point skill, one row per forecast week on a shared colour scale

Usage: make_report.py [--out ehd_report.html]
"""
import argparse
import base64
import io
import json
import os

import numpy as np
import xarray as xr
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt            # noqa: E402

ROOT = "/home/tanpv/Progs/claude/HW_EC_10"
OUTD = f"{ROOT}/outputs/downscale"
STATIC = f"{ROOT}/outputs/Cli/static_fields_0.25.nc"
WEEKS = ["W1", "W2", "W3", "W4", "W5", "W6", "ALL"]
TARGETS = {"POD": (">", 0.60), "POFD": ("<", 0.20), "FAR": ("<", 0.20),
           "CSI": (">", 0.40), "FBI": ("~", 0.75)}


def fig_to_img(fig):
    buf = io.BytesIO()
    fig.savefig(buf, format="png", dpi=110, bbox_inches="tight")
    plt.close(fig)
    return ("data:image/png;base64,"
            + base64.b64encode(buf.getvalue()).decode())


def load(name):
    p = os.path.join(OUTD, name)
    return json.load(open(p)) if os.path.exists(p) else None


def table(rows, header, cls=""):
    h = "".join(f"<th>{c}</th>" for c in header)
    b = "".join("<tr>" + "".join(f"<td>{c}</td>" for c in r) + "</tr>" for r in rows)
    return f'<table class="{cls}"><thead><tr>{h}</tr></thead><tbody>{b}</tbody></table>'


def fmt(v, target=None, digits=3):
    if v is None or (isinstance(v, float) and not np.isfinite(v)):
        return "&ndash;"
    s = f"{v:.{digits}f}"
    if target:
        op, lim = target
        ok = (v > lim) if op == ">" else (v < lim) if op == "<" else abs(v - lim) < 0.15
        s = f'<span class="{"ok" if ok else "bad"}">{s}</span>'
    return s


def skill_maps(pred_files, static):
    """One row of maps per forecast week, shared colour scale, RMSE per grid point."""
    obs = xr.open_dataset(f"{ROOT}/outputs/OBS_4EC15/tx_obs_aligned.nc")
    land = (np.isfinite(static.vn.values) & (static.vn.values > 0))
    li = np.where(land.ravel())[0]
    o = obs.tx_obs.values.reshape(20, 104, 46, -1)[..., li]
    lat, lon = obs.lat.values, obs.lon.values
    bounds = [(0, 11), (11, 18), (18, 25), (25, 32), (32, 39), (39, 46)]
    imgs = {}
    for name, path in pred_files.items():
        if not os.path.exists(path):
            continue
        p = np.load(path)
        fig, axes = plt.subplots(1, 6, figsize=(19, 3.4), constrained_layout=True)
        rms = []
        for (l0, l1) in bounds:
            d = p[:, :, l0:l1] - o[:, :, l0:l1]
            rms.append(np.sqrt(np.nanmean(d ** 2, axis=(0, 1, 2))))
        vmax = float(np.nanpercentile(np.concatenate(rms), 98))
        for k, (ax, r) in enumerate(zip(axes, rms)):
            grid = np.full(land.size, np.nan)
            grid[li] = r
            im = ax.pcolormesh(lon, lat, grid.reshape(land.shape),
                               vmin=0, vmax=vmax, cmap="YlOrRd", shading="auto")
            ax.set_title(f"W{k+1}", fontsize=11)
            ax.set_aspect(1.0)
            ax.set_xticks([]); ax.set_yticks([])
        fig.colorbar(im, ax=axes, shrink=0.85, label="RMSE (°C)")
        imgs[name] = fig_to_img(fig)
    return imgs


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default=f"{ROOT}/ehd_report.html")
    args = ap.parse_args()
    static = xr.open_dataset(STATIC)

    base = load("baselines.json")
    ai = load("ai_lgbm.json")
    ehd_b = load("ehd_baselines.json")
    ehd_c = load("ehd_classifier.json")
    ehd_f = load("ehd_full.json")
    ehd_r = load("ehd_full_real.json")      # native 1.5 deg REAL MARS data
    ehd_rp = load("ehd_full_real_pcs.json") # native 1.5 deg REAL + wide-domain EOF/PC

    parts = []

    # --- downscaling skill
    if base or ai:
        rows = []
        src = dict(base or {})
        if ai:
            src["lgbm"] = ai.get("lgbm")
        for m, r in src.items():
            if not isinstance(r, dict):
                continue
            rows.append([m] + [fmt(r.get(w, {}).get("RMSE"), digits=2) for w in WEEKS]
                        + [fmt(r.get("ALL", {}).get("MAE"), digits=2),
                           fmt(r.get("ALL", {}).get("CORR"), digits=3)])
        parts.append("<h2>1. Kỹ năng downscaling Tx (RMSE theo tuần, °C)</h2>"
                     + table(rows, ["phương pháp"] + WEEKS + ["MAE (ALL)", "CORR (ALL)"]))

    # --- event scores
    for title, blob, key in [("2. Điểm sự kiện EHD — các baseline", ehd_b, None),
                             ("3. Phân loại chỉ dùng Tx", ehd_c, "classifier"),
                             ("4. Mô hình đủ predictor (proxy 1.5°)", ehd_f, "full"),
                             ("5. Mô hình đủ predictor — SỐ LIỆU THẬT 1.5°",
                              ehd_r, "full"),
                             ("6. SỐ LIỆU THẬT 1.5° + EOF/PC miền rộng",
                              ehd_rp, "full")]:
        if not blob:
            continue
        items = blob.items() if key is None else [(key, blob[key])]
        for lvl in ("gridpoint", "regional_EHDE"):
            rows = []
            for nm, r in items:
                if not isinstance(r, dict) or lvl not in r:
                    continue
                for w in WEEKS:
                    v = r[lvl].get(w)
                    if not v:
                        continue
                    rows.append([nm, w, fmt(v["base"]), fmt(v["POD"], TARGETS["POD"]),
                                 fmt(v["POFD"], TARGETS["POFD"]),
                                 fmt(v["FAR"], TARGETS["FAR"]),
                                 fmt(v["FBI"], TARGETS["FBI"]),
                                 fmt(v["CSI"], TARGETS["CSI"])])
            if rows:
                parts.append(f"<h3>{title} — {lvl}</h3>"
                             + table(rows, ["mô hình", "tuần", "base", "POD", "POFD",
                                            "FAR", "FBI", "CSI"]))
        parts.insert(len(parts), "")

    # --- feature importance (standard: render for EVERY model that emits it)
    def fi_chart(name, gain):
        top = list(gain.items())[:25]
        fig, ax = plt.subplots(figsize=(7, 6.5))
        ax.barh([k for k, _ in top][::-1], [v for _, v in top][::-1], color="#c2410c")
        ax.set_xlabel("normalized gain")
        ax.set_title(f"Feature importance — {name}")
        return f'<h3>{name}</h3><img src="{fig_to_img(fig)}">'

    fi_sources = []
    if ehd_rp and "feature_gain" in ehd_rp.get("full", {}):
        fi_sources.append(("SỐ LIỆU THẬT 1.5° + EOF/PC miền rộng",
                           ehd_rp["full"]["feature_gain"]))
    if ehd_r and "feature_gain" in ehd_r.get("full", {}):
        fi_sources.append(("EHD đủ predictor — SỐ LIỆU THẬT 1.5°",
                           ehd_r["full"]["feature_gain"]))
    if ehd_f and "feature_gain" in ehd_f.get("full", {}):
        fi_sources.append(("EHD đủ predictor (proxy 1.5°)", ehd_f["full"]["feature_gain"]))
    if ai and isinstance(ai.get("lgbm"), dict) and "feature_gain" in ai:
        fi_sources.append(("downscaling LightGBM", ai["feature_gain"]))
    if ehd_c and "feature_gain" in ehd_c:
        fi_sources.append(("phân loại chỉ dùng Tx", ehd_c["feature_gain"]))
    if fi_sources:
        parts.append("<h2>5. Đóng góp của các đặc trưng</h2>"
                     + "".join(fi_chart(n, g) for n, g in fi_sources))

    # --- maps
    imgs = skill_maps({"BCSD": f"{OUTD}/pred/pred_bcsd.npy",
                       "LightGBM": f"{OUTD}/pred/pred_lgbm.npy"}, static)
    for nm, src in imgs.items():
        parts.append(f"<h3>Bản đồ RMSE theo tuần — {nm}</h3><img src='{src}'>")

    css = """
    body{font-family:system-ui,Segoe UI,Arial,sans-serif;margin:32px;max-width:1400px;
         color:#1f2328;line-height:1.5}
    h1{border-bottom:3px solid #c2410c;padding-bottom:8px}
    h2{margin-top:34px;color:#7c2d12}
    table{border-collapse:collapse;margin:10px 0 22px;font-size:14px}
    th,td{border:1px solid #d0d7de;padding:5px 10px;text-align:right}
    th{background:#f6f8fa;font-weight:600}
    td:first-child,th:first-child{text-align:left}
    tbody tr:nth-child(even){background:#fafbfc}
    .ok{color:#137333;font-weight:600}.bad{color:#b3261e}
    img{max-width:100%;border:1px solid #d0d7de;border-radius:6px;margin:8px 0}
    .note{background:#fff8e1;border-left:4px solid #f59e0b;padding:12px 16px;margin:18px 0}
    """
    html = f"""<!doctype html><html lang="vi"><head><meta charset="utf-8">
<title>HW_EC_10 — Dự báo EHD từ S2S ECMWF 1.5°</title><style>{css}</style></head><body>
<h1>HW_EC_10 — Downscaling và dự báo ngày nóng cực đoan từ S2S ECMWF 1.5°</h1>
<div class="note"><b>Chỉ tiêu mục tiêu:</b> POD &gt; 0,6 · POFD &lt; 0,2 · FAR &lt; 0,2 ·
CSI &gt; 0,4 · FBI ~ 0,75. Giá trị <span class="ok">xanh</span> là đạt,
<span class="bad">đỏ</span> là chưa đạt.<br>
<b>Lưu ý:</b> FAR = 1 − POD/FBI là một đẳng thức, nên FAR không độc lập với POD và FBI.
Kết quả dưới đây dựng trên proxy 1.5° (trung bình bộ EC 0.25° lên lưới 1.5° thật);
sẽ chạy lại trên số liệu 1.5° gốc khi tải xong.<br>
<b>Kiểm định:</b> k-fold 5 × 4 năm; mọi khí hậu, ngưỡng và chuẩn hoá đều dựng
riêng từ các năm training của từng fold.</div>
{''.join(parts)}
</body></html>"""
    with open(args.out, "w") as fh:
        fh.write(html)
    print(f"wrote {args.out} ({os.path.getsize(args.out)/1e6:.2f} MB)")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())