#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
draw_vietnam_nc_maps.py

Draw operational Vietnam forecast maps from a NetCDF file.

Main change in this version
---------------------------
1) No geopandas/shapely/fiona/pyogrio dependency. The program reads the
   Vietnam7.shp polygon file directly.
2) Input/output paths are set in the USER SETTINGS block below. You can run:

       python draw_vietnam_nc_maps.py

   without passing command-line arguments.

Expected NetCDF variables
-------------------------
    t2_fcs(lon, lat, lead)
    t2_ano(lon, lat, lead)
    pr_fcs(lon, lat, lead)
    pr_ano(lon, lat, lead)

Required Python packages
------------------------
    numpy, matplotlib, xarray, netCDF4 or h5netcdf
"""

from __future__ import annotations

import argparse
import re
import struct
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Iterable, List, Optional, Sequence, Tuple

import numpy as np

# Use a non-interactive backend so the script runs safely on HPC/login nodes.
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.cm import ScalarMappable
from matplotlib.colors import BoundaryNorm
from matplotlib.collections import LineCollection
from matplotlib.path import Path as MplPath

try:
    import xarray as xr
except ImportError as exc:  # pragma: no cover
    raise SystemExit(
        "ERROR: xarray is required. Install with: pip install xarray netCDF4"
    ) from exc


# =============================================================================
# USER SETTINGS: edit this block when input/output paths change
# =============================================================================

# Input NetCDF file.
INPFILE = "/home/tanpv/Progs/562_VNU/2026/Dr_SFC_GD2/outputs/2021/01/01/Ope_Pr_T2_and_Ano.nc"

# Initial month. Accepted formats: "04-2026", "4-2026", "2026-04", "2026-4".
INIT_MONTH = "01-2021"

# Shapefile of the 7 Vietnam climate regions. The .dbf, .shx, .prj files should
# stay in the same directory as this .shp file.
SHPFILE = "/home/tanpv/GrADS_SCR/VN_SubRegs/Vietnam7.shp"

# Output directory and output file settings.
OUTDIR = "./maps"
OUTPUT_PREFIX = "Dr_SFC_GD2_maps"
OUTPUT_FORMAT = "png"       # png, pdf, jpg, jpeg, tif, tiff
OUTPUT_DPI = 300

# Plot switches.
CLIP_TO_VIETNAM = True       # True: mask grid cells outside Vietnam7 polygons
DRAW_SEA_LABEL = True        # True: add "Biển Đông" label
DRAW_GRID = False            # True: draw light latitude-longitude grid

# Optional manual map extent. Use None for automatic extent from shapefile.
# Format: (lonmin, lonmax, latmin, latmax)
MAP_EXTENT = None
# Example:
# MAP_EXTENT = (101.0, 111.0, 7.5, 24.5)

# Figure size in inches.
FIG_WIDTH = 13.5
FIG_HEIGHT = 18.0

# If True and command-line arguments are supplied, they can override the settings
# above. If False, the program always uses only the USER SETTINGS block.
ALLOW_COMMAND_LINE_OVERRIDE = True


# =============================================================================
# Plot specifications
# =============================================================================

@dataclass(frozen=True)
class ProductSpec:
    varname: str
    title: str
    cbar_label: str
    levels: Sequence[float]
    cmap: str
    extend: str = "both"


PRODUCTS: Tuple[ProductSpec, ...] = (
    ProductSpec(
        varname="t2_fcs",
        title="Monthly 2-m Temperature (°C) forecast initialized at {init_month}",
        cbar_label="°C",
        levels=(8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30, 32),
        cmap="turbo",
        extend="both",
    ),
    ProductSpec(
        varname="t2_ano",
        title="Monthly 2-m Temperature Anomaly (°C) forecast initialized at {init_month}",
        cbar_label="°C",
        levels=(-2.0, -1.5, -1.2, -0.8, -0.4, -0.2, 0.0, 0.2, 0.4, 0.8, 1.2, 1.5, 2.0),
        cmap="RdYlGn_r",
        extend="both",
    ),
    ProductSpec(
        varname="pr_fcs",
        title="Monthly Rainfall Amount (mm/month) forecast initialized at {init_month}",
        cbar_label="mm/month",
        levels=(0, 5, 10, 20, 30, 50, 70, 100, 130, 150, 180, 210, 240, 280, 320, 360, 400, 450, 500),
        cmap="YlGnBu",
        extend="max",
    ),
    ProductSpec(
        varname="pr_ano",
        title="Monthly Relative Rainfall Anomaly (%) forecast initialized at {init_month}",
        cbar_label="%",
        levels=(-100, -80, -60, -40, -30, -20, -10, -5, 0, 5, 10, 20, 30, 40, 60, 80, 100),
        cmap="RdYlBu",
        extend="both",
    ),
)


@dataclass
class RuntimeConfig:
    inpfile: str
    init_month: str
    shpfile: str
    outdir: str
    prefix: str
    output_format: str
    dpi: int
    clip: bool
    sea_label: bool
    extent: Optional[Tuple[float, float, float, float]]
    fig_width: float
    fig_height: float
    draw_grid: bool


@dataclass
class RegionGeometry:
    """Minimal geometry container read from a polygon shapefile."""

    rings: List[np.ndarray]          # each ring is an Nx2 array: columns lon, lat
    bounds: Tuple[float, float, float, float]  # minx, miny, maxx, maxy


KEY_VALUE_ALIASES = {
    "Inpfile": "--inpfile",
    "inpfile": "--inpfile",
    "INPFILE": "--inpfile",
    "InitMonth": "--init-month",
    "initmonth": "--init-month",
    "initial_month": "--init-month",
    "InitialMonth": "--init-month",
    "Shpfile": "--shpfile",
    "shpfile": "--shpfile",
    "ShapeFile": "--shpfile",
    "shapefile": "--shpfile",
    "Outdir": "--outdir",
    "outdir": "--outdir",
    "Prefix": "--prefix",
    "prefix": "--prefix",
    "DPI": "--dpi",
    "dpi": "--dpi",
    "Format": "--format",
    "format": "--format",
    "Clip": "--clip",
    "clip": "--clip",
    "SeaLabel": "--sea-label",
    "sea_label": "--sea-label",
    "Extent": "--extent",
    "extent": "--extent",
}


def str_to_bool(text: str) -> bool:
    s = str(text).strip().lower()
    if s in {"1", "true", "t", "yes", "y", "co", "có"}:
        return True
    if s in {"0", "false", "f", "no", "n", "khong", "không"}:
        return False
    raise ValueError(f"Cannot parse boolean value: {text}")


def parse_extent_text(text: Optional[str]) -> Optional[Tuple[float, float, float, float]]:
    if text is None or str(text).strip() == "":
        return None
    vals = [float(x.strip()) for x in str(text).split(",")]
    if len(vals) != 4:
        raise ValueError("Extent must have format lonmin,lonmax,latmin,latmax")
    lonmin, lonmax, latmin, latmax = vals
    if lonmin >= lonmax or latmin >= latmax:
        raise ValueError("Invalid extent: lonmin < lonmax and latmin < latmax are required")
    return lonmin, lonmax, latmin, latmax


def normalize_key_value_args(argv: Sequence[str]) -> List[str]:
    """Convert key=value arguments to argparse-compatible arguments."""
    out: List[str] = []
    for item in argv:
        if item.startswith("--") or "=" not in item:
            out.append(item)
            continue
        key, value = item.split("=", 1)
        mapped = KEY_VALUE_ALIASES.get(key)
        if mapped is None:
            out.append(item)
        else:
            out.extend([mapped, value])
    return out


def config_from_user_settings() -> RuntimeConfig:
    return RuntimeConfig(
        inpfile=INPFILE,
        init_month=INIT_MONTH,
        shpfile=SHPFILE,
        outdir=OUTDIR,
        prefix=OUTPUT_PREFIX,
        output_format=OUTPUT_FORMAT,
        dpi=int(OUTPUT_DPI),
        clip=bool(CLIP_TO_VIETNAM),
        sea_label=bool(DRAW_SEA_LABEL),
        extent=MAP_EXTENT,
        fig_width=float(FIG_WIDTH),
        fig_height=float(FIG_HEIGHT),
        draw_grid=bool(DRAW_GRID),
    )


def parse_args_to_config(argv: Optional[Sequence[str]] = None) -> RuntimeConfig:
    """
    Build runtime config.

    With no command-line arguments, this uses the USER SETTINGS block. Command-line
    arguments remain available only for convenience/backward compatibility.
    """
    cfg = config_from_user_settings()

    raw_argv = list(sys.argv[1:] if argv is None else argv)
    if not raw_argv or not ALLOW_COMMAND_LINE_OVERRIDE:
        return cfg

    argv2 = normalize_key_value_args(raw_argv)
    parser = argparse.ArgumentParser(
        description="Draw Vietnam forecast maps from a NetCDF file and Vietnam7 shapefile.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("--inpfile", default=cfg.inpfile, help="Input NetCDF file")
    parser.add_argument("--init-month", default=cfg.init_month, help="Initial month, e.g. 04-2026 or 2026-04")
    parser.add_argument("--shpfile", default=cfg.shpfile, help="Path to VN_SubRegs/Vietnam7.shp")
    parser.add_argument("--outdir", default=cfg.outdir, help="Output directory")
    parser.add_argument("--prefix", default=cfg.prefix, help="Output file prefix")
    parser.add_argument(
        "--format",
        default=cfg.output_format,
        choices=("png", "pdf", "jpg", "jpeg", "tif", "tiff"),
        help="Output figure format",
    )
    parser.add_argument("--dpi", type=int, default=cfg.dpi, help="Output resolution for raster formats")
    parser.add_argument("--clip", default="yes" if cfg.clip else "no", help="yes/no: mask outside Vietnam")
    parser.add_argument("--sea-label", default="yes" if cfg.sea_label else "no", help="yes/no: add Biển Đông label")
    parser.add_argument("--extent", default=None, help="lonmin,lonmax,latmin,latmax")
    parser.add_argument("--fig-width", type=float, default=cfg.fig_width, help="Figure width in inches")
    parser.add_argument("--fig-height", type=float, default=cfg.fig_height, help="Figure height in inches")
    parser.add_argument("--draw-grid", default="yes" if cfg.draw_grid else "no", help="yes/no: draw grid")
    ns = parser.parse_args(argv2)

    return RuntimeConfig(
        inpfile=ns.inpfile,
        init_month=ns.init_month,
        shpfile=ns.shpfile,
        outdir=ns.outdir,
        prefix=ns.prefix,
        output_format=ns.format,
        dpi=ns.dpi,
        clip=str_to_bool(ns.clip),
        sea_label=str_to_bool(ns.sea_label),
        extent=parse_extent_text(ns.extent) if ns.extent is not None else cfg.extent,
        fig_width=ns.fig_width,
        fig_height=ns.fig_height,
        draw_grid=str_to_bool(ns.draw_grid),
    )


# =============================================================================
# Time utilities
# =============================================================================


def parse_initial_month(text: str) -> Tuple[int, int, str, str]:
    """
    Parse initial month.

    Returns
    -------
    year, month, label_for_title, yyyymm_for_filename
    """
    s = text.strip()
    m1 = re.fullmatch(r"(\d{1,2})[-/](\d{4})", s)
    m2 = re.fullmatch(r"(\d{4})[-/](\d{1,2})", s)
    if m1:
        month = int(m1.group(1))
        year = int(m1.group(2))
    elif m2:
        year = int(m2.group(1))
        month = int(m2.group(2))
    else:
        raise ValueError("INIT_MONTH must have format MM-YYYY, M-YYYY, YYYY-MM, or YYYY-M")
    if not 1 <= month <= 12:
        raise ValueError(f"Invalid month in INIT_MONTH: {text}")
    return year, month, f"{month:02d}-{year:04d}", f"{year:04d}{month:02d}"


def add_months(year: int, month: int, offset: int) -> Tuple[int, int]:
    """Add offset months to year/month. Offset=1 gives the following month."""
    total = year * 12 + (month - 1) + offset
    return total // 12, total % 12 + 1


def month_label(year: int, month: int) -> str:
    """Return target-month label such as 5-2026, matching the supplied sample."""
    return f"{month}-{year}"


# =============================================================================
# NetCDF utilities
# =============================================================================


def centers_to_edges(coord: np.ndarray) -> np.ndarray:
    """Convert 1-D grid centers to cell edges."""
    coord = np.asarray(coord, dtype=float)
    if coord.ndim != 1 or coord.size < 2:
        raise ValueError("Coordinate must be a 1-D array with at least two points")
    diff = np.diff(coord)
    edges = np.empty(coord.size + 1, dtype=float)
    edges[1:-1] = coord[:-1] + diff / 2.0
    edges[0] = coord[0] - diff[0] / 2.0
    edges[-1] = coord[-1] + diff[-1] / 2.0
    return edges


def open_dataset(inpfile: Path) -> "xr.Dataset":
    if not inpfile.exists():
        raise FileNotFoundError(f"Input NetCDF file not found: {inpfile}")
    try:
        return xr.open_dataset(inpfile, decode_times=True, mask_and_scale=True)
    except Exception:
        # Some operational files may have non-standard time encodings. For plotting these maps,
        # lon/lat/lead are sufficient.
        return xr.open_dataset(inpfile, decode_times=False, mask_and_scale=True)


def get_1d_coord(ds: "xr.Dataset", name: str) -> np.ndarray:
    if name not in ds:
        raise KeyError(f"Coordinate '{name}' not found in NetCDF file")
    arr = np.asarray(ds[name].values, dtype=float)
    if arr.ndim != 1:
        raise ValueError(f"Coordinate '{name}' must be 1-D")
    return arr


def get_lead_values(ds: "xr.Dataset") -> np.ndarray:
    if "lead" not in ds:
        raise KeyError("Coordinate 'lead' not found in NetCDF file")
    vals = np.asarray(ds["lead"].values)
    if vals.ndim != 1:
        raise ValueError("Coordinate 'lead' must be 1-D")
    return vals.astype(int)


def extract_2d_field(ds: "xr.Dataset", varname: str, lead_index: int) -> np.ndarray:
    if varname not in ds:
        raise KeyError(f"Variable '{varname}' not found in NetCDF file")
    da = ds[varname]
    required_dims = {"lon", "lat", "lead"}
    if not required_dims.issubset(set(da.dims)):
        raise ValueError(
            f"Variable '{varname}' must contain dimensions lon, lat, lead. Found: {da.dims}"
        )
    field = da.isel(lead=lead_index).transpose("lat", "lon").values.astype(float)
    field = np.where(np.isclose(field, -99.0), np.nan, field)
    return field


# =============================================================================
# Minimal shapefile reader and geometry utilities
# =============================================================================


def _ring_signed_area(ring: np.ndarray) -> float:
    """Signed area of a closed or open polygon ring."""
    if ring.shape[0] < 3:
        return 0.0
    x = ring[:, 0]
    y = ring[:, 1]
    if not (np.isclose(x[0], x[-1]) and np.isclose(y[0], y[-1])):
        x = np.r_[x, x[0]]
        y = np.r_[y, y[0]]
    return float(0.5 * np.sum(x[:-1] * y[1:] - x[1:] * y[:-1]))


def read_polygon_shapefile(shpfile: Path) -> RegionGeometry:
    """
    Read polygon rings from an ESRI .shp file without geopandas/shapely.

    Supported shape types: Polygon, PolygonZ, PolygonM, PolyLine variants.
    Vietnam7.shp should normally be Polygon.
    """
    if not shpfile.exists():
        raise FileNotFoundError(f"Shapefile not found: {shpfile}")

    rings: List[np.ndarray] = []
    with shpfile.open("rb") as f:
        header = f.read(100)
        if len(header) != 100:
            raise ValueError(f"Invalid shapefile header: {shpfile}")
        file_code = struct.unpack(">i", header[0:4])[0]
        if file_code != 9994:
            raise ValueError(f"Invalid shapefile file code in {shpfile}: {file_code}")
        shp_type = struct.unpack("<i", header[32:36])[0]
        # Header bounds are xmin, ymin, xmax, ymax.
        header_bounds = struct.unpack("<4d", header[36:68])

        while True:
            rec_header = f.read(8)
            if not rec_header:
                break
            if len(rec_header) != 8:
                raise ValueError(f"Corrupt shapefile record header in {shpfile}")
            _rec_no, content_len_16bit = struct.unpack(">2i", rec_header)
            content = f.read(content_len_16bit * 2)
            if len(content) != content_len_16bit * 2:
                raise ValueError(f"Corrupt shapefile record content in {shpfile}")
            if len(content) < 4:
                continue
            rec_shape_type = struct.unpack("<i", content[0:4])[0]
            if rec_shape_type == 0:
                continue
            # Polygon/PolygonZ/PolygonM and PolyLine variants share the same initial layout.
            if rec_shape_type not in {3, 5, 13, 15, 23, 25}:
                raise ValueError(
                    f"Unsupported shapefile geometry type {rec_shape_type}. "
                    "Expected Polygon/PolygonZ/PolygonM or PolyLine."
                )
            if len(content) < 44:
                continue
            num_parts, num_points = struct.unpack("<2i", content[36:44])
            parts_offset = 44
            points_offset = parts_offset + 4 * num_parts
            points_nbytes = 16 * num_points
            if len(content) < points_offset + points_nbytes:
                raise ValueError(f"Corrupt shapefile points in {shpfile}")
            parts = list(struct.unpack(f"<{num_parts}i", content[parts_offset:points_offset]))
            points_flat = struct.unpack(f"<{num_points * 2}d", content[points_offset:points_offset + points_nbytes])
            pts = np.asarray(points_flat, dtype=float).reshape(num_points, 2)
            parts_end = parts[1:] + [num_points]
            for start, end in zip(parts, parts_end):
                ring = pts[start:end]
                if ring.shape[0] >= 2:
                    rings.append(ring.copy())

    if not rings:
        raise ValueError(f"No polygon/polyline rings found in shapefile: {shpfile}")

    all_points = np.vstack(rings)
    bounds = (
        float(np.nanmin(all_points[:, 0])),
        float(np.nanmin(all_points[:, 1])),
        float(np.nanmax(all_points[:, 0])),
        float(np.nanmax(all_points[:, 1])),
    )
    # Use header bounds only if point-derived bounds look invalid.
    if not np.all(np.isfinite(bounds)):
        bounds = tuple(float(x) for x in header_bounds)  # type: ignore[assignment]
    return RegionGeometry(rings=rings, bounds=bounds)


def compute_outer_segments(rings: Sequence[np.ndarray], ndigits: int = 7) -> List[np.ndarray]:
    """
    Approximate national boundary by keeping line segments that appear only once.

    Internal boundaries between adjacent climate regions are usually duplicated in
    opposite directions. Removing duplicated segments leaves the exterior outline.
    """
    counts: Dict[Tuple[Tuple[float, float], Tuple[float, float]], int] = {}
    original: Dict[Tuple[Tuple[float, float], Tuple[float, float]], np.ndarray] = {}

    for ring in rings:
        if ring.shape[0] < 2:
            continue
        pts = ring
        if not np.allclose(pts[0], pts[-1]):
            pts = np.vstack([pts, pts[0]])
        for p1, p2 in zip(pts[:-1], pts[1:]):
            a = (round(float(p1[0]), ndigits), round(float(p1[1]), ndigits))
            b = (round(float(p2[0]), ndigits), round(float(p2[1]), ndigits))
            if a == b:
                continue
            key = (a, b) if a <= b else (b, a)
            counts[key] = counts.get(key, 0) + 1
            original.setdefault(key, np.asarray([[p1[0], p1[1]], [p2[0], p2[1]]], dtype=float))

    return [original[key] for key, count in counts.items() if count == 1]


def make_country_mask(regions: RegionGeometry, lon2d: np.ndarray, lat2d: np.ndarray) -> np.ndarray:
    """
    Return True for grid centers inside at least one Vietnam7 polygon ring.

    This avoids shapely/geopandas. For Vietnam climate-region polygons, holes are
    not expected to materially affect the map. If the shapefile contains holes,
    they may be included in the mask.
    """
    points = np.column_stack([lon2d.ravel(), lat2d.ravel()])
    inside = np.zeros(points.shape[0], dtype=bool)
    for ring in regions.rings:
        if ring.shape[0] < 3:
            continue
        path = MplPath(ring, closed=True)
        inside |= path.contains_points(points, radius=1.0e-10)
    return inside.reshape(lon2d.shape)


# =============================================================================
# Plotting utilities
# =============================================================================


def automatic_extent(bounds: Tuple[float, float, float, float]) -> Tuple[float, float, float, float]:
    minx, miny, maxx, maxy = bounds
    lon_margin = max(0.3, (maxx - minx) * 0.05)
    lat_margin = max(0.3, (maxy - miny) * 0.03)
    return minx - lon_margin, maxx + lon_margin, miny - lat_margin, maxy + lat_margin


def apply_common_axis_style(
    ax: plt.Axes,
    extent: Tuple[float, float, float, float],
    show_ylabels: bool,
    draw_grid: bool,
) -> None:
    lonmin, lonmax, latmin, latmax = extent
    ax.set_xlim(lonmin, lonmax)
    ax.set_ylim(latmin, latmax)
    ax.set_aspect("equal", adjustable="box")

    xticks = [104, 108]
    yticks = list(range(8, 25, 2))
    ax.set_xticks([x for x in xticks if lonmin <= x <= lonmax])
    ax.set_yticks([y for y in yticks if latmin <= y <= latmax])
    ax.set_xticklabels([f"{int(x)}E" for x in ax.get_xticks()], fontsize=7)
    if show_ylabels:
        ax.set_yticklabels([f"{int(y)}N" for y in ax.get_yticks()], fontsize=7)
    else:
        ax.set_yticklabels([])
        ax.tick_params(axis="y", length=0)
    ax.tick_params(axis="x", labelsize=7, length=2, pad=1)
    ax.tick_params(axis="y", labelsize=7, length=2, pad=1)
    for spine in ax.spines.values():
        spine.set_linewidth(0.8)
    if draw_grid:
        ax.grid(True, linewidth=0.25, alpha=0.35)


def draw_boundaries(ax: plt.Axes, regions: RegionGeometry, outer_segments: Sequence[np.ndarray]) -> None:
    # Climate-region boundaries: all polygon rings, thin lines.
    region_lines = [ring for ring in regions.rings if ring.shape[0] >= 2]
    if region_lines:
        ax.add_collection(LineCollection(region_lines, colors="0.20", linewidths=0.45, zorder=5))
    # Approximate national boundary: external segments, thicker lines.
    if outer_segments:
        ax.add_collection(LineCollection(list(outer_segments), colors="0.05", linewidths=0.85, zorder=6))


def draw_sea_label(ax: plt.Axes, extent: Tuple[float, float, float, float]) -> None:
    lonmin, lonmax, latmin, latmax = extent
    x = lonmin + 0.82 * (lonmax - lonmin)
    y = latmin + 0.48 * (latmax - latmin)
    ax.text(
        x,
        y,
        "Biển Đông",
        fontsize=6.5,
        rotation=-65,
        ha="center",
        va="center",
        color="0.25",
        zorder=7,
    )


def make_output_filename(prefix: str, init_yyyymm: str, fmt: str) -> str:
    return f"{prefix}_init_{init_yyyymm}.{fmt}"


# =============================================================================
# Main plotting routine
# =============================================================================


def plot_maps(cfg: RuntimeConfig) -> Path:
    inpfile = Path(cfg.inpfile).expanduser().resolve()
    shpfile = Path(cfg.shpfile).expanduser().resolve()
    outdir = Path(cfg.outdir).expanduser().resolve()
    outdir.mkdir(parents=True, exist_ok=True)

    init_year, init_month, init_label, init_yyyymm = parse_initial_month(cfg.init_month)
    ds = open_dataset(inpfile)
    regions = read_polygon_shapefile(shpfile)
    outer_segments = compute_outer_segments(regions.rings)

    for spec in PRODUCTS:
        if spec.varname not in ds:
            raise KeyError(f"Required variable '{spec.varname}' is missing from {inpfile}")

    lon = get_1d_coord(ds, "lon")
    lat = get_1d_coord(ds, "lat")
    lead_values = get_lead_values(ds)
    if lead_values.size < 1:
        raise ValueError("No lead times found in NetCDF file")

    lon_edges = centers_to_edges(lon)
    lat_edges = centers_to_edges(lat)
    lon2d, lat2d = np.meshgrid(lon, lat)

    country_mask = None
    if cfg.clip:
        country_mask = make_country_mask(regions, lon2d, lat2d)

    extent = cfg.extent if cfg.extent is not None else automatic_extent(regions.bounds)
    target_labels = [month_label(*add_months(init_year, init_month, int(lead))) for lead in lead_values]

    nlead = int(lead_values.size)
    if nlead != 6:
        print(f"WARNING: expected 6 lead times, found {nlead}. Plotting all available leads.", file=sys.stderr)

    fig = plt.figure(figsize=(cfg.fig_width, cfg.fig_height))
    gs = fig.add_gridspec(
        nrows=8,
        ncols=nlead + 1,
        width_ratios=[1.0] * nlead + [0.06],
        height_ratios=[0.12, 1.0, 0.12, 1.0, 0.12, 1.0, 0.12, 1.0],
        hspace=0.28,
        wspace=0.06,
    )

    for row, spec in enumerate(PRODUCTS):
        title_ax = fig.add_subplot(gs[2 * row, :nlead])
        title_ax.axis("off")
        title_ax.text(
            0.5,
            0.3,
            spec.title.format(init_month=init_label),
            ha="center",
            va="center",
            fontsize=10,
            fontweight="semibold",
        )

        cmap = plt.get_cmap(spec.cmap)
        norm = BoundaryNorm(spec.levels, ncolors=cmap.N, clip=False)

        for ilead in range(nlead):
            ax = fig.add_subplot(gs[2 * row + 1, ilead])

            field = extract_2d_field(ds, spec.varname, ilead)
            if country_mask is not None:
                field = np.where(country_mask, field, np.nan)

            ax.pcolormesh(
                lon_edges,
                lat_edges,
                field,
                cmap=cmap,
                norm=norm,
                shading="auto",
                zorder=1,
            )
            draw_boundaries(ax, regions, outer_segments)
            if cfg.sea_label:
                draw_sea_label(ax, extent)
            apply_common_axis_style(ax, extent, show_ylabels=(ilead == 0), draw_grid=cfg.draw_grid)
            ax.set_title(target_labels[ilead], fontsize=8, pad=2)

        cax = fig.add_subplot(gs[2 * row + 1, nlead])
        sm = ScalarMappable(norm=norm, cmap=cmap)
        sm.set_array([])
        cbar = fig.colorbar(sm, cax=cax, orientation="vertical", extend=spec.extend)
        cbar.set_label(spec.cbar_label, fontsize=7)
        cbar.ax.tick_params(labelsize=6, length=2, pad=1)

    outfile = outdir / make_output_filename(cfg.prefix, init_yyyymm, cfg.output_format)
    fig.savefig(outfile, dpi=cfg.dpi, bbox_inches="tight")
    plt.close(fig)
    ds.close()
    return outfile


def main(argv: Optional[Sequence[str]] = None) -> int:
    try:
        cfg = parse_args_to_config(argv)
        outfile = plot_maps(cfg)
    except Exception as exc:
        print(f"ERROR: {exc}", file=sys.stderr)
        return 2
    print(f"Created: {outfile}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())