from __future__ import annotations

import os
from typing import Optional, Sequence, Tuple

import numpy as np
import xarray as xr
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.colors import BoundaryNorm, ListedColormap
import cartopy.crs as ccrs
import cartopy.feature as cfeature
from scipy.ndimage import gaussian_filter, generic_filter


# =========================================================
# Global style controls
# =========================================================
TITLE_FONTSIZE = 12
LABEL_FONTSIZE = 10
TICK_FONTSIZE = 9
LEGEND_FONTSIZE = 8

CBAR_PAD = 0.04
CBAR_SHRINK = 0.78
CBAR_ASPECT = 32

DEFAULT_FIGSIZE_MAP = (10, 6)
DEFAULT_FIGSIZE_TS = (10, 4.5)
DEFAULT_FIGSIZE_CURVE = (6.5, 6)


# =========================================================
# Basic validators / filesystem helpers
# =========================================================
def _require_dims(da: xr.DataArray, dims: Sequence[str], name: str) -> None:
    missing = [d for d in dims if d not in da.dims]
    if missing:
        raise ValueError(f"{name} is missing required dims: {missing}. Found dims={da.dims}")


def _ensure_output_dir(path: Optional[str]) -> None:
    if path is None:
        return
    directory = os.path.dirname(path)
    if directory:
        os.makedirs(directory, exist_ok=True)


# =========================================================
# Cross-dateline helpers
# =========================================================
def _normalize_lon_to_360(lon: float) -> float:
    return float(lon) % 360.0


def _normalize_lon_to_180(lon: float) -> float:
    lon180 = ((float(lon) + 180.0) % 360.0) - 180.0
    if np.isclose(lon180, -180.0) and float(lon) > 0:
        return 180.0
    return lon180


def _extent_crosses_dateline(extent: Optional[Sequence[float]]) -> bool:
    """
    extent = [lon_min, lon_max, lat_min, lat_max]
    """
    if extent is None:
        return False
    if len(extent) != 4:
        raise ValueError(f"extent must be [lon_min, lon_max, lat_min, lat_max], got {extent!r}")

    lon_min, lon_max = float(extent[0]), float(extent[1])
    lon_min_360 = _normalize_lon_to_360(lon_min)
    lon_max_360 = _normalize_lon_to_360(lon_max)
    return lon_max_360 < lon_min_360


def _selected_extent_width_360(extent: Sequence[float]) -> float:
    """
    Width of the selected longitude arc in degrees, following the arc from
    lon_min eastward to lon_max.
    """
    if len(extent) != 4:
        raise ValueError(f"extent must be [lon_min, lon_max, lat_min, lat_max], got {extent!r}")

    lon_min = _normalize_lon_to_360(extent[0])
    lon_max = _normalize_lon_to_360(extent[1])

    if lon_max >= lon_min:
        return lon_max - lon_min
    return (lon_max + 360.0) - lon_min


def _selected_extent_midpoint(extent: Sequence[float]) -> float:
    """
    Midpoint longitude of the selected arc, returned in [-180, 180].
    Works for both wrapped and non-wrapped extents.
    """
    if len(extent) != 4:
        raise ValueError(f"extent must be [lon_min, lon_max, lat_min, lat_max], got {extent!r}")

    lon_min = _normalize_lon_to_360(extent[0])
    width = _selected_extent_width_360(extent)

    mid = (lon_min + 0.5 * width) % 360.0
    return _normalize_lon_to_180(mid)


def _choose_central_longitude(extent: Optional[Sequence[float]]) -> float:
    """
    Choose a map center that is portable across regions.

    - Near-global domains stay centered at 0
    - Regional domains are centered on the midpoint of the selected arc
    """
    if extent is None:
        return 0.0

    width = _selected_extent_width_360(extent)

    if width >= 359.0:
        return 0.0

    return _selected_extent_midpoint(extent)


def _to_plot_lon_convention(lon_values: np.ndarray, *, central_longitude: float) -> np.ndarray:
    """
    Convert lon values into the plotting frame centered on central_longitude.

    Resulting longitudes fall in:
        [central_longitude - 180, central_longitude + 180)
    """
    lon_values = np.asarray(lon_values, dtype=float)
    return ((lon_values - central_longitude + 180.0) % 360.0) + central_longitude - 180.0


def _normalize_extent_for_plot(
    extent: Optional[Sequence[float]],
    *,
    central_longitude: float,
):
    """
    Convert extent longitudes into the plotting frame centered on central_longitude.

    For wrapped extents, preserve the selected arc as one continuous interval.
    """
    if extent is None:
        return None

    if len(extent) != 4:
        raise ValueError(f"extent must be [lon_min, lon_max, lat_min, lat_max], got {extent!r}")

    lon_min, lon_max, lat_min, lat_max = [float(v) for v in extent]

    lon_min_360 = _normalize_lon_to_360(lon_min)
    lon_max_360 = _normalize_lon_to_360(lon_max)

    if lon_max_360 < lon_min_360:
        lon_max_360 += 360.0

    lon_min_plot = ((lon_min_360 - central_longitude + 180.0) % 360.0) + central_longitude - 180.0
    lon_max_plot = lon_min_plot + (lon_max_360 - lon_min_360)

    return [lon_min_plot, lon_max_plot, lat_min, lat_max]


