# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Measure your noise floor and your tuning sweep.

The benchmark sorts 2 million random integers (NumPy). (1) The identical run
10 times: the A/A spread. (2) A "tuning sweep" of 30 identical configurations
(a dummy parameter), best picked, then re-measured 5 times. (3) A against B
(identical code) in blocks and interleaved, with a CPU-heavy background load
started halfway through.
"""
import argparse, json, multiprocessing as mp, statistics, time
import numpy as np

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--size", type=int, default=2_000_000, help="Integers to sort per run.")
p.add_argument("--reps", type=int, default=5, help="Sorts per run (the run reports their median).")
p.add_argument("--sweep", type=int, default=30, help="Configurations in the dummy sweep.")
p.add_argument("--pairs", type=int, default=20, help="A/B runs of each in the order test.")
p.add_argument("--load-procs", type=int, default=8, help="Background busy processes in the order test.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

data = np.random.default_rng(7).integers(0, 1 << 40, args.size)
def run(_cfg=None):
    ts = []
    for _ in range(args.reps):
        a = data.copy(); t = time.perf_counter(); a.sort(); ts.append(time.perf_counter() - t)
    return statistics.median(ts) * 1000

def burn(stop):
    x = 0
    while not stop.is_set(): x += 1

run()                                                                      # warm up
aa = [run() for _ in range(10)]
spread = (max(aa) - min(aa)) / statistics.mean(aa)

sweep = [{"config": i, "ms": run(i)} for i in range(args.sweep)]
best = min(sweep, key=lambda r: r["ms"]); rerun = [run(best["config"]) for _ in range(5)]
sweep_mean = statistics.mean(r["ms"] for r in sweep)

def order_test(interleaved):
    stop = mp.Event(); procs = []; a, b = [], []
    seq = ["A", "B"] * args.pairs if interleaved else ["A"] * args.pairs + ["B"] * args.pairs
    for i, who in enumerate(seq):
        if i == len(seq) // 2:                                             # background load starts halfway through
            procs = [mp.Process(target=burn, args=(stop,)) for _ in range(args.load_procs)]; [q.start() for q in procs]
        (a if who == "A" else b).append(run())
    stop.set(); [q.join() for q in procs]
    return (statistics.mean(b) - statistics.mean(a)) / statistics.mean(a)

if __name__ == "__main__":
    blocked, inter = order_test(False), order_test(True)
    summary = {"aa_spread": spread, "aa_min": min(aa), "aa_max": max(aa), "sweep_best": best["ms"], "sweep_mean": sweep_mean,
               "best_gain": (sweep_mean - best["ms"]) / sweep_mean, "rerun_mean": statistics.mean(rerun),
               "rerun_gain": (sweep_mean - statistics.mean(rerun)) / sweep_mean, "blocked_diff": blocked, "interleaved_diff": inter}
    summary.update({k + "_pct": v * 100 for k, v in list(summary.items()) if k in ("aa_spread", "best_gain", "rerun_gain", "blocked_diff", "interleaved_diff")})
    print(f"A/A: 10 identical runs {min(aa):.1f}-{max(aa):.1f} ms, spread {spread:.1%}")
    print(f"sweep: best of {args.sweep} identical configs {best['ms']:.1f} ms ({summary['best_gain']:.1%} better than their mean); re-measured {summary['rerun_mean']:.1f} ms ({summary['rerun_gain']:.1%})")
    print(f"identical A and B under a load that starts halfway: blocked order says B is {blocked:+.1%}, interleaved says {inter:+.1%}")
    if args.json:
        with open(args.json, "w") as f:
            json.dump({"summary": summary, "tables": {"aa": [{"run": i + 1, "ms": v} for i, v in enumerate(aa)], "sweep": sweep}}, f)