thermograph/grading.py

"""Turn a raw daily-weather record into day-of-year climatology and letter grades.

For a given day of the year we build the reference distribution from every
historical day whose day-of-year is within +/- `HALF_WINDOW` days of it (wrapping
around the year end). An observed value is then placed on that distribution as an
empirical percentile and mapped to a human-readable grade.
"""
import datetime
import warnings

import numpy as np
import polars as pl

HALF_WINDOW = 7  # +/- 7 days -> a 15-day seasonal window
RAIN_THRESHOLD = 0.01  # inches; a day with < this much precip counts as "dry"
DSR_LOOKBACK_MIN_DAYS = 14  # min history required before a window to seed dry streaks

# Metrics graded on the diverging temperature scale (percentile -> TEMP_BANDS):
# the two air temps, the combined "feels like" (heat index / wind chill) plus its
# separate apparent high/low sides (fmax/fmin — the calendar recombines them
# client-side around a user-chosen comfort temperature), wind speed and wind
# gust. Each is a single daily scalar placed on its own ±7-day historical
# distribution, so calm/low reads as the cool (blue) end and windy/high as the
# hot (red) end — same coloring + categories as temperature.
TEMP_METRICS = ("tmax", "tmin", "feels", "fmax", "fmin", "humid", "wind", "gust")
# All metrics that get a percentile summary (temperature-like + precipitation).
CLIMO_METRICS = TEMP_METRICS + ("precip",)

# Percentile -> (grade label, css class) for temperature. Higher percentile = warmer.
# 9 symmetric tiers around the middle 40-60% "Normal", escalating to "Near Record"
# at both edges. Each entry's number is the tier's LOWER bound; a tier spans
# [lower, next-lower) — i.e. lower-inclusive, upper-exclusive (a p90 day is "Very
# High", not "High"). Boundaries sit on multiples of 5 (plus the 1/99 record edges).
TEMP_BANDS = [
    (99, "Near Record", "rec-hot"),      # >=99  extreme high (danger)
    (90, "Very High", "very-hot"),       # 90-99
    (75, "High", "hot"),                 # 75-90
    (60, "Above Normal", "warm"),        # 60-75
    (40, "Normal", "normal"),            # 40-60 (the middle)
    (25, "Below Normal", "cool"),        # 25-40
    (10, "Low", "cold"),                 # 10-25
    (1, "Very Low", "very-cold"),        # 1-10
    (0, "Near Record", "rec-cold"),      # <1    extreme low (danger)
]

# Precipitation is graded ONLY among days that actually rained (>= RAIN_THRESHOLD)
# in the seasonal window — a "rain percentile". Rain is one-directional (heavier =
# more extreme), so these 9 tiers are sequential light->heavy, using the SAME cut
# points as temperature. Dry days are handled separately (the "dry" class, colored
# by dry streak in the UI). Same [lower, upper) convention as TEMP_BANDS.
RAIN_BANDS = [
    (99, "Extreme", "wet-9"),       # heaviest rain for the season (darkest)
    (90, "Very Heavy", "wet-8"),    # 90-99
    (75, "Heavy", "wet-7"),         # 75-90
    (60, "Mod–Heavy", "wet-6"),     # 60-75
    (40, "Moderate", "wet-5"),      # 40-60
    (25, "Light–Mod", "wet-4"),     # 25-40
    (10, "Light", "wet-3"),         # 10-25
    (1, "Very Light", "wet-2"),     # 1-10
    (0, "Trace", "wet-1"),          # <1    lightest measurable rain
]


def _ladder_from(bands, bottom_lo=None):
    """Derive the single-day detail view's tier ladder from a band table, so the
    two can never drift apart: each tier as (class, label, printable percentile
    range, lower-bound pct, upper-bound pct). The top tier is open-ended (>99);
    the bottom runs down from the 1st percentile — ``bottom_lo`` marks its lower
    bound (None for temperature; 0 for rain, whose lightest tier bottoms out at
    the smallest measured rain day)."""
    top_thr, top_label, top_css = bands[0]
    out = [(top_css, top_label, f">{top_thr}%", top_thr, None)]
    for (thr, label, css), (prev_thr, _, _) in zip(bands[1:-1], bands[:-2]):
        out.append((css, label, f"{thr}–{prev_thr}%", thr, prev_thr))
    edge = bands[-2][0]
    out.append((bands[-1][2], bands[-1][1], f"<{edge}%", bottom_lo, edge))
    return out


