# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Barriers multiply variance.

E[max]/E[mean] of n task durations against 1 + c*sqrt(2 ln n), a BSP job with
and without a barrier every round, and the same max under a heavy tail.
"""
import argparse, json
import numpy as np

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--c", type=float, default=0.3, help="Coefficient of variation of a task's duration (mean 1).")
p.add_argument("--task-counts", type=lambda s: [int(v) for v in s.split(",")], default=[10, 100, 1000, 10000],
               help="Values of n (tasks per stage), comma-separated.")
p.add_argument("--trials", type=int, default=2000, help="Simulated stages per value of n.")
p.add_argument("--rounds", type=int, default=30, help="BSP supersteps.")
p.add_argument("--workers", type=int, default=1000, help="BSP workers.")
p.add_argument("--pareto-alpha", type=float, default=2.5, help="Tail index of the heavy-tailed durations.")
p.add_argument("--seed", type=int, default=0, help="Random seed.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
rng = np.random.default_rng(args.seed)
c = args.c

rows = []
for n in args.task_counts:
    d = np.clip(rng.normal(1.0, c, (args.trials, n)), 0.05, None)   # trials of n tasks
    sim, formula = d.max(axis=1).mean(), 1 + c * np.sqrt(2 * np.log(n))
    rows.append({"n": n, "simulated": float(sim), "formula": float(formula)})
    print(f"n = {n:>6,}   simulated {sim:.2f}   formula {formula:.2f}")

# BSP: every round waits for its slowest worker, versus no barrier
d = np.clip(rng.normal(1.0, c, (args.rounds, args.workers)), 0.05, None)
bsp = d.max(axis=1).sum()               # every round waits for its slowest worker
async_ = d.sum(axis=0).max()            # each worker runs its rounds back to back
print(f"\n{args.rounds} rounds x {args.workers} workers   with barriers {bsp:.1f}   without {async_:.1f}   "
      f"ratio {bsp / async_:.2f}x")

# heavy tail: a Pareto-like distribution instead of the normal
h = rng.pareto(args.pareto_alpha, (args.trials, args.workers)) + 1
heavy = h.max(axis=1).mean() / h.mean()
print(f"heavy-tail max/mean ({args.workers} tasks): {heavy:.1f}")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {
            "max_over_mean": rows,
            "bsp": [{"with_barriers": float(bsp), "without": float(async_), "ratio": float(bsp / async_),
                     "heavy_tail_max_over_mean": float(heavy)}],
        }}, f)