"""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 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", *SEASONS)                 # payload slice keys

# Per-metric descriptor — display label, overall-score weight, and the
# (positive-drift, negative-drift) direction words — in ONE table so a scored
# metric is defined in a single place. Temps / humidity / feels-like dominate the
# weighting; wet bulb medium; wind + precip lightest. 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",  "weight": 1.5,  "dir": ("warmer", "cooler")},
    "tmin":    {"label": "Low temp",   "weight": 1.5,  "dir": ("warmer", "cooler")},
    "feels":   {"label": "Feels like", "weight": 2.0,  "dir": ("warmer", "cooler")},
    "humid":   {"label": "Humidity",   "weight": 2.0,  "dir": ("muggier", "drier")},
    "wetbulb": {"label": "Wet bulb",   "weight": 1.25, "dir": ("warmer", "cooler")},
    "wind":    {"label": "Wind",       "weight": 0.5,  "dir": ("windier", "calmer")},
    "gust":    {"label": "Gusts",      "weight": 0.5,  "dir": ("gustier", "calmer")},
    "precip":  {"label": "Precip",     "weight": 0.75, "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,
            "weight": METRICS[metric]["weight"], "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}",
        "weight": METRICS[metric]["weight"], "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:
    """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: 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
    wsum = sum(e["weight"] for e in present)
    mad = sum(e["weight"] * e["mad"] for e in present) / wsum
    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,
    }