_TEMP_LADDER = _ladder_from(TEMP_BANDS)
# Rain tiers are on the rain-day-only percentile scale (see _grade_precip).
_RAIN_LADDER = _ladder_from(RAIN_BANDS, bottom_lo=0)


def _band(pct: float, bands) -> tuple[str, str]:
    # The top tier is strict (pct > its threshold): "Near Record" high means
    # strictly beyond the 99th percentile — the top <1% — mirroring the
    # strictly-below-1st bottom tier (its lower neighbor already catches pct >= 1).
    # Everything in between stays lower-inclusive, upper-exclusive.
    top_thr, top_label, top_css = bands[0]
    if pct > top_thr:
        return top_label, top_css
    for threshold, label, css in bands[1:]:
        if pct >= threshold:
            return label, css
    return bands[-1][1], bands[-1][2]


def _as_date(d) -> datetime.date:
    """Coerce a date-ish value (``date``, ``datetime``, or ISO string) to a plain
    ``datetime.date`` — the one date type the grading layer works in."""
    if isinstance(d, str):
        return datetime.date.fromisoformat(d[:10])
    if isinstance(d, datetime.datetime):
        return d.date()
    return d


def _finite(col: pl.Series) -> np.ndarray:
    """Finite float values of a polars Series as an ndarray, with nulls and NaN
    dropped — the frame→numpy bridge every percentile routine grades on."""
    a = np.asarray(col.to_numpy(), dtype="float64")
    return a[~np.isnan(a)]


def window_mask(doys: np.ndarray, target_doy: int, half: int = HALF_WINDOW) -> np.ndarray:
    diff = np.abs(doys.astype(int) - int(target_doy))
    circular = np.minimum(diff, 366 - diff)
    return circular <= half


def empirical_percentile(samples: np.ndarray, value) -> float | None:
    """Mid-rank percentile of `value` within `samples` (handles ties correctly)."""
    n = samples.size
    if n == 0 or value is None or (isinstance(value, float) and np.isnan(value)):
        return None
    less = int(np.sum(samples < value))
    equal = int(np.sum(samples == value))
    return round(100.0 * (less + 0.5 * equal) / n, 1)


def climatology(df: pl.DataFrame, target_doy: int) -> dict:
    """Summarize the +/-7 day historical distribution for one day of the year."""
    doys = df["doy"].to_numpy()
    sub = df.filter(window_mask(doys, target_doy))
    years = sub["date"].dt.year()
    out = {
        "target_doy": int(target_doy),
        "n_samples": int(len(sub)),
        "year_range": [int(years.min()), int(years.max())] if len(sub) else None,
    }
    for var in CLIMO_METRICS:
        if var not in sub.columns:
            out[var] = None
            continue
        v = _finite(sub[var])
        if v.size == 0:
            out[var] = None
            continue
        p10, p40, p50, p60, p90 = np.percentile(v, [10, 40, 50, 60, 90])
        out[var] = {
            "min": round(float(np.min(v)), 2),
            "p10": round(float(p10), 1),
            "p40": round(float(p40), 1),   # Normal band is now 40-60 (see TEMP_BANDS)
            "p50": round(float(p50), 1),
            "p60": round(float(p60), 1),
            "p90": round(float(p90), 1),
            "max": round(float(np.max(v)), 2),
            "mean": round(float(np.mean(v)), 1),
        }
    return out


def _band_stats(samples: np.ndarray) -> dict | None:
    if samples.size == 0:
        return None
    # Percentiles for the chart's nested "normal" fan, matching the 9 tier bounds:
    # p40-p60 is the Normal band; p25/p75, p10/p90 and p1/p99 mark the successive
    # Below/Above Normal, Low/High, Very Low/High and Near-Record edges.
    p1, p10, p25, p40, p50, p60, p75, p90, p99 = np.percentile(
        samples, [1, 10, 25, 40, 50, 60, 75, 90, 99]
    )
    return {
        "p1": round(float(p1), 1),
        "p10": round(float(p10), 1),
        "p25": round(float(p25), 1),
        "p40": round(float(p40), 1),
        "p50": round(float(p50), 1),
        "p60": round(float(p60), 1),
        "p75": round(float(p75), 1),
        "p90": round(float(p90), 1),
        "p99": round(float(p99), 1),
    }


