"""Value-drift check: quantify how much the migration moved the data, per variable
and by grade band, by comparing the two live history sources head-to-head.

Open-Meteo is kept as a dormant fallback, so both sources are still fetchable: this
pulls NASA POWER (the new primary) and the Open-Meteo archive (ERA5) for a sample of
cells and reports, per variable:
  - mean-absolute-difference / bias / max — the RAW value drift, and
  - grade-band divergence — the share of (day, metric) whose graded band actually
    changes between sources.
Since Open-Meteo IS ERA5, this doubles as a fidelity check for the ERA5 seed.

    python drift_check.py [--limit N] [--days N]   # on a networked box

Note a real property of the grading model worth reading the two numbers together:
grades are percentiles within each source's OWN distribution, so a uniform bias
between sources shifts MAD but leaves grades unchanged. High MAD with low grade
divergence means "different absolute values, same story" — usually fine. High grade
divergence is the signal that actually matters.

The comparison logic (compare_values / grade_band_divergence) is unit-tested; the
fetch is not.
"""
import datetime
import sys

import polars as pl

from data import cities
from data import climate
from data import grading
from data import grid

# Variables compared for raw drift (feels is derived; precip/wind/gust/humid optional).
_METRICS = ("tmax", "tmin", "precip", "wind", "gust", "humid", "feels")


def compare_values(a: pl.DataFrame, b: pl.DataFrame, metrics=_METRICS) -> dict:
    """Per-variable raw-drift stats over the dates the two frames share: n aligned
    days, mean-absolute-difference, mean bias (a - b), and max absolute difference.
    Days where either source is null for a metric are dropped from that metric."""
    j = a.join(b, on="date", how="inner", suffix="_b")
    out = {}
    for m in metrics:
        if m not in a.columns or m not in b.columns:
            continue
        d = j.select((pl.col(m) - pl.col(f"{m}_b")).alias("d")).drop_nulls()
        if d.height == 0:
            out[m] = {"n": 0, "mad": None, "bias": None, "max_abs": None}
            continue
        diff = d["d"]
        out[m] = {
            "n": d.height,
            "mad": round(float(diff.abs().mean()), 3),
            "bias": round(float(diff.mean()), 3),
            "max_abs": round(float(diff.abs().max()), 3),
        }
    return out


def grade_band_divergence(a: pl.DataFrame, b: pl.DataFrame, start, end) -> dict:
    """Per-metric share of days whose graded band differs between the two sources over
    [start, end]. Each source's days are graded against its OWN ±7-day climatology
    (grading.grade_range), so this measures whether the migration changes the story a
    day tells, not just its raw value."""
    ga = {d["date"]: d for d in grading.grade_range(a, start, end)}
    gb = {d["date"]: d for d in grading.grade_range(b, start, end)}
    per: dict[str, list[int]] = {}
    for date in ga.keys() & gb.keys():
        da, db = ga[date], gb[date]
        for m in grading.CLIMO_METRICS:
            va, vb = da.get(m), db.get(m)
            if va and vb:
                slot = per.setdefault(m, [0, 0])   # [compared, differ]
                slot[0] += 1
                slot[1] += int(va["g"] != vb["g"])
    return {m: {"compared": c, "differ": d, "pct": round(100.0 * d / c, 1) if c else None}
            for m, (c, d) in sorted(per.items())}


def check_cell(cell: dict, days: int) -> dict:
    """Fetch both sources for a cell and compare them over the most recent `days`."""
    nasa = climate._fetch_history_nasa(cell)      # new primary (+ Meteostat gusts)
    om = climate._fetch_history(cell)             # Open-Meteo archive (ERA5)
    end = min(nasa["date"].max(), om["date"].max())
    start = end - datetime.timedelta(days=days)
    return {"values": compare_values(nasa, om),
            "grades": grade_band_divergence(nasa, om, start, end)}


def _mean(xs: list) -> "float | None":
    xs = [x for x in xs if x is not None]
    return round(sum(xs) / len(xs), 3) if xs else None


def main(limit: "int | None" = None, days: int = 365) -> None:
    todo = cities.all_cities()
    if limit:
        todo = todo[:limit]
    mad_by_metric: dict[str, list] = {}
    grade_by_metric: dict[str, list] = {}
    checked = failed = 0
    for i, c in enumerate(todo, 1):
        cell = grid.snap(c["lat"], c["lon"])
        try:
            rep = check_cell(cell, days)
        except Exception as e:  # noqa: BLE001 - one cell's outage shouldn't abort the sweep
            failed += 1
            print(f"[{i}/{len(todo)}] FAILED {c['slug']}: {e}")
            continue
        checked += 1
        vals = ", ".join(f"{m} mad={s['mad']}" for m, s in rep["values"].items()
                         if s.get("mad") is not None)
        grds = ", ".join(f"{m} {s['pct']}%" for m, s in rep["grades"].items())
        print(f"[{i}/{len(todo)}] {c['slug']} ({cell['id']})")
        print(f"    values: {vals}")
        print(f"    grade drift: {grds}")
        for m, s in rep["values"].items():
            mad_by_metric.setdefault(m, []).append(s.get("mad"))
        for m, s in rep["grades"].items():
            grade_by_metric.setdefault(m, []).append(s.get("pct"))

    print(f"\n=== summary over {checked} cells (failed={failed}) ===")
    print("mean MAD (raw drift):")
    for m in _METRICS:
        if m in mad_by_metric:
            print(f"  {m:7s} {_mean(mad_by_metric[m])}")
    print("mean grade-band divergence (share of days a grade changes):")
    for m in sorted(grade_by_metric):
        print(f"  {m:7s} {_mean(grade_by_metric[m])}%")


if __name__ == "__main__":
    argv = sys.argv[1:]
    lim = int(argv[argv.index("--limit") + 1]) if "--limit" in argv else None
    dys = int(argv[argv.index("--days") + 1]) if "--days" in argv else 365
    main(limit=lim, days=dys)