#!/usr/bin/env python3
"""Compute and plot the Stage-1 monthly climatologies for Dr_SFC."""

from __future__ import annotations

import argparse
import html
import logging
import math
import sys
from datetime import datetime, timezone
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.colors import BoundaryNorm
import numpy as np
import shapefile
from scipy.io import netcdf_file


FILL = np.float32(-99.0)
RAIN_LEVELS = np.array([5, 10, 15, 20, 30, 50, 70, 100, 150, 200, 250, 300, 350, 400, 500], dtype=float)
RCMS = ("RegCM", "clWRF")
LEADS = range(1, 7)
MONTHS = range(1, 13)


def configure_log(path: Path) -> logging.Logger:
    path.parent.mkdir(parents=True, exist_ok=True)
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)s %(message)s",
        handlers=[logging.FileHandler(path, encoding="utf-8"), logging.StreamHandler(sys.stdout)],
    )
    return logging.getLogger("dr_sfc.climatology")


def as_lon_lat_time(ds, name: str) -> np.ndarray:
    var = ds.variables[name]
    dims = tuple(var.dimensions)
    aliases = {"lon": "lon", "longitude": "lon", "lat": "lat", "latitude": "lat", "mon": "time", "time": "time"}
    canonical = tuple(aliases.get(d.lower(), d.lower()) for d in dims)
    if set(canonical) != {"lon", "lat", "time"}:
        raise ValueError(f"{name}: unsupported dimensions {dims}")
    raw = var[:]
    data = np.ma.filled(raw, np.nan).astype(np.float64)
    data[np.isclose(data, -99.0)] = np.nan
    order = [canonical.index(k) for k in ("lon", "lat", "time")]
    return np.transpose(data, order)


def monthly_mean(data: np.ndarray) -> np.ndarray:
    if data.shape[-1] != 480:
        raise ValueError(f"Expected 480 months, found {data.shape[-1]}")
    out = np.full(data.shape[:2] + (12,), np.nan, dtype=np.float64)
    with np.errstate(invalid="ignore"):
        for month0 in range(12):
            out[..., month0] = np.nanmean(data[..., month0::12], axis=-1)
    return out


def load_mask(path: Path, lon: np.ndarray, lat: np.ndarray) -> np.ndarray:
    # This mask is NetCDF3 classic. scipy avoids a site HDF5-close defect.
    with netcdf_file(path, "r", mmap=False) as ds:
        reg = ds.variables["reg"]
        dims = tuple(d.decode() if isinstance(d, bytes) else d for d in reg.dimensions)
        dims = tuple(d.lower() for d in dims)
        arr = np.array(reg[:], copy=True)
        if dims == ("lat", "lon"):
            arr = arr.T
        elif dims != ("lon", "lat"):
            raise ValueError(f"reg: unsupported dimensions {reg.dimensions}")
        mlon = np.array(ds.variables["lon"][:], copy=True)
        mlat = np.array(ds.variables["lat"][:], copy=True)
    if arr.shape != (lon.size, lat.size) or not np.allclose(mlon, lon) or not np.allclose(mlat, lat):
        raise ValueError("Mask grid does not match input grid")
    return (arr >= 1) & (arr <= 7)


def compute(input_root: Path, mask_file: Path, log: logging.Logger):
    model = {rcm: {"pr": [], "t2": []} for rcm in RCMS}
    obs_ref = None
    lon_ref = lat_ref = None
    source_files = []
    for lead in LEADS:
        for rcm in RCMS:
            path = input_root / f"lead{lead}_{rcm}_01_BL_Test_VN.nc"
            if not path.is_file():
                raise FileNotFoundError(path)
            log.info("Reading lead=%d RCM=%s: %s", lead, rcm, path)
            source_files.append(str(path))
            with netcdf_file(path, "r", mmap=False) as ds:
                lon = np.asarray(ds.variables["lon"][:], dtype=np.float32)
                lat = np.asarray(ds.variables["lat"][:], dtype=np.float32)
                if lon_ref is None:
                    lon_ref, lat_ref = lon, lat
                elif not (np.array_equal(lon, lon_ref) and np.array_equal(lat, lat_ref)):
                    raise ValueError(f"Grid mismatch in {path}")
                current_obs = {v: monthly_mean(as_lon_lat_time(ds, f"{v}_o")) for v in ("pr", "t2")}
                if obs_ref is None:
                    obs_ref = current_obs
                else:
                    for v in ("pr", "t2"):
                        if not np.allclose(obs_ref[v], current_obs[v], equal_nan=True):
                            raise ValueError(f"Observed {v}_o differs in {path}")
                for v in ("pr", "t2"):
                    model[rcm][v].append(monthly_mean(as_lon_lat_time(ds, v)))
                    log.info("Computed %s climatology for lead=%d RCM=%s", v, lead, rcm)
    mask = load_mask(mask_file, lon_ref, lat_ref)
    for v in ("pr", "t2"):
        obs_ref[v] = np.where(mask[..., None], obs_ref[v], np.nan)
        for rcm in RCMS:
            model[rcm][v] = np.where(
                mask[..., None, None], np.stack(model[rcm][v], axis=2), np.nan
            )
    return lon_ref, lat_ref, obs_ref, model, source_files


