thermograph/data/scoring.py

"""Climate-shift scoring: how far a cell's recent record (the last ``RECENT_YEARS``
years) has drifted from its full multi-decade baseline.

For each metric and each percentile category ``q`` we take the recent-years value
at that percentile and ask where it lands in the baseline distribution. The gap
between where it lands and ``q`` — in percentile points, so it is unit-free and
comparable across metrics — is the divergence. A metric's score is the mean
absolute divergence over the categories, mapped to 0-100; its sign (bias) says
which way it drifted. Metrics are computed per meteorological season and the
seasonal differentials are averaged into an annual view, then averaged equally
across all metrics into one overall score.

Baseline is the ENTIRE record, including the recent years — the same all-years
climatology every other view grades against. The recent window's ~13% overlap
mildly attenuates the divergence, uniformly for every cell; it is flagged in the
payload (``baseline_overlaps_recent``).
"""
import numpy as np
import polars as pl

from data import grading

RECENT_YEARS = 6
QS = (10, 25, 50, 75, 90)                     # the scored percentile categories
SEASONS = {"djf": (12, 1, 2), "mam": (3, 4, 5),
           "jja": (6, 7, 8), "son": (9, 10, 11)}
SLICES = ("annual", *SEASONS)                 # payload slice keys

# Per-metric descriptor — display label and the (positive-drift, negative-drift)
# direction words — in ONE table so a scored metric is defined in a single place.
# The overall score weights every metric equally (see _overall). wetbulb is derived
# at the read boundary (climate.py); the rest are raw daily columns. Insertion order
# is the canonical compute order (the frontend renders in its own display order).
METRICS = {
    "tmax":    {"label": "High temp",  "dir": ("warmer", "cooler")},
    "tmin":    {"label": "Low temp",   "dir": ("warmer", "cooler")},
    "feels":   {"label": "Feels like", "dir": ("warmer", "cooler")},
    "humid":   {"label": "Humidity",   "dir": ("muggier", "drier")},
    "wetbulb": {"label": "Wet bulb",   "dir": ("warmer", "cooler")},
    "wind":    {"label": "Wind",       "dir": ("windier", "calmer")},
    "gust":    {"label": "Gusts",      "dir": ("gustier", "calmer")},
    "precip":  {"label": "Precip",     "dir": ("wetter", "drier")},
}
SCORE_METRICS = tuple(METRICS)

MIN_BASELINE_YEARS = 25   # below this, no scoring at all (payload carries a reason)
MIN_SLICE_SAMPLES = 300   # recent-slice floor (~540 expected per 6-yr season)
MIN_WET_DAYS = 30         # recent wet-day floor for the precip amount component
SATURATION = 15.0         # mean |divergence| (pct points) that maps to score 100

# A day whose peak wet bulb reaches this counts as a heat-stress ("wet-bulb") day
# rather than a normal one — 26 °C is the onset of serious heat stress for
# sustained exertion. Reported as a recent-vs-baseline share alongside the score.
WETBULB_STRESS_F = 78.8   # 26 °C

# score threshold -> (label, css for positive bias, css for negative bias). The
# css classes are the existing 9-step temperature scale, so no new color tokens
# are needed — the frontend's TIER_COLORS resolves them for free. Lower-bound
# inclusive, checked high-to-low.
TIERS = [
    (80, "Extreme shift", "rec-hot", "rec-cold"),
    (60, "Strong shift", "very-hot", "very-cold"),
    (35, "Notable shift", "hot", "cold"),
    (15, "Mild shift", "warm", "cool"),
    (0, "Steady", "normal", "normal"),
]


def _minus_years(d, years: int):
    """`d` shifted back `years` calendar years, clamping Feb 29 to Feb 28."""
    try:
        return d.replace(year=d.year - years)
    except ValueError:
        return d.replace(year=d.year - years, day=28)


def recent_baseline_split(df: pl.DataFrame):
    """(recent, baseline): recent = the last ``RECENT_YEARS`` years, baseline = the
    full record (the recent years included — see module docstring)."""
    latest = df["date"].max()
    cutoff = _minus_years(latest, RECENT_YEARS)
    return df.filter(pl.col("date") > cutoff), df


def _slice(df: pl.DataFrame, season: str) -> pl.DataFrame:
    """The pooled days of one meteorological season (DJF wraps the year end, but the
    samples are pooled across all years so the boundary is irrelevant)."""
    return df.filter(pl.col("date").dt.month().is_in(list(SEASONS[season])))