def _grade_value(samples: np.ndarray, value, bands) -> dict | None:
    """Grade a temperature value by its empirical percentile in the window."""
    pct = empirical_percentile(samples, value)
    if pct is None:
        return None
    label, css = _band(pct, bands)
    return {"value": round(float(value), 2), "percentile": pct, "grade": label, "class": css}


def _grade_precip(samples: np.ndarray, value) -> dict | None:
    """Grade precipitation by its "rain percentile" — the rank of the day's rainfall
    among *rain days only* (>= RAIN_THRESHOLD) in the window. Dry days get the "dry"
    class with no percentile (the UI colors them by dry streak instead)."""
    if value is None or (isinstance(value, float) and np.isnan(value)):
        return None
    value = float(value)
    if value < RAIN_THRESHOLD:
        return {"value": round(value, 2), "percentile": None, "grade": "Dry", "class": "dry"}
    rain = samples[samples >= RAIN_THRESHOLD]
    pct = empirical_percentile(rain, value)
    if pct is None:            # no historical rain days in window (extremely rare)
        pct = 100.0
    label, css = _band(pct, RAIN_BANDS)
    return {"value": round(value, 2), "percentile": pct, "grade": label, "class": css}


def dry_streaks(dates, precips) -> dict[str, int]:
    """Map each date (ISO string) to days since the last measurable rain, walking a
    chronological precip series. Missing precip counts as a dry day. `dates` is an
    iterable of ``datetime.date`` (a polars Date column's ``.to_list()``)."""
    out: dict[str, int] = {}
    streak = 0
    for d, p in zip(dates, precips):
        wet = p is not None and not (isinstance(p, float) and np.isnan(p)) and p >= RAIN_THRESHOLD
        streak = 0 if wet else streak + 1
        out[_as_date(d).isoformat()] = streak
    return out


def grade_range(df: pl.DataFrame, start, end) -> list[dict]:
    """Grade every historical day in [start, end] against its own ±7-day window.

    Powers the calendar view. Reuses one window per day-of-year across the whole
    range, so a 2-year span computes at most ~366 windows (not one per day). Each
    day is returned in a compact shape: value (v), percentile (pct), css class (c),
    grade label (g).
    """
    start, end = _as_date(start), _as_date(end)
    full = df.sort("date")
    doys_all = full["doy"].to_numpy()
    cols = {v: np.asarray(full[v].to_numpy(), dtype="float64")
            for v in CLIMO_METRICS if v in full.columns}
    masks: dict[int, np.ndarray] = {}
    cache: dict[tuple[int, str], np.ndarray] = {}

    def samples(doy: int, var: str) -> np.ndarray:
        key = (doy, var)
        if key not in cache:
            mask = masks.get(doy)
            if mask is None:
                mask = masks[doy] = window_mask(doys_all, doy)
            arr = cols[var][mask]
            cache[key] = arr[~np.isnan(arr)]
        return cache[key]

    # Days since last measurable rain. Computed over the ENTIRE record that precedes
    # the window — not just the window, nor a fixed N-day buffer — so the streak on
    # the first shown day is exact even when a dry spell straddles the window start.
    # A fixed 14-day lookback would be the bare minimum but still undercounts longer
    # droughts (the record has 25-day dry streaks); using the full history (already
    # cached per cell) is both correct for any streak length and free.
    if full["date"].min() > start - datetime.timedelta(days=DSR_LOOKBACK_MIN_DAYS):
        # Should never happen (history starts in 1980); guards against a future
        # change that trims history and would silently truncate streaks.
        warnings.warn("dry-streak lookback shorter than the 14-day minimum", stacklevel=2)
    dsr_map = dry_streaks(full["date"].to_list(), cols["precip"])

    sub = full.filter((pl.col("date") >= start) & (pl.col("date") <= end))
    out = []
    for row in sub.iter_rows(named=True):
        doy = int(row["doy"])
        date = row["date"].isoformat()
        rec = {"date": date, "dsr": dsr_map.get(date)}
        for var in TEMP_METRICS:
            if var not in cols:
                rec[var] = None
                continue
            g = _grade_value(samples(doy, var), row[var], TEMP_BANDS)
            rec[var] = {"v": g["value"], "pct": g["percentile"], "c": g["class"], "g": g["grade"]} if g else None
        gp = _grade_precip(samples(doy, "precip"), row["precip"])
        rec["precip"] = {"v": gp["value"], "pct": gp["percentile"], "c": gp["class"], "g": gp["grade"]} if gp else None
        out.append(rec)
    return out


