# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy", "scipy"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Fit a USL curve to something you run.

A parallel workload: N worker processes each scan their own 64 MB array
(a memory-bandwidth-bound sum) for a fixed time; throughput is the total GB
scanned per second. Measured at N = 1, 2, 3, 4, 6, 8; the Universal
Scalability Law is fitted to those points, and then throughput is measured
beyond them (N = 12, 16, 24) to test the fit's prediction.
"""
import argparse, json, multiprocessing as mp, time
import numpy as np
from scipy.optimize import curve_fit

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--fit-n", default="1,2,3,4,6,8", help="Worker counts used for the fit.")
p.add_argument("--test-n", default="12,16,24", help="Worker counts beyond the fit.")
p.add_argument("--seconds", type=float, default=3.0, help="Seconds per measurement.")
p.add_argument("--mb", type=int, default=64, help="Array per worker (MB).")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

def worker(ready, seconds, out):
    a = np.ones(args.mb * 2**20 // 8); a.sum(); n = 0
    ready.wait()                                                       # every worker is up and has its array
    stop_at = time.time() + seconds
    while time.time() < stop_at: a.sum(); n += 1
    out.put(n * a.nbytes)

def measure(n):
    out = mp.Queue(); ready = mp.Barrier(n)
    ps = [mp.Process(target=worker, args=(ready, args.seconds, out)) for _ in range(n)]
    for q in ps: q.start()
    total = sum(out.get() for _ in ps)
    for q in ps: q.join()
    return total / args.seconds / 1e9

if __name__ == "__main__":
    fit_n = [int(x) for x in args.fit_n.split(",")]; test_n = [int(x) for x in args.test_n.split(",")]
    X = {n: measure(n) for n in fit_n + test_n}
    N = np.array(fit_n, dtype=float); C = np.array([X[n] for n in fit_n]) / X[1]
    usl = lambda n, a, b: n / (1 + a * (n - 1) + b * n * (n - 1))
    (a, b), _ = curve_fit(usl, N, C, bounds=(0, 1))
    peak = float(np.sqrt((1 - a) / b)) if b > 0 else None
    rows = [{"N": n, "gb_s": X[n], "relative": X[n] / X[1], "fit": float(usl(n, a, b)), "used_for_fit": "yes" if n in fit_n else "no (test)"} for n in fit_n + test_n]
    summary = {"alpha": float(a), "beta": float(b), "peak_n": peak, "best_n": max(X, key=X.get), "best_gb_s": max(X.values())}
    print(f"alpha={a:.4f} beta={b:.6f} " + (f"peak at N*={peak:.0f}" if peak else "no peak (beta = 0)"))
    for r in rows: print(f"N={r['N']:>3}  {r['gb_s']:6.1f} GB/s  relative {r['relative']:5.2f}  USL fit {r['fit']:5.2f}  ({r['used_for_fit']})")
    if args.json:
        with open(args.json, "w") as f:
            json.dump({"summary": summary, "tables": {"throughput": rows}}, f)