def write_netcdf(path: Path, lon, lat, obs, model, rcm, sources, log):
    path.parent.mkdir(parents=True, exist_ok=True)
    with netcdf_file(path, "w", version=1) as ds:
        ds.createDimension("lon", lon.size)
        ds.createDimension("lat", lat.size)
        ds.createDimension("lead", 6)
        ds.createDimension("month", 12)
        ds.title = f"Dr_SFC {rcm} monthly climatology (1982-2021)"
        ds.rcm = rcm
        ds.history = f"Created {datetime.now(timezone.utc).isoformat()} by climatology_pipeline.py"
        ds.source_files = "\n".join(s for s in sources if f"_{rcm}_" in s)
        for name, vals, units, long_name in (
            ("lon", lon, "degrees_east", "longitude"), ("lat", lat, "degrees_north", "latitude"),
            ("lead", np.arange(1, 7), "month", "forecast lead"), ("month", np.arange(1, 13), "1", "calendar month"),
        ):
            var = ds.createVariable(name, "f", (name,))
            var[:] = vals
            var.units, var.long_name = units, long_name
        definitions = (
            (f"pr_cli_{rcm}", model["pr"], ("lon", "lat", "lead", "month"), "mm/month", f"{rcm} monthly precipitation climatology"),
            (f"t2_cli_{rcm}", model["t2"], ("lon", "lat", "lead", "month"), "degree_Celsius", f"{rcm} monthly mean temperature climatology"),
            ("pr_o_cli", obs["pr"], ("lon", "lat", "month"), "mm/month", "Observed monthly precipitation climatology"),
            ("t2_o_cli", obs["t2"], ("lon", "lat", "month"), "degree_Celsius", "Observed monthly mean temperature climatology"),
        )
        for name, values, dims, units, long_name in definitions:
            var = ds.createVariable(name, "f", dims)
            var._FillValue = FILL
            var.missing_value = FILL
            var.units, var.long_name = units, long_name
            var[:] = np.where(np.isfinite(values), values, FILL).astype(np.float32)
    log.info("Wrote %s", path)


def shape_segments(path: Path):
    reader = shapefile.Reader(str(path))
    segments = []
    for shape in reader.shapes():
        points = np.asarray(shape.points)
        stops = list(shape.parts) + [len(points)]
        segments.extend(points[stops[i]:stops[i + 1]] for i in range(len(stops) - 1))
    return segments


def plot_sheet(values, lon, lat, boundary, levels, cmap, title, unit, output, log):
    fig, axes = plt.subplots(2, 6, figsize=(11.69, 8.27), constrained_layout=False)
    fig.subplots_adjust(left=.045, right=.985, bottom=.13, top=.91, wspace=.18, hspace=.20)
    norm = BoundaryNorm(levels, plt.get_cmap(cmap).N, clip=True)
    mesh = None
    for i, ax in enumerate(axes.flat):
        mesh = ax.pcolormesh(lon, lat, values[..., i].T, cmap=cmap, norm=norm, shading="nearest", rasterized=True)
        for seg in boundary:
            ax.plot(seg[:, 0], seg[:, 1], color="black", linewidth=.35)
        ax.set_title(f"Tháng {i + 1}", fontsize=8)
        ax.set_xlim(101.125, 110.875); ax.set_ylim(7.625, 24.375)
        ax.set_xticks([102, 104, 106, 108, 110]); ax.set_yticks([8, 12, 16, 20, 24])
        ax.tick_params(labelsize=5, length=2)
        ax.grid(color="0.65", linewidth=.25, linestyle=":")
        if i < 6: ax.set_xticklabels([])
        if i % 6: ax.set_yticklabels([])
    fig.suptitle(title, fontsize=13, fontweight="bold")
    cax = fig.add_axes([.08, .055, .84, .025])
    cb = fig.colorbar(mesh, cax=cax, orientation="horizontal", ticks=levels, extend="both")
    cb.ax.tick_params(labelsize=6, rotation=45 if len(levels) > 16 else 0)
    cb.set_label(unit, fontsize=8)
    output.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(output, dpi=180, metadata={"Title": title})
    plt.close(fig)
    log.info("Rendered %s", output)


