#!/usr/bin/env python3
"""Plot operational precipitation/2-m temperature forecasts and anomalies.

This reproduces the layout and plotting logic of Plot_Pr_T2_fcs_anomaly.gs:
  * two output figures: Pr_forecast and T2_forecast;
  * each figure has 2 rows x 6 lead months;
  * row 1 = forecast, row 2 = anomaly;
  * target month = initial month + lead (1..6);
  * filled fields are drawn grid-cell by grid-cell (GrADS ``gxout grfill`` style),
    not by contour shading/interpolation;
  * Vietnam national outline and the seven climate-region boundaries are drawn
    from VN_SubRegs/Vietnam7.shp.

Expected NetCDF variables:
  pr_fcs(lon, lat, lead), t2_fcs(lon, lat, lead),
  pr_ano(lon, lat, lead), t2_ano(lon, lat, lead), lon, lat, lead.
"""

from __future__ import annotations

import argparse
import calendar
import os
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.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) to mask the grid to Vietnam. "
        f"Import failed: {exc}"
    )


# -----------------------------------------------------------------------------
# Plot settings copied from the GrADS script as closely as practical.
# -----------------------------------------------------------------------------
PR_LEVELS = np.array(
    [5, 10, 20, 30, 50, 70, 100, 120, 140, 160, 180, 200,
     240, 280, 320, 350, 400, 450, 500], dtype=float
)
PR_ANO_LEVELS = np.array(
    [-100, -80, -60, -40, -30, -20, -10, -5, 0, 5, 10, 20, 30, 40, 60, 80, 100],
    dtype=float,
)
T2_LEVELS = np.array([10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30, 32], dtype=float)
T2_ANO_LEVELS = np.array(
    [-2, -1.5, -1, -0.8, -0.6, -0.4, -0.2, 0, 0.2, 0.4, 0.6, 0.8, 1, 1.5, 2],
    dtype=float,
)

# Similar visual progressions to the GrADS color.gs settings.
CMAP_PR = plt.get_cmap("YlGnBu")
CMAP_PR_ANO = plt.get_cmap("RdYlBu")       # dry/red -> near-zero/white -> wet/blue
CMAP_T2 = plt.get_cmap("RdYlGn_r")         # cool/green -> warm/red
CMAP_T2_ANO = plt.get_cmap("RdYlGn_r")     # negative/green -> positive/red


@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 6-month Pr/T2 forecasts and anomalies in the layout of Plot_Pr_T2_fcs_anomaly.gs."
    )
    p.add_argument("--year", type=int, required=True, help="Operational initial year, e.g. 2026")
    p.add_argument("--month", type=int, required=True, choices=range(1, 13), metavar="1..12",
                   help="Operational initial month, e.g. 7")
    p.add_argument("--day", type=int, required=True, help="Operational initial day, e.g. 1")
    p.add_argument("--input", required=True, help="Input NetCDF file Ope_Pr_T2_and_Ano.nc")
    p.add_argument("--output-prefix", required=True,
                   help="Prefix exactly like GrADS Dat_out; e.g. /.../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("--pr-ano-scale", type=float, default=100.0,
                   help=("Scale applied to pr_ano before plotting. Default 100 reproduces the GrADS "
                         "expression pr_ano*100 and the Relative Rainfall Anomaly (%%) panel. "
                         "Use 1 if pr_ano is already stored in percent."))
    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)
    # GrADS gxout grfill displays the NetCDF grid cells as they are.  Therefore
    # no additional polygon masking is applied by default; missing values already
    # present in the NetCDF remain missing.  This is important near the coastline,
    # where a centre-point polygon mask can incorrectly remove valid coastal cells.
    p.add_argument("--mask-country", action="store_true",
                   help=("Optionally mask grid-cell centres outside the union of Vietnam7 polygons. "
                         "Default is OFF to reproduce GrADS grfill and retain coastal grid cells."))
    # Kept only for backward compatibility with older calling scripts.
    p.add_argument("--no-mask", action="store_true", help=argparse.SUPPRESS)
    return p.parse_args()


def add_months(year: int, month: int, delta: int) -> tuple[int, int]:
    """Return (year, month) after adding delta months."""
    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)

    # Vietnam7 is expected to be geographic lon/lat. Catch projected shapefiles early.
    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 (EPSG:4326-like). "
            f"Detected bounds: {bounds}. Reproject it to lon/lat first."
        )
    return RegionGeometry(parts=parts, union=country, bounds=bounds)


