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

10 million Zipf-distributed events over 1 million keys into a Count-Min
sketch with w = 2,719 and d = 5 (epsilon = 0.001, delta = 1%). Relative error
for the top 10 keys and for keys with true count near 10. Then Misra-Gries
and Space-Saving with 1,000 counters.
"""
import argparse, json
from collections import Counter
import numpy as np

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--events", type=int, default=10_000_000, help="Events in the stream.")
p.add_argument("--keys", type=int, default=1_000_000, help="Key space.")
p.add_argument("--zipf", type=float, default=1.3, help="Zipf exponent.")
p.add_argument("--width", type=int, default=2719, help="Count-Min width w.")
p.add_argument("--depth", type=int, default=5, help="Count-Min depth d.")
p.add_argument("--counters", type=int, default=1000, help="Counters for Misra-Gries and Space-Saving.")
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)
stream = (rng.zipf(args.zipf, args.events) % args.keys).astype(np.int64)
true = np.bincount(stream, minlength=args.keys)
seeds = rng.integers(1, 2**31 - 1, (args.depth, 2))
cms = np.zeros((args.depth, args.width), dtype=np.int64)
cols = [((a * stream + b) % 2147483647) % args.width for a, b in seeds]
for r in range(args.depth): np.add.at(cms[r], cols[r], 1)
def est(keys): return np.min([cms[r][((a * keys + b) % 2147483647) % args.width] for r, (a, b) in enumerate(seeds)], axis=0)
eps_n = np.e / args.width * args.events
top = np.argsort(-true)[:10]
tail = np.where((true >= 8) & (true <= 12))[0][:1000]
top_err = (est(top) - true[top]) / true[top]; tail_err = (est(tail) - true[tail]) / true[tail]
print(f"bound e/w x N = {eps_n:,.0f}; max overestimate seen {int((est(np.arange(args.keys)) - true).max()):,}")
print(f"top 10 keys: median relative error {np.median(top_err):.3%}; keys with ~10 events: median relative error {np.median(tail_err):,.0%}")

def misra_gries(s, k):
    c = {}
    for x in s:
        if x in c: c[x] += 1
        elif len(c) < k: c[x] = 1
        else:
            for y in list(c):
                c[y] -= 1
                if c[y] == 0: del c[y]
    return c
def space_saving(s, k):
    c = {}
    for x in s:
        if x in c: c[x] += 1
        elif len(c) < k: c[x] = 1
        else:
            y = min(c, key=c.get); c[x] = c.pop(y) + 1
    return c
sub = stream[: 1_000_000].tolist(); tsub = Counter(sub); N = len(sub)
heavy = {k for k, v in tsub.items() if v > N / (args.counters + 1)}
res = []
for name, fn in (("Misra-Gries", misra_gries), ("Space-Saving", space_saving)):
    c = fn(sub, args.counters)
    top_true = [k for k, _ in tsub.most_common(10)]; top_est = sorted(c, key=c.get, reverse=True)[:10]
    res.append({"algorithm": name, "heavy_hitters": len(heavy), "heavy_found": len(heavy & set(c)),
                "top10_same_order": "yes" if top_true == top_est else "no"})
    print(f"{name:<12} (on the first 1M events): {res[-1]['heavy_found']} of {len(heavy)} keys above N/{args.counters + 1} found; "
          f"top-10 order identical: {res[-1]['top10_same_order']}")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"heavy": res, "summary": [{"bound": eps_n, "top_err_pct": float(np.median(top_err)) * 100,
                   "tail_err_pct": float(np.median(tail_err)) * 100}]}}, f)