#!/usr/bin/env python3
"""Compare MangoHud frame logs from CachyOS and NixOS (see MIGRATION.md).

    tools/bench-compare                 every log in ~/benchmarks
    tools/bench-compare DIR|FILE ...    the logs given

Logs are grouped by the OS MangoHud wrote into each file's header, so logs
from both systems can sit in one folder. For each run it shows the average
FPS, the 1% and 0.1% lows (the FPS of the slowest 1% and 0.1% of frames) and
the longest frame; then each system's average over its runs, and how far
NixOS is from CachyOS.

Record with log_interval=0 so every frame is logged; with sampled logs the
lows are only estimates.
"""
import csv
import statistics
import sys
from pathlib import Path


def read_log(path):
    """Return (os name, frame times in ms) from one MangoHud CSV log."""
    rows = list(csv.reader(path.open(newline="", errors="replace")))
    os_name = "unknown"
    frametimes = None
    for i, row in enumerate(rows):
        cols = [c.strip().lower() for c in row]
        if cols[:1] == ["os"] and i + 1 < len(rows) and rows[i + 1]:
            os_name = rows[i + 1][0].strip() or os_name
        elif "frametime" in cols:
            col = cols.index("frametime")
            frametimes = []
            for data in rows[i + 1 :]:
                try:
                    frametimes.append(float(data[col]))
                except (IndexError, ValueError):
                    continue
            break
    if not frametimes:
        return os_name, []
    # Older MangoHud versions log microseconds, newer ones milliseconds.
    if statistics.median(frametimes) > 200:
        frametimes = [t / 1000 for t in frametimes]
    return os_name, [t for t in frametimes if t > 0]


def low(frametimes, fraction):
    """FPS of the slowest `fraction` of frames (averaged, as MangoHud does)."""
    worst = sorted(frametimes, reverse=True)
    n = max(1, int(len(worst) * fraction))
    return 1000 / statistics.fmean(worst[:n])


def stats(frametimes):
    return {
        "avg": 1000 / statistics.fmean(frametimes),
        "1%": low(frametimes, 0.01),
        "0.1%": low(frametimes, 0.001),
        "max ms": max(frametimes),
    }


def short_os(name):
    lower = name.lower()
    if "cachyos" in lower:
        return "CachyOS"
    if "nixos" in lower:
        return "NixOS"
    return name


def main(args):
    targets = [Path(a) for a in args] or [Path.home() / "benchmarks"]
    files = []
    for target in targets:
        if target.is_dir():
            files += sorted(target.rglob("*.csv"))
        elif target.exists():
            files.append(target)
        else:
            sys.exit(f"bench-compare: {target} doesn't exist")
    files = [f for f in files if not f.stem.endswith("_summary")]
    if not files:
        sys.exit("bench-compare: no MangoHud logs found")

    keys = ["avg", "1%", "0.1%", "max ms"]
    header = f"{'':38}" + "".join(f"{k:>9}" for k in keys)
    groups = {}
    print(header)
    for f in files:
        os_name, frametimes = read_log(f)
        if len(frametimes) < 100:
            print(f"{f.name[:38]:38}  skipped: {len(frametimes)} frames")
            continue
        s = stats(frametimes)
        system = short_os(os_name)
        groups.setdefault(system, []).append(s)
        label = f"{system}  {f.name}"[:38]
        print(f"{label:38}" + "".join(f"{s[k]:9.1f}" for k in keys))

    print()
    print(f"{'average of runs':38}" + "".join(f"{k:>9}" for k in keys))
    means = {}
    for system, runs in groups.items():
        means[system] = {k: statistics.fmean(r[k] for r in runs) for k in keys}
        label = f"{system} ({len(runs)} runs)"
        print(f"{label:38}" + "".join(f"{means[system][k]:9.1f}" for k in keys))

    if "CachyOS" in means and "NixOS" in means:
        a, b = means["CachyOS"], means["NixOS"]
        diff = "".join(f"{(b[k] - a[k]) / a[k] * 100:+8.1f}%" for k in keys)
        print(f"{'NixOS vs CachyOS':38}{diff}")
        print("(higher is better for FPS columns, lower is better for max ms)")
        # Run-to-run noise on the same system, to judge whether a gap is real.
        for system, runs in groups.items():
            if len(runs) > 1:
                spread = statistics.pstdev(r["avg"] for r in runs)
                print(f"{system}: average FPS varies ±{spread:.1f} between runs")


if __name__ == "__main__":
    main(sys.argv[1:])