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

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"
SUN_THRESHOLD = 0.5    # `sun` fraction >= this counts as a "sunny" day (>=50% of daylight)
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 among days with ANY rain (> 0) in the seasonal window — a
# "rain percentile". Rain is one-directional (heavier = more extreme), so these 8
# tiers are sequential light->heavy, using the SAME cut points as temperature. Dry
# days (no rain at all) are handled separately (the "dry" class, colored by dry
# streak in the UI). Same [lower, upper) convention as TEMP_BANDS. The eight tiers
# fill the wet-2..wet-9 colour ramp with no gap.
RAIN_BANDS = [
    (99, "Extreme", "wet-9"),       # >99   heaviest rain for the season (darkest)
    (95, "Severe", "wet-8"),        # 95-99  (top half of the old Very Heavy)
    (90, "Very Heavy", "wet-7"),    # 90-95  (lower half)
    (60, "Heavy", "wet-6"),         # 60-90
    (40, "Typical", "wet-5"),       # 40-60
    (25, "Brisk", "wet-4"),         # 25-40
    (10, "Light", "wet-3"),         # 10-25
    (0, "Trace", "wet-2"),          # <10    the 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 pct_ordinal(pct) -> str:
    """A percentile as a display ordinal: 66.4 -> '66th', 99.6 -> '99th'.

    Floored into 1..99 on purpose. An empirical percentile is a rank against the
    sample, so rounding can land on 100 (or 0), and "100th percentile" reads as a
    measurement error rather than "as extreme as it has ever been". The band label
    ("Near Record") carries the how-extreme part.

    Canonical for every surface — the frontend mirrors it as pctOrd() in
    shared.js, so the Day page, the calendar tooltip, the chart, the city pages
    and the homepage strip all say the same thing about the same reading.
    """
    try:
        # floor(x + 0.5), NOT round(): Python's round() is half-to-even, so it
        # gives 16 for 16.5 while JavaScript's Math.round gives 17 — the same
        # reading would then read "16th" server-side and "17th" client-side.
        n = min(99, max(1, math.floor(float(pct) + 0.5)))
    except (TypeError, ValueError):
        return "—"
    suffix = "th" if 10 <= n % 100 <= 20 else {1: "st", 2: "nd", 3: "rd"}.get(n % 10, "th")
    return f"{n}{suffix}"


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 _window(df: pl.DataFrame, target_doy: int) -> pl.DataFrame:
    """The ±7-day seasonal sub-frame around a day-of-year — the reference sample
    every per-day-of-year summary (climatology, day detail, grading) draws from."""
    return df.filter(window_mask(df["doy"].to_numpy(), target_doy))


def _window_meta(sub: pl.DataFrame) -> dict:
    """Sample size and the [min, max] calendar years covered by a window sub-frame."""
    years = sub["date"].dt.year()
    return {
        "n_samples": int(len(sub)),
        "year_range": [int(years.min()), int(years.max())] if len(sub) else None,
    }


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."""
    sub = _window(df, target_doy)
    out = {"target_doy": int(target_doy), **_window_meta(sub)}
    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),
        }
    # Plain-language "insight" frequencies over the same window: the share of days
    # that see measurable rain, and the share that are sunny. Powers the friendly
    # summary sentences on the home page.
    out["rain_freq"] = _rain_frac(sub)
    out["sun_freq"] = _sun_frac(sub)
    return out


def _rain_frac(df: pl.DataFrame) -> float | None:
    """Fraction of days with measurable rain (>= RAIN_THRESHOLD), or None if no
    usable precip data."""
    if "precip" not in df.columns:
        return None
    v = _finite(df["precip"])
    if v.size == 0:
        return None
    return round(float(np.mean(v >= RAIN_THRESHOLD)), 3)


def _sun_frac(df: pl.DataFrame) -> float | None:
    """Fraction of days that are sunny (`sun` fraction >= SUN_THRESHOLD), or None
    if no usable sunshine data (e.g. NASA/MET-sourced cells)."""
    if "sun" not in df.columns:
        return None
    v = _finite(df["sun"])
    if v.size == 0:
        return None
    return round(float(np.mean(v >= SUN_THRESHOLD)), 3)


# Days per calendar month (non-leap); used to turn a rain fraction into an
# average count of rainy days for the month.
_DAYS_IN_MONTH = (31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31)


def month_climate(df: pl.DataFrame) -> dict:
    """Per calendar month (1..12), a compact climatological summary across the full
    history: typical high/low (median), the share of days that see rain / are
    sunny, and the implied average count of rainy days. Feeds the calendar view's
    plain-language insight strip and cross-month ("Nx more likely") comparisons."""
    if "date" not in df.columns:
        return {}
    months = df["date"].dt.month().to_numpy()
    out: dict = {}
    for m in range(1, 13):
        sub = df.filter(pl.Series(months == m))
        if sub.is_empty():
            continue
        tmax = _finite(sub["tmax"]) if "tmax" in sub.columns else np.array([])
        tmin = _finite(sub["tmin"]) if "tmin" in sub.columns else np.array([])
        rain_freq = _rain_frac(sub)
        entry = {
            "n_samples": int(len(sub)),
            "tmax_med": round(float(np.median(tmax)), 1) if tmax.size else None,
            "tmin_med": round(float(np.median(tmin)), 1) if tmin.size else None,
            "rain_freq": rain_freq,
            "sun_freq": _sun_frac(sub),
            "rainy_days": (round(rain_freq * _DAYS_IN_MONTH[m - 1], 1)
                           if rain_freq is not None else None),
        }
        out[str(m)] = entry
    return out


def all_time_records(df: pl.DataFrame) -> dict:
    """All-time record high/low (and the date each occurred) per metric across the
    full archive — the raw material for the records page and the city teaser."""
    out: dict = {}
    for var in CLIMO_METRICS:
        if var not in df.columns:
            continue
        sub = df.select(["date", var]).drop_nulls(var)
        if sub.is_empty():
            continue
        hi = sub.row(int(sub[var].arg_max()), named=True)
        lo = sub.row(int(sub[var].arg_min()), named=True)
        out[var] = {
            "max": round(float(hi[var]), 2),
            "max_date": hi["date"].isoformat() if hasattr(hi["date"], "isoformat") else str(hi["date"]),
            "min": round(float(lo[var]), 2),
            "min_date": lo["date"].isoformat() if hasattr(lo["date"], "isoformat") else str(lo["date"]),
        }
    return out


def longest_dry_streak(df: pl.DataFrame) -> tuple[int, str | None]:
    """Longest run of consecutive days without any rain, and the ISO date the streak
    began. A dry day has no rain at all (precip <= 0); null precip counts as dry. The
    threshold matches _grade_precip's dry/rain split — any measurable rain, however
    slight, is a rain day that breaks the streak (not the 0.01" rain-frequency line)."""
    if "precip" not in df.columns:
        return (0, None)
    d = df.select(["date", "precip"]).sort("date")
    dates, precips = d["date"].to_list(), d["precip"].to_list()
    best_len, best_start = 0, None
    cur_len, cur_start = 0, None
    for dt, p in zip(dates, precips):
        wet = p is not None and not (isinstance(p, float) and np.isnan(p)) and p > 0
        if wet:
            cur_len, cur_start = 0, None
        else:
            if cur_len == 0:
                cur_start = dt
            cur_len += 1
            if cur_len > best_len:
                best_len, best_start = cur_len, cur_start
    start = best_start.isoformat() if best_start and hasattr(best_start, "isoformat") else (
        str(best_start) if best_start else None)
    return (best_len, start)


def _band_stats(samples: np.ndarray) -> dict | None:
    if samples.size == 0:
        return None
    # Percentiles for the chart's nested "normal" fan. 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. p95 additionally splits the rain fan's top
    # region into Very Heavy (90-95) and Severe (95-99); unused by the temperature fan.
    p1, p10, p25, p40, p50, p60, p75, p90, p95, p99 = np.percentile(
        samples, [1, 10, 25, 40, 50, 60, 75, 90, 95, 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),
        "p95": round(float(p95), 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* (any measurable rain, > 0) in the window. A day with no
    rain at all gets the "dry" class with no percentile (the UI colors it by dry
    streak instead); any rain, however slight, is at least a Trace day."""
    if value is None or (isinstance(value, float) and np.isnan(value)):
        return None
    value = float(value)
    if value <= 0:
        return {"value": round(value, 2), "percentile": None, "grade": "Dry", "class": "dry"}
    rain = samples[samples > 0]
    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 day with any rain, walking a
    chronological precip series. Any measurable rain (precip > 0) resets the count, so
    this matches _grade_precip's dry/rain split; 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 > 0
        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)}


# The percentile marks the rain ladder needs — every band's lower threshold except
# 0 (the bottom tier bottoms out at the smallest rain day). Derived from RAIN_BANDS
# so adding or splitting a tier can't leave this behind.
_RAIN_MARKS = sorted({thr for thr, _, _ in RAIN_BANDS if thr > 0})


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 > 0]           # any measurable rain (dry == no rain)
    n = int(samples.size)
    tiers = []
    if rain.size:
        marks = {m: round(float(np.percentile(rain, m)), 2) for m in _RAIN_MARKS}
        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": "none", "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 = _window(df, doy)

    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,
        **_window_meta(sub),
        "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
    sub = _window(df, 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