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

Ten servers with very different request counts (10,000 down to 100) draw
latencies from log-normal distributions. Compares (a) the average of the ten
p99 values, (b) the true p99 over all samples, and (c) the p99 read from the
merged bucket counts of Prometheus-style histograms.
"""
import argparse, json
import numpy as np

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--seed", type=int, default=1, help="Random seed.")
p.add_argument("--buckets", default="5,10,25,50,75,100,150,200,300,500,750,1000,2500,5000", help="Bucket upper bounds (ms).")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
rng = np.random.default_rng(args.seed)

counts = [10_000, 5_000, 3_000, 2_000, 1_000, 800, 500, 300, 200, 100]
servers = [rng.lognormal(np.log(20 + 15 * i), 0.4 + 0.08 * i, n) for i, n in enumerate(counts)]   # small servers are slower and wider
bounds = np.array([float(b) for b in args.buckets.split(",")] + [np.inf])
hist = sum(np.histogram(s, bins=np.concatenate([[0], bounds]))[0] for s in servers)               # merge = add the counts

def hist_quantile(q, bounds, counts):
    """Prometheus histogram_quantile: find the bucket, interpolate linearly inside it."""
    cum = np.cumsum(counts); rank = q * cum[-1]; i = int(np.searchsorted(cum, rank))
    lo = 0 if i == 0 else bounds[i - 1]; hi = bounds[i]; below = 0 if i == 0 else cum[i - 1]
    return float(lo + (hi - lo) * (rank - below) / counts[i]), float(lo), float(hi)

per = [float(np.percentile(s, 99)) for s in servers]
true = float(np.percentile(np.concatenate(servers), 99)); merged, lo, hi = hist_quantile(0.99, bounds, hist)
rows = [{"method": "(a) average of the ten p99s", "p99_ms": float(np.mean(per))},
        {"method": "(b) true p99 over all samples", "p99_ms": true},
        {"method": "(c) p99 from merged buckets", "p99_ms": merged}]
for r in rows: r["err_pct"] = (r["p99_ms"] / true - 1) * 100
shards = [{"server": i + 1, "requests": n, "p99_ms": v} for i, (n, v) in enumerate(zip(counts, per))]
summary = {"bucket_lo": lo, "bucket_hi": hi}
for r in rows: print(f"{r['method']:<32} {r['p99_ms']:7.1f} ms  ({r['err_pct']:+.1f}%)")
print(f"the true p99 falls in the bucket {lo:g}-{hi:g} ms")

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