"""Row math for the Video Benchmark Space.

Pure functions that take ``list[dict]`` rows (a subset of each
benchmark's results schema — projected down to the benchmark's
``keep_keys`` in ``schema.py``) and return JSON-serializable values. The Gradio handlers
in ``app.py`` call these to build the Results table heatmap, the
Leaderboards bar chart, and the three Compare-tab plots.

Nothing here imports Gradio, pandas, or plotly: the module is purely
arithmetic so it stays trivially testable and reusable from notebooks.
"""
from __future__ import annotations

from typing import Any, Iterable


Row = dict[str, Any]
Scale = dict[str, Any]  # {"min": float, "max": float, "lower": bool}


# --------------------------------------------------------------------------- #
# Helpers
# --------------------------------------------------------------------------- #
def _is_num(v: Any) -> bool:
    """True for finite ints/floats; False for ``None``, NaN, and bools.

    Bools are explicitly rejected so accidental flag columns can't sneak
    into numeric aggregations.
    """
    if v is None:
        return False
    if isinstance(v, bool):
        return False
    if isinstance(v, (int, float)):
        # NaN is the only float that's not equal to itself.
        return v == v
    return False


# --------------------------------------------------------------------------- #
# Per-metric normalization (used by the Results table color ramp and the
# Leaderboards weighted score).
# --------------------------------------------------------------------------- #
def normalize_scale(rows: list[Row], col: dict[str, Any]) -> Scale:
    """Return ``{min, max, lower}`` for a metric column.

    Used to map raw values into [0, 1] before feeding the rdylgn color
    ramp: ``t = (v - min) / (max - min); if lower: t = 1 - t``. When all
    values are missing, returns a degenerate scale tagged with
    ``empty=True`` so callers can fall back to mid-gray.
    """
    vals = [r.get(col["key"]) for r in rows]
    nums = [v for v in vals if _is_num(v)]
    lower = bool(col.get("lower"))
    if not nums:
        return {"min": 0.0, "max": 0.0, "lower": lower, "empty": True}
    lo = min(nums)
    hi = max(nums)
    return {"min": lo, "max": hi, "lower": lower}


def _normalize_value(v: Any, scale: Scale) -> float:
    """Map ``v`` into [0, 1] using a scale from :func:`normalize_scale`.

    Returns ``0.5`` for empty scales, ``0.0`` for non-numeric values, and
    inverts the result when the column is polarity-``lower`` so that
    ``1.0`` always means "best on this axis".
    """
    if scale.get("empty"):
        return 0.5
    if not _is_num(v):
        return 0.0
    lo = scale["min"]
    hi = scale["max"]
    if hi == lo:
        return 0.5
    t = (v - lo) / (hi - lo)
    return (1 - t) if scale["lower"] else t


def dedupe_latest(rows: list[Row], config_keys: Iterable[str]) -> list[Row]:
    """Collapse re-runs of the same config down to a single best row.

    Workers append a fresh row every time a configuration is benchmarked,
    so the Hub can hold several rows that are identical across every
    ``config_keys`` field and differ only in their metrics. Grouping by
    that key tuple, we keep the most complete row (fewest ``None`` / NaN
    fields), so an entry with missing items never wins over one without.
    Ties fall back to the most samples (``num_samples``) and then the
    latest ``created_at`` (ISO-8601 strings sort chronologically).
    """
    keys = tuple(config_keys)

    def rank(r: Row) -> tuple[int, int, str]:
        missing = sum(v is None or (isinstance(v, float) and v != v)
                      for v in r.values())
        n = r.get("num_samples")
        return (-missing,
                n if isinstance(n, int) and not isinstance(n, bool) else -1,
                str(r.get("created_at") or ""))

    best: dict[tuple, Row] = {}
    order: list[tuple] = []
    for r in rows:
        k = tuple(r.get(c) for c in keys)
        if k not in best:
            order.append(k)
            best[k] = r
        elif rank(r) >= rank(best[k]):
            best[k] = r
    return [best[k] for k in order]


