#!/usr/bin/env python3
"""Verify ECMWF S2S extreme-heat forecasts over Vietnam (2003--2022).

Outputs:
  ehd_verification.nc   gridded W1--W6 verification metrics
  ehd_verification.html self-contained report with all requested maps/diagrams

The recomputed 90th-percentile and EHD-threshold products are read from
outputs/Cli/Qp_tx_mod.nc and outputs/Cli/Qp_tx_obs.nc.
"""
from __future__ import annotations

import argparse
import base64
import glob
import io
import os
from datetime import date
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from netCDF4 import Dataset
import numpy as np

WEEKS = ((0, 11), (11, 18), (18, 25), (25, 32), (32, 39), (39, 46))
BASE = Path(__file__).resolve().parent
DATA = "/home/tanpv/Data/QG.25.196"


def noleap_doy(d: date) -> int:
    """Zero-based day in a 365-day calendar; Feb 29 shares Feb 28."""
    n = d.timetuple().tm_yday - 1
    if d.month == 2 and d.day == 29:
        return 58
    if d.month > 2 and d.year % 4 == 0 and (d.year % 100 or d.year % 400 == 0):
        n -= 1
    return n


def safe_div(a, b):
    return np.divide(a, b, out=np.full_like(a, np.nan, dtype="f8"), where=b != 0)


def init_acc(shape):
    z = lambda: np.zeros(shape, dtype="f8")
    return {
        "n": z(), "sum": z(), "abs": z(), "sq": z(),
        "w_n": z(), "w_sum": z(), "w_abs": z(), "w_sq": z(),
        "x_n": z(), "x_sum": z(), "x_abs": z(), "x_sq": z(),
        "c_n": z(), "c_sum": z(), "c_abs": z(), "c_sq": z(),
        "hit": z(), "miss": z(), "fa": z(), "cn": z(),
    }


def add_errors(a, prefix, err, valid):
    e = np.where(valid, err, 0.0)
    # Daily errors arrive as (lead, lat, lon); reduced weekly values are 2-D.
    axis = 0 if valid.ndim == 3 else None
    a[prefix + "n"] += valid.sum(axis=axis) if axis is not None else valid
    a[prefix + "sum"] += e.sum(axis=axis) if axis is not None else e
    a[prefix + "abs"] += np.abs(e).sum(axis=axis) if axis is not None else np.abs(e)
    a[prefix + "sq"] += (e * e).sum(axis=axis) if axis is not None else e * e


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


