# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Run your own: config dump, seed, sweep, raw results.

One command that (1) records versions, kernel, CPU model and every non-default
setting into config.txt; (2) generates the data from a fixed seed; (3) runs the
benchmark (random lookups into an in-memory array of keys) at sizes 1x, 3x,
10x and 30x; (4) writes every raw per-run result into results.csv; all of it
for 5 seeds. Prints the per-key cost per size with the seed-to-seed spread.
"""
import argparse, csv, json, os, platform, statistics, subprocess, sys, time
import numpy as np

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--base-mb", type=float, default=8, help="Data size at 1x (MB).")
p.add_argument("--sizes", default="1,3,10,30", help="Size multipliers.")
p.add_argument("--seeds", type=int, default=5, help="Seeds to run.")
p.add_argument("--lookups", type=int, default=10_000_000, help="Random lookups per run.")
p.add_argument("--out", default="run-your-own-out", help="Directory for config.txt and results.csv.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

defaults = {"base_mb": 8, "sizes": "1,3,10,30", "seeds": 5, "lookups": 10_000_000}
def sh(cmd):
    try: return subprocess.run(cmd, capture_output=True, text=True, timeout=10).stdout.strip()
    except Exception: return "unavailable"
os.makedirs(args.out, exist_ok=True)
cpu = sh(["sysctl", "-n", "machdep.cpu.brand_string"]) if sys.platform == "darwin" else \
      next((l.split(":", 1)[1].strip() for l in open("/proc/cpuinfo") if l.startswith("model name")), platform.processor())
config = {"python": sys.version.split()[0], "numpy": np.__version__, "os": platform.platform(), "kernel": platform.release(),
          "cpu": cpu, "cores": os.cpu_count(), "command": " ".join(sys.argv),
          "non_default_settings": {k: v for k, v in vars(args).items() if k in defaults and v != defaults[k]} or "none"}
with open(os.path.join(args.out, "config.txt"), "w") as f:
    for k, v in config.items(): f.write(f"{k}: {v}\n")

rows = []
for seed in range(args.seeds):
    rng = np.random.default_rng(seed)
    for m in map(float, args.sizes.split(",")):
        n = int(args.base_mb * m * 1e6 / 8)
        keys = rng.integers(0, 1 << 62, n)                              # the data set, from the seed
        idx = rng.integers(0, n, args.lookups)
        keys[idx[:1000]].sum()                                          # warm up
        t = time.perf_counter(); s = keys[idx].sum(); el = time.perf_counter() - t
        rows.append({"seed": seed, "size_x": m, "size_mb": n * 8 / 1e6, "ns_per_key": el / args.lookups * 1e9})
with open(os.path.join(args.out, "results.csv"), "w", newline="") as f:
    w = csv.DictWriter(f, fieldnames=list(rows[0])); w.writeheader(); w.writerows(rows)

sizes = []
for m in sorted({r["size_x"] for r in rows}):
    xs = [r["ns_per_key"] for r in rows if r["size_x"] == m]
    sizes.append({"size_x": m, "size_mb": next(r["size_mb"] for r in rows if r["size_x"] == m), "mean_ns": statistics.mean(xs),
                  "min_ns": min(xs), "max_ns": max(xs), "spread_pct": (max(xs) - min(xs)) / statistics.mean(xs) * 100})
for i, s in enumerate(sizes):
    s["step_x"] = s["mean_ns"] / sizes[i - 1]["mean_ns"] if i else 1.0
    print(f"{s['size_x']:>4g}x ({s['size_mb']:6.0f} MB): {s['mean_ns']:5.2f} ns per key  (seeds {s['min_ns']:.2f}-{s['max_ns']:.2f}, spread {s['spread_pct']:.1f}%)  x{s['step_x']:.2f} vs previous size")
print(f"wrote {args.out}/config.txt and {args.out}/results.csv ({len(rows)} raw rows)")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"config": config, "tables": {"sizes": sizes, "runs": rows}}, f)