# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Conflict grows with the square of concurrency.

L transactions alive at once, each touching k random keys out of N: the number
of pairs that share a key, against L²k²/2N, for uniform and Zipf-skewed key
choice.
"""
import argparse, json
import numpy as np

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--n", type=int, default=1_000_000, help="Keys N.")
p.add_argument("--k", type=int, default=10, help="Keys per transaction k.")
p.add_argument("--ls", default="50,100,200,500", help="Concurrent transactions L.")
p.add_argument("--zipf", type=float, default=1.2, help="Zipf exponent for the skewed run.")
p.add_argument("--trials", type=int, default=200, help="Trials per point.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

def conflicting_pairs(L, zipf_a=None):
    rng = np.random.default_rng(1); total = 0
    for _ in range(args.trials):
        keys = [set(rng.integers(0, args.n, args.k)) if zipf_a is None else set((rng.zipf(zipf_a, args.k) - 1) % args.n) for _ in range(L)]
        for i in range(L):
            for j in range(i + 1, L):
                if keys[i] & keys[j]: total += 1
    return total / args.trials
rows = []
for L in map(int, args.ls.split(",")):
    rows.append({"L": L, "uniform": conflicting_pairs(L), "formula": L * L * args.k * args.k / (2 * args.n), "zipf": conflicting_pairs(L, args.zipf)})
for i, r in enumerate(rows): r["uniform_vs_prev"] = r["uniform"] / rows[i - 1]["uniform"] if i and rows[i - 1]["uniform"] else None
for r in rows: print(f"L = {r['L']:>3}: uniform {r['uniform']:8.2f} pairs (formula {r['formula']:6.2f})   Zipf {r['zipf']:9.1f} pairs")

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