#!/usr/bin/env python3
"""Plot drought-occurrence probability maps from Dr_Prob.nc.

This script reads the probability variables
    mild, mord, seve, extr
from Dr_Prob.nc and produces 12 figures in total:
    4 drought levels x 3 timescales (1, 3, 6 months).
Each figure contains 6 lead-month panels arranged in one row.

Main characteristics
--------------------
* GrADS-like grid-cell filling using pcolormesh (grfill style).
* Vietnam outline and 7 climate-region boundaries from Vietnam7.shp.
* "Biển Đông" label on every panel.
* Title format:
    Probability of 3-Month <level> Drought Occurrence initialized at dd-mm-yyyy
* White-to-dark-red probability color scale from 0 to 100 percent.
"""

from __future__ import annotations

import argparse
import calendar
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.cm import ScalarMappable
from matplotlib.colors import BoundaryNorm, LinearSegmentedColormap
from matplotlib.ticker import FuncFormatter, FixedLocator
import numpy as np
import shapefile  # pyshp
import xarray as xr

try:
    import shapely
    from shapely.geometry import shape as shapely_shape
    from shapely.ops import unary_union
except Exception as exc:  # pragma: no cover
    raise SystemExit(
        "This script requires shapely (>=2 recommended). "
        f"Import failed: {exc}"
    )


# -----------------------------------------------------------------------------
# Probability color scale: 0-100 % from white to dark red.
# -----------------------------------------------------------------------------
PROB_LEVELS = np.arange(0.0, 110.0, 10.0)  # 0, 10, ..., 100
PROB_CMAP = LinearSegmentedColormap.from_list(
    "white_to_darkred",
    [
        "#ffffff",
        "#fff5f0",
        "#fee0d2",
        "#fcbba1",
        "#fc9272",
        "#fb6a4a",
        "#ef3b2c",
        "#cb181d",
        "#a50f15",
        "#67000d",
    ],
    N=256,
)
PROB_NORM = BoundaryNorm(PROB_LEVELS, ncolors=PROB_CMAP.N, clip=True)

LEVEL_LABELS = {
    "mild": "Mild",
    "mord": "Moderate",
    "seve": "Severe",
    "extr": "Extreme Severe",
}


@dataclass(frozen=True)
class RegionGeometry:
    parts: list
    union: object
    bounds: tuple[float, float, float, float]


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="Plot drought-occurrence probabilities from Dr_Prob.nc."
    )
    p.add_argument("--year", type=int, required=True, help="Initial year, e.g. 2026")
    p.add_argument("--month", type=int, required=True, choices=range(1, 13), metavar="1..12",
                   help="Initial month, e.g. 7")
    p.add_argument("--day", type=int, required=True, help="Initial day, e.g. 1")
    p.add_argument("--input", required=True, help="Input NetCDF file, e.g. Dr_Prob.nc")
    p.add_argument("--output-dir", required=True,
                   help="Output directory, e.g. /home/.../PDF/20260701")
    p.add_argument("--shape", required=True,
                   help="Climate-region shapefile, e.g. VN_SubRegs/Vietnam7.shp")
    p.add_argument("--nlead", type=int, default=6, help="Number of lead months to plot (default: 6)")
    p.add_argument("--lon-min", type=float, default=101.0)
    p.add_argument("--lon-max", type=float, default=111.0)
    p.add_argument("--lat-min", type=float, default=7.5)
    p.add_argument("--lat-max", type=float, default=24.0)
    p.add_argument("--dpi", type=int, default=200)
    p.add_argument("--mask-country", action="store_true",
                   help=("Optionally mask grid-cell centres outside the union of Vietnam7 polygons. "
                         "Default is OFF to preserve the original grid cells like GrADS grfill."))
    return p.parse_args()


def add_months(year: int, month: int, delta: int) -> tuple[int, int]:
    idx = year * 12 + (month - 1) + delta
    return idx // 12, idx % 12 + 1


def _repair_geom(g):
    if g.is_empty:
        return g
    if not g.is_valid:
        try:
            return shapely.make_valid(g)
        except Exception:
            return g.buffer(0)
    return g


def load_regions(shp_path: str) -> RegionGeometry:
    shp = Path(shp_path)
    if shp.suffix.lower() != ".shp":
        shp = shp.with_suffix(".shp")
    if not shp.exists():
        raise FileNotFoundError(f"Shapefile not found: {shp}")

    reader = shapefile.Reader(str(shp))
    parts = []
    for rec_shape in reader.shapes():
        g = _repair_geom(shapely_shape(rec_shape.__geo_interface__))
        if not g.is_empty:
            parts.append(g)

    if not parts:
        raise ValueError(f"No polygon geometry found in {shp}")

    country = _repair_geom(unary_union(parts))
    bounds = tuple(float(v) for v in country.bounds)

    minx, miny, maxx, maxy = bounds
    if not (90 <= minx <= 120 and 90 <= maxx <= 130 and 0 <= miny <= 30 and 0 <= maxy <= 35):
        raise ValueError(
            "Vietnam7.shp does not appear to be in geographic lon/lat coordinates. "
            f"Detected bounds: {bounds}. Reproject it to lon/lat first."
        )
    return RegionGeometry(parts=parts, union=country, bounds=bounds)


