# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy", "numba"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Compute a compilation break-even.

An interpreted Python loop against the same loop compiled with numba: the
compile cost, the saving per tuple, and the data volume at which compiling pays.
"""
import argparse, json, statistics, time
import numpy as np
from numba import njit

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--rows", type=int, default=1_000_000, help="Tuples per timed run.")
p.add_argument("--expressions", type=lambda s: s.split(","), default=["simple", "complex"],
               help="Which loop bodies to time: simple (one compare) and/or complex (five more arithmetic operations).")
p.add_argument("--repeats", type=int, default=3, help="Runs per timing; the median is reported.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()


def py_sum(a, thr):
    s = 0.0
    for v in a:
        if v > thr: s += v
    return s

def py_complex(a, thr):
    s = 0.0
    for v in a:
        w = (v * 1.5 + 0.25) * v - v / 3.0             # five arithmetic operations
        if w > thr: s += w
    return s

kernels = {"simple": py_sum, "complex": py_complex}


def median_time(fn, a):
    runs = []
    for _ in range(args.repeats):
        t = time.perf_counter(); fn(a, 0.5); runs.append(time.perf_counter() - t)
    return statistics.median(runs)


a = np.random.default_rng(0).random(args.rows)

# numba's own start-up (loading LLVM and its type system) is paid once per process, on the first
# compile; time it on a throwaway function so each expression below shows its compile cost alone
t = time.perf_counter(); njit(lambda v: v + 1)(1.0); startup = time.perf_counter() - t
print(f"numba start-up (once per process): {startup:.2f} s")
results, curve = [], []
for name in args.expressions:
    py_fn = kernels[name]
    nb_fn = njit(py_fn)                                 # a fresh compiled copy
    t = time.perf_counter(); nb_fn(a, 0.5); compile_cost = time.perf_counter() - t   # first call = compile + run
    py = median_time(py_fn, a) / len(a)
    nb = median_time(nb_fn, a) / len(a)
    saving = py - nb
    break_even, break_even_cold = compile_cost / saving, (startup + compile_cost) / saving
    results.append({"expression": name, "startup_s": startup, "compile_s": compile_cost, "interpreted_ns": py * 1e9,
                    "compiled_ns": nb * 1e9, "saving_ns": saving * 1e9, "break_even_tuples": break_even,
                    "break_even_cold_tuples": break_even_cold})
    print(f"{name:>8}: compile {compile_cost:.2f} s · saving {saving * 1e9:.1f} ns/tuple · "
          f"break-even {break_even / 1e6:.1f} M tuples ({break_even_cold / 1e6:.1f} M counting start-up)")
    for n in (10 ** k for k in range(3, 11)):           # total time to process n tuples each way
        curve.append({"expression": name, "tuples": n, "interpreted_s": py * n, "compiled_s": compile_cost + nb * n,
                      "compiled_cold_s": startup + compile_cost + nb * n})

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"results": results, "curve": curve}}, f)