# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Mean, percentile averaging and fan-out.

One million latencies: 99% around 5 ms, 1% around 700 ms. Prints the mean and
percentiles, the average of ten unequal shards' p99 against the global p99,
and the share of fan-out requests (the maximum of 100 draws) slower than the
single-call p99.
"""
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="Latencies to draw.")
p.add_argument("--fanout", type=int, default=100, help="Calls per fanned-out request.")
p.add_argument("--slow-share", type=float, default=0.01, help="Share of slow calls.")
p.add_argument("--seed", type=int, default=1, 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)
def draw(n):
    slow = rng.random(n) < args.slow_share
    return np.where(slow, rng.normal(700, 80, n), rng.normal(5, 1, n)).clip(0.5)
x = draw(args.n)
p50, p99, p999 = np.percentile(x, [50, 99, 99.9])
stats = [{"statistic": k, "ms": float(v)} for k, v in (("mean", x.mean()), ("p50", p50), ("p99", p99), ("p99.9", p999))]

# ten unequal shards, and the slow calls are not spread evenly across them (they lean to the later shards)
cuts = np.sort(rng.choice(np.arange(1, args.n), 9, replace=False))
skewed = np.split(x[np.argsort(rng.random(args.n) + (x > 100) * rng.random(args.n))], cuts)
shard_rows = [{"shard": i + 1, "rows": len(s), "p99_ms": float(np.percentile(s, 99))} for i, s in enumerate(skewed)]
avg_p99 = float(np.mean([r["p99_ms"] for r in shard_rows]))

fan = np.concatenate([draw(10_000 * args.fanout).reshape(10_000, args.fanout).max(axis=1) for _ in range(10)])
share = float((fan > p99).mean()); theory = 1 - (1 - args.slow_share) ** args.fanout

summary = {"mean": float(x.mean()), "p50": float(p50), "p99": float(p99), "p999": float(p999), "avg_shard_p99": avg_p99,
           "fan_share": share, "fan_theory": theory, "fan_p50": float(np.percentile(fan, 50))}
print(f"mean {x.mean():.1f} ms  p50 {p50:.1f}  p99 {p99:.1f}  p99.9 {p999:.1f}")
print(f"average of ten shard p99s {avg_p99:.1f} ms against the global p99 {p99:.1f} ms")
print(f"fan-out of {args.fanout}: {share:.3f} of requests slower than the single-call p99 (1 - 0.99^{args.fanout} = {theory:.3f}); median request {summary['fan_p50']:.0f} ms")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"summary": summary, "tables": {"stats": stats, "shards": shard_rows}}, f)