def _prepare_lon_lat_for_plot(
    da2d: xr.DataArray,
    *,
    central_longitude: float,
) -> xr.DataArray:
    """
    Normalize longitude coordinates to match the intended map view and sort longitude.
    """
    _require_dims(da2d, ("lat", "lon"), "da2d")

    lon_vals = np.asarray(da2d["lon"].values, dtype=float)
    lon_new = _to_plot_lon_convention(lon_vals, central_longitude=central_longitude)

    return da2d.assign_coords(lon=lon_new).sortby("lon")


def _create_geo_axes(
    *,
    extent: Optional[Sequence[float]] = None,
    central_longitude: Optional[float] = None,
):
    if central_longitude is None:
        central_longitude = _choose_central_longitude(extent)
    ax = plt.axes(projection=ccrs.PlateCarree(central_longitude=central_longitude))
    return ax


# =========================================================
# Helper colormap
# =========================================================
def _build_tercile_palette() -> Tuple[ListedColormap, BoundaryNorm, np.ndarray]:
    """
    Build a 15-bin palette:
      1..5   below-normal   (light -> dark brown)
      6..10  near-normal    (light -> dark gray)
      11..15 above-normal   (light -> dark green)
    """
    colors = [
        "#f6e3cf", "#e8c39b", "#d39a63", "#b56f35", "#7a4518",
        "#f0f0f0", "#d9d9d9", "#bdbdbd", "#8c8c8c", "#525252",
        "#e5f5e0", "#a1d99b", "#74c476", "#31a354", "#006d2c",
    ]
    cmap = ListedColormap(colors, name="tercile_dominant")
    levels = np.arange(0.5, 16.5, 1.0)
    norm = BoundaryNorm(levels, cmap.N)
    return cmap, norm, levels


# =========================================================
# Category filtering helpers
# =========================================================
def apply_majority_filter_category(
    cat2d: np.ndarray,
    *,
    size: int = 3,
    valid_mask: Optional[np.ndarray] = None,
) -> np.ndarray:
    """
    Majority filter for dominant category only.

    Parameters
    ----------
    cat2d : np.ndarray
        2D array of dominant category indices (typically 0, 1, 2).
    size : int
        Neighborhood size for majority filtering.
    valid_mask : np.ndarray, optional
        Boolean mask of cells that are valid for category filtering.

    Returns
    -------
    np.ndarray
        Smoothed category array with NaNs preserved outside valid_mask.
    """
    arr = np.asarray(cat2d, dtype=float)

    if valid_mask is None:
        valid_mask = np.isfinite(arr)
    else:
        valid_mask = np.asarray(valid_mask, dtype=bool)

    work = np.where(valid_mask, arr, np.nan)

    def _majority_func(values):
        vals = values[np.isfinite(values)]
        if vals.size == 0:
            return np.nan
        vals_int = vals.astype(int)
        counts = np.bincount(vals_int)
        return float(np.argmax(counts))

    filtered = generic_filter(
        work,
        function=_majority_func,
        size=size,
        mode="nearest",
    )

    filtered[~valid_mask] = np.nan
    return filtered


# =========================================================
# Generic helpers
# =========================================================
def _prepare_2d_field(
    da: xr.DataArray,
    *,
    time_index: int = 0,
    sample_dim: str = "time",
) -> xr.DataArray:
    out = da
    if sample_dim in out.dims:
        out = out.isel({sample_dim: time_index})
    _require_dims(out, ("lat", "lon"), "2D field")
    return out.transpose("lat", "lon")


def _add_base_map(
    ax,
    extent: Optional[Sequence[float]] = None,
    *,
    central_longitude: float = 0.0,
) -> None:
    ax.coastlines(linewidth=0.9)
    ax.add_feature(cfeature.BORDERS, linewidth=0.5)

    if extent is not None:
        extent_plot = _normalize_extent_for_plot(extent, central_longitude=central_longitude)
        ax.set_extent(extent_plot, crs=ccrs.PlateCarree())


def _style_small_legend(ax, ncol: int = 1) -> None:
    ax.legend(
        fontsize=LEGEND_FONTSIZE,
        frameon=False,
        ncol=ncol,
        handlelength=1.4,
        handletextpad=0.5,
        borderpad=0.2,
        labelspacing=0.3,
        columnspacing=0.7,
        loc="best",
    )


def _add_map_colorbar(fig, ax, mappable, label: str):
    cbar = fig.colorbar(
        mappable,
        ax=ax,
        orientation="horizontal",
        pad=CBAR_PAD,
        shrink=CBAR_SHRINK,
        aspect=CBAR_ASPECT,
    )
    cbar.set_label(label, fontsize=LABEL_FONTSIZE)
    cbar.ax.tick_params(labelsize=TICK_FONTSIZE)
    return cbar


def _default_years(n: int) -> np.ndarray:
    if n == 30:
        return np.arange(1991, 2021)
    return np.arange(1991, 1991 + n)


