"""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

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"]
    return entry


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 headline bias comes only from the temperature-direction metrics so the
    overall reads 'warmer / cooler' rather than being muddied by wind/precip."""
    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, bias)
    direction = "warmer" if bias >= 0 else "cooler"
    return {"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}"}


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()

    slices = {}
    for key in SLICES:
        base_s, rec_s = _slice(baseline, key), _slice(recent, key)
        entries = {}
        for m in SCORE_METRICS:
            if m not in history.columns:
                entries[m] = _null_entry(m, "not available for this location")
                continue
            entries[m] = _metric_entry(m, grading._finite(base_s[m]), grading._finite(rec_s[m]))
        slices[key] = {"overall": _overall(entries), "metrics": entries}

    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,
        "slices": slices,
    }