def divergence(base: np.ndarray, rec: np.ndarray) -> dict | None:
    """Core math for one metric within one slice. For each category ``q`` in
    ``QS``: the recent value at that percentile, where it lands in the baseline,
    and the signed gap ``d = landed − q``. Returns per-category ``d`` plus the
    mean-absolute-divergence (``mad``, drives the score) and mean signed
    divergence (``bias``, drives the direction). ``None`` when the recent slice is
    too thin to trust."""
    if rec.size < MIN_SLICE_SAMPLES or base.size == 0:
        return None
    per_q, deltas = [], []
    for q in QS:
        pct = grading.empirical_percentile(base, float(np.percentile(rec, q)))
        if pct is None:
            continue
        per_q.append({"q": q, "d": round(pct - q, 1)})
        deltas.append(pct - q)
    if not per_q:
        return None
    a = np.asarray(deltas)
    return {"per_q": per_q,
            "mad": round(float(np.mean(np.abs(a))), 1),
            "bias": round(float(np.mean(a)), 1)}


def precip_divergence(base: np.ndarray, rec: np.ndarray) -> dict | None:
    """Precip drift, handling its zero-inflation as two parts: how the wet-day
    *amounts* shifted (percentile divergence over rain days only) and how the
    wet-day *frequency* shifted (points of the wet-day share). They average into
    one mad/bias; the frequency alone carries it when rain days are too few for a
    stable amount distribution. (The wet-day share for the card's freq read-out is
    computed separately by ``_freq_for``.)"""
    if rec.size < MIN_SLICE_SAMPLES or base.size == 0:
        return None
    thr = grading.RAIN_THRESHOLD
    freq_d = (float(np.mean(rec >= thr)) - float(np.mean(base >= thr))) * 100.0
    wet_rec = rec[rec >= thr]
    amount = divergence(base[base >= thr], wet_rec) if wet_rec.size >= MIN_WET_DAYS else None
    if amount is not None:
        return {"per_q": amount["per_q"],
                "mad": round(0.5 * amount["mad"] + 0.5 * abs(freq_d), 1),
                "bias": round(0.5 * amount["bias"] + 0.5 * freq_d, 1)}
    return {"per_q": [], "mad": round(abs(freq_d), 1), "bias": round(freq_d, 1)}


def _freq_for(metric: str, base: np.ndarray, rec: np.ndarray) -> dict | None:
    """The day-count read-out a metric card carries, or ``None``. Precip: the
    wet-day share; wet bulb: the heat-stress-day share. Both recent-vs-baseline in
    percentage points, over the whole slice — counts, not distribution shifts."""
    if metric == "precip":
        return _day_share(base, rec, grading.RAIN_THRESHOLD)
    if metric == "wetbulb":
        return _day_share(base, rec, WETBULB_STRESS_F, {"threshold_f": WETBULB_STRESS_F})
    return None


def _day_share(base: np.ndarray, rec: np.ndarray, thr: float, extra: dict | None = None) -> dict:
    """Share of days at or above ``thr`` — recent (``f6``) vs baseline (``f45``) and
    the gap (``d``), in percentage points."""
    f6 = float(np.mean(rec >= thr)) * 100.0
    f45 = float(np.mean(base >= thr)) * 100.0
    return {"f6": round(f6, 1), "f45": round(f45, 1), "d": round(f6 - f45, 1), **(extra or {})}


def score_of(mad: float) -> int:
    """Mean-absolute-divergence (percentile points) -> 0-100 score. 0 = matches
    the baseline; 100 = a drift of ``SATURATION`` points or more."""
    return int(round(100.0 * min(mad / SATURATION, 1.0)))


def tier_of(score: int, bias: float):
    """(label, css class) for a score, the class picked from the warm or cool
    ladder by the sign of ``bias``."""
    for thr, label, hot, cold in TIERS:
        if score >= thr:
            return label, (hot if bias >= 0 else cold)
    return TIERS[-1][1], TIERS[-1][2]


def _direction(metric: str, bias: float) -> str:
    up, down = METRICS[metric]["dir"]
    return up if bias >= 0 else down


def _null_entry(metric: str, reason: str) -> dict:
    return {"key": metric, "label": METRICS[metric]["label"], "score": None,
            "reason": reason}


def _build_entry(metric: str, mad: float, bias: float, per_q: list,
                 n_recent: int, n_base: int, freq: dict | None) -> dict:
    """Assemble a scored metric entry from its divergence magnitude/direction — the
    one place the entry shape is defined, shared by the seasonal and annual paths."""
    score = score_of(mad)
    tier, css = tier_of(score, bias)
    direction = _direction(metric, bias)
    entry = {
        "key": metric, "label": METRICS[metric]["label"],
        "score": score, "mad": round(mad, 1), "bias": round(bias, 1),
        "tier": tier, "class": css, "direction": direction,
        "grade": tier if score < 15 else f"{tier}, {direction}",
        "per_q": per_q,
        "n_recent": n_recent, "n_base": n_base,
    }
    if freq is not None:
        entry["freq"] = freq
    return entry


