thermograph/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 pooled
into an annual view, then weighted (temps / humidity / feels-like heaviest) 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

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", "djf", "mam", "jja", "son")

# All scored metrics. wetbulb is derived at the read boundary (climate.py); the
# rest are raw daily columns. Order here is the canonical compute order; the
# frontend renders in its own display order (Precip first).
SCORE_METRICS = ("tmax", "tmin", "feels", "humid", "wetbulb", "wind", "gust", "precip")

# Temps / humidity / feels-like dominate; wet bulb medium; wind + precip lightest.
WEIGHTS = {"feels": 2.0, "humid": 2.0, "tmax": 1.5, "tmin": 1.5,
           "wetbulb": 1.25, "precip": 0.75, "wind": 0.5, "gust": 0.5}

# Metrics whose signed drift feeds the headline "warmer / cooler" (so wind and
# precip signs never cancel the temperature story in the overall bias).
TEMP_DIR_METRICS = ("tmax", "tmin", "feels", "wetbulb")

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

METRIC_LABELS = {
    "tmax": "High temp", "tmin": "Low temp", "feels": "Feels like",
    "humid": "Humidity", "wetbulb": "Wet bulb", "wind": "Wind",
    "gust": "Gusts", "precip": "Precip",
}

# (word for a positive drift, word for a negative drift).
DIRECTION_WORDS = {
    "tmax": ("warmer", "cooler"), "tmin": ("warmer", "cooler"),
    "feels": ("warmer", "cooler"), "wetbulb": ("warmer", "cooler"),
    "humid": ("muggier", "drier"), "wind": ("windier", "calmer"),
    "gust": ("gustier", "calmer"), "precip": ("wetter", "drier"),
}

# 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, key: str) -> pl.DataFrame:
    """A slice of the record: the whole thing for ``"annual"``, else 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)."""
    if key == "annual":
        return df
    return df.filter(pl.col("date").dt.month().is_in(list(SEASONS[key])))


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 detail 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:
        v6 = float(np.percentile(rec, q))
        pct = grading.empirical_percentile(base, v6)
        if pct is None:
            continue
        per_q.append({"q": q, "v6": round(v6, 2), "pct": pct, "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."""
    if rec.size < MIN_SLICE_SAMPLES or base.size == 0:
        return None
    thr = grading.RAIN_THRESHOLD
    f6 = float(np.mean(rec >= thr)) * 100.0
    f45 = float(np.mean(base >= thr)) * 100.0
    freq_d = f6 - f45
    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:
        mad = 0.5 * amount["mad"] + 0.5 * abs(freq_d)
        bias = 0.5 * amount["bias"] + 0.5 * freq_d
    else:
        mad, bias = abs(freq_d), freq_d
    return {"per_q": amount["per_q"] if amount else [],
            "mad": round(mad, 1), "bias": round(bias, 1),
            "freq": {"f6": round(f6, 1), "f45": round(f45, 1), "d": round(freq_d, 1)}}


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 = DIRECTION_WORDS[metric]
    return up if bias >= 0 else down


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


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")
    score = score_of(div["mad"])
    tier, css = tier_of(score, div["bias"])
    direction = _direction(metric, div["bias"])
    entry = {
        "key": metric, "label": METRIC_LABELS[metric],
        "score": score, "mad": div["mad"], "bias": div["bias"],
        "tier": tier, "class": css, "direction": direction,
        "grade": tier if score < 15 else f"{tier}, {direction}",
        "weight": WEIGHTS[metric], "per_q": div["per_q"],
        "n_recent": int(rec.size), "n_base": int(base.size),
    }
    if "freq" in div:
        entry["freq"] = div["freq"]
    if metric == "wetbulb":
        # How often the peak wet bulb reaches heat-stress levels — a "wet-bulb day"
        # vs a normal one — recent share vs the full record.
        entry["freq"] = _stress_freq(base, rec, WETBULB_STRESS_F)
    return entry


def _stress_freq(base: np.ndarray, rec: np.ndarray, thr: float) -> dict:
    """Share of days at or above ``thr`` — recent vs baseline, in percentage
    points. The rest of the days are 'normal'."""
    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), "threshold_f": thr}


def _overall(entries: dict) -> dict | None:
    """Weighted roll-up of the present metrics in one slice. Weights renormalize
    over whatever is present (a metric missing for the source drops out cleanly).

    The total reads as a direction-agnostic NET CHANGE: it is colored by magnitude
    on the intensity ramp and labeled by tier alone, not 'warmer/cooler' — the
    per-metric cards still carry direction. A signed ``bias`` (from the
    temperature-direction metrics) stays in the payload for reference."""
    present = [e for e in entries.values() if e.get("score") is not None]
    if not present:
        return None
    wsum = sum(e["weight"] for e in present)
    mad = sum(e["weight"] * e["mad"] for e in present) / wsum
    dirs = [e for e in present if e["key"] in TEMP_DIR_METRICS]
    bias = (sum(e["weight"] * e["bias"] for e in dirs) / sum(e["weight"] for e in dirs)
            if dirs else 0.0)
    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), "bias": round(bias, 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)
    score = score_of(mad)
    tier, css = tier_of(score, bias)
    direction = _direction(metric, bias)
    per_q = []
    for q in QS:
        ds = [pq["d"] for e in present for pq in e.get("per_q", []) if pq["q"] == q]
        if ds:
            per_q.append({"q": q, "d": round(sum(ds) / len(ds), 1)})
    entry = {
        "key": metric, "label": METRIC_LABELS[metric],
        "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}",
        "weight": WEIGHTS[metric], "per_q": per_q,
        "n_recent": sum(e["n_recent"] for e in present),
        "n_base": sum(e["n_base"] for e in present),
    }
    if metric == "precip":
        div = precip_divergence(grading._finite(baseline["precip"]), grading._finite(recent["precip"]))
        if div:
            entry["freq"] = div["freq"]
    elif metric == "wetbulb":
        entry["freq"] = _stress_freq(grading._finite(baseline["wetbulb"]),
                                     grading._finite(recent["wetbulb"]), WETBULB_STRESS_F)
    return entry


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 key in SEASONS:
        base_s, rec_s = _slice(baseline, key), _slice(recent, key)
        entries = {m: _slice_entry(m, history, base_s, rec_s) for m in SCORE_METRICS}
        slices[key] = {"overall": _overall(entries), "metrics": entries}

    # …then the annual headline is the average of those seasonal differentials.
    annual = {}
    for m in SCORE_METRICS:
        seasonal = [slices[k]["metrics"][m] for k in SEASONS]
        annual[m] = _annual_entry(m, seasonal, history, baseline, recent)
    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,
    }