thermograph/tests/data/test_scoring.py

"""Divergence-scoring tests: synthetic records with known shifts drive known
scores. The recent window is always the last 6 years OF THE SAME frame, so a
"shift" here means perturbing only those trailing years."""
import datetime

import numpy as np
import polars as pl
import pytest

from data import scoring

ALL_METRICS = ("tmax", "tmin", "feels", "humid", "wetbulb", "wind", "gust", "precip")


def _years_before(d: datetime.date, years: int) -> datetime.date:
    try:
        return d.replace(year=d.year - years)
    except ValueError:
        return d.replace(year=d.year - years, day=28)


def make_frame(years: int = 45, seed: int = 3, drop: tuple = (), end: str | None = None):
    """A full multi-metric daily record, every metric i.i.d. across years around a
    seasonal cycle — so the last 6 years are statistically identical to the rest
    until a test perturbs them. Columns in ``drop`` are omitted (to exercise the
    missing-metric path)."""
    end_ts = datetime.date.fromisoformat(end) if end else datetime.date(2026, 7, 11)
    start = _years_before(end_ts, years)
    dates = [start + datetime.timedelta(days=i) for i in range((end_ts - start).days + 1)]
    n = len(dates)
    rng = np.random.default_rng(seed)
    doy = np.array([d.timetuple().tm_yday for d in dates])
    seasonal = 55 + 30 * np.sin((doy - 100) / 366.0 * 2 * np.pi)
    tmax = seasonal + rng.normal(0, 8, n)
    tmin = tmax - 15 + rng.normal(0, 3, n)
    cols = {
        "date": dates,
        "tmax": np.round(tmax, 1),
        "tmin": np.round(tmin, 1),
        "feels": np.round(tmax + rng.normal(0, 2, n), 1),
        "humid": np.round(np.clip(12 + rng.normal(0, 3, n), 1, None), 1),
        "wetbulb": np.round(tmax - 12 + rng.normal(0, 2, n), 1),
        "wind": np.round(np.clip(8 + rng.normal(0, 3, n), 0, None), 1),
        "gust": np.round(np.clip(16 + rng.normal(0, 5, n), 0, None), 1),
        "precip": np.where(rng.random(n) < 0.3, np.round(rng.gamma(1.5, 0.2, n), 2), 0.0),
    }
    for d in drop:
        cols.pop(d)
    df = pl.DataFrame(cols)
    return df.with_columns(pl.col("date").dt.ordinal_day().cast(pl.Int16).alias("doy"))


def _recent_mask(df: pl.DataFrame) -> np.ndarray:
    cutoff = scoring._minus_years(df["date"].max(), scoring.RECENT_YEARS)
    return (df["date"] > cutoff).to_numpy()


def shift_metric(df, metric, amount, months=None):
    """Add ``amount`` to ``metric`` for the recent-6-years rows (optionally only
    within ``months``) — the perturbation the scorer should detect."""
    mask = _recent_mask(df)
    if months is not None:
        mask = mask & df["date"].dt.month().is_in(list(months)).to_numpy()
    vals = df[metric].to_numpy().copy()
    vals[mask] = vals[mask] + amount
    return df.with_columns(pl.Series(metric, np.round(vals, 2)))


# --- baseline / no-shift ------------------------------------------------------

def test_no_shift_is_steady():
    out = scoring.build_scores(make_frame())
    ann = out["slices"]["annual"]
    assert ann["overall"]["score"] < 15               # Steady
    for k in ("tmax", "tmin", "feels", "humid"):
        m = ann["metrics"][k]
        assert m["score"] is not None
        assert abs(m["bias"]) <= 5                     # within sampling noise
        assert m["mad"] <= 5


def test_payload_metadata():
    out = scoring.build_scores(make_frame(years=45))
    assert out["recent_years"] == 6
    assert out["baseline_overlaps_recent"] is True
    assert out["n_recent"] < out["n_baseline"]
    assert set(out["slices"]) == set(scoring.SLICES)
    assert out["recent_range"][1] == out["baseline_range"][1]


# --- known directional shift --------------------------------------------------

def test_warming_shift_scores_up_and_warmer():
    base = scoring.build_scores(make_frame())["slices"]["annual"]["metrics"]["tmax"]
    hot = scoring.build_scores(shift_metric(make_frame(), "tmax", 9.0))
    m = hot["slices"]["annual"]["metrics"]["tmax"]
    assert m["score"] >= 30 and m["score"] > base["score"]
    assert m["bias"] > 0
    assert m["direction"] == "warmer"
    assert "warmer" in m["grade"]
    assert all(q["d"] > 0 for q in m["per_q"])         # every category shifted up
    # The overall total is a direction-agnostic net change — no warmer/cooler word.
    ov = hot["slices"]["annual"]["overall"]
    assert ov["descriptor"] == "net change"
    assert "—" not in ov["grade"] and "warmer" not in ov["grade"]
    assert ov["score"] > 0 and "direction" not in ov


