"""Observed Tx climatology, the Tx90p threshold, and the EHD index.

Conventions enforced here (see CLAUDE.md hard constraints):
  * EHD day  iff  Tx >= min(Tx90p, abs_degC).      threshold_rule="min" is PRODUCTION.
  * Tx90p is the calendar-day `percentile`-th percentile of observed Tx over a baseline
    period, pooled across years within a CENTERED calendar-day window
    (percentile_window_days, odd). This climatology window is NOT a forecast running
    window; it is only used to build the fixed climatology.
  * The gridded obs calendar is 365_day (29 Feb removed), so day-of-year is 1..365 and an
    obs array of length 365*nyears reshapes cleanly to (nyears, 365, ...).

All functions take explicit `baseline_years` so the caller (cross-validation) can fit the
climatology WITHOUT the held-out year. No function here silently uses all years.
"""
from __future__ import annotations
import numpy as np
import xarray as xr


def load_gridded_obs(path: str, var: str = "tasmax", fill: float = -9999.0) -> xr.DataArray:
    """Open gridded obs Tx as a DataArray (time, lat, lon) in degC with NaN for fill."""
    time_coder = xr.coders.CFDatetimeCoder(use_cftime=True)
    ds = xr.open_dataset(path, decode_times=time_coder)
    da = ds[var].astype("float64")
    da = da.where(da != fill)
    da = da.where(np.isfinite(da))
    return da


def _years_doy(da: xr.DataArray) -> tuple[np.ndarray, np.ndarray]:
    years = np.array([t.year for t in da["time"].values])
    # 365_day calendar: day-of-year is position within the year block (1..365).
    doy = np.empty(years.shape, dtype=int)
    for y in np.unique(years):
        m = years == y
        doy[m] = np.arange(1, m.sum() + 1)
    return years, doy


def tx90p_climatology(
    da: xr.DataArray,
    baseline_years: list[int],
    percentile: int = 90,
    window_days: int = 15,
) -> xr.DataArray:
    """Calendar-day percentile climatology on a 365-day year.

    Returns DataArray (doy=1..365, lat, lon). For each day-of-year d, pools obs from all
    baseline years within a centered window of `window_days` calendar days (wrapping the
    year boundary) and takes the `percentile`-th percentile per grid cell.
    """
    assert window_days % 2 == 1, "window_days must be odd"
    half = window_days // 2
    years, doy = _years_doy(da)
    sel = np.isin(years, baseline_years)
    vals = da.values[sel]                    # (nsel_time, lat, lon)
    d_sel = doy[sel]
    nlat, nlon = vals.shape[1], vals.shape[2]

    # index days by doy for fast window gathering
    by_doy = {d: np.where(d_sel == d)[0] for d in range(1, 366)}
    out = np.full((365, nlat, nlon), np.nan, dtype="float64")
    for d in range(1, 366):
        offs = [((d - 1 + k) % 365) + 1 for k in range(-half, half + 1)]
        idx = np.concatenate([by_doy[o] for o in offs]) if offs else np.array([], int)
        if idx.size == 0:
            continue
        block = vals[idx]                    # (n, lat, lon)
        out[d - 1] = np.nanpercentile(block, percentile, axis=0)
    return xr.DataArray(
        out,
        coords={"doy": np.arange(1, 366), "lat": da["lat"], "lon": da["lon"]},
        dims=("doy", "lat", "lon"),
        name=f"tx{percentile}p",
    )


def ehd_threshold(tx90p: xr.DataArray, abs_degC: float = 35.0, rule: str = "min") -> xr.DataArray:
    """Combine the percentile climatology with the absolute cap.

    rule="min" -> threshold = min(Tx90p, abs_degC)   [PRODUCTION, user-fixed]
    rule="max" -> threshold = max(Tx90p, abs_degC)   [DIAGNOSTIC comparison only]
    """
    if rule == "min":
        thr = xr.where(tx90p < abs_degC, tx90p, abs_degC)
    elif rule == "max":
        thr = xr.where(tx90p > abs_degC, tx90p, abs_degC)
    else:
        raise ValueError(f"unknown rule {rule!r}")
    return thr.rename(f"ehd_threshold_{rule}")


def ehd_flags(tx_daily: xr.DataArray, threshold_by_doy: xr.DataArray, doy: np.ndarray) -> xr.DataArray:
    """Boolean EHD flags for a daily Tx array given per-doy thresholds.

    tx_daily: (time, lat, lon) daily Tx in degC; doy: 1..365 per time step.
    """
    thr = threshold_by_doy.sel(doy=xr.DataArray(doy, dims="time")).values
    return xr.DataArray(
        (tx_daily.values >= thr),
        coords=tx_daily.coords, dims=tx_daily.dims, name="ehd",
    ).where(np.isfinite(tx_daily))