#!/usr/bin/env python3
"""Build a self-contained EHD report with W1--W6 arranged by metric."""

from __future__ import annotations

import argparse
import base64
import io
from html import escape
from pathlib import Path

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

BASE = Path(__file__).resolve().parent


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",
}

SIGNED = {
    "tx_me_daily", "tx_me_weekly_mean", "tx_me_weekly_max",
    "ehd_count_me", "ets",
}


def load_subregions(path):
    """Read polygon parts and label positions from a shapefile/DBF basename."""
    reader = shapefile.Reader(path)
    name_index = next(
        (i for i, field in enumerate(reader.fields[1:]) if field.name.lower() == "name"),
        None,
    )
    regions = []
    for shape, record in zip(reader.shapes(), reader.records()):
        stops = list(shape.parts) + [len(shape.points)]
        parts = [np.asarray(shape.points[stops[i]:stops[i + 1]])
                 for i in range(len(stops) - 1)]
        xmin, ymin, xmax, ymax = shape.bbox
        label = str(record[name_index]) if name_index is not None else ""
        label_dx = {"S5": 0.80, "S6": 0.35}.get(label, 0.0)
        x = (xmin + xmax) / 2 + label_dx
        y = (ymin + ymax) / 2
        regions.append((parts, label, (x, y)))
    return regions


def draw_subregions(ax, regions):
    """Overlay sub-region outlines and identifiers without obscuring map values."""
    for parts, label, (x, y) in regions:
        for points in parts:
            ax.plot(points[:, 0], points[:, 1], color="black", linewidth=.65,
                    solid_joinstyle="round", zorder=3)
        if label:
            ax.text(x, y, label, ha="center", va="center", fontsize=10,
                    fontweight="bold", color="black", zorder=4,
                    bbox={"boxstyle": "round,pad=.12", "facecolor": "white",
                          "edgecolor": "none", "alpha": .68})


def row_image(values, lon, lat, title, signed, subregions):
    """Draw all six weeks using one color scale for this metric."""
    finite = values[np.isfinite(values)]
    if signed:
        limit = np.nanpercentile(np.abs(finite), 98) if finite.size else 1.0
        limit = limit if limit > 0 else 1.0
        vmin, vmax, cmap = -limit, limit, "RdBu_r"
    else:
        vmax = np.nanpercentile(finite, 98) if finite.size else 1.0
        vmax = vmax if vmax > 0 else 1.0
        vmin, cmap = 0.0, "viridis"

    fig, axes = plt.subplots(1, 6, figsize=(18, 6.35), sharex=True, sharey=True)
    image = None
    for wi, ax in enumerate(axes):
        image = ax.pcolormesh(
            lon, lat, values[wi], shading="auto", cmap=cmap,
            vmin=vmin, vmax=vmax,
        )
        draw_subregions(ax, subregions)
        ax.set_title(f"W{wi + 1}", fontweight="bold")
        ax.set_xlabel("Longitude (°E)")
        ax.set_aspect("equal", adjustable="box")
        if wi == 0:
            ax.set_ylabel("Latitude (°N)")
        ax.tick_params(labelsize=8)
    fig.suptitle(title, fontsize=15, fontweight="bold", y=.99)
    fig.colorbar(image, ax=axes, orientation="horizontal", fraction=.07,
                 pad=.16, aspect=55, label=title)
    fig.subplots_adjust(left=.045, right=.992, bottom=.23, top=.87, wspace=.02)
    buf = io.BytesIO()
    fig.savefig(buf, format="png", dpi=135, bbox_inches="tight")
    plt.close(fig)
    return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", default=str(BASE / "outputs" / "ehd_verification.nc"))
    parser.add_argument("--output", default=str(BASE / "ehd_verification_metric_rows.html"))
    parser.add_argument(
        "--subregions",
        default="/home/tanpv/GrADS_SCR/VN_SubRegs1/Vietnam7.dbf",
        help="Vietnam sub-region .dbf or .shp path (companion files are required)",
    )
    args = parser.parse_args()
    subregions = load_subregions(args.subregions)

    with Dataset(args.input) as data:
        missing = [name for name in TITLES if name not in data.variables]
        if missing:
            raise KeyError(f"Missing verification variables: {missing}")
        if data["week"].shape != (6,):
            raise ValueError(f"Expected six forecast weeks, got {data['week'].shape}")
        lon = np.asarray(data["lon"][:])
        lat = np.asarray(data["lat"][:])
        paired = getattr(data, "paired_initializations", "unknown")
        rows = []
        for name, title in TITLES.items():
            values = np.ma.filled(data[name][:], np.nan)
            rows.append((title, row_image(
                values, lon, lat, title, name in SIGNED, subregions,
            )))

    html = ["""<!doctype html><html><head><meta charset="utf-8">
<meta name="viewport" content="width=device-width,initial-scale=1">
<title>EHD verification — metrics by lead time</title>
<style>
body{font:15px system-ui,sans-serif;margin:0;color:#18212b;background:#f5f7fa}
header,main{max-width:1900px;margin:auto;padding:20px}
h1{margin-bottom:.35rem}.intro{margin-top:0;color:#53606d}
.row{background:#fff;border:1px solid #ccd4df;border-radius:8px;margin:0 0 18px;
 padding:10px;box-shadow:0 1px 3px #17202a12;overflow-x:auto}
.row img{display:block;width:100%;min-width:1250px;height:auto}
.note{color:#53606d;font-size:13px}
</style></head><body><header>
<h1>ECMWF S2S extreme-heat verification over Vietnam</h1>""",
        f'<p class="intro">Period: 2003–2022; paired initializations: {paired}. '
        'Each row shows W1–W6 for one metric. All six maps in a row share that '
        'metric’s color scale. Black outlines mark the seven Vietnam sub-regions (S1–S7).</p></header><main>']
    for title, uri in rows:
        html.append(f'<section class="row"><img alt="{escape(title)}: W1 to W6" src="{uri}"></section>')
    html.append("""<p class="note">Color limits are calculated independently for each metric
from the 98th percentile across all six weeks. Signed metrics use a symmetric diverging
scale; nonnegative metrics use a sequential scale.</p></main></body></html>""")

    with open(args.output, "w", encoding="utf-8") as out:
        out.write("".join(html))
    print(f"Wrote {args.output} with {len(rows)} metric rows")


if __name__ == "__main__":
    main()