def test_cooling_shift_reads_cooler():
    out = scoring.build_scores(shift_metric(make_frame(), "tmin", -9.0))
    m = out["slices"]["annual"]["metrics"]["tmin"]
    assert m["bias"] < 0
    assert m["direction"] == "cooler"
    assert m["class"] in ("cool", "cold", "very-cold", "rec-cold")


# --- seasonal isolation -------------------------------------------------------

def test_summer_only_shift_isolated_to_jja():
    out = scoring.build_scores(shift_metric(make_frame(), "tmax", 12.0, months=(6, 7, 8)))
    jja = out["slices"]["jja"]["metrics"]["tmax"]
    djf = out["slices"]["djf"]["metrics"]["tmax"]
    assert jja["mad"] > djf["mad"]
    assert jja["score"] > djf["score"]
    assert jja["bias"] > 0


# --- precip zero-inflation ----------------------------------------------------

def test_precip_frequency_shift():
    df = make_frame()
    mask = _recent_mask(df)
    rng = np.random.default_rng(99)
    p = df["precip"].to_numpy().copy()
    # Double the wet-day frequency in the recent window at similar amounts.
    extra = mask & (rng.random(len(p)) < 0.35) & (p < scoring.grading.RAIN_THRESHOLD)
    p[extra] = 0.12
    df = df.with_columns(pl.Series("precip", np.round(p, 2)))
    m = scoring.build_scores(df)["slices"]["annual"]["metrics"]["precip"]
    assert m["freq"]["d"] > 8                           # wetter days more often
    assert m["direction"] == "wetter"
    assert m["bias"] > 0


# --- wet-bulb heat-stress-day frequency ---------------------------------------

def test_wetbulb_stress_frequency_present():
    out = scoring.build_scores(make_frame())
    assert out["wetbulb_stress_f"] == scoring.WETBULB_STRESS_F
    freq = out["slices"]["annual"]["metrics"]["wetbulb"]["freq"]
    assert freq["threshold_f"] == scoring.WETBULB_STRESS_F
    assert 0 <= freq["f6"] <= 100 and 0 <= freq["f45"] <= 100
    assert freq["d"] == pytest.approx(freq["f6"] - freq["f45"], abs=0.05)


def test_wetbulb_stress_frequency_rises_with_heat():
    # A large recent wet-bulb shift pushes summer days over the stress threshold.
    out = scoring.build_scores(shift_metric(make_frame(), "wetbulb", 30.0))
    freq = out["slices"]["annual"]["metrics"]["wetbulb"]["freq"]
    assert freq["f6"] > freq["f45"] and freq["d"] > 0


# --- overall aggregation ------------------------------------------------------

def test_overall_is_equal_average_of_metric_mads():
    # Every present metric counts the same — no per-metric weighting.
    entries = {
        "a": {"score": 20, "mad": 3.0},
        "b": {"score": 80, "mad": 12.0},
        "c": {"score": None},              # missing -> dropped from the average
    }
    ov = scoring._overall(entries)
    assert ov["mad"] == 7.5                                # (3.0 + 12.0) / 2
    assert ov["score"] == scoring.score_of(7.5) == 50


def test_missing_gust_column_nulls_out_but_overall_survives():
    out = scoring.build_scores(make_frame(drop=("gust",)))
    ann = out["slices"]["annual"]
    g = ann["metrics"]["gust"]
    assert g["score"] is None and "not available" in g["reason"]
    assert ann["overall"]["score"] is not None          # averaged over present metrics


def test_all_null_metric_reports_not_enough_data():
    df = make_frame().with_columns(pl.lit(None, dtype=pl.Float64).alias("humid"))
    m = scoring.build_scores(df)["slices"]["annual"]["metrics"]["humid"]
    assert m["score"] is None
    assert "not enough" in m["reason"]


# --- guards -------------------------------------------------------------------

def test_short_record_is_unavailable():
    out = scoring.build_scores(make_frame(years=10))
    assert "unavailable" in out
    assert "slices" not in out


# --- unit helpers -------------------------------------------------------------

def test_score_and_tier_mapping():
    assert scoring.score_of(0) == 0
    assert scoring.score_of(scoring.SATURATION) == 100
    assert scoring.score_of(scoring.SATURATION * 2) == 100   # clamped
    assert scoring.tier_of(90, 5)[1] == "rec-hot"
    assert scoring.tier_of(90, -5)[1] == "rec-cold"
    assert scoring.tier_of(5, 1)[0] == "Steady"


def test_divergence_none_below_min_samples():
    base = np.linspace(0, 100, 5000)
    rec = np.linspace(0, 100, scoring.MIN_SLICE_SAMPLES - 1)
    assert scoring.divergence(base, rec) is None