"""Small plotting helpers for the CAMulator timescales tutorial."""

import cartopy.crs as ccrs
import numpy as np


def map_axis(fig, nrows, ncols, index, region=None, extent=None, central_longitude=180):
    """Add a map axis with coastlines. region = {'lat': (min, max), 'lon': (min, max)}, longitudes in degrees east
    (negative = west; lon min > max crosses the date line). (The older extent = (lon_min, lon_max, lat_min, lat_max) in
    0-360 also works.)"""
    if region is not None:
        lo0, lo1 = region["lon"]
        if lo1 < lo0:
            lo1 += 360
        extent = (lo0, lo1, *region["lat"])
        central_longitude = (lo0 + lo1) / 2
    ax = fig.add_subplot(nrows, ncols, index, projection=ccrs.PlateCarree(central_longitude=central_longitude))
    if extent is not None:
        ax.set_extent(extent, crs=ccrs.PlateCarree())
    ax.coastlines(linewidth=0.6)
    return ax


DATA_CRS = ccrs.PlateCarree()  # pass as transform= when plotting lat/lon data on a map axis


def sinlat_axis(ax, ticks=(-60, -30, -10, 0, 10, 30, 60)):
    """Label an x axis that holds sin(latitude) with latitudes, so each band takes the room its area does."""
    ax.set_xticks(np.sin(np.deg2rad(ticks)))
    ax.set_xticklabels([f'{abs(t)}°{"S" if t < 0 else "N" if t > 0 else ""}' for t in ticks])
    ax.set_xlabel("Latitude (spaced by area)")


def sinlat(lat):
    return np.sin(np.deg2rad(lat))


def polar_axis(fig, nrows, ncols, index, lat_min=50):
    """Add a north-polar stereographic map axis with coastlines, showing latitudes north of lat_min."""
    ax = fig.add_subplot(nrows, ncols, index, projection=ccrs.NorthPolarStereo())
    ax.set_extent([-180, 180, lat_min, 90], crs=ccrs.PlateCarree())
    ax.coastlines(linewidth=0.6)
    return ax


# Plain-language names, units and reading notes for variables participants can choose (everyday units, see
# core.EVERYDAY_UNITS).
VARIABLES = {
    "FLUT": ("OLR", "W/m²", "outgoing longwave radiation; lower = more deep cloud"),
    "PRECT": ("Rain", "mm/day", None),
    "TREFHT": ("2-m temperature", "K", None),
    "TS": ("Surface temperature", "K", None),
    "PS": ("Surface pressure", "hPa", None),
    "U10": ("10-m wind speed", "m/s", None),
    "westerly_stress": ("Surface wind stress", "mN/m²", "positive = westerly"),
    "rain": ("Rain", "mm/day", None),
    "T": ("2-m temperature", "K", None),
    "vapour": ("Column water vapour", "mm", None),
}


def name(variable):
    """Short display name, e.g. 'OLR', 'Rain'."""
    return VARIABLES[variable][0] if variable in VARIABLES else variable


def units(variable):
    return VARIABLES[variable][1] if variable in VARIABLES else ""


def label(variable):
    """Axis / colour-bar label: name, units and a reading note, e.g. 'OLR (W/m²; ... lower = more deep cloud)'."""
    if variable not in VARIABLES:
        return variable
    short, unit, note = VARIABLES[variable]
    return f"{short} ({unit}; {note})" if note else f"{short} ({unit})"


def symmetric_limit(field, quantile=0.98):
    """A colour limit that covers almost all of a field's values, the same on both sides of zero."""
    return float(np.nanquantile(np.abs(np.asarray(field)), quantile)) or 1.0


def km_axis(ax, circumference_km):
    """Add a top x axis showing the wavelength in km for an axis of zonal wavenumbers."""
    top = ax.secondary_xaxis("top", functions=(lambda k: circumference_km / np.maximum(k, 1e-9),
                                               lambda km: circumference_km / np.maximum(km, 1e-9)))
    top.set_xlabel("Wavelength (km)")
    return top


def region_text(region):
    """{'lat': (-5, 5), 'lon': (35, 45)} -> '5°S-5°N, 35-45°E' (for titles)."""
    def lat(v):
        return f"{abs(v):g}°{'S' if v < 0 else 'N'}" if v else "0°"
    def lon(v):
        return f"{abs(v):g}°{'W' if v < 0 else 'E'}" if v else "0°"
    la0, la1 = region["lat"]
    text = f"{lat(la0)}-{lat(la1)}"
    if "lon" in region:
        text += f", {lon(region['lon'][0])}-{lon(region['lon'][1])}"
    return text