def _temp_ladder(samples: np.ndarray) -> dict | None:
    """Value at each temperature tier boundary within the ±7-day window."""
    if samples.size == 0:
        return None
    marks = {m: round(float(np.percentile(samples, m)), 1) for m in (1, 10, 25, 40, 50, 60, 75, 90, 99)}
    tiers = [
        {"c": c, "label": label, "range": rng,
         "lo": marks[lo] if lo is not None else None,
         "hi": marks[hi] if hi is not None else None}
        for c, label, rng, lo, hi in _TEMP_LADDER
    ]
    return {"tiers": tiers, "median": marks[50],
            "min": round(float(samples.min()), 1), "max": round(float(samples.max()), 1)}


def _precip_ladder(samples: np.ndarray) -> dict | None:
    """Value at each rain-day tier boundary, plus how often the window is dry."""
    if samples.size == 0:
        return None
    rain = samples[samples >= RAIN_THRESHOLD]
    n = int(samples.size)
    tiers = []
    if rain.size:
        marks = {m: round(float(np.percentile(rain, m)), 2) for m in (1, 10, 25, 40, 60, 75, 90, 99)}
        rmin = round(float(rain.min()), 2)
        tiers = [
            {"c": c, "label": label, "range": rng,
             "lo": rmin if lo == 0 else marks[lo],
             "hi": marks[hi] if hi is not None else None}
            for c, label, rng, lo, hi in _RAIN_LADDER
        ]
    tiers.append({"c": "dry", "label": "Dry", "range": f"< {RAIN_THRESHOLD}\"", "lo": 0.0, "hi": None})
    return {"tiers": tiers, "dry_pct": round(100.0 * (n - rain.size) / n, 1),
            "rain_days": int(rain.size),
            "min": round(float(samples.min()), 2), "max": round(float(samples.max()), 2)}


def day_detail(df: pl.DataFrame, date, obs: dict | None) -> dict:
    """Full percentile breakdown for one day: the value at every tier boundary in
    its own ±7-day window, plus where the observed values (if any) land.

    Powers the single-day detail page. `obs` may be None when the date isn't yet
    in the record — then only the climatological ladders are returned."""
    d = _as_date(date)
    doy = d.timetuple().tm_yday
    sub = df.filter(window_mask(df["doy"].to_numpy(), doy))
    years = sub["date"].dt.year()

    metrics = {}
    for var in TEMP_METRICS:
        if var not in sub.columns:
            continue
        vals = _finite(sub[var])
        metrics[var] = {
            "ladder": _temp_ladder(vals),
            "obs": _grade_value(vals, obs.get(var) if obs else None, TEMP_BANDS),
        }
    precip = _finite(sub["precip"])
    metrics["precip"] = {
        "ladder": _precip_ladder(precip),
        "obs": _grade_precip(precip, obs.get("precip") if obs else None),
    }
    return {
        "date": d.isoformat(),
        "doy": doy,
        "n_samples": int(len(sub)),
        "year_range": [int(years.min()), int(years.max())] if len(sub) else None,
        "metrics": metrics,
    }


def grade_day(df: pl.DataFrame, date, obs: dict) -> dict:
    """Grade one observed day against its own day-of-year +/-7 window."""
    d = _as_date(date)
    doy = d.timetuple().tm_yday
    doys = df["doy"].to_numpy()
    sub = df.filter(window_mask(doys, doy))

    result = {"date": d.isoformat(), "doy": doy}
    for var in TEMP_METRICS:
        result[var] = (
            _grade_value(_finite(sub[var]), obs.get(var), TEMP_BANDS)
            if var in sub.columns else None
        )
    result["precip"] = _grade_precip(_finite(sub["precip"]), obs.get("precip"))

    # Per-day normal band (this day-of-year's ±7 window) so the trend chart can
    # draw the "normal" envelope each actual value is compared against.
    result["normals"] = {
        var: _band_stats(_finite(sub[var]))
        for var in CLIMO_METRICS if var in sub.columns
    }

    # A single "departure" score: how far the day strayed from the median (50th pct),
    # taking the most extreme of high/low. 0 = perfectly normal, 50 = record extreme.
    departures = [
        abs(result[v]["percentile"] - 50)
        for v in ("tmax", "tmin")
        if result[v] is not None
    ]
    result["departure"] = round(max(departures), 1) if departures else None
    return result