#!/usr/bin/env python3
"""the benchmark suite behind `make bench`, `make bench-save` and `make bench-check`.

usage: tests/bench/suite.py run [--zy PATH] [--out FILE]
       tests/bench/suite.py check [--zy PATH] [--baseline FILE] [--out FILE]

it wraps tests/bench/bench.sh, which times the kernels of tests/programs/bench.zc
(collatz, sieve, digits, fib) plain and o3 on each backend it can build, and adds the
256-word pointer sweep of tests/bench/sweep.zc, checked (plain column) and unbounded, and
the two array kernels of tests/bench/arrays.zc (a matmul whose hot loop is proved, a
gather that keeps its check), timed here on build/bin/zy the same way. interp rows are dropped unless INTERP=1. before
timing, ZYTRACE=1 on tests/programs/o3.zc must show at least one compiled o3 region, and
jit o3 collatz must be under half of jit plain collatz, or it refuses.
selftest checks the verdict logic on doctored tables in a second, without timing: the slope between 1 and N repetitions, the
median of TRIALS slopes. bench.sh itself is run as it is (it needs bc).

one row per backend and kernel. columns: plain ms and o3 ms per repetition; checks, the
number of bounds checks the compiler keeps in the row's plain function (the lines of
`zy --explain-checks` with that fn; o3 and unbounded code is never listed, so it counts
0), or '-' when zy has no --explain-checks; peak heap words and gc mark time, which the
runtime does not expose: '-'.

check compares with the baseline (tests/bench.baseline): each of the plain and o3
columns fails on its own when it is more than TOLERANCE (10%) slower and also more than
MIN_MS (0.1 ms) slower, so sub-ms cells don't flake; rows missing from a run don't fail, and a row fails
when both sides have a checks count and it changed (more is a lost proof; fewer means the
row no longer measures what it did, so re-save the baseline on purpose). when anything looks slower the whole
measurement is repeated, up to RETRIES more times, and each cell keeps its best figure,
so one noisy run does not fail the check.
"""
import shutil
import platform
import argparse
import json
import os
import statistics
import subprocess
import sys
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent.parent
TRIALS = int(os.environ.get("TRIALS", "5"))
MIN_S = float(os.environ.get("MIN_S", "0.4"))
TOLERANCE = float(os.environ.get("TOLERANCE", "0.10"))
MIN_MS = float(os.environ.get("MIN_MS", "0.1"))
INTERP = os.environ.get("INTERP", "") == "1"
RETRIES = int(os.environ.get("RETRIES", "1"))

# the rows this script times itself: name, source, kernel letter (bench.sh has the rest)
SWEEPS = [
    ("zy:sweep.checked", "tests/bench/sweep.zc", "c"),
    ("zy:sweep.unbounded", "tests/bench/sweep.zc", "u"),
    ("zy:arrays.matmul", "tests/bench/arrays.zc", "m"),
    ("zy:arrays.gather", "tests/bench/arrays.zc", "g"),
]
# where each row's checks are counted: the source and the function (bench.sh rows are
# backend:kernel, all from tests/programs/bench.zc)
CHECK_FNS = {
    "collatz": ("tests/programs/bench.zc", "csum"),
    "sieve": ("tests/programs/bench.zc", "sieve"),
    "digits": ("tests/programs/bench.zc", "digits"),
    "fib": ("tests/programs/bench.zc", "fib"),
    "sweep.checked": ("tests/bench/sweep.zc", "sweep"),
    "sweep.unbounded": ("tests/bench/sweep.zc", "sweepu"),
    "arrays.matmul": ("tests/bench/arrays.zc", "matmul"),
    "arrays.gather": ("tests/bench/arrays.zc", "gather"),
}
COLUMNS = ["program", "plain_ms", "o3_ms", "checks", "peak_heap_words", "gc_mark_ms"]
KERNELS = ("collatz", "sieve", "digits", "fib")