def iter_polygons(geom) -> Iterable:
    """Yield Polygon objects from Polygon/MultiPolygon/GeometryCollection."""
    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, zorder=7):
    for poly in iter_polygons(geom):
        x, y = poly.exterior.xy
        ax.plot(x, y, color="black", linewidth=linewidth, zorder=zorder)
        for ring in poly.interiors:
            x, y = ring.xy
            ax.plot(x, y, color="black", linewidth=linewidth, zorder=zorder)


def draw_boundaries(ax, regions: RegionGeometry):
    # Seven climate-region boundaries.
    for g in regions.parts:
        draw_polygon_boundary(ax, g, linewidth=0.45, zorder=7)
    # National outer boundary, slightly thicker.
    draw_polygon_boundary(ax, regions.union, linewidth=0.9, 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 get_2d(ds: xr.Dataset, varname: str, 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]
    if "lead" not in da.dims:
        raise ValueError(f"Variable '{varname}' has no 'lead' dimension: {da.dims}")
    da = da.isel(lead=lead_index)
    if not {"lat", "lon"}.issubset(da.dims):
        raise ValueError(f"Variable '{varname}' must contain lat/lon dimensions; got {da.dims}")
    da = da.transpose("lat", "lon")
    arr = np.asarray(da.values, dtype=float)

    # xarray normally decodes _FillValue/missing_value to NaN. This is an extra guard.
    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 discrete_norm(levels, cmap):
    # GrADS grfill assigns one color to each value interval. BoundaryNorm gives
    # pcolormesh the same discrete-bin behavior without spatial interpolation.
    return BoundaryNorm(levels, ncolors=cmap.N, clip=False)


def centers_to_edges(coord: np.ndarray) -> np.ndarray:
    """Convert 1-D grid-cell centres to cell edges for GrADS-like grfill.

    GrADS ``gxout grfill`` paints the full grid box associated with each grid
    point.  Matplotlib's pcolormesh reproduces that behavior when it is given
    cell edges rather than centre coordinates.  This routine also works for a
    non-uniform 1-D coordinate.
    """
    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 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([104, 108]))
    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=7, length=2.5, pad=1.5)
    for spine in ax.spines.values():
        spine.set_linewidth(0.8)
        spine.set_color("black")


def plot_panel(ax, lon, lat, z, levels, cmap, args, regions, month_label, show_y):
    setup_axis(ax, args, show_y)

    # GrADS equivalent: set gxout grfill
    # ----------------------------------
    # Do NOT use contourf here. contourf constructs filled contour polygons and
    # therefore interpolates spatially between grid points; near a coastline or
    # missing-value boundary this can suppress/reshape valid edge-grid information.
    # pcolormesh with explicit cell edges paints every source grid box directly.
    lon_edge = centers_to_edges(lon)
    lat_edge = centers_to_edges(lat)
    gf = ax.pcolormesh(
        lon_edge, lat_edge, z,
        cmap=cmap,
        norm=discrete_norm(levels, cmap),
        shading="flat",
        edgecolors="none",
        linewidth=0.0,
        antialiased=False,
        zorder=1,
    )

    # The original GrADS program plots a contour overlay after grfill
    # (Con_pr, Con_pr_ano, Con_t2, Con_t2_ano). Keep that overlay here; only the
    # color filling itself is grid-cell based.
    ax.contour(
        lon, lat, z, levels=levels, colors="black",
        linewidths=0.20, alpha=0.55, zorder=3
    )
    draw_boundaries(ax, regions)

    # Label the East Sea (Biển Đông), matching the placement/style of the
    # reference GrADS map.  Coordinates are in lon/lat degrees.
    ax.text(
        109.15, 17.45, "Biển Đông",
        rotation=-55,
        rotation_mode="anchor",
        ha="center", va="center",
        fontsize=7.5, fontweight="bold", color="black",
        zorder=9,
    )

    ax.set_title(month_label, fontsize=8.5, pad=3)
    return gf


def format_colorbar(cb, levels, fontsize=6.7):
    cb.set_ticks(levels)
    cb.ax.tick_params(labelsize=fontsize, length=2, pad=1.5)
    cb.outline.set_linewidth(0.7)