def _infer_time_axis_values(scores: xr.DataArray, years: Optional[Sequence[int]] = None) -> np.ndarray:
    """
    Prefer user-supplied years, then DataArray time coordinate if available,
    otherwise fall back to a default synthetic year axis.
    """
    sample_dim = scores.dims[0]
    n = scores.sizes[sample_dim]

    if years is not None:
        years = np.asarray(years)
        if len(years) != n:
            raise ValueError("Length of years must match number of time samples.")
        return years

    if sample_dim in scores.coords:
        vals = np.asarray(scores[sample_dim].values)
        if len(vals) == n:
            return vals

    return _default_years(n)


def _compute_auto_levels(
    data: np.ndarray,
    *,
    symmetric: bool = False,
    default: Tuple[float, float] = (-1.0, 1.0),
    n_levels: int = 21,
) -> np.ndarray:
    finite = np.asarray(data[np.isfinite(data)], dtype=float)
    if finite.size == 0:
        return np.linspace(default[0], default[1], n_levels)

    vmin = float(np.nanmin(finite))
    vmax = float(np.nanmax(finite))

    if symmetric or (vmin < 0 < vmax):
        vmaxa = max(abs(vmin), abs(vmax))
        if np.isclose(vmaxa, 0.0):
            vmaxa = 1.0
        return np.linspace(-vmaxa, vmaxa, n_levels)

    if np.isclose(vmin, vmax):
        pad = 1.0 if np.isclose(vmin, 0.0) else abs(vmin) * 0.1 + 1e-3
        vmin -= pad
        vmax += pad

    return np.linspace(vmin, vmax, n_levels)


def _add_gridlines(ax, crs=ccrs.PlateCarree()):
    gl = ax.gridlines(
        crs=crs,
        draw_labels=True,
        linewidth=0.4,
        color="gray",
        alpha=0.5,
        linestyle="--",
    )
    gl.top_labels = False
    gl.right_labels = False
    gl.xlabel_style = {"size": 8}
    gl.ylabel_style = {"size": 8}
    return gl


def _save_or_close(fig, output_path: Optional[str]) -> None:
    if output_path:
        _ensure_output_dir(output_path)
        fig.savefig(output_path, dpi=150, bbox_inches="tight")
    plt.close(fig)


def _create_map_figure(
    *,
    figsize: Tuple[float, float],
    extent: Optional[Sequence[float]],
    central_longitude: Optional[float],
):
    if central_longitude is None:
        central_longitude = _choose_central_longitude(extent)

    fig = plt.figure(figsize=figsize)
    ax = _create_geo_axes(extent=extent, central_longitude=central_longitude)
    _add_base_map(ax, extent=extent, central_longitude=central_longitude)
    _add_gridlines(ax)
    return fig, ax, central_longitude


def _add_cyclic_if_needed(
    da2d: xr.DataArray,
    *,
    central_longitude: float,
    extent: Optional[Sequence[float]] = None,
) -> xr.DataArray:
    """
    Add a cyclic longitude point for wrapped single-panel maps so contourf
    does not leave a seam at the wrap boundary.

    Uses a manual cyclic append because transformed plotting longitudes may
    not be perfectly equally spaced.
    """
    _require_dims(da2d, ("lat", "lon"), "da2d")

    if extent is not None and not _extent_crosses_dateline(extent):
        return da2d

    lon = np.asarray(da2d["lon"].values, dtype=float)
    lat = np.asarray(da2d["lat"].values, dtype=float)
    vals = np.asarray(da2d.values, dtype=float)

    if lon.size < 2:
        return da2d

    dlon = np.diff(lon)
    dlon = dlon[np.isfinite(dlon)]
    dlon = dlon[np.abs(dlon) < 180.0]

    if dlon.size == 0:
        return da2d

    dlon_mean = float(np.median(dlon))

    vals_cyc = np.concatenate([vals, vals[:, :1]], axis=1)
    lon_cyc = np.concatenate([lon, [lon[-1] + dlon_mean]])

    return xr.DataArray(
        vals_cyc,
        dims=("lat", "lon"),
        coords={"lat": lat, "lon": lon_cyc},
        name=da2d.name,
        attrs=da2d.attrs,
    )


def _smooth_2d_masked(data2d: np.ndarray, sigma: float) -> np.ndarray:
    """
    Gaussian smooth a 2D field while preserving NaNs.

    Uses weighted smoothing so masked / NaN regions do not contaminate
    nearby valid values.
    """
    arr = np.asarray(data2d, dtype=float)

    if sigma is None or sigma <= 0:
        return arr

    valid = np.isfinite(arr).astype(float)
    filled = np.where(np.isfinite(arr), arr, 0.0)

    smooth_num = gaussian_filter(filled * valid, sigma=sigma, mode="nearest")
    smooth_den = gaussian_filter(valid, sigma=sigma, mode="nearest")

    out = np.full_like(arr, np.nan, dtype=float)
    mask = smooth_den > 1e-12
    out[mask] = smooth_num[mask] / smooth_den[mask]

    return out