def composite_rank(rows: list[Row], metric_cols: list[dict[str, Any]]) -> list[Row]:
    """Sort rows by their composite score across the visible metric columns.

    Score is the unweighted average of ``(1 - normalized_value)`` across
    ``metric_cols``, so *lower* = better. Sort is stable: ties fall back
    to original input order so repeated calls on the same input produce
    the same ordering.
    """
    if not rows:
        return []
    if not metric_cols:
        return list(rows)

    scales = [(c, normalize_scale(rows, c)) for c in metric_cols]

    def score(r: Row) -> float:
        s = 0.0
        n = 0
        for c, sc in scales:
            v = r.get(c["key"])
            s += (1 - _normalize_value(v, sc))
            n += 1
        return s / n if n else 0.0

    indexed = list(enumerate(rows))
    indexed.sort(key=lambda pair: (score(pair[1]), pair[0]))
    return [r for _, r in indexed]


# --------------------------------------------------------------------------- #
# Filtering
# --------------------------------------------------------------------------- #
def filter_rows(rows: list[Row], filters: dict[str, Iterable[str]] | None) -> list[Row]:
    """Keep rows whose column values match the Results-tab chip selections.

    ``filters`` maps a column key to the set of allowed values. Values
    are always compared as strings (the CheckboxGroups emit strings,
    even for numeric columns like ``g`` / ``crf``), so row values are
    coerced via ``str()`` before lookup. An empty / missing entry for a
    key disables that filter.
    """
    if not filters:
        return list(rows)
    cleaned: dict[str, set[str]] = {}
    for k, vs in filters.items():
        if not vs:
            continue
        s = {str(v) for v in vs}
        if s:
            cleaned[k] = s
    if not cleaned:
        return list(rows)
    out: list[Row] = []
    for r in rows:
        keep = True
        for k, allowed in cleaned.items():
            if str(r.get(k)) not in allowed:
                keep = False
                break
        if keep:
            out.append(r)
    return out


def metric_scales(rows: list[Row], cols: list[dict[str, Any]]) -> dict[str, Scale]:
    """Return ``{column_key: scale}`` for every metric column in ``cols``.

    Pre-computing scales once per render lets the heatmap colorize each
    cell without re-walking the rows for each column.
    """
    return {c["key"]: normalize_scale(rows, c) for c in cols if c.get("metric")}


# --------------------------------------------------------------------------- #
# Leaderboards
# --------------------------------------------------------------------------- #
def leaderboard(
    rows: list[Row], ts: str, cat: str, *,
    axes: list[str],
    cats: dict[str, Any],
    group_keys: tuple[str, ...],
    columns: list[dict[str, Any]],
) -> dict[str, Any]:
    """Rank configs for ``cat`` at access pattern ``ts`` using caller-supplied config.

    Keeps rows matching ``ts``, groups by ``group_keys`` (averaging each axis
    across repeats), normalizes per axis, and scores by the weighted sum of
    ``(1 - normalized)`` from ``cats[cat]['weights']`` (lower = better). Returns
    ``{axes, axis_keys, items}`` with each ``items[i].values`` oriented so 1.0 =
    best. Raises ``ValueError`` for an unknown category and ``RuntimeError`` if an
    axis is missing from ``columns``.
    """
    if cat not in cats:
        raise ValueError(f"unknown leaderboard category: {cat!r}")
    weights = cats[cat]["weights"]
    cols_by_key = {c["key"]: c for c in columns}
    axis_cols = [cols_by_key.get(k) for k in axes]
    if any(c is None for c in axis_cols):
        missing = [k for k, c in zip(axes, axis_cols) if c is None]
        raise RuntimeError(f"missing column metadata for axes: {missing}")

    scoped = [r for r in rows if r.get("timestamps_mode") == ts]

    groups: dict[tuple, list[Row]] = {}
    order: list[tuple] = []
    for r in scoped:
        key = tuple(r.get(k) for k in group_keys)
        if key not in groups:
            groups[key] = []
            order.append(key)
        groups[key].append(r)

    aggregated: list[Row] = []
    for key in order:
        rs = groups[key]
        ag: Row = dict(rs[0])
        for c in axis_cols:
            vals = [r.get(c["key"]) for r in rs]
            nums = [v for v in vals if _is_num(v)]
            ag[c["key"]] = sum(nums) / len(nums) if nums else None
        aggregated.append(ag)

    if not aggregated:
        return {"axes": [c["short"] or c["label"] for c in axis_cols],
                "axis_keys": axes, "items": []}

    scales = [normalize_scale(aggregated, c) for c in axis_cols]

    scored: list[tuple[float, int, Row, list[float]]] = []
    for idx, row in enumerate(aggregated):
        sum_w = 0.0
        wsum = 0.0
        values: list[float] = []
        for i, c in enumerate(axis_cols):
            w = weights.get(c["key"], 1)
            n = _normalize_value(row.get(c["key"]), scales[i])
            values.append(n)
            sum_w += w * (1 - n)
            wsum += w
        score = sum_w / wsum if wsum else 0.0
        scored.append((score, idx, row, values))

    scored.sort(key=lambda t: (t[0], t[1]))
    items = [
        {"row": _slim_row(r, group_keys, axes), "values": vals, "score": s}
        for s, _, r, vals in scored
    ]
    return {
        "axes": [c["short"] or c["label"] for c in axis_cols],
        "axis_keys": axes,
        "items": items,
    }