def make_figure(
    ds: xr.Dataset,
    lon: np.ndarray,
    lat: np.ndarray,
    regions: RegionGeometry,
    args: argparse.Namespace,
    var_top: str,
    var_bottom: str,
    levels_top: np.ndarray,
    levels_bottom: np.ndarray,
    cmap_top,
    cmap_bottom,
    title_top: str,
    title_bottom: str,
    out_suffix: str,
    bottom_scale: float = 1.0,
):
    nlead = args.nlead
    fig = plt.figure(figsize=(14.5, 8.0))
    gs = fig.add_gridspec(
        2, nlead + 1,
        width_ratios=[1.0] * nlead + [0.055],
        left=0.055, right=0.955, bottom=0.065, top=0.935,
        wspace=0.06, hspace=0.25,
    )

    axes = [[fig.add_subplot(gs[r, c]) for c in range(nlead)] for r in range(2)]
    cax_top = fig.add_subplot(gs[0, -1])
    cax_bottom = fig.add_subplot(gs[1, -1])

    last_cf_top = None
    last_cf_bottom = None
    for i in range(nlead):
        ty, tm = add_months(args.year, args.month, i + 1)
        target_label = f"{tm}-{ty}"

        z_top = get_2d(ds, var_top, i)
        z_bottom = get_2d(ds, var_bottom, i) * bottom_scale
        # Match GrADS grfill: by default keep all valid NetCDF grid cells and rely
        # on the file's own missing values.  Polygon centre-point masking is optional
        # because it can drop legitimate coastal/border cells.
        if args.mask_country and not args.no_mask:
            z_top = mask_to_country(z_top, lon, lat, regions.union)
            z_bottom = mask_to_country(z_bottom, lon, lat, regions.union)

        last_cf_top = plot_panel(
            axes[0][i], lon, lat, z_top, levels_top, cmap_top,
            args, regions, target_label, show_y=(i == 0),
        )
        last_cf_bottom = plot_panel(
            axes[1][i], lon, lat, z_bottom, levels_bottom, cmap_bottom,
            args, regions, target_label, show_y=(i == 0),
        )

    init_label = f"{args.month:02d}-{args.year}"
    fig.text(0.50, 0.965, f"{title_top} forecast initialized at {init_label}",
             ha="center", va="center", fontsize=10.5)
    fig.text(0.50, 0.505, f"{title_bottom} forecast initialized at {init_label}",
             ha="center", va="center", fontsize=10.5)

    cb1 = fig.colorbar(last_cf_top, cax=cax_top, orientation="vertical", extend="both")
    cb2 = fig.colorbar(last_cf_bottom, cax=cax_bottom, orientation="vertical", extend="both")
    format_colorbar(cb1, levels_top)
    format_colorbar(cb2, levels_bottom)

    out = Path(f"{args.output_prefix}{out_suffix}")
    out.parent.mkdir(parents=True, exist_ok=True)
    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"):
        if coord not in ds:
            raise KeyError(f"Coordinate '{coord}' not found in input NetCDF")
    for v in ("pr_fcs", "pr_ano", "t2_fcs", "t2_ano"):
        if v not in ds:
            raise KeyError(f"Required variable '{v}' not found in input NetCDF")
    nlead_file = int(ds.sizes.get("lead", 0))
    if nlead_file < args.nlead:
        raise ValueError(f"Input contains only {nlead_file} lead(s), but --nlead={args.nlead}")


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

    # Validate day without changing the month+lead logic used for target labels.
    _, 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)

        # Ensure increasing coordinates for stable contour plotting.
        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)

        make_figure(
            ds, lon, lat, regions, args,
            var_top="pr_fcs", var_bottom="pr_ano",
            levels_top=PR_LEVELS, levels_bottom=PR_ANO_LEVELS,
            cmap_top=CMAP_PR, cmap_bottom=CMAP_PR_ANO,
            title_top="Monthly Rainfall Amount (mm/month)",
            title_bottom="Monthly Relative Rainfall Anomaly (%)",
            out_suffix="Pr_forecast.png",
            bottom_scale=args.pr_ano_scale,
        )

        make_figure(
            ds, lon, lat, regions, args,
            var_top="t2_fcs", var_bottom="t2_ano",
            levels_top=T2_LEVELS, levels_bottom=T2_ANO_LEVELS,
            cmap_top=CMAP_T2, cmap_bottom=CMAP_T2_ANO,
            title_top="Monthly 2-m Temperature (°C)",
            title_bottom="Monthly 2-m Temperature Anomaly (°C)",
            out_suffix="T2_forecast.png",
            bottom_scale=1.0,
        )

    return 0


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