# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Find the best batch size on your machine.

8-byte values through a pipeline of three operators (multiply, add, compare),
each reading the whole batch and writing a new one, across a sweep of batch sizes.
"""
import argparse, json, platform, statistics, subprocess, time
from pathlib import Path
import numpy as np

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--values", type=int, default=64_000_000, help="8-byte values pushed through the pipeline.")
p.add_argument("--batch-sizes", type=lambda s: [int(v) for v in s.split(",")],
               default=[64, 256, 1024, 4096, 16384, 131072, 1_048_576, 16_777_216],
               help="Batch sizes to try (values per batch), comma-separated.")
p.add_argument("--repeats", type=int, default=3, help="Runs per batch size; the median time is reported.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()


def cache_sizes():
    """L1d and L2 in bytes, from the OS (sysctl on macOS, sysfs on Linux)."""
    if platform.system() == "Darwin":
        def q(key):
            out = subprocess.run(["sysctl", "-n", key], capture_output=True, text=True).stdout.strip()
            return int(out) if out.isdigit() else 0
        # Apple silicon reports the performance cores as perflevel0
        return {"L1d": q("hw.perflevel0.l1dcachesize") or q("hw.l1dcachesize"),
                "L2": q("hw.perflevel0.l2cachesize") or q("hw.l2cachesize")}
    sizes = {}
    for d in sorted(Path("/sys/devices/system/cpu/cpu0/cache").glob("index*")):
        level, kind = (d / "level").read_text().strip(), (d / "type").read_text().strip()
        raw = (d / "size").read_text().strip()
        n = int(raw[:-1]) * (1024 if raw.endswith("K") else 1024 * 1024 if raw.endswith("M") else 1)
        if level == "1" and kind == "Data": sizes["L1d"] = n
        if level == "2": sizes["L2"] = n
    return sizes


x = np.random.default_rng(0).random(args.values)

def run(size):
    acc = 0
    for i in range(0, len(x), size):
        b = x[i:i+size]
        acc += np.count_nonzero((b * 2.0 + 1.0) > 2.0)   # 3 operators, 2 temporary arrays
    return acc

rows = []
for size in args.batch_sizes:
    runs = []
    for _ in range(args.repeats):
        t = time.perf_counter(); run(size); runs.append(time.perf_counter() - t)
    rows.append({"batch_size": size, "seconds": statistics.median(runs)})
    print(f"{size:>12,} values  {rows[-1]['seconds']:7.3f} s")

best = min(rows, key=lambda r: r["seconds"])
print(f"\nfastest: {best['batch_size']:,} values ({best['batch_size'] * 8 / 1024:,.0f} KB per column)")
caches = [{"cache": k, "bytes": v, "values": v // 8} for k, v in cache_sizes().items() if v]
for c in caches:
    print(f"{c['cache']:>3} = {c['bytes'] / 1024:,.0f} KB  ->  {c['values']:,} values of 8 bytes")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"batches": rows, "caches": caches, "summary": [{
            "best_batch": best["batch_size"], "best_bytes": best["batch_size"] * 8, "best_s": best["seconds"],
            "smallest_batch": rows[0]["batch_size"], "smallest_s": rows[0]["seconds"]}]}}, f)