def iter_polygons(geom) -> Iterable:
    if geom.is_empty:
        return
    gt = geom.geom_type
    if gt == "Polygon":
        yield geom
    elif gt in ("MultiPolygon", "GeometryCollection"):
        for child in geom.geoms:
            yield from iter_polygons(child)


def draw_polygon_boundary(ax, geom, linewidth=0.55, color="black", zorder=7):
    for poly in iter_polygons(geom):
        x, y = poly.exterior.xy
        ax.plot(x, y, color=color, linewidth=linewidth, zorder=zorder)
        for ring in poly.interiors:
            x, y = ring.xy
            ax.plot(x, y, color=color, linewidth=linewidth, zorder=zorder)


def draw_boundaries(ax, regions: RegionGeometry):
    for g in regions.parts:
        draw_polygon_boundary(ax, g, linewidth=0.35, color="#666666", zorder=7)
    draw_polygon_boundary(ax, regions.union, linewidth=0.80, color="black", zorder=8)


def mask_to_country(z: np.ndarray, lon: np.ndarray, lat: np.ndarray, country_geom) -> np.ndarray:
    xx, yy = np.meshgrid(lon, lat)
    try:
        inside = shapely.intersects_xy(country_geom, xx, yy)
    except Exception:
        inside = shapely.contains_xy(country_geom, xx, yy)
    return np.where(inside, z, np.nan)


def centers_to_edges(coord: np.ndarray) -> np.ndarray:
    c = np.asarray(coord, dtype=float)
    if c.ndim != 1 or c.size < 2:
        raise ValueError("Coordinate must be a 1-D array with at least two points")
    mid = 0.5 * (c[:-1] + c[1:])
    edges = np.empty(c.size + 1, dtype=float)
    edges[1:-1] = mid
    edges[0] = c[0] - 0.5 * (c[1] - c[0])
    edges[-1] = c[-1] + 0.5 * (c[-1] - c[-2])
    return edges


def get_2d(ds: xr.Dataset, varname: str, ts_index: int, lead_index: int) -> np.ndarray:
    if varname not in ds:
        raise KeyError(f"Variable '{varname}' not found. Available: {list(ds.data_vars)}")
    da = ds[varname]
    needed = {"lon", "lat", "timescale", "lead"}
    if not needed.issubset(da.dims):
        raise ValueError(f"Variable '{varname}' must contain dims {needed}; got {da.dims}")

    da = da.isel(timescale=ts_index, lead=lead_index)
    da = da.transpose("lat", "lon")
    arr = np.asarray(da.values, dtype=float)

    for attr in ("_FillValue", "missing_value"):
        val = da.attrs.get(attr)
        if val is not None and np.isfinite(val):
            arr[np.isclose(arr, float(val))] = np.nan
    arr[~np.isfinite(arr)] = np.nan
    return arr


def lon_formatter(x, pos=None):
    return f"{abs(x):g}°{'E' if x >= 0 else 'W'}"


def lat_formatter(y, pos=None):
    return f"{abs(y):g}°{'N' if y >= 0 else 'S'}"


def setup_axis(ax, args, show_y: bool):
    ax.set_xlim(args.lon_min, args.lon_max)
    ax.set_ylim(args.lat_min, args.lat_max)
    ax.set_aspect("equal", adjustable="box")
    ax.xaxis.set_major_locator(FixedLocator([102, 104, 106, 108, 110]))
    ax.xaxis.set_major_formatter(FuncFormatter(lon_formatter))
    ax.yaxis.set_major_locator(FixedLocator(np.arange(8, 25, 2)))
    ax.yaxis.set_major_formatter(FuncFormatter(lat_formatter))
    if not show_y:
        ax.set_yticklabels([])
    ax.tick_params(axis="both", labelsize=6.8, length=2.5, pad=1.5)
    for spine in ax.spines.values():
        spine.set_linewidth(0.75)
        spine.set_color("black")


def plot_panel(ax, lon, lat, z, args, regions, month_label, show_y):
    setup_axis(ax, args, show_y)
    ax.set_facecolor("white")

    lon_edge = centers_to_edges(lon)
    lat_edge = centers_to_edges(lat)
    pm = ax.pcolormesh(
        lon_edge, lat_edge, z,
        cmap=PROB_CMAP,
        norm=PROB_NORM,
        shading="flat",
        edgecolors="none",
        linewidth=0.0,
        antialiased=False,
        zorder=1,
    )

    # Optional thin contour hints to help distinguish broad probability zones.
    try:
        ax.contour(
            lon, lat, z, levels=[20, 40, 60, 80], colors="#6e1f1f",
            linewidths=0.25, alpha=0.45, zorder=3
        )
    except Exception:
        pass

    draw_boundaries(ax, regions)
    ax.text(
        109.10, 17.45, "Biển Đông",
        rotation=-55,
        rotation_mode="anchor",
        ha="center", va="center",
        fontsize=6.8, fontstyle="italic", color="#4f4f4f",
        zorder=10, clip_on=False,
    )
    ax.set_title(month_label, fontsize=8.1, pad=3)
    return pm


