#!/usr/bin/env python3
"""
Summarize a results.txt produced by bench.sh into a per-label table.

Usage: summarize.py [results.txt]

Prints median and min-max of every metric for every label, and flags
cross-label deltas against the ~5 percent noise floor of this machine.
"""

import re
import sys
from statistics import median

NOISE_FLOOR_PCT = 5.0

RE_CORPUS = re.compile(
    r"^RESULT corpus label=(\S+) build=\S+ target=(\d+) chars=\d+ "
    r"prompt_n=(\d+) rep=\d+ prompt_ms=([\d.]+) tps=([\d.]+)")
RE_LB = re.compile(
    r"^LBRESULT label=(\S+) build=\S+ test=(\S+) tps=([\d.]+)")


def main():
    path = sys.argv[1] if len(sys.argv) > 1 else "/root/bench/results.txt"
    # data[metric][label] = list of tps
    data = {}
    order = []
    labels = []
    for line in open(path):
        m = RE_CORPUS.match(line)
        if m:
            label, target, _pn, _ms, tps = m.groups()
            metric = "pp%s-corpus" % target
        else:
            m = RE_LB.match(line)
            if not m:
                continue
            label, test, tps = m.groups()
            metric = "%s-llama-bench" % test
        if metric not in data:
            data[metric] = {}
            order.append(metric)
        data[metric].setdefault(label, []).append(float(tps))
        if label not in labels:
            labels.append(label)

    if not data:
        print("no parseable results in %s" % path)
        return 1

    def sort_key(m):
        n = re.search(r"(\d+)", m)
        return (0 if "corpus" in m else 1, int(n.group(1)) if n else 0)
    order.sort(key=sort_key)

    w = max(len(l) for l in labels) + 2
    head = "metric".ljust(20) + "".join(l.ljust(max(w, 24)) for l in labels)
    print("")
    print("t/s, median (min-max), n samples")
    print(head)
    print("-" * len(head))
    for metric in order:
        row = metric.ljust(20)
        for label in labels:
            vals = data[metric].get(label)
            if not vals:
                row += "-".ljust(max(w, 24))
            else:
                cell = "%.1f (%.1f-%.1f) n=%d" % (
                    median(vals), min(vals), max(vals), len(vals))
                row += cell.ljust(max(w, 24))
        print(row)

    if len(labels) >= 2:
        base = labels[0]
        print("")
        print("deltas vs %s (noise floor %.0f%%, anything under it is UNPROVEN)"
              % (base, NOISE_FLOOR_PCT))
        for other in labels[1:]:
            print("  %s vs %s:" % (other, base))
            for metric in order:
                a = data[metric].get(base)
                b = data[metric].get(other)
                if not a or not b:
                    continue
                ma, mb = median(a), median(b)
                pct = (mb - ma) / ma * 100.0
                verdict = "SIGNIFICANT" if abs(pct) >= NOISE_FLOOR_PCT else "unproven"
                print("    %-20s %+6.1f%%  %s" % (metric, pct, verdict))
    return 0


if __name__ == "__main__":
    sys.exit(main())