#!/usr/bin/env python3
# Time tests/bench/sieve and tests/bench/trees on C and on zc.
# The number is the slope between 1 repetition and N: that drops process
# startup and, for zc, the time spent compiling the benchmark. It does not
# bootstrap the compiler. Each cell is the median of TRIALS slopes, with the
# min and the max as the spread.
#
#   python3 tests/bench/run.py            # everything
#   python3 tests/bench/run.py check      # checksums only
#   python3 tests/bench/run.py sieve      # one program
#   TRIALS=5 GAP=0.8 python3 tests/bench/run.py
import os
import shutil
import subprocess
import sys
import tempfile
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent.parent
BENCH = Path(__file__).resolve().parent
TRIALS = int(os.environ.get("TRIALS", "5"))
GAP = float(os.environ.get("GAP", "0.8"))
NMAX = int(os.environ.get("NMAX", "200000"))
SIEVE_N = b"9592"   # primes below 100000
TREE_N = b"8191"    # nodes in a full tree of depth 12


def log(msg):
    print(msg, file=sys.stderr, flush=True)


def run(cmd, stdin=None, env=None, cwd=None):
    e = os.environ.copy()
    if env:
        e.update(env)
    t0 = time.perf_counter()
    p = subprocess.run(cmd, input=stdin, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=e, cwd=cwd)
    return time.perf_counter() - t0, p.returncode, p.stdout, p.stderr


def pin(cmd):
    if shutil.which("taskset"):
        return ["taskset", "-c", "0"] + list(cmd)
    return list(cmd)


def need(cmd, cwd=None):
    p = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, cwd=cwd)
    if p.returncode != 0:
        sys.stderr.write(p.stderr.decode(errors="replace"))
        raise SystemExit(f"command failed: {' '.join(cmd)}")
    return p


def machine_lines():
    lines = []
    u = subprocess.check_output(["uname", "-srm"], text=True).strip()
    lines.append(u)
    model = ""
    if Path("/proc/cpuinfo").exists():
        for line in Path("/proc/cpuinfo").read_text().splitlines():
            if line.startswith("model name"):
                model = line.split(":", 1)[1].strip()
                break
    if model:
        lines.append(model)
    lines.append("gcc " + subprocess.check_output(["gcc", "-dumpversion"], text=True).strip())
    if shutil.which("clang"):
        lines.append(subprocess.check_output(["clang", "--version"], text=True).splitlines()[0])
    if shutil.which("qemu-aarch64"):
        lines.append(subprocess.check_output(["qemu-aarch64", "--version"], text=True).splitlines()[0])
    if shutil.which("llvm-config"):
        lines.append("llvm " + subprocess.check_output(["llvm-config", "--version"], text=True).strip())
    lines.append("pinned to cpu 0 with taskset" if shutil.which("taskset") else "not pinned")
    return lines


