#!/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
all labels pooled): 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"}
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 baseline_for(summary, rows, label, scene, baseline):
    if (label, baseline, scene) in summary:
        return summary[(label, baseline, scene)], label
    pooled = [r for r in rows if r["variant"] == baseline and r["scene"] == scene]
    if not pooled:
        return None, None
    s = summarise([dict(r, label="*") for r in pooled])
    return s.get(("*", baseline, scene)), "*"


def table(summary, rows, baseline=None, noise=None):
    spreads = noise_spread(summary, noise) if noise else None
    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:
            if var == baseline:
                cells += ["—", "", "", "baseline"]
            else:
                b, blab = baseline_for(summary, rows, lab, scene, baseline)
                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")
    a = ap.parse_args(argv)
    rows = load_rows(a.file)
    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())