def _plot_map_field(
    da2d: xr.DataArray,
    *,
    title: str,
    cmap,
    levels: Sequence[float],
    extent: Optional[Sequence[float]],
    figsize: Tuple[float, float],
    output_path: Optional[str],
    central_longitude: Optional[float],
    colorbar_label: str,
    norm=None,
    extend: str = "both",
    smooth_sigma: float = 0.0,
):
    fig, ax, central_longitude = _create_map_figure(
        figsize=figsize,
        extent=extent,
        central_longitude=central_longitude,
    )

    da_plot = _prepare_lon_lat_for_plot(da2d, central_longitude=central_longitude)
    da_plot = _add_cyclic_if_needed(
        da_plot,
        central_longitude=central_longitude,
        extent=extent,
    )

    data_smoothed = _smooth_2d_masked(da_plot.values, sigma=smooth_sigma)
    data = np.ma.masked_invalid(data_smoothed)

    if isinstance(cmap, str):
        cmap = plt.get_cmap(cmap).copy()
        cmap.set_bad(color="white")
    elif hasattr(cmap, "copy"):
        cmap = cmap.copy()
        if hasattr(cmap, "set_bad"):
            cmap.set_bad(color="white")

    cf = ax.contourf(
        da_plot["lon"].values,
        da_plot["lat"].values,
        data,
        levels=levels,
        cmap=cmap,
        norm=norm,
        extend=extend,
        transform=ccrs.PlateCarree(),
    )

    _add_map_colorbar(fig, ax, cf, colorbar_label)
    ax.set_title(title, fontsize=TITLE_FONTSIZE, fontweight="bold")
    plt.tight_layout()
    _save_or_close(fig, output_path)


# =========================================================
# Tercile code conversion
# =========================================================
def _compute_dominant_tercile_codes(
    prob_da: xr.DataArray,
    time_index: int = 0,
    *,
    category_dim: str = "category",
    sample_dim: str = "time",
) -> xr.DataArray:
    arr = prob_da
    if sample_dim in arr.dims:
        arr = arr.isel({sample_dim: time_index})
    _require_dims(arr, (category_dim, "lat", "lon"), "prob_da")
    arr = arr.transpose(category_dim, "lat", "lon")

    probs = np.asarray(arr.values, dtype=float)
    valid_any = np.any(np.isfinite(probs), axis=0)
    probs_filled = np.where(np.isfinite(probs), probs, -np.inf)

    dom_cat = np.argmax(probs_filled, axis=0)
    dom_prob = np.max(probs_filled, axis=0)

    shade = np.full(dom_prob.shape, np.nan, dtype=float)
    bins = [1.0 / 3.0, 0.45, 0.55, 0.65, 0.75, np.inf]
    for i in range(5):
        mask = (dom_prob >= bins[i]) & (dom_prob < bins[i + 1])
        shade[mask] = i + 1

    codes = np.full(dom_cat.shape, np.nan, dtype=float)
    for cat in range(3):
        mask = dom_cat == cat
        codes[mask] = shade[mask] + cat * 5

    codes[~valid_any] = np.nan

    return xr.DataArray(
        codes,
        dims=("lat", "lon"),
        coords={"lat": arr["lat"], "lon": arr["lon"]},
        name="dominant_tercile_code",
    )


# =========================================================
# Deterministic forecast map
# =========================================================
def plot_deterministic_forecast(
    forecast: xr.DataArray,
    *,
    time_index: int = 0,
    title: str = "Deterministic Forecast",
    cmap: str = "viridis",
    levels: Optional[Sequence[float]] = None,
    extent: Optional[Sequence[float]] = None,
    figsize: Tuple[float, float] = DEFAULT_FIGSIZE_MAP,
    output_path: Optional[str] = None,
    central_longitude: Optional[float] = None,
    smooth_sigma: float = 0.0,
):
    da2d = _prepare_2d_field(forecast, time_index=time_index)
    if levels is None:
        levels = _compute_auto_levels(da2d.values, symmetric=False)

    _plot_map_field(
        da2d,
        title=title,
        cmap=cmap,
        levels=levels,
        extent=extent,
        figsize=figsize,
        output_path=output_path,
        central_longitude=central_longitude,
        colorbar_label=forecast.name if forecast.name else "forecast",
        norm=None,
        extend="both",
        smooth_sigma=smooth_sigma,
    )