def _slim_row(r: Row, group_keys: tuple[str, ...], axes: list[str]) -> Row:
    keep = (*group_keys, "timestamps_mode", "repo_id", *axes)
    return {k: r.get(k) for k in keep if k in r}


# --------------------------------------------------------------------------- #
# Compare-tab aggregations
# --------------------------------------------------------------------------- #
def compare_bar(rows: list[Row], metric: str, group_by: str) -> list[dict[str, Any]]:
    """Mean of ``metric`` grouped by ``group_by``, sorted descending by value.

    Returns ``[{"k": label, "v": mean}, ...]``. The ``lerobot/`` prefix is
    stripped from dataset labels so x-axis ticks stay readable when
    ``group_by="repo_id"``.
    """
    g: dict[str, list[float]] = {}
    order: list[str] = []
    for r in rows:
        v = r.get(metric)
        if not _is_num(v):
            continue
        k = str(r.get(group_by))
        if k not in g:
            g[k] = []
            order.append(k)
        g[k].append(v)
    out = [
        {"k": k.replace("lerobot/", ""), "v": sum(g[k]) / len(g[k])}
        for k in order
        if g[k]
    ]
    out.sort(key=lambda d: d["v"], reverse=True)
    return out


def compare_scatter(
    rows: list[Row], *, x: str, y: str, label_keys: tuple[str, ...] = ("g", "crf"),
) -> list[dict[str, Any]]:
    """One ``{x, y, c, label}`` point per row, skipping non-numeric ``x``/``y``."""
    pts: list[dict[str, Any]] = []
    for r in rows:
        xv = r.get(x)
        yv = r.get(y)
        if not _is_num(xv) or not _is_num(yv):
            continue
        label = " ".join(f"{k}={r.get(k)}" for k in label_keys)
        pts.append({"x": xv, "y": yv, "c": r.get("vcodec"), "label": label})
    return pts


def compare_stacked(rows: list[Row]) -> dict[str, Any]:
    """Aggregate decoding latency per codec, broken down by access pattern.

    Returns ``{codecs, modes, data}`` where ``data[i].segments`` carries
    the per-mean-decode-time mean for each access pattern. The codec list
    is sorted by total stacked height ascending so the fastest codec sits
    at the left of the chart.
    """
    codecs: list[str] = []
    seen_c: set[str] = set()
    modes: list[str] = []
    seen_m: set[str] = set()
    for r in rows:
        c = r.get("vcodec")
        if c is not None and c not in seen_c:
            seen_c.add(c)
            codecs.append(c)
        m = r.get("timestamps_mode")
        if m is not None and m not in seen_m:
            seen_m.add(m)
            modes.append(m)

    data: list[dict[str, Any]] = []
    for codec in codecs:
        segs: list[dict[str, Any]] = []
        for mode in modes:
            vals = [
                r["median_load_time_video_ms"]
                for r in rows
                if r.get("vcodec") == codec
                and r.get("timestamps_mode") == mode
                and _is_num(r.get("median_load_time_video_ms"))
            ]
            v = sum(vals) / len(vals) if vals else 0.0
            segs.append({"key": mode, "v": v})
        total = sum(s["v"] for s in segs)
        data.append({"k": codec, "segments": segs, "total": total})

    data.sort(key=lambda d: d["total"])
    return {"codecs": codecs, "modes": modes, "data": data}