def make_html(path: Path, figures, nc_paths, log_path: Path):
    cards = "\n".join(
        f'<section><h2>{html.escape(title)}</h2><p>{html.escape(note)}</p><a href="figures/{img.name}"><img src="figures/{img.name}" alt="{html.escape(title)}"></a></section>'
        for title, note, img in figures
    )
    content = f"""<!doctype html><html lang="vi"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width">
<title>Dr_SFC – Trung bình khí hậu tháng</title><style>body{{font:16px system-ui;margin:auto;max-width:1200px;padding:24px;color:#18212b}}img{{width:100%;border:1px solid #ccd3da}}section{{margin:2rem 0}}code{{background:#eef1f4;padding:.15rem .3rem}}</style></head>
<body><h1>Dr_SFC – Trung bình khí hậu tháng (1982–2021)</h1><p>Kết quả tính từ 12 file NetCDF thật (2 RCM × 6 lead). Khí hậu mô hình được tính riêng cho RegCM và clWRF; vùng ngoài Việt Nam được mask.</p>
<p>NetCDF RegCM: <code>{html.escape(str(nc_paths["RegCM"]))}</code><br>NetCDF clWRF: <code>{html.escape(str(nc_paths["clWRF"]))}</code><br>Log: <code>{html.escape(str(log_path))}</code></p>{cards}</body></html>"""
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(content, encoding="utf-8")


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--project-root", type=Path, default=Path(__file__).resolve().parent)
    p.add_argument("--input-root", type=Path, default=Path("/home/tanpv/Data/562_VNU/0.25_VN/For_BC"))
    p.add_argument("--mask-file", type=Path, default=Path("/home/tanpv/Progs/Dr_VN_Ope/DATA/Mask_VN_0.25_Corrected.nc"))
    p.add_argument("--shape-file", type=Path, default=Path("/home/tanpv/GrADS_SCR/Vietnam7.shp"))
    args = p.parse_args()
    if not args.shape_file.exists():
        fallback = Path("/home/tanpv/GrADS_SCR/VN_SubRegs1/Vietnam7.shp")
        if fallback.exists(): args.shape_file = fallback
    stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    log_path = args.project_root / "logs" / f"climatology_{stamp}.log"
    log = configure_log(log_path)
    log.info("Starting climatology pipeline")
    lon, lat, obs, model, sources = compute(args.input_root, args.mask_file, log)
    nc_paths = {}
    for rcm in RCMS:
        nc_paths[rcm] = args.project_root / "Clim" / f"Monthly_Clim_{rcm}.nc"
        write_netcdf(nc_paths[rcm], lon, lat, obs, model[rcm], rcm, sources, log)
    boundary = shape_segments(args.shape_file)
    all_t = np.concatenate([obs["t2"].ravel()] + [model[rcm]["t2"].ravel() for rcm in RCMS])
    tlevels = np.arange(math.floor(np.nanmin(all_t)), math.ceil(np.nanmax(all_t)) + 1, 1)
    fig_dir = args.project_root / "outputs" / "Clim" / "figures"
    figures = []
    for var, unit, levels, cmap in (("t2", "°C", tlevels, "turbo"), ("pr", "mm/tháng", RAIN_LEVELS, "YlGnBu")):
        out = fig_dir / f"Observed_Clim_{var}_m.png"
        title = f"Khí hậu quan trắc – {'nhiệt độ' if var == 't2' else 'lượng mưa'}"
        plot_sheet(obs[var], lon, lat, boundary, levels, cmap, title, unit, out, log)
        figures.append((title, "12 tháng; cùng một thang màu giữa các panel.", out))
        for rcm in RCMS:
            for lead in LEADS:
                out = fig_dir / f"{rcm}_Clim_{var}_m_lead_{lead}.png"
                title = f"Khí hậu {rcm} – {'nhiệt độ' if var == 't2' else 'lượng mưa'} – lead {lead}"
                plot_sheet(model[rcm][var][..., lead - 1, :], lon, lat, boundary, levels, cmap, title, unit, out, log)
                figures.append((title, f"{rcm}; 12 tháng; cùng một thang màu.", out))
    html_path = args.project_root / "outputs" / "Clim" / "index.html"
    make_html(html_path, figures, nc_paths, log_path)
    log.info("Wrote HTML report %s", html_path)
    log.info("Pipeline completed successfully")


if __name__ == "__main__":
    main()