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

Measures the machine's two ceilings (peak memory bandwidth from a large array
copy, peak compute from a float64 matrix multiply), then the achieved GB/s
and FLOP/s of a scan (sum), a filter (sum where x > t) and a hash aggregation
(group-by sum over 1,000 keys) on 50 million float64 values. Places each on the roofline by its arithmetic intensity (FLOPs per byte).
"""
import argparse, json, time
import numpy as np

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--n", type=int, default=50_000_000, help="Values per operator (float64).")
p.add_argument("--repeats", type=int, default=5, help="Runs per measurement (best is kept).")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

def best(fn):
    ts = []
    for _ in range(args.repeats):
        t = time.perf_counter(); fn(); ts.append(time.perf_counter() - t)
    return min(ts)
rng = np.random.default_rng(0)
x = rng.random(args.n); keys = rng.integers(0, 1000, args.n); dst = np.empty_like(x)
bw_peak = 2 * x.nbytes / best(lambda: np.copyto(dst, x)) / 1e9      # read + write
m = rng.random((2048, 2048)); flops_peak = 2 * 2048**3 / best(lambda: m @ m) / 1e9

ops = [("scan: sum(x)", lambda: x.sum(), x.nbytes, args.n),
       ("filter: sum(x[x > 0.5])", lambda: x[x > 0.5].sum(), x.nbytes + args.n // 2 * 8, 1.5 * args.n),
       ("hash aggregation: group-by sum, 1,000 keys", lambda: np.bincount(keys, weights=x, minlength=1000), x.nbytes + keys.nbytes, args.n)]
rows = []
for name, fn, nbytes, flops in ops:
    t = best(fn)
    ai = flops / nbytes; gbs = nbytes / t / 1e9; gflops = flops / t / 1e9
    roof = min(flops_peak, ai * bw_peak)
    rows.append({"operator": name, "intensity": ai, "gb_s": gbs, "gflop_s": gflops, "roof_gflop_s": roof, "pct_of_roof": gflops / roof * 100,
                 "bound": "memory" if ai * bw_peak < flops_peak else "compute", "pct_of_bandwidth": gbs / bw_peak * 100})
ridge = flops_peak / bw_peak
summary = {"bw_peak": bw_peak, "flops_peak": flops_peak, "ridge": ridge}
print(f"ceilings: memory {bw_peak:.1f} GB/s, compute {flops_peak:.0f} GFLOP/s -> ridge point {ridge:.1f} FLOPs per byte")
for r in rows: print(f"{r['operator']:<44} {r['intensity']:.3f} FLOP/B  {r['gb_s']:6.1f} GB/s ({r['pct_of_bandwidth']:5.1f}% of bandwidth)  {r['gflop_s']:6.2f} GFLOP/s  -> {r['bound']}-bound side")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"summary": summary, "tables": {"operators": rows}}, f)