def make_figure(ds: xr.Dataset, lon: np.ndarray, lat: np.ndarray, regions: RegionGeometry,
                args: argparse.Namespace, varname: str, level_label: str,
                ts_index: int, ts_value: int):
    nlead = args.nlead
    fig = plt.figure(figsize=(12.9, 4.45))
    gs = fig.add_gridspec(
        1, nlead,
        left=0.042, right=0.988, bottom=0.165, top=0.84,
        wspace=0.018, hspace=0.0,
    )
    axes = [fig.add_subplot(gs[0, c]) for c in range(nlead)]

    last_pm = None
    for i in range(nlead):
        ty, tm = add_months(args.year, args.month, i + 1)
        target_label = f"{tm}-{ty}"
        z = get_2d(ds, varname, ts_index, i)
        if args.mask_country:
            z = mask_to_country(z, lon, lat, regions.union)
        z = np.clip(z, 0.0, 100.0)

        last_pm = plot_panel(axes[i], lon, lat, z, args, regions, target_label, show_y=(i == 0))

    init_label = f"{args.day:02d}-{args.month:02d}-{args.year:04d}"
    fig.text(
        0.50, 0.905,
        f"Probability of {ts_value}-Month {level_label} Drought Occurrence initialized at {init_label}",
        ha="center", va="center", fontsize=10.3
    )

    cax = fig.add_axes([0.24, 0.075, 0.52, 0.042])
    sm = ScalarMappable(norm=PROB_NORM, cmap=PROB_CMAP)
    sm.set_array([])
    cb = fig.colorbar(sm, cax=cax, orientation="horizontal", extend="neither")
    cb.set_ticks(PROB_LEVELS)
    cb.ax.tick_params(labelsize=7, length=2, pad=1.5)
    cb.outline.set_linewidth(0.7)
    cb.ax.set_xlabel("Probability (%)", fontsize=7.8, labelpad=2)

    out_dir = Path(args.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    out = out_dir / f"Dr_Prob_{varname}_{ts_value}month.png"
    fig.savefig(out, dpi=args.dpi, bbox_inches="tight", facecolor="white")
    plt.close(fig)
    print(f"Wrote: {out}")


def validate_input(ds: xr.Dataset, args: argparse.Namespace):
    for coord in ("lon", "lat", "timescale", "lead"):
        if coord not in ds:
            raise KeyError(f"Coordinate '{coord}' not found in input NetCDF")
    for v in ("mild", "mord", "seve", "extr"):
        if v not in ds:
            raise KeyError(f"Required variable '{v}' not found in input NetCDF")

    nlead_file = int(ds.sizes.get("lead", 0))
    nts_file = int(ds.sizes.get("timescale", 0))
    if nlead_file < args.nlead:
        raise ValueError(f"Input contains only {nlead_file} lead(s), but --nlead={args.nlead}")
    if nts_file < 1:
        raise ValueError("No timescale found in input NetCDF")


def main() -> int:
    args = parse_args()

    _, ndays = calendar.monthrange(args.year, args.month)
    if not 1 <= args.day <= ndays:
        raise ValueError(f"Invalid initial date: {args.year}-{args.month:02d}-{args.day:02d}")

    inp = Path(args.input)
    if not inp.exists():
        raise FileNotFoundError(f"Input NetCDF not found: {inp}")

    regions = load_regions(args.shape)

    with xr.open_dataset(inp, mask_and_scale=True, decode_times=True) as ds:
        validate_input(ds, args)
        lon = np.asarray(ds["lon"].values, dtype=float)
        lat = np.asarray(ds["lat"].values, dtype=float)

        if np.any(np.diff(lon) < 0):
            ds = ds.sortby("lon")
            lon = np.asarray(ds["lon"].values, dtype=float)
        if np.any(np.diff(lat) < 0):
            ds = ds.sortby("lat")
            lat = np.asarray(ds["lat"].values, dtype=float)

        ts_values = np.asarray(ds["timescale"].values)
        ts_values = [int(round(float(v))) for v in ts_values]

        for varname, label in LEVEL_LABELS.items():
            for its, ts in enumerate(ts_values):
                make_figure(ds, lon, lat, regions, args, varname, label, its, ts)

    return 0


if __name__ == "__main__":
    try:
        sys.exit(main())
    except Exception as exc:
        print(f"ERROR: {exc}", file=sys.stderr)
        sys.exit(2)