"""Fixed climatological P90 and EHD threshold (LOCKED definitions).

P90 is computed ONCE from observed Tx over the reference period (2003-2022 by
default), per grid cell and per calendar day, pooling all reference years'
values in a centred cyclic 31-day window (d-15 .. d+15) that wraps across the
year boundary. It is NOT monthly, seasonal, annual or fold-specific.

Leap-year / 29 February handling (documented, testable):
    Every date is mapped to a FIXED 365-entry calendar-day index via its
    (month, day) on a non-leap calendar. 29 February is pooled with 28
    February (index 59). Consequently the climatology has exactly 365 day
    indices in every year, leap or not, and 29-Feb observations contribute to
    the late-February / early-March windows just like 28-Feb observations.
"""
import numpy as np

# Cumulative days before each month on a NON-leap (365-day) calendar.
_CUM_NOLEAP = np.array([0, 31, 59, 90, 120, 151, 181, 212, 243, 273, 304, 334])
DAYS_IN_YEAR = 365


def day_of_year_noleap(dates):
    """Map datetime64 dates to a fixed calendar-day index in 1..365.

    29 February is mapped to 59 (the index of 28 February). All other dates
    use their non-leap day-of-year, so e.g. 1 March is always index 60
    regardless of leap year.
    """
    dates = np.asarray(dates, dtype="datetime64[D]")
    months = dates.astype("datetime64[M]")
    month = (months.astype(int) % 12) + 1
    day = (dates - months).astype("timedelta64[D]").astype(int) + 1
    doy = _CUM_NOLEAP[month - 1] + day
    is_feb29 = (month == 2) & (day == 29)
    doy = np.where(is_feb29, 59, doy)
    return doy.astype(int)


def _years(dates):
    d = np.asarray(dates, dtype="datetime64[Y]").astype(int) + 1970
    return d


def p90_climatology(tx, dates, ref_years=(2003, 2022), window=15, q=90):
    """Fixed per-cell, per-calendar-day P90 climatology.

    Parameters
    ----------
    tx : ndarray, shape (T, ...)
        Daily observed Tx. Leading axis is time; remaining axes are spatial.
    dates : array of datetime64, shape (T,)
    ref_years : (int, int)
        Inclusive reference-period bounds (default 2003-2022).
    window : int
        Half-width of the centred cyclic window in days (default 15 -> 31 days).
    q : float
        Percentile (default 90).

    Returns
    -------
    p90 : ndarray, shape (365, ...)
        P90 threshold per calendar-day index (1..365 -> rows 0..364).
    """
    tx = np.asarray(tx, dtype="float64")
    doy = day_of_year_noleap(dates)
    yr = _years(dates)
    sel = (yr >= ref_years[0]) & (yr <= ref_years[1])
    tx_ref = tx[sel]
    doy_ref = doy[sel]

    spatial = tx.shape[1:]
    out = np.full((DAYS_IN_YEAR,) + spatial, np.nan, dtype="float64")
    for d in range(1, DAYS_IN_YEAR + 1):
        offs = ((np.arange(d - window, d + window + 1) - 1) % DAYS_IN_YEAR) + 1
        mask = np.isin(doy_ref, offs)
        vals = tx_ref[mask]
        if vals.shape[0] == 0:
            continue
        with np.errstate(all="ignore"):
            out[d - 1] = np.nanpercentile(vals, q, axis=0)
    return out


def ehd_threshold(p90, cap=35.0):
    """EHD_THRESHOLD(s,d) = min(P90(s,d), cap). min, never max. cap=35.0 degC."""
    return np.minimum(np.asarray(p90, dtype="float64"), cap)


def threshold_for_dates(thr_clim, dates):
    """Expand a (365, ...) climatological field onto a (T, ...) date series."""
    idx = day_of_year_noleap(dates) - 1
    return np.asarray(thr_clim)[idx]