# =========================================================
# Dominant tercile map
# =========================================================
def plot_tercile_dominant_forecast(
    probabilities: xr.DataArray,
    *,
    time_index: int = 0,
    title: str = "Dominant Tercile Forecast",
    extent: Optional[Sequence[float]] = None,
    figsize: Tuple[float, float] = DEFAULT_FIGSIZE_MAP,
    output_path: Optional[str] = None,
    min_prob_percent: float = 40.0,
    central_longitude: Optional[float] = None,
    category_dim: str = "category",
    sample_dim: str = "time",
    majority_filter_size: int = 3,
    instability_margin_percent: float = 5.0,
):
    """
    Plot dominant tercile forecast.

    Only probabilities >= min_prob_percent are shown.
    Lower values are masked out.

    This plot intentionally remains unsmoothed because it is a discrete
    categorical map; smoothing would create artificial category artifacts.

    Parameters
    ----------
    majority_filter_size : int
        Neighborhood size for majority filtering of dominant category only.
        Set <= 1 to disable.
    instability_margin_percent : float
        If the difference between the top two category probabilities is smaller
        than this threshold, the cell is masked as unstable.
    """
    arr = probabilities
    if sample_dim in arr.dims:
        arr = arr.isel({sample_dim: time_index})
    _require_dims(arr, (category_dim, "lat", "lon"), "probabilities")
    arr = arr.transpose(category_dim, "lat", "lon")

    probs = np.asarray(arr.values, dtype=float)
    valid_any = np.any(np.isfinite(probs), axis=0)
    probs_filled = np.where(np.isfinite(probs), probs, -np.inf)

    dom_cat_raw = np.argmax(probs_filled, axis=0)
    dom_prob = np.max(probs_filled, axis=0)

    finite_dom = dom_prob[np.isfinite(dom_prob)]
    if finite_dom.size > 0 and np.nanmax(finite_dom) <= 1.0:
        dom_prob_pct = dom_prob * 100.0
        probs_pct = probs_filled * 100.0
    else:
        dom_prob_pct = dom_prob.copy()
        probs_pct = probs_filled.copy()

    first_bin = float(min_prob_percent)
    if first_bin >= 45.0:
        raise ValueError("min_prob_percent must be less than 45 to preserve tercile shading bins.")

    weak_mask = dom_prob_pct < first_bin

    # instability mask: top two probabilities too close to each other
    sorted_probs_pct = np.sort(probs_pct, axis=0)

    top = sorted_probs_pct[-1]
    second = sorted_probs_pct[-2]

    # Safe margin computation (no warnings)
    valid = np.isfinite(top) & np.isfinite(second)

    margin_pct = np.full_like(top, np.nan, dtype=float)
    margin_pct[valid] = top[valid] - second[valid]

    # Only evaluate instability where valid
    unstable_mask = np.zeros_like(top, dtype=bool)
    unstable_mask[valid] = margin_pct[valid] < float(instability_margin_percent)

    shade = np.full(dom_prob.shape, np.nan, dtype=float)
    bins = [first_bin, 45.0, 55.0, 65.0, 75.0, np.inf]
    for i in range(5):
        mask = (dom_prob_pct >= bins[i]) & (dom_prob_pct < bins[i + 1])
        shade[mask] = i + 1

    valid_cat_mask = valid_any & (~weak_mask) & (~unstable_mask)

    dom_cat_filtered = np.where(valid_cat_mask, dom_cat_raw, np.nan).astype(float)
    if majority_filter_size is not None and majority_filter_size > 1:
        dom_cat_filtered = apply_majority_filter_category(
            dom_cat_filtered,
            size=int(majority_filter_size),
            valid_mask=valid_cat_mask,
        )

    codes = np.full(dom_cat_raw.shape, np.nan, dtype=float)
    n_categories = probs.shape[0]
    for cat in range(n_categories):
        mask = dom_cat_filtered == cat
        codes[mask] = shade[mask] + cat * 5

    codes[~valid_any] = np.nan
    codes[weak_mask] = np.nan
    codes[unstable_mask] = np.nan

    dom = xr.DataArray(
        codes,
        dims=("lat", "lon"),
        coords={"lat": arr["lat"], "lon": arr["lon"]},
        name="dominant_tercile_code",
    )

    cmap, norm, levels = _build_tercile_palette()

    fig, ax, central_longitude = _create_map_figure(
        figsize=figsize,
        extent=extent,
        central_longitude=central_longitude,
    )
    dom_plot = _prepare_lon_lat_for_plot(dom, central_longitude=central_longitude)
    dom_plot = _add_cyclic_if_needed(
        dom_plot,
        central_longitude=central_longitude,
        extent=extent,
    )
    data = np.ma.masked_invalid(dom_plot.values)

    ax.contourf(
        dom_plot["lon"].values,
        dom_plot["lat"].values,
        data,
        levels=levels,
        cmap=cmap,
        norm=norm,
        extend="neither",
        transform=ccrs.PlateCarree(),
    )

    ax.set_title(title, fontsize=TITLE_FONTSIZE, fontweight="bold")

    # -----------------------------------------------------
    # Three separate legends: Below / Near / Above
    # -----------------------------------------------------
    fig.subplots_adjust(bottom=0.16)

    cax1 = fig.add_axes([0.18, 0.07, 0.18, 0.028])
    cax2 = fig.add_axes([0.41, 0.07, 0.18, 0.028])
    cax3 = fig.add_axes([0.64, 0.07, 0.18, 0.028])

    boundaries = [first_bin, 45.0, 55.0, 65.0, 75.0, 100.0]
    ticks = [first_bin, 45.0, 55.0, 65.0, 75.0]
    norm5 = BoundaryNorm(boundaries, 5)

    # Below-normal
    cmap_below = ListedColormap(cmap.colors[0:5])
    sm1 = plt.cm.ScalarMappable(cmap=cmap_below, norm=norm5)
    sm1.set_array([])
    cbar1 = fig.colorbar(sm1, cax=cax1, orientation="horizontal", ticks=ticks)
    cbar1.ax.tick_params(labelsize=7, length=2, pad=1)
    cbar1.set_label("Below", fontsize=8, labelpad=1)

    # Near-normal
    cmap_near = ListedColormap(cmap.colors[5:10])
    sm2 = plt.cm.ScalarMappable(cmap=cmap_near, norm=norm5)
    sm2.set_array([])
    cbar2 = fig.colorbar(sm2, cax=cax2, orientation="horizontal", ticks=ticks)
    cbar2.ax.tick_params(labelsize=7, length=2, pad=1)
    cbar2.set_label("Near", fontsize=8, labelpad=1)

    # Above-normal
    cmap_above = ListedColormap(cmap.colors[10:15])
    sm3 = plt.cm.ScalarMappable(cmap=cmap_above, norm=norm5)
    sm3.set_array([])
    cbar3 = fig.colorbar(sm3, cax=cax3, orientation="horizontal", ticks=ticks)
    cbar3.ax.tick_params(labelsize=7, length=2, pad=1)
    cbar3.set_label("Above", fontsize=8, labelpad=1)

    _save_or_close(fig, output_path)


