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

A power-law graph of 200,000 vertices split across 8 workers by vertex
(v % 8) and by edge (hash of the edge). 10 PageRank supersteps; each worker's
messages and processing time are measured per superstep (workers run one
after another, each timed on its own share). The sum of the slowest worker's
times against the sum of the average; then a combiner that sums messages per
destination; then PageRank at 10, 20 and 30 iterations, and connected
components, which must run until no label changes.
"""
import argparse, json, time
import numpy as np

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--n", type=int, default=200_000, help="Vertices.")
p.add_argument("--m", type=int, default=5, help="Edges per new vertex.")
p.add_argument("--workers", type=int, default=8, help="Workers.")
p.add_argument("--steps", type=int, default=10, help="Supersteps.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

rng = np.random.default_rng(0); W = args.workers
targets = np.empty(args.n * args.m * 2, dtype=np.int64); src, dst = [], []; filled = 0
for v in range(args.m, args.n):
    picks = targets[rng.integers(0, filled, args.m)] if filled else np.arange(args.m)
    src.append(np.full(args.m, v)); dst.append(picks); targets[filled:filled + args.m] = picks; targets[filled + args.m:filled + 2 * args.m] = v; filled += 2 * args.m
s = np.concatenate(src); d = np.concatenate(dst); s, d = np.concatenate([s, d]), np.concatenate([d, s])
outdeg = np.bincount(s, minlength=args.n)
def partition(mode):
    if mode == "vertex": return s % W                                    # the worker that owns the source sends its messages
    h = (s * 0x9E3779B1 + d * 0x85EBCA77) & 0xFFFFFFFF; return h % W     # edge-cut-free: edges spread by hash
def run(mode, combiner):
    part = partition(mode); rank = np.full(args.n, 1 / args.n); slow, avg, msgs = 0.0, 0.0, []
    for _ in range(args.steps):
        times, counts, incoming = [], [], np.zeros(args.n)
        for w in range(W):
            mask = part == w; t = time.perf_counter()
            ss, dd = s[mask], d[mask]; contrib = rank[ss] / outdeg[ss]
            if combiner:                                                  # sum per destination before sending
                u, inv = np.unique(dd, return_inverse=True); summed = np.bincount(inv, weights=contrib); n_msg = len(u); np.add.at(incoming, u, summed)
            else: n_msg = len(dd); np.add.at(incoming, dd, contrib)
            times.append(time.perf_counter() - t); counts.append(n_msg)
        rank = 0.15 / args.n + 0.85 * incoming
        slow += max(times); avg += sum(times) / W; msgs.append(counts)
    return {"partitioning": f"by {mode}", "combiner": "yes" if combiner else "no", "slowest_sum_s": slow, "average_sum_s": avg,
            "straggler_ratio": slow / avg, "max_worker_msgs": int(max(max(c) for c in msgs)), "mean_worker_msgs": float(np.mean(msgs))}, rank
rows = []
for mode in ("vertex", "edge"):
    for comb in (False, True):
        r, _ = run(mode, comb); rows.append(r)
def pagerank(iters):
    rank = np.full(args.n, 1 / args.n)
    for _ in range(iters):
        inc = np.bincount(d, weights=rank[s] / outdeg[s], minlength=args.n); rank = 0.15 / args.n + 0.85 * inc
    return rank
r10, r20, r30 = pagerank(10), pagerank(20), pagerank(30)
label = np.arange(args.n); rounds = 0
while True:
    new = label.copy(); np.minimum.at(new, d, label[s]); rounds += 1
    if (new == label).all(): break
    label = new
summary = {"pr_10_vs_30_l1": float(np.abs(r10 - r30).sum()), "pr_20_vs_30_l1": float(np.abs(r20 - r30).sum()), "cc_rounds": rounds,
           "top10_same_10_vs_30": "yes" if set(np.argsort(-r10)[:10]) == set(np.argsort(-r30)[:10]) else "no"}
for r in rows: print(f"{r['partitioning']:<10} combiner {r['combiner']:<3}: slowest worker summed {r['slowest_sum_s']:.3f} s vs average {r['average_sum_s']:.3f} s ({r['straggler_ratio']:.2f}x); busiest worker {r['max_worker_msgs']:,} messages vs mean {r['mean_worker_msgs']:,.0f}")
print(f"PageRank: L1 difference 10 vs 30 iterations {summary['pr_10_vs_30_l1']:.2e}, 20 vs 30 {summary['pr_20_vs_30_l1']:.2e}; same top 10: {summary['top10_same_10_vs_30']}; connected components needed {rounds} rounds")

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