#!/usr/bin/env python3
"""Summarise bench.py session rows and apply the BENCH-MATRIX win rule.

    analyze.py [--file ~/onion-bench/results/session.jsonl] [--label B2]
               [--baseline C0 --noise B0] [--include-invalid] [--exclude-backup]

Groups windows by (label, variant, scene). Each session contributes its median p99 and
median 1 % low (frametime-based); a group's value is the median of its session medians
and its spread is max - min of those. With --baseline/--noise, every non-baseline group
is compared with the baseline variant on the same scene: same label if present, else the
same block (a tagged label 'B3.dh-off' looks in 'B3', 'B3.baseline', then all 'B3.*'
pooled), else all labels pooled. A tagged session of the baseline variant itself is
compared with 'B3' / 'B3.baseline' when one exists: win = median p99 improves by more than the noise label's p99 spread
AND 1 % low does not worsen by more than its 1 % low spread; loss = p99 worsens beyond
the spread or 1 % low drops beyond it; otherwise tie (ties keep upstream defaults).
Rows flagged crashed / timeout / no_data / no_markers / low_n / vsync_suspect are
excluded unless --include-invalid. Stdlib only.
"""
import argparse
import json
import statistics
import sys
from pathlib import Path

INVALID = {"crashed", "timeout", "no_data", "no_markers", "low_n", "vsync_suspect",
           "shader_not_loaded"}
DEFAULT_FILE = Path.home() / "onion-bench" / "results" / "session.jsonl"


def load_rows(path):
    rows = []
    for line in Path(path).read_text().splitlines():
        line = line.strip()
        if line:
            rows.append(json.loads(line))
    return [r for r in rows if r.get("kind", "window") == "window"]


def is_valid(row, exclude_backup=False):
    flags = set(row.get("flags") or [])
    if not row.get("metrics") or flags & INVALID:
        return False
    return not (exclude_backup and "backup_in_window" in flags)


def summarise(rows):
    """{(label, variant, scene): stats} from valid window rows."""
    by_session = {}
    for r in rows:
        key = (r.get("label"), r["variant"], r["scene"])
        by_session.setdefault(key, {}).setdefault(r["session"], []).append(r["metrics"])
    out = {}
    for key, sessions in by_session.items():
        p99s = [statistics.median(m["p99_ms"] for m in ms) for ms in sessions.values()]
        lows = [statistics.median(m["one_low_fps"] for m in ms) for ms in sessions.values()]
        fps = [statistics.median(m["avg_fps"] for m in ms) for ms in sessions.values()]
        out[key] = {
            "sessions": len(sessions),
            "windows": sum(len(ms) for ms in sessions.values()),
            "p99": statistics.median(p99s), "p99_spread": max(p99s) - min(p99s),
            "low": statistics.median(lows), "low_spread": max(lows) - min(lows),
            "fps": statistics.median(fps),
        }
    return out


def noise_spread(summary, noise_label):
    """(p99 spread, 1 % low spread) of the noise block, worst case across its groups."""
    groups = [v for (lab, _, _), v in summary.items() if lab == noise_label]
    if not groups:
        return None
    return max(g["p99_spread"] for g in groups), max(g["low_spread"] for g in groups)


def verdict(base, cand, spread_p99, spread_low):
    gain = base["p99"] - cand["p99"]          # > 0: candidate's p99 is lower (better)
    low_drop = base["low"] - cand["low"]      # > 0: candidate's 1 % low is worse
    if gain > spread_p99 and low_drop <= spread_low:
        return "win"
    if gain < -spread_p99 or low_drop > spread_low:
        return "loss"
    return "tie"


def block_of(label):
    """'B3.dh-off' -> 'B3'; untagged labels (and None) are their own block."""
    return label.split(".", 1)[0] if label else label


def _pooled(rows, labels, baseline, scene, name):
    pooled = [r for r in rows if r["variant"] == baseline and r["scene"] == scene
              and (labels is None or r.get("label") in labels)]
    if not pooled:
        return None
    return summarise([dict(r, label=name) for r in pooled]).get((name, baseline, scene))