def bench_sh():
    """runs tests/bench/bench.sh and returns its rows: (backend:kernel, plain ms, o3 ms)"""
    r = subprocess.run(["bash", str(ROOT / "tests/bench/bench.sh")], cwd=ROOT, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    if r.returncode != 0:
        sys.exit("bench: tests/bench/bench.sh failed:\n" + r.stdout[-2000:] + r.stderr[-2000:])
    rows = []
    for line in r.stdout.splitlines():
        w = line.split()
        # backend names have spaces ("x86-64 jit"): the kernel is the fourth field from the end
        if len(w) >= 5 and w[-4] in KERNELS and w[-1].endswith("x"):
            backend = "-".join(w[:-4])
            rows.append((backend + ":" + w[-4], float(w[-3]), float(w[-2])))
    if not INTERP:
        rows = [r for r in rows if not r[0].startswith("interp:")]
    if not rows:
        sys.exit("bench: no rows in the output of tests/bench/bench.sh:\n" + r.stdout[-2000:])
    return rows


def run_once(zy, src, letter, reps):
    data = src + b"`" + letter.encode() + str(reps).encode()
    t = time.perf_counter()
    r = subprocess.run([zy], input=data, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    dt = time.perf_counter() - t
    if r.returncode != 0:
        sys.exit("bench: %s %s failed: %s" % (zy, letter, r.stderr.decode(errors="replace")[:300]))
    return dt


def slope_ms(zy, path, letter):
    src = (ROOT / path).read_bytes()
    n = 4
    while True:
        one = min(run_once(zy, src, letter, 1) for _ in range(2))
        many = run_once(zy, src, letter, n)
        if many - one >= MIN_S or n >= 1 << 16:
            break
        n *= 4
    slopes = []
    for _ in range(TRIALS):
        one = run_once(zy, src, letter, 1)
        many = run_once(zy, src, letter, n)
        slopes.append((many - one) / (n - 1) * 1000)
    return statistics.median(slopes)


_explained = {}


def explain(zy, path):
    """the kept-check counts per function of path from zy --explain-checks, or None when
    this zy has no --explain-checks (or its output is not format v1)"""
    if (zy, path) not in _explained:
        counts = None
        try:
            r = subprocess.run([zy, "--explain-checks", path], input=(ROOT / path).read_bytes(),
                               stdout=subprocess.PIPE, stderr=subprocess.PIPE, timeout=60)
            lines = r.stdout.decode(errors="replace").splitlines()
            if r.returncode == 0 and lines and json.loads(lines[0]).get("version") == 1:
                counts = {}
                for l in lines[1:]:
                    fn = json.loads(l).get("fn")
                    counts[fn] = counts.get(fn, 0) + 1
        except (OSError, ValueError, subprocess.TimeoutExpired):
            counts = None
        _explained[(zy, path)] = counts
    return _explained[(zy, path)]


def checks_of(zy, program):
    """the checks column of a row: its function's kept checks, or '-'"""
    where = CHECK_FNS.get(program.split(":", 1)[-1])
    if not where:
        return "-"
    counts = explain(zy, where[0])
    return "-" if counts is None else str(counts.get(where[1], 0))


def fmt_row(cells):
    return "%-26s %10s %10s %8s %16s %11s" % tuple(cells)


def num(x):
    return "-" if x is None else "%.3f" % x


def write_table(path, rows):
    lines = ["# tinyvm bench: ms per repetition, from tests/bench/bench.sh and tests/bench/suite.py; '-' where a figure is not available",
             fmt_row(COLUMNS)]
    for r in rows:
        lines.append(fmt_row([r["program"], num(r["plain"]), num(r["o3"]), r["checks"], r["peak"], r["gc"]]))
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text("\n".join(lines) + "\n")


def read_table(path):
    rows = {}
    f = lambda x: None if x == "-" else float(x)
    for line in path.read_text().splitlines():
        if not line.strip() or line.startswith("#") or line.split()[0] == "program":
            continue
        w = line.split()
        rows[w[0]] = {"program": w[0], "plain": f(w[1]), "o3": f(w[2]), "checks": w[3], "peak": w[4], "gc": w[5]}
    return rows


def region_guard(zy):
    """refuses to go on when the jit compiles none of the o3 regions of tests/programs/o3.zc"""
    src = (ROOT / "tests/programs/o3.zc").read_bytes()
    env = dict(os.environ, ZYTRACE="1")
    r = subprocess.run([zy], input=src, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env)
    err = r.stderr.decode(errors="replace").splitlines()
    done = [l for l in err if l.startswith("o3: ") and " bytes of " in l or " compiled by llvm" in l]
    stay = [l for l in err if l.startswith("o3: ") and "stays on the region interpreter" in l]
    print("bench: o3 regions in o3.zc: %d compiled, %d on the region interpreter" % (len(done), len(stay)), file=sys.stderr, flush=True)
    if not done:
        sys.exit("bench: REFUSING: %s compiled no o3 regions (ZYTRACE=1 on tests/programs/o3.zc); "
                 "the numbers would be the region interpreter's. rebuild zy with the jit and unset ZYO3INTERP" % zy)
    return len(done), len(stay)


def o3_guard(rows):
    """jit o3 collatz must be under half of jit plain collatz (needs no interp rows)"""
    d = {r["program"]: r for r in rows}
    j = d.get("x86-64-jit:collatz")
    if j is None:
        return
    if j["o3"] is None or j["plain"] is None:
        sys.exit("bench: REFUSING: x86-64-jit:collatz has no plain or o3 figure")
    if j["o3"] * 2 >= j["plain"]:
        sys.exit("bench: REFUSING: jit o3 collatz %.2f ms is not under half of jit plain %.2f ms; o3 regions are not compiled" % (j["o3"], j["plain"]))


def run_all(zy):
    region_guard(zy)
    rows = []
    for name, plain, o3 in bench_sh():
        print("bench: %-26s %10.3f %10.3f ms" % (name, plain, o3), file=sys.stderr, flush=True)
        rows.append({"program": name, "plain": plain, "o3": o3, "checks": checks_of(zy, name), "peak": "-", "gc": "-"})
    for name, path, letter in SWEEPS:
        ms = slope_ms(zy, path, letter)
        print("bench: %-26s %10.3f ms" % (name, ms), file=sys.stderr, flush=True)
        rows.append({"program": name, "plain": ms, "o3": None, "checks": checks_of(zy, name), "peak": "-", "gc": "-"})
    o3_guard(rows)
    return rows


def slower(now, base):
    return now is not None and base is not None and base > 0 and now > base * (1 + TOLERANCE) and now - base > MIN_MS


def pct(now, base):
    return "-" if now is None or base is None or base <= 0 else "%+.1f%%" % ((now - base) / base * 100)


def expected_backends():
    """the backends bench.sh should report on this machine, found the same way it finds them"""
    want = ["zy"]
    if platform.system() == "Linux" and platform.machine() == "x86_64":
        want.append("x86-64-jit")
    if os.environ.get("LLVM_CONFIG") or shutil.which("llvm-config") or shutil.which("llvm-config-19"):
        want.append("llvm")
    if INTERP:
        want.append("interp")
    return want


def require_rows(rows):
    """bench-save refuses a table that is missing a backend this machine has"""
    have = {r["program"].split(":")[0] for r in rows}
    missing = [b for b in expected_backends() if b not in have]
    if missing:
        sys.exit("bench: REFUSING to save: no rows for %s (a backend failed to build?)" % ", ".join(missing))


def main():
    ap = argparse.ArgumentParser(description="tinyvm benchmark suite")
    ap.add_argument("mode", choices=["run", "check", "selftest"])
    ap.add_argument("--rows", help="check: compare this saved table instead of measuring")
    ap.add_argument("--zy", default=str(ROOT / "build/bin/zy"))
    ap.add_argument("--out", default=str(ROOT / "build/bench.txt"))
    ap.add_argument("--baseline", default=str(ROOT / "tests/bench.baseline"))
    a = ap.parse_args()
    if a.mode == "selftest":
        return selftest()
    rows = list(read_table(Path(a.rows)).values()) if a.rows else run_all(a.zy)
    if a.mode == "run":
        if os.environ.get("BENCH_SAVE") == "1":
            require_rows(rows)
        write_table(Path(a.out), rows)
        print(Path(a.out).read_text(), end="")
        return 0
    base_path = Path(a.baseline)
    if not base_path.exists():
        sys.exit("bench: no baseline at %s (run make bench-save)" % base_path)
    base = read_table(base_path)
    tries = 0
    while not a.rows and tries < RETRIES and any(slower(r[c], base[r["program"]][c]) for r in rows if r["program"] in base for c in ("plain", "o3")):
        tries += 1
        print("bench: something looks slower; measuring again (%d of %d)" % (tries, RETRIES), file=sys.stderr, flush=True)
        again = {r["program"]: r for r in run_all(a.zy)}
        for r in rows:
            g = again.get(r["program"])
            if g:
                for c in ("plain", "o3"):
                    if r[c] is not None and g[c] is not None:
                        r[c] = min(r[c], g[c])
    bad = 0
    print("%-26s %9s %9s %8s %9s %9s %8s %6s %6s  %s" % ("program", "base pl", "now pl", "change", "base o3", "now o3", "change", "b chk", "n chk", "verdict"))
    names = set()
    for r in rows:
        names.add(r["program"])
        b = base.get(r["program"])
        if not b:
            print("%-26s %9s %9s %8s %9s %9s %8s %6s %6s  new" % (r["program"], "-", num(r["plain"]), "-", "-", num(r["o3"]), "-", "-", r["checks"]))
            continue
        why = [c for c in ("plain", "o3") if slower(r[c], b[c])]
        verdict = ", ".join(c + " SLOWER" for c in why)
        if b["checks"] != "-" and r["checks"] != "-" and int(r["checks"]) != int(b["checks"]):
            more = int(r["checks"]) > int(b["checks"])
            verdict = (verdict + ", " if verdict else "") + ("MORE CHECKS" if more else "FEWER CHECKS")
        if verdict:
            bad += 1
        print("%-26s %9s %9s %8s %9s %9s %8s %6s %6s  %s" % (r["program"], num(b["plain"]), num(r["plain"]), pct(r["plain"], b["plain"]),
              num(b["o3"]), num(r["o3"]), pct(r["o3"], b["o3"]), b["checks"], r["checks"], verdict or "ok"))
    for name in base:
        if name not in names:
            print("%-26s %9s %9s  gone" % (name, num(base[name]["plain"]), "-"))
    write_table(Path(a.out), rows)
    print("bench-check: %s (%d of %d rows over the limit; a column fails when over %d%% and over %g ms slower)" % ("FAIL" if bad else "ok", bad, len(rows), round(TOLERANCE * 100), MIN_MS))
    return 1 if bad else 0


def selftest():
    """no timing: runs check on a doctored table against a doctored baseline and expects the verdicts"""
    import tempfile
    t = Path(tempfile.mkdtemp())
    base = [{"program": n, "plain": p, "o3": o, "checks": "-", "peak": "-", "gc": "-"} for n, p, o in
            [("a:big", 10.0, 2.0), ("a:tiny", 0.5, 0.4), ("a:gone", 1.0, 1.0)]]
    base.append({"program": "a:chk", "plain": 1.0, "o3": None, "checks": "3", "peak": "-", "gc": "-"})
    write_table(t / "base", base)
    def case(name, now, want):
        write_table(t / "now", now)
        r = subprocess.run([sys.executable, __file__, "check", "--rows", str(t / "now"), "--baseline", str(t / "base"), "--out", str(t / "out")],
                           stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
        ok = (r.returncode != 0) == want
        print("selftest: %-36s exit %d, want %s: %s" % (name, r.returncode, "fail" if want else "ok", "pass" if ok else "WRONG"))
        if not ok:
            print(r.stdout)
        return ok
    row = lambda n, p, o, k="-": {"program": n, "plain": p, "o3": o, "checks": k, "peak": "-", "gc": "-"}
    two = lambda k: [row("a:big", 10.0, 2.0), row("a:tiny", 0.5, 0.4), row("a:chk", 1.0, None, k)]
    good = [
        case("same numbers", [row("a:big", 10.0, 2.0), row("a:tiny", 0.5, 0.4)], False),
        case("big plain 20% slower", [row("a:big", 12.0, 2.0), row("a:tiny", 0.5, 0.4)], True),
        case("big o3 +0.3 ms (15%)", [row("a:big", 10.0, 2.3), row("a:tiny", 0.5, 0.4)], True),
        case("tiny 30% but +0.12 ms", [row("a:big", 10.0, 2.0), row("a:tiny", 0.62, 0.4)], True),
        case("tiny 20% but only +0.08 ms", [row("a:big", 10.0, 2.0), row("a:tiny", 0.5, 0.48)], False),
        case("big +0.5 ms but only 5%", [row("a:big", 10.5, 2.0), row("a:tiny", 0.5, 0.4)], False),
        case("rows missing (no interp)", [row("a:big", 10.0, 2.0)], False),
        case("checks the same", two("3"), False),
        case("checks rose 3 -> 4", two("4"), True),
        case("checks fell 3 -> 0", two("0"), True),
        case("checks '-' now (no --explain-checks)", two("-"), False),
    ]
    print("selftest: %s" % ("ok" if all(good) else "FAIL"))
    return 0 if all(good) else 1


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