# =========================================================
# Generic skill map
# =========================================================
def plot_skill_map(
    skill_da: xr.DataArray,
    *,
    title: str,
    cmap: str = "RdBu_r",
    levels: Optional[Sequence[float]] = None,
    extent: Optional[Sequence[float]] = None,
    figsize: Tuple[float, float] = DEFAULT_FIGSIZE_MAP,
    output_path: Optional[str] = None,
    central_longitude: Optional[float] = None,
    smooth_sigma: float = 0.0,
):
    da2d = skill_da.transpose("lat", "lon")
    if levels is None:
        levels = _compute_auto_levels(da2d.values, symmetric=False)

    _plot_map_field(
        da2d,
        title=title,
        cmap=cmap,
        levels=levels,
        extent=extent,
        figsize=figsize,
        output_path=output_path,
        central_longitude=central_longitude,
        colorbar_label=skill_da.name if skill_da.name else "skill",
        norm=None,
        extend="both",
        smooth_sigma=smooth_sigma,
    )


# =========================================================
# ROC curves
# =========================================================
def plot_roc_curves(
    roc_dict: dict,
    *,
    title: str = "ROC Curves",
    output_path: Optional[str] = None,
):
    fig = plt.figure(figsize=DEFAULT_FIGSIZE_CURVE)
    ax = plt.gca()

    colors = {"below": "saddlebrown", "near": "gray", "above": "green"}

    for label, d in roc_dict.items():
        fpr = np.asarray(d["fpr"], dtype=float)
        tpr = np.asarray(d["tpr"], dtype=float)

        valid = np.isfinite(fpr) & np.isfinite(tpr)
        fpr = fpr[valid]
        tpr = tpr[valid]

        auc_val = np.nan
        if fpr.size >= 2:
            order = np.argsort(fpr)
            fpr_sorted = fpr[order]
            tpr_sorted = tpr[order]
            auc_val = np.trapz(tpr_sorted, fpr_sorted)

        legend_label = f"{label} (AUC={auc_val:.2f})" if np.isfinite(auc_val) else f"{label} (AUC=NaN)"

        ax.plot(
            fpr,
            tpr,
            label=legend_label,
            linewidth=1.8,
            color=colors.get(label, None),
        )

    ax.plot([0, 1], [0, 1], linestyle="--", color="black", linewidth=0.9)
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    ax.set_xlabel("False Positive Rate", fontsize=LABEL_FONTSIZE)
    ax.set_ylabel("True Positive Rate", fontsize=LABEL_FONTSIZE)
    ax.set_title(title, fontsize=TITLE_FONTSIZE, fontweight="bold")
    ax.tick_params(labelsize=TICK_FONTSIZE)

    _style_small_legend(ax, ncol=1)
    ax.grid(True, alpha=0.3)

    plt.tight_layout()
    _save_or_close(fig, output_path)


# =========================================================
# Reliability diagrams
# =========================================================
def plot_reliability_diagrams(
    rel_dict: dict,
    *,
    title: str = "Reliability",
    output_path: Optional[str] = None,
):
    fig = plt.figure(figsize=DEFAULT_FIGSIZE_CURVE)
    ax = plt.gca()

    colors = {"below": "saddlebrown", "near": "gray", "above": "green"}

    for label, d in rel_dict.items():
        x = np.asarray(d["forecast_bin_mean"], dtype=float)
        y = np.asarray(d["observed_frequency"], dtype=float)
        valid = np.isfinite(x) & np.isfinite(y)

        ax.plot(
            x[valid],
            y[valid],
            marker="o",
            markersize=4,
            linewidth=1.8,
            label=label,
            color=colors.get(label, None),
        )

    ax.plot([0, 1], [0, 1], linestyle="--", color="black", linewidth=0.9)
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    ax.set_xlabel("Forecast Probability", fontsize=LABEL_FONTSIZE)
    ax.set_ylabel("Observed Frequency", fontsize=LABEL_FONTSIZE)
    ax.set_title(title, fontsize=TITLE_FONTSIZE, fontweight="bold")
    ax.tick_params(labelsize=TICK_FONTSIZE)

    _style_small_legend(ax, ncol=1)
    ax.grid(True, alpha=0.3)

    plt.tight_layout()
    _save_or_close(fig, output_path)