def map_figure(arr, lon, lat, title, cmap="RdBu_r", nonnegative=False):
    fig, ax = plt.subplots(figsize=(5.2, 5.5))
    finite = arr[np.isfinite(arr)]
    if finite.size:
        if nonnegative:
            lo, hi = 0, np.nanpercentile(finite, 98)
            if hi <= 0: hi = 1
        else:
            hi = np.nanpercentile(np.abs(finite), 98)
            if hi <= 0: hi = 1
            lo = -hi
    else:
        lo, hi = (0, 1) if nonnegative else (-1, 1)
    im = ax.pcolormesh(lon, lat, arr, shading="auto", cmap=cmap, vmin=lo, vmax=hi)
    ax.set(xlabel="Longitude (°E)", ylabel="Latitude (°N)", title=title)
    ax.set_aspect("equal", adjustable="box")
    fig.colorbar(im, ax=ax, shrink=.78)
    fig.tight_layout()
    return png_uri(fig)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", default=DATA)
    ap.add_argument("--output-dir", default=str(BASE / "outputs"))
    ap.add_argument("--threshold-dir", default=str(BASE / "outputs" / "Cli"))
    ap.add_argument("--html-output", default=str(BASE / "ehd_verification.html"))
    args = ap.parse_args()
    os.makedirs(args.output_dir, exist_ok=True)

    with Dataset(f"{args.threshold_dir}/Qp_tx_mod.nc") as d:
        qmod = np.ma.filled(d["tx_thres"][:], np.nan).astype("f4")
        qmod_lon, qmod_lat = np.asarray(d["lon"][:]), np.asarray(d["lat"][:])
    with Dataset(f"{args.threshold_dir}/Qp_tx_obs.nc") as d:
        qobs = np.ma.filled(d["tx_thres"][:], np.nan).astype("f4")
        qobs_lon, qobs_lat = np.asarray(d["lon"][:]), np.asarray(d["lat"][:])
    with Dataset(f"{args.data}/OBS/Mask_VN_S2S_EC_0.25.nc") as d:
        lon, lat = d["lon"][:], d["lat"][:]
        mask = ~np.ma.getmaskarray(d["vn"][:]) & np.isfinite(np.ma.filled(d["vn"][:], np.nan))
    expected_mod = (365, 46, len(lat), len(lon))
    expected_obs = (365, len(lat), len(lon))
    if qmod.shape != expected_mod or qobs.shape != expected_obs:
        raise ValueError(f"Threshold shapes must be {expected_mod} and {expected_obs}; got {qmod.shape} and {qobs.shape}")
    if not (np.allclose(qmod_lon, lon) and np.allclose(qmod_lat, lat) and
            np.allclose(qobs_lon, lon) and np.allclose(qobs_lat, lat)):
        raise ValueError("Threshold coordinates do not match the Vietnam mask grid")

    shape = mask.shape
    acc = [init_acc(shape) for _ in WEEKS]
    # ROC histograms use continuous forecast severity; reliability uses the
    # deterministic forecast as probabilities 0 and 1 (two bins).
    roc_edges = np.linspace(-15, 15, 121)
    pos_hist = np.zeros((6, len(roc_edges) - 1) + shape, dtype="u4")
    neg_hist = np.zeros_like(pos_hist)
    rel_fc = np.zeros((6, 2), dtype="f8")
    rel_ob = np.zeros((6, 2), dtype="f8")
    paired = 0

    files = sorted(glob.glob(f"{args.data}/EC_Combined/*/EC_all_*.nc"))
    for k, mod_path in enumerate(files, 1):
        stamp = os.path.basename(mod_path)[7:17]
        obs_path = f"{args.data}/OBS_4EC/{stamp[:4]}/Obs_{stamp}.nc"
        if not os.path.exists(obs_path):
            continue
        d0 = date.fromisoformat(stamp)
        doys = np.array([noleap_doy(date.fromordinal(d0.toordinal() + i)) for i in range(46)])
        with Dataset(mod_path) as dm, Dataset(obs_path) as do:
            txm = np.ma.filled(dm["mx2t6"][:46], np.nan).astype("f4")
            txo = np.ma.filled(do["tasmax"][:46], np.nan).astype("f4")
        thm = qmod[doys, np.arange(46)]
        tho = qobs[doys]
        sm, so = txm - thm, txo - tho
        vm = np.isfinite(txm) & np.isfinite(thm) & mask
        vo = np.isfinite(txo) & np.isfinite(tho) & mask
        valid = vm & vo
        em, eo = sm >= 0, so >= 0
        for wi, (i0, i1) in enumerate(WEEKS):
            a, v = acc[wi], valid[i0:i1]
            err = txm[i0:i1] - txo[i0:i1]
            add_errors(a, "", err, v)
            count = v.sum(0)
            good = count > 0
            mmean = np.divide(np.where(v, txm[i0:i1], 0).sum(0), count, out=np.zeros(shape), where=good)
            omean = np.divide(np.where(v, txo[i0:i1], 0).sum(0), count, out=np.zeros(shape), where=good)
            add_errors(a, "w_", mmean - omean, good)
            mm = np.max(np.where(v, txm[i0:i1], -np.inf), axis=0)
            om = np.max(np.where(v, txo[i0:i1], -np.inf), axis=0)
            add_errors(a, "x_", np.where(good, mm - om, 0.0), good)
            cm = np.where(v, em[i0:i1], 0).sum(0)
            co = np.where(v, eo[i0:i1], 0).sum(0)
            add_errors(a, "c_", cm - co, good)
            fm, ob = em[i0:i1], eo[i0:i1]
            a["hit"] += (v & fm & ob).sum(0)
            a["miss"] += (v & ~fm & ob).sum(0)
            a["fa"] += (v & fm & ~ob).sum(0)
            a["cn"] += (v & ~fm & ~ob).sum(0)
            # Global two-bin reliability counts.
            for b, forecast in enumerate((~fm, fm)):
                use = v & forecast
                rel_fc[wi, b] += use.sum()
                rel_ob[wi, b] += (use & ob).sum()
            # Per-grid severity histograms for approximate ROC/AUC.
            vals = sm[i0:i1]
            bins = np.clip(np.searchsorted(roc_edges, vals, side="right") - 1, 0, len(roc_edges)-2)
            for t in range(i1-i0):
                yy, xx = np.where(v[t])
                np.add.at(pos_hist[wi], (bins[t, yy, xx], yy, xx), ob[t, yy, xx].astype("u4"))
                np.add.at(neg_hist[wi], (bins[t, yy, xx], yy, xx), (~ob[t, yy, xx]).astype("u4"))
        paired += 1
        if k % 100 == 0:
            print(f"Processed {k}/{len(files)} initializations", flush=True)

    names = []
    fields = []
    for wi, a in enumerate(acc):
        out = {}
        for label, p in (("daily", ""), ("weekly_mean", "w_"), ("weekly_max", "x_"), ("ehd_count", "c_")):
            out[f"tx_me_{label}" if label != "ehd_count" else "ehd_count_me"] = safe_div(a[p+"sum"], a[p+"n"])
            out[f"tx_mae_{label}" if label != "ehd_count" else "ehd_count_mae"] = safe_div(a[p+"abs"], a[p+"n"])
            out[f"tx_rmse_{label}" if label != "ehd_count" else "ehd_count_rmse"] = np.sqrt(safe_div(a[p+"sq"], a[p+"n"]))
        h, m, f, c = a["hit"], a["miss"], a["fa"], a["cn"]
        total = h + m + f + c
        out.update(pod=safe_div(h, h+m), pofd=safe_div(f, f+c), far=safe_div(f, h+f),
                   fbi=safe_div(h+f, h+m), csi=safe_div(h, h+m+f))
        random_hits = safe_div((h+m)*(h+f), total)
        out["ets"] = safe_div(h-random_hits, h+m+f-random_hits)
        # Histogram AUC via trapezoidal integration of threshold ROC curve.
        tp = np.cumsum(pos_hist[wi, ::-1], axis=0)
        fp = np.cumsum(neg_hist[wi, ::-1], axis=0)
        tpr = safe_div(tp, tp[-1]); fpr = safe_div(fp, fp[-1])
        out["auc"] = np.trapz(tpr, fpr, axis=0)
        for n, x in out.items():
            x = np.where(mask, x, np.nan).astype("f4")
            if wi == 0: names.append(n)
            fields.append((wi, n, x))

    nc_path = os.path.join(args.output_dir, "ehd_verification.nc")
    with Dataset(nc_path, "w", format="NETCDF4") as d:
        d.createDimension("week", 6); d.createDimension("lat", len(lat)); d.createDimension("lon", len(lon))
        d.createVariable("week", "i1", ("week",))[:] = np.arange(1, 7)
        d.createVariable("lat", "f4", ("lat",))[:] = lat
        d.createVariable("lon", "f4", ("lon",))[:] = lon
        d.setncattr("paired_initializations", paired)
        d.setncattr("threshold_definition", "max(day-of-year 90th percentile, 35 degC)")
        for n in names:
            v = d.createVariable(n, "f4", ("week", "lat", "lon"), zlib=True, fill_value=-9999.)
            stack = np.stack([x for w, nn, x in fields if nn == n])
            v[:] = np.ma.masked_invalid(stack)

    titles = {
        "tx_me_daily":"Tx daily ME (°C)", "tx_mae_daily":"Tx daily MAE (°C)", "tx_rmse_daily":"Tx daily RMSE (°C)",
        "tx_me_weekly_mean":"Weekly-mean Tx ME (°C)", "tx_mae_weekly_mean":"Weekly-mean Tx MAE (°C)", "tx_rmse_weekly_mean":"Weekly-mean Tx RMSE (°C)",
        "tx_me_weekly_max":"Weekly-max Tx ME (°C)", "tx_mae_weekly_max":"Weekly-max Tx MAE (°C)", "tx_rmse_weekly_max":"Weekly-max Tx RMSE (°C)",
        "ehd_count_me":"EHD-count ME (days)", "ehd_count_mae":"EHD-count MAE (days)", "ehd_count_rmse":"EHD-count RMSE (days)",
        "pod":"POD", "pofd":"POFD", "far":"FAR", "fbi":"Frequency Bias Index", "csi":"CSI", "ets":"ETS", "auc":"ROC AUC"
    }
    html = ["<!doctype html><html><head><meta charset='utf-8'><title>EHD verification</title>",
            "<style>body{font:15px system-ui;max-width:1180px;margin:auto;padding:24px;color:#18212b}select{font-size:16px;padding:6px}.grid{display:grid;grid-template-columns:repeat(auto-fit,minmax(430px,1fr));gap:14px}.card{border:1px solid #ccd;padding:8px}.card img{width:100%}small{color:#53606d}</style></head><body>",
            "<h1>ECMWF S2S extreme-heat verification over Vietnam</h1>",
            f"<p>Period: 2003–2022; paired initializations: {paired}. Threshold: max(model/observed day-of-year Q90, 35 °C). All calculations and maps use the supplied Vietnam mask.</p>",
            "<p><label>Forecast week: <select id='week' onchange='showWeek(this.value)'>" + "".join(f"<option value='{i}'>W{i}</option>" for i in range(1,7)) + "</select></label></p>"]
    by_week = {i: [] for i in range(6)}
    for wi, n, x in fields:
        nonneg = n not in {"tx_me_daily","tx_me_weekly_mean","tx_me_weekly_max","ehd_count_me","ets"}
        by_week[wi].append((titles[n], map_figure(x, lon, lat, f"W{wi+1}: {titles[n]}", "viridis" if nonneg else "RdBu_r", nonneg)))
    for wi in range(6):
        display = "grid" if wi == 0 else "none"
        html.append(f"<section class='grid wk' id='w{wi+1}' style='display:{display}'>")
        for title, uri in by_week[wi]: html.append(f"<div class='card'><img alt='{title}' src='{uri}'></div>")
        # Reliability diagram for deterministic probabilities 0/1.
        obsfreq = safe_div(rel_ob[wi], rel_fc[wi])
        fig, ax = plt.subplots(figsize=(5,4)); ax.plot([0,1],[0,1],'k--',label='Perfect')
        ax.plot([0,1], obsfreq, 'o-', label='EC deterministic'); ax.set(xlim=(-.03,1.03),ylim=(-.03,1.03),xlabel='Forecast probability',ylabel='Observed frequency',title=f'W{wi+1}: Reliability diagram'); ax.grid(alpha=.25); ax.legend()
        html.append(f"<div class='card'><img alt='Reliability diagram' src='{png_uri(fig)}'><small>Binary deterministic forecasts supply probability bins 0 and 1.</small></div></section>")
    html += ["<h2>Method notes</h2><p>Daily metrics pool matched lead days. Weekly-mean and weekly-maximum errors are first reduced within each initialization, then verified across initializations. EHD-count errors compare weekly counts. POD, POFD, FAR, FBI, CSI and ETS use daily contingency tables. ROC/AUC uses continuous model severity as the discrimination score (121 bins from −15 to 15 °C); reliability uses the only probabilities available from a deterministic product, 0 and 1.</p>",
             "<script>function showWeek(n){document.querySelectorAll('.wk').forEach(x=>x.style.display='none');document.getElementById('w'+n).style.display='grid'}</script></body></html>"]
    with open(args.html_output, "w", encoding="utf-8") as f: f.write("".join(html))
    print(f"Wrote {nc_path} and {args.html_output} from {paired} initializations")


if __name__ == "__main__":
    main()