#!/usr/bin/env python3
"""Shared numerical and map helpers for the manuscript figure scripts."""

from __future__ import annotations

import cartopy.crs as ccrs
from cartopy.mpl.ticker import LatitudeFormatter, LongitudeFormatter
from cartopy.util import add_cyclic_point
import numpy as np
from scipy.ndimage import gaussian_filter


GLOBAL_SEASONS = (("Winter", "DJF", "JJA"), ("Summer", "JJA", "DJF"))
STATES = (("high_composite", "High BAM"), ("low_composite", "Low BAM"))


def finite_percentile(arrays, percentile: float, ignore_zeros: bool = False) -> float:
    pieces = [np.abs(value[np.isfinite(value)]).ravel() for value in arrays]
    if ignore_zeros:
        pieces = [piece[piece > 0] for piece in pieces]
    combined = np.concatenate(pieces) if pieces else np.asarray([])
    return float(np.percentile(combined, percentile)) if combined.size else 1.0


def smooth_nan(field: np.ndarray, sigma: float, periodic_longitude: bool = False) -> np.ndarray:
    valid = np.isfinite(field)
    mode = ("nearest", "wrap") if periodic_longitude else "nearest"
    numerator = gaussian_filter(np.where(valid, field, 0.0), sigma=sigma, mode=mode)
    denominator = gaussian_filter(valid.astype(float), sigma=sigma, mode=mode)
    result = np.full(field.shape, np.nan, dtype=float)
    reliable = denominator >= 0.5
    result[reliable] = numerator[reliable] / denominator[reliable]
    return result


def draw_rectangular_map(axis, lon, lat, color_field, contour_field, norm, cmap,
                         contour_levels, title, contour_color="black",
                         contour_linewidth=0.55, coastline_alpha=0.5,
                         contour_fmt="%g"):
    color_cyclic, lon_cyclic = add_cyclic_point(color_field, coord=lon)
    contour_cyclic, _ = add_cyclic_point(contour_field, coord=lon)
    mesh = axis.pcolormesh(
        lon_cyclic, lat, color_cyclic, transform=ccrs.PlateCarree(),
        cmap=cmap, norm=norm, shading="auto", rasterized=True,
    )
    if np.any(np.isfinite(contour_cyclic)):
        contours = axis.contour(
            lon_cyclic, lat, contour_cyclic, levels=contour_levels,
            transform=ccrs.PlateCarree(), colors=contour_color,
            linewidths=contour_linewidth,
        )
        #axis.clabel(contours, fmt=contour_fmt, fontsize=10, inline=True)
    axis.coastlines(linewidth=0.55, alpha=coastline_alpha)
    axis.set_global()
    axis.set_xticks(np.arange(-180, 181, 60), crs=ccrs.PlateCarree())
    axis.set_yticks(np.arange(-90, 91, 30), crs=ccrs.PlateCarree())
    axis.xaxis.set_major_formatter(LongitudeFormatter(dateline_direction_label=True))
    axis.yaxis.set_major_formatter(LatitudeFormatter())
    axis.tick_params(labelsize=10)
    axis.set_title(title, fontsize=10)
    return mesh