# =========================================================
# CCA loading map
# =========================================================
def plot_cca_loading_map(
    loadings: xr.DataArray,
    *,
    mode: int = 1,
    title: Optional[str] = None,
    cmap: str = "RdBu_r",
    levels: Optional[Sequence[float]] = None,
    extent: Optional[Sequence[float]] = None,
    figsize: Tuple[float, float] = DEFAULT_FIGSIZE_MAP,
    output_path: Optional[str] = None,
    central_longitude: Optional[float] = None,
    smooth_sigma: float = 0.0,
):
    """
    Plot a single CCA loading map.

    Expected dims:
        (lat, lon, mode)
    """
    if "mode" not in loadings.dims:
        raise ValueError("Loading DataArray must have a 'mode' dimension.")

    da = loadings.sel(mode=mode).transpose("lat", "lon")

    if levels is None:
        levels = _compute_auto_levels(da.values, symmetric=True)

    _plot_map_field(
        da,
        title=title or f"{loadings.name if loadings.name else 'CCA Loading'} - Mode {mode}",
        cmap=cmap,
        levels=levels,
        extent=extent,
        figsize=figsize,
        output_path=output_path,
        central_longitude=central_longitude,
        colorbar_label=loadings.name if loadings.name else "CCA loading",
        norm=None,
        extend="both",
        smooth_sigma=smooth_sigma,
    )