def build(td):
    td = Path(td)
    log("building vms")
    need(["make", "-s", "build/core/genome.bin"], cwd=ROOT)
    cc = os.environ.get("CC") or subprocess.check_output(["sh", "src/tools/findcc.sh"], cwd=ROOT, text=True).strip()
    core = ["-Isrc/core", "-Ibuild/core/map", "--embed-dir=build/core"]
    zy = td / "zy"
    need([cc, "-O3", "-w", "-DJTYPE=int", "src/core/core.c", "src/core/io_linux.c", *core, "-o", str(zy)], cwd=ROOT)
    vms = {"interp": [str(zy)]}
    if shutil.which("llvm-config"):
        cflags = subprocess.check_output(["llvm-config", "--cflags"], text=True).split()
        ld = subprocess.check_output(["llvm-config", "--ldflags", "--libs"], text=True).split()
        zyl = td / "zyl"
        need([cc, "-O3", "-w", "-DJTYPE=int", "-DZYLLVM", *cflags, "src/core/core.c", "src/core/io_linux.c", *core, "-o", str(zyl), *ld], cwd=ROOT)
        vms["llvm"] = [str(zyl)]
    if shutil.which("aarch64-linux-gnu-gcc") and shutil.which("qemu-aarch64"):
        zyj = td / "zyj"
        need(["aarch64-linux-gnu-gcc", "-O3", "-w", "-DJTYPE=int", "-static", "src/core/core.c", "src/core/io_linux.c", *core, "-o", str(zyj)], cwd=ROOT)
        vms["arm64"] = ["qemu-aarch64", str(zyj)]
    log("building C")
    bins = {}
    for comp, tag in (("gcc", "gcc"), ("clang", "clang")):
        if comp == "clang" and not shutil.which("clang"):
            continue
        for lto in (False, True):
            name = tag + ("-lto" if lto else "")
            flags = [comp, "-O3", "-march=native"] + (["-flto"] if lto else [])
            sw = td / f"sieve-{name}"
            tw = td / f"trees-{name}"
            need(flags + ["-o", str(sw), str(BENCH / "sieve.c")])
            need(flags + ["-o", str(tw), str(BENCH / "trees.c")])
            bins[f"sieve-{name}"] = sw
            bins[f"trees-{name}"] = tw
    # PGO: train and use the same output path so the .gcda name matches.
    if shutil.which("gcc"):
        for src, key, arg in (("sieve.c", "sieve-gcc-pgo", ["word", "80"]), ("trees.c", "trees-gcc-pgo", ["malloc", "40"])):
            out = td / key
            prof = td / f"prof-{key}"
            prof.mkdir()
            base = ["gcc", "-O3", "-march=native", "-flto", "-o", str(out), str(BENCH / src)]
            need(base + [f"-fprofile-generate={prof}"])
            need([str(out), *arg])
            p = subprocess.run(base + [f"-fprofile-use={prof}", "-fprofile-correction"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            if p.returncode != 0 or b"profile count data file not found" in p.stderr:
                log(f"PGO skipped for {key}: {p.stderr.decode(errors='replace')[:200]}")
            else:
                bins[key] = out
                log(f"PGO applied for {key}")
    return vms, bins


def slope(invoke, expected, trials=TRIALS):
    def once(n):
        cmd, stdin, env = invoke(n)
        dt, rc, out, err = run(pin(cmd), stdin, env)
        if rc != 0 or out.strip() != expected:
            raise SystemExit(
                f"bad run n={n} rc={rc}\ncmd={cmd}\nout={out[:200]!r}\nerr={err[:400]!r}"
            )
        return dt

    dt1 = once(1)
    n = 16
    per = 0.0
    while n <= NMAX:
        dtn = once(n)
        per = (dtn - dt1) / (n - 1)
        gap = dtn - dt1
        if per > 0 and gap >= GAP:
            break
        if per > 0 and gap >= 0.05:
            need_n = int(GAP / per) + 1
            need_n = min(NMAX, max(need_n, n))
            if need_n != n:
                dtn = once(need_n)
                per = (dtn - dt1) / (need_n - 1)
                n = need_n
                # a noisy calibrating run can undershoot the gap. grow until
                # the kernel time actually dominates, or we hit the cap.
                while n < NMAX and per > 0 and (per * (n - 1)) < GAP:
                    n = min(NMAX, max(n + 1, int(GAP / per) + 1))
                    dtn = once(n)
                    per = (dtn - dt1) / (n - 1)
            break
        nxt = n * 4
        if nxt > NMAX:
            n = NMAX
            break
        n = nxt
    samples = []
    for i in range(trials):
        a = once(1)
        b = once(n)
        ms = (b - a) / (n - 1) * 1000.0
        samples.append(ms)
        log(f"    trial {i + 1}/{trials}: {ms:.4f} ms  (n={n})")
    samples.sort()
    mid = samples[len(samples) // 2]
    return {"median": mid, "min": samples[0], "max": samples[-1], "n": n, "samples": samples}


def zc_call(vm, src, which, env=None):
    def invoke(n):
        return vm, src + b"`" + f"{which}{n}".encode(), env
    return invoke


def c_call(binary, mode):
    def invoke(n):
        return [str(binary), mode, str(n)], None, None
    return invoke


def show(title, rows):
    print()
    print(title)
    print(f"{'path':<52} {'median ms':>10} {'min':>10} {'max':>10} {'reps':>8}")
    for name, r in rows:
        print(f"{name:<52} {r['median']:10.4f} {r['min']:10.4f} {r['max']:10.4f} {r['n']:8d}")
    print("samples:")
    for name, r in rows:
        nums = " ".join(f"{s:.4f}" for s in r["samples"])
        print(f"  {name}: {nums}")
    sys.stdout.flush()


def main():
    which = sys.argv[1] if len(sys.argv) > 1 else "all"
    os.chdir(ROOT)
    for line in machine_lines():
        print(line)
    print(f"trials {TRIALS}, gap {GAP}s, slope of 1 vs N, median of the trials")
    print("arm64 is qemu-user on this host, not native arm64 hardware")
    sys.stdout.flush()
    td = tempfile.mkdtemp(prefix="zcbench-")
    log(f"work dir {td}")
    vms, bins = build(td)
    sieve_src = (BENCH / "sieve.zc").read_bytes()
    tree_src = (BENCH / "trees.zc").read_bytes()

    log("checksums")
    checks = [
        (pin(vms["interp"]), sieve_src + b"`a1", SIEVE_N, None),
        (pin(vms["interp"]), sieve_src + b"`b1", SIEVE_N, None),
        (pin(vms["interp"]), sieve_src + b"`c1", SIEVE_N, None),
        (pin(vms["interp"]), tree_src + b"`a2", TREE_N, None),
        (pin(vms["interp"]), tree_src + b"`b1", TREE_N, None),
        (pin(vms["interp"]), tree_src + b"`c1", TREE_N, None),
        (pin([str(bins["sieve-gcc"]), "word", "1"]), None, SIEVE_N, None),
        (pin([str(bins["sieve-gcc"]), "byte", "1"]), None, SIEVE_N, None),
        (pin([str(bins["trees-gcc"]), "malloc", "1"]), None, TREE_N, None),
        (pin([str(bins["trees-gcc"]), "arena", "1"]), None, TREE_N, None),
    ]
    if "llvm" in vms:
        checks.append((pin(vms["llvm"]), sieve_src + b"`c1", SIEVE_N, {"ZYHOT": "1"}))
    if "arm64" in vms:
        checks.append((pin(vms["arm64"]), sieve_src + b"`c1", SIEVE_N, None))
    for cmd, stdin, expect, env in checks:
        dt, rc, out, err = run(cmd, stdin, env)
        if rc != 0 or out.strip() != expect:
            raise SystemExit(f"checksum failed {cmd} rc={rc} out={out!r} err={err[:300]!r}")
        log(f"  ok {expect.decode()} in {dt:.2f}s")
    if which == "check":
        print("checksums ok")
        return

    sieve_rows = []
    tree_rows = []

    def add(rows, name, invoke, expect):
        log(f"timing {name}")
        rows.append((name, slope(invoke, expect)))
        r = rows[-1][1]
        log(f"  median {r['median']:.4f} ms  spread {r['min']:.4f}..{r['max']:.4f}")

    if which in ("all", "sieve"):
        for key, label, mode in (
            ("sieve-gcc", "gcc -O3 -march=native, word flags", "word"),
            ("sieve-gcc-lto", "gcc -O3 -march=native -flto, word flags", "word"),
            ("sieve-gcc-pgo", "gcc -O3 -march=native -flto PGO, word flags", "word"),
            ("sieve-clang", "clang -O3 -march=native, word flags", "word"),
            ("sieve-clang-lto", "clang -O3 -march=native -flto, word flags", "word"),
            ("sieve-gcc-lto", "gcc -O3 -march=native -flto, byte flags", "byte"),
        ):
            if key not in bins:
                continue
            add(sieve_rows, label, c_call(bins[key], mode), SIEVE_N)
        add(sieve_rows, "zc interp, plain, bounds-checked int[?]", zc_call(vms["interp"], sieve_src, "a"), SIEVE_N)
        add(sieve_rows, "zc interp, o3, bounds-checked int[?]", zc_call(vms["interp"], sieve_src, "b"), SIEVE_N)
        add(sieve_rows, "zc interp, o3, raw *int", zc_call(vms["interp"], sieve_src, "c"), SIEVE_N)
        if "llvm" in vms:
            hot = {"ZYHOT": "1"}
            add(sieve_rows, "zc llvm O3, ZYHOT=1, o3 raw *int", zc_call(vms["llvm"], sieve_src, "c", hot), SIEVE_N)
            add(sieve_rows, "zc llvm O3, ZYHOT=1, o3 bounds-checked", zc_call(vms["llvm"], sieve_src, "b", hot), SIEVE_N)
        if "arm64" in vms:
            add(sieve_rows, "zc arm64 jit, qemu-user, o3 raw *int", zc_call(vms["arm64"], sieve_src, "c"), SIEVE_N)
        show("sieve: primes below 100000, checksum 9592", sieve_rows)

    if which in ("all", "trees"):
        for key, label, mode in (
            ("trees-gcc", "gcc -O3 -march=native, malloc", "malloc"),
            ("trees-gcc-lto", "gcc -O3 -march=native -flto, malloc", "malloc"),
            ("trees-gcc-pgo", "gcc -O3 -march=native -flto PGO, malloc", "malloc"),
            ("trees-clang", "clang -O3 -march=native, malloc", "malloc"),
            ("trees-clang-lto", "clang -O3 -march=native -flto, malloc", "malloc"),
            ("trees-gcc-lto", "gcc -O3 -march=native -flto, arena", "arena"),
            ("trees-clang-lto", "clang -O3 -march=native -flto, arena", "arena"),
        ):
            if key not in bins:
                continue
            add(tree_rows, label, c_call(bins[key], mode), TREE_N)
        add(tree_rows, "zc interp, plain new/free", zc_call(vms["interp"], tree_src, "a"), TREE_N)
        add(tree_rows, "zc interp, o3 new/free", zc_call(vms["interp"], tree_src, "b"), TREE_N)
        add(tree_rows, "zc interp, o3 arena", zc_call(vms["interp"], tree_src, "c"), TREE_N)
        if "llvm" in vms:
            hot = {"ZYHOT": "1"}
            add(tree_rows, "zc llvm O3, ZYHOT=1, plain new/free", zc_call(vms["llvm"], tree_src, "a", hot), TREE_N)
            add(tree_rows, "zc llvm O3, ZYHOT=1, o3 new/free", zc_call(vms["llvm"], tree_src, "b", hot), TREE_N)
            add(tree_rows, "zc llvm O3, ZYHOT=1, o3 arena", zc_call(vms["llvm"], tree_src, "c", hot), TREE_N)
        if "arm64" in vms:
            add(tree_rows, "zc arm64 jit, qemu-user, plain new/free", zc_call(vms["arm64"], tree_src, "a"), TREE_N)
            add(tree_rows, "zc arm64 jit, qemu-user, o3 arena", zc_call(vms["arm64"], tree_src, "c"), TREE_N)
        show("trees: depth 12, checksum 8191", tree_rows)

    print(f"\nwork dir left at {td}")


if __name__ == "__main__":
    main()