def _metric_entry(metric: str, base: np.ndarray, rec: np.ndarray) -> dict:
    div = precip_divergence(base, rec) if metric == "precip" else divergence(base, rec)
    if div is None:
        return _null_entry(metric, "not enough recent data")
    return _build_entry(metric, div["mad"], div["bias"], div["per_q"],
                        int(rec.size), int(base.size), _freq_for(metric, base, rec))


def _overall(entries: dict) -> dict | None:
    """Equal-weight roll-up of the present metrics in one slice: every metric (and,
    through each metric's own mean, every percentile) counts the same. Metrics
    missing for the source simply drop out of the average.

    The total reads as a direction-agnostic NET CHANGE: colored by magnitude on the
    intensity ramp and labeled by tier alone, not 'warmer/cooler' — the per-metric
    cards carry direction."""
    present = [e for e in entries.values() if e.get("score") is not None]
    if not present:
        return None
    mad = sum(e["mad"] for e in present) / len(present)
    score = score_of(mad)
    tier, css = tier_of(score, 1.0)   # magnitude only — net change, not a direction
    return {"score": score, "mad": round(mad, 1), "tier": tier, "class": css,
            "grade": tier, "descriptor": "net change"}


def _slice_entry(metric: str, history: pl.DataFrame, base_s: pl.DataFrame,
                 rec_s: pl.DataFrame) -> dict:
    if metric not in history.columns:
        return _null_entry(metric, "not available for this location")
    return _metric_entry(metric, grading._finite(base_s[metric]), grading._finite(rec_s[metric]))


def _annual_entry(metric: str, seasonal: list, history: pl.DataFrame,
                  baseline: pl.DataFrame, recent: pl.DataFrame) -> dict:
    """The metric's headline score, built as the AVERAGE of its seasonal
    differentials rather than from an annually-pooled distribution. Pooling all
    days together widens the reference spread and hides a shift confined to one
    season (a hot-summer trend barely moves the all-year percentiles); averaging
    the four seasons' divergences keeps it visible. Frequency read-outs (precip wet
    days, wet-bulb heat-stress days) are day counts, so those stay pooled over the
    whole year."""
    if metric not in history.columns:
        return _null_entry(metric, "not available for this location")
    present = [e for e in seasonal if e.get("score") is not None]
    if not present:
        return _null_entry(metric, "not enough recent data")
    mad = sum(e["mad"] for e in present) / len(present)
    bias = sum(e["bias"] for e in present) / len(present)
    per_q = []
    for q in QS:
        ds = [pq["d"] for e in present for pq in e["per_q"] if pq["q"] == q]
        if ds:
            per_q.append({"q": q, "d": round(sum(ds) / len(ds), 1)})
    freq = None
    if metric in ("precip", "wetbulb"):
        freq = _freq_for(metric, grading._finite(baseline[metric]), grading._finite(recent[metric]))
    return _build_entry(metric, mad, bias, per_q,
                        sum(e["n_recent"] for e in present),
                        sum(e["n_base"] for e in present), freq)


def build_scores(history: pl.DataFrame) -> dict:
    """Full score payload for a cell's history frame. Returns an ``{"unavailable":
    reason}`` shape when the record is too short to score."""
    years = history["date"].dt.year()
    span = int(years.max() - years.min() + 1)
    if span < MIN_BASELINE_YEARS:
        return {"unavailable": f"Only {span} years of record here; climate drift "
                               f"needs at least {MIN_BASELINE_YEARS}."}
    recent, baseline = recent_baseline_split(history)
    latest = history["date"].max()

    # Each season scored on its own distribution first…
    slices = {}
    for season in SEASONS:
        base_s, rec_s = _slice(baseline, season), _slice(recent, season)
        entries = {m: _slice_entry(m, history, base_s, rec_s) for m in SCORE_METRICS}
        slices[season] = {"overall": _overall(entries), "metrics": entries}

    # …then the annual headline is the average of those seasonal differentials.
    annual = {m: _annual_entry(m, [slices[s]["metrics"][m] for s in SEASONS],
                               history, baseline, recent) for m in SCORE_METRICS}
    slices = {"annual": {"overall": _overall(annual), "metrics": annual}, **slices}

    return {
        "recent_years": RECENT_YEARS,
        "recent_range": [recent["date"].min().isoformat(), latest.isoformat()],
        "baseline_range": [history["date"].min().isoformat(), latest.isoformat()],
        "n_recent": int(recent.height),
        "n_baseline": int(history.height),
        "baseline_overlaps_recent": True,
        "wetbulb_stress_f": WETBULB_STRESS_F,
        "slices": slices,
    }