# =========================================================
# Single CCA score set
# =========================================================
def plot_cca_scores_by_mode(
    scores: xr.DataArray,
    *,
    years: Optional[Sequence[int]] = None,
    prefix: str = "CCA Scores",
    output_dir: Optional[str] = None,
):
    """
    Generate one plot per mode.

    Expected dims:
        (time, mode)
    """
    if "mode" not in scores.dims:
        raise ValueError("Scores must have 'mode' dimension.")

    xvals = _infer_time_axis_values(scores, years=years)
    modes = scores["mode"].values

    for m in modes:
        fig = plt.figure(figsize=(9, 4.5))
        ax = plt.gca()

        y = scores.sel(mode=m).values

        ax.plot(
            xvals,
            y,
            linewidth=2.0,
            color="black",
        )

        ax.axhline(0, color="black", linewidth=0.8)
        ax.set_xlabel("Year", fontsize=LABEL_FONTSIZE)
        ax.set_ylabel("CCA Variate", fontsize=LABEL_FONTSIZE)
        ax.set_title(f"{prefix} - Mode {m}", fontsize=TITLE_FONTSIZE, fontweight="bold")
        ax.tick_params(labelsize=TICK_FONTSIZE)

        step = max(1, len(xvals) // 8)
        ax.set_xticks(xvals[::step])

        ax.grid(True, alpha=0.25)
        plt.tight_layout()

        if output_dir:
            os.makedirs(output_dir, exist_ok=True)
            fname = f"{prefix.lower().replace(' ', '_')}_mode{m}.png"
            fig.savefig(os.path.join(output_dir, fname), dpi=150, bbox_inches="tight")

        plt.close(fig)


def plot_cca_scores(
    scores: xr.DataArray,
    *,
    modes: Optional[Sequence[int]] = None,
    years: Optional[Sequence[int]] = None,
    title: Optional[str] = None,
    ylabel: Optional[str] = None,
    figsize: Tuple[float, float] = DEFAULT_FIGSIZE_TS,
    output_path: Optional[str] = None,
):
    """
    Plot one set of CCA variate time series.

    Expected dims:
        (time, mode)
    """
    if "mode" not in scores.dims:
        raise ValueError("Scores DataArray must have a 'mode' dimension.")

    xvals = _infer_time_axis_values(scores, years=years)

    if modes is None:
        modes = scores["mode"].values.tolist()

    fig = plt.figure(figsize=figsize)
    ax = plt.gca()

    for m in modes:
        ax.plot(
            xvals,
            scores.sel(mode=m).values,
            linewidth=1.8,
            label=f"Mode {m}",
        )

    ax.axhline(0, color="black", linewidth=0.8)
    ax.set_xlabel("Year", fontsize=LABEL_FONTSIZE)
    ax.set_ylabel(
        ylabel if ylabel is not None else (scores.name if scores.name else "CCA score"),
        fontsize=LABEL_FONTSIZE,
    )

    if title is None:
        title = scores.name if scores.name else "CCA Variate Time Series"
    ax.set_title(title, fontsize=TITLE_FONTSIZE, fontweight="bold")
    ax.tick_params(labelsize=TICK_FONTSIZE)

    step = max(1, len(xvals) // 8)
    ax.set_xticks(xvals[::step])

    _style_small_legend(ax, ncol=min(len(modes), 3))
    ax.grid(True, alpha=0.25)

    plt.tight_layout()
    _save_or_close(fig, output_path)


# =========================================================
# Combined X and Y CCA score plots by mode
# =========================================================
def plot_combined_cca_scores_by_mode(
    x_scores: xr.DataArray,
    y_scores: xr.DataArray,
    *,
    years: Optional[Sequence[int]] = None,
    canonical_corrs: Optional[Sequence[float]] = None,
    output_dir: Optional[str] = None,
):
    """
    Generate one X/Y combined plot per mode.

    Parameters
    ----------
    x_scores, y_scores : xr.DataArray
        Expected dims: (time, mode)
    years : sequence, optional
        X-axis year values
    canonical_corrs : sequence, optional
        Canonical correlations by mode. If provided, adds (r=...) to titles.
    output_dir : str, optional
        Directory to save plots
    """
    if "mode" not in x_scores.dims or "mode" not in y_scores.dims:
        raise ValueError("Both x_scores and y_scores must have a 'mode' dimension.")

    n = x_scores.sizes[x_scores.dims[0]]
    if y_scores.sizes[y_scores.dims[0]] != n:
        raise ValueError("x_scores and y_scores must have the same number of samples.")

    xvals = _infer_time_axis_values(x_scores, years=years)
    modes = x_scores["mode"].values

    if canonical_corrs is not None:
        canonical_corrs = np.asarray(canonical_corrs, dtype=float)

    for m in modes:
        fig = plt.figure(figsize=(9, 4.5))
        ax = plt.gca()

        ax.plot(
            xvals,
            x_scores.sel(mode=m).values,
            linewidth=2.0,
            linestyle="-",
            label="X",
        )

        ax.plot(
            xvals,
            y_scores.sel(mode=m).values,
            linewidth=2.0,
            linestyle="--",
            label="Y",
        )

        corr_text = ""
        if canonical_corrs is not None:
            idx = int(m) - 1
            if 0 <= idx < len(canonical_corrs) and np.isfinite(canonical_corrs[idx]):
                corr_text = f" (r={float(canonical_corrs[idx]):.2f})"

        ax.axhline(0, color="black", linewidth=0.8)
        ax.set_xlabel("Year", fontsize=LABEL_FONTSIZE)
        ax.set_ylabel("CCA Variate", fontsize=LABEL_FONTSIZE)
        ax.set_title(f"CCA Variates - Mode {m}{corr_text}", fontsize=TITLE_FONTSIZE, fontweight="bold")
        ax.tick_params(labelsize=TICK_FONTSIZE)

        step = max(1, len(xvals) // 8)
        ax.set_xticks(xvals[::step])

        _style_small_legend(ax, ncol=2)
        ax.grid(True, alpha=0.25)

        plt.tight_layout()

        if output_dir:
            os.makedirs(output_dir, exist_ok=True)
            fname = f"cca_xy_mode{m}.png"
            fig.savefig(os.path.join(output_dir, fname), dpi=150, bbox_inches="tight")

        plt.close(fig)
        
# =========================================================
# Combined X and Y CCA score plot
# =========================================================
def plot_combined_cca_scores(
    x_scores: xr.DataArray,
    y_scores: xr.DataArray,
    *,
    modes: Optional[Sequence[int]] = None,
    years: Optional[Sequence[int]] = None,
    title: str = "CCA Variates (X & Y)",
    figsize: Tuple[float, float] = (10, 5),
    output_path: Optional[str] = None,
):
    """
    Plot X and Y CCA time series together on the same axes.

    X modes are solid lines.
    Y modes are dashed lines.

    Expected dims for both:
        (time, mode)
    """
    if "mode" not in x_scores.dims or "mode" not in y_scores.dims:
        raise ValueError("Both x_scores and y_scores must have a 'mode' dimension.")

    n = x_scores.sizes[x_scores.dims[0]]
    if y_scores.sizes[y_scores.dims[0]] != n:
        raise ValueError("x_scores and y_scores must have the same number of samples.")

    xvals = _infer_time_axis_values(x_scores, years=years)

    if modes is None:
        n_modes = min(x_scores.sizes["mode"], y_scores.sizes["mode"])
        modes = list(range(1, n_modes + 1))

    fig = plt.figure(figsize=figsize)
    ax = plt.gca()

    for m in modes:
        ax.plot(
            xvals,
            x_scores.sel(mode=m).values,
            linewidth=1.8,
            linestyle="-",
            label=f"X Mode {m}",
        )

    for m in modes:
        ax.plot(
            xvals,
            y_scores.sel(mode=m).values,
            linewidth=1.8,
            linestyle="--",
            label=f"Y Mode {m}",
        )

    ax.axhline(0, color="black", linewidth=0.8)
    ax.set_xlabel("Year", fontsize=LABEL_FONTSIZE)
    ax.set_ylabel("CCA Variate", fontsize=LABEL_FONTSIZE)
    ax.set_title(title, fontsize=TITLE_FONTSIZE, fontweight="bold")
    ax.tick_params(labelsize=TICK_FONTSIZE)

    step = max(1, len(xvals) // 8)
    ax.set_xticks(xvals[::step])

    _style_small_legend(ax, ncol=1)
    ax.grid(True, alpha=0.25)

    plt.tight_layout()
    _save_or_close(fig, output_path)