def baseline_for(summary, rows, label, scene, baseline):
    """(baseline stats, where-label) for one candidate group, or (None, None).

    Order: the candidate's own label; then its block - the untagged block label
    ('B3', for tagged candidates), '<block>.baseline', then every '<block>.*' label pooled; then all labels
    pooled. A tagged session of the baseline variant itself (e.g. 'B3.dh-off' running
    C1 against baseline C1) is compared only with 'B3' or 'B3.baseline'."""
    block = block_of(label)
    tagged = bool(label) and block != label
    order = [label] + ([block] if tagged else []) + ([f"{block}.baseline"] if block else [])
    for lab in order:
        if (lab, baseline, scene) in summary:
            return summary[(lab, baseline, scene)], lab
    if block:
        labels = {r.get("label") for r in rows
                  if r.get("label") == block or str(r.get("label")).startswith(block + ".")}
        b = _pooled(rows, labels, baseline, scene, f"{block}.*")
        if b is not None:
            return b, f"{block}.*"
    b = _pooled(rows, None, baseline, scene, "*")
    return (b, "*") if b is not None else (None, None)


def self_baseline(summary, label, scene, baseline):
    """Baseline for a tagged group that runs the baseline variant itself, or None
    (then the group *is* the baseline)."""
    block = block_of(label)
    if not label or block == label:
        return None, None
    for lab in (block, f"{block}.baseline"):
        if lab != label and (lab, baseline, scene) in summary:
            return summary[(lab, baseline, scene)], lab
    return None, None


def table(summary, rows, baseline=None, noise=None):
    spreads = noise_spread(summary, noise) if noise else None
    full = summarise(rows)  # baselines may live in labels filtered out of the table
    hdr = ["label", "variant", "scene", "sess", "win#", "p99 ms", "p99 spr", "1%low",
           "1%low spr", "avg fps"]
    if baseline:
        hdr += ["vs", "Δp99 ms", "Δ1%low", "verdict"]
    out = ["| " + " | ".join(hdr) + " |", "|" + "---|" * len(hdr)]
    for key in sorted(summary, key=lambda k: (str(k[0]), k[2], k[1])):
        lab, var, scene = key
        s = summary[key]
        cells = [str(lab), var, scene, str(s["sessions"]), str(s["windows"]),
                 f"{s['p99']:.2f}", f"{s['p99_spread']:.2f}", f"{s['low']:.1f}",
                 f"{s['low_spread']:.1f}", f"{s['fps']:.0f}"]
        if baseline:
            b, blab = (self_baseline(full, lab, scene, baseline) if var == baseline
                       else baseline_for(full, rows, lab, scene, baseline))
            if var == baseline and b is None:
                cells += ["—", "", "", "baseline"]
            else:
                if b is None:
                    cells += ["—", "", "", "no baseline"]
                elif spreads is None:
                    cells += [f"{baseline}@{blab}", f"{b['p99'] - s['p99']:+.2f}",
                              f"{s['low'] - b['low']:+.1f}", "no noise data"]
                else:
                    cells += [f"{baseline}@{blab}", f"{b['p99'] - s['p99']:+.2f}",
                              f"{s['low'] - b['low']:+.1f}", verdict(b, s, *spreads)]
        out.append("| " + " | ".join(cells) + " |")
    if baseline and noise:
        out.append("")
        out.append(f"noise ({noise}): " + (
            f"p99 spread {spreads[0]:.2f} ms, 1%low spread {spreads[1]:.1f} fps"
            if spreads else "no rows - verdicts unavailable"))
    return "\n".join(out)


def main(argv=None):
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("--file", default=str(DEFAULT_FILE))
    ap.add_argument("--label", action="append", help="only these labels (repeatable); "
                    "the noise label is always kept")
    ap.add_argument("--baseline", help="baseline variant, e.g. C0")
    ap.add_argument("--noise", help="label whose spread is the win threshold, e.g. B0")
    ap.add_argument("--include-invalid", action="store_true")
    ap.add_argument("--exclude-backup", action="store_true")
    ap.add_argument("--since", help="only rows with ts >= this ISO time, e.g. 2026-10-16T22:24")
    a = ap.parse_args(argv)
    rows = load_rows(a.file)
    if a.since:
        rows = [r for r in rows if r.get("ts", "") >= a.since]
    total = len(rows)
    if not a.include_invalid:
        rows = [r for r in rows if is_valid(r, a.exclude_backup)]
    warn = sum(1 for r in rows if "backup_in_window" in (r.get("flags") or []))
    shown = rows
    if a.label:
        keep = set(a.label) | ({a.noise} if a.noise else set())
        shown = [r for r in rows if r.get("label") in keep]
    if not shown:
        print("no valid rows", file=sys.stderr)
        return 1
    print(table(summarise(shown), rows, a.baseline, a.noise))
    print(f"\n{len(rows)}/{total} windows valid; {warn} with backup_in_window (warn only)")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())