# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Digest collisions, and a Bloom filter used for data.

Part A: the birthday formula, by shrinking SHA-256 to 32 bits: the share of
200 runs of 100,000 random digests that contain a collision, against
1 − exp(−n²/2^33). Part B: a Bloom filter (8 million bits, 4 hashes) used to
drop "duplicates" from 1,000,000 distinct keys, and the same filter used as a
negative oracle in front of an exact set.
"""
import argparse, hashlib, json, math, random

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--n", type=int, default=100_000, help="Part A: digests per run.")
p.add_argument("--bits", type=int, default=32, help="Part A: digest width.")
p.add_argument("--runs", type=int, default=200, help="Part A: runs.")
p.add_argument("--keys", type=int, default=1_000_000, help="Part B: distinct keys.")
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()
random.seed(args.seed)

def collides(n, bits):
    seen = set()
    for _ in range(n):
        d = int.from_bytes(hashlib.sha256(random.randbytes(16)).digest(), "big") >> (256 - bits)
        if d in seen: return True
        seen.add(d)
    return False
measured = sum(collides(args.n, args.bits) for _ in range(args.runs)) / args.runs
formula = 1 - math.exp(-args.n * args.n / 2 ** (args.bits + 1))

class Bloom:
    def __init__(self, m_bits, k): self.m, self.k, self.bits = m_bits, k, bytearray(m_bits // 8 + 1)
    def _idx(self, key):
        h = hashlib.sha256(key.encode()).digest()
        return [int.from_bytes(h[4 * i:4 * i + 4], "big") % self.m for i in range(self.k)]
    def add(self, key):
        for i in self._idx(key): self.bits[i // 8] |= 1 << (i % 8)
    def might(self, key): return all(self.bits[i // 8] >> (i % 8) & 1 for i in self._idx(key))
b, kept_naive = Bloom(8_000_000, 4), 0
for i in range(args.keys):
    key = f"event-{i}"
    if not b.might(key): b.add(key); kept_naive += 1
fill = sum(bin(x).count("1") for x in b.bits) / b.m
b2, exact, kept_oracle, lookups = Bloom(8_000_000, 4), set(), 0, 0
for i in range(args.keys):
    key = f"event-{i}"
    if b2.might(key):
        lookups += 1
        if key in exact: continue                                       # only the exact set may say "seen"
    b2.add(key); exact.add(key); kept_oracle += 1
summary = {"collision_measured": measured, "collision_formula": formula, "kept_naive": kept_naive, "lost_naive": args.keys - kept_naive,
           "lost_pct": (args.keys - kept_naive) / args.keys * 100, "fill_pct": fill * 100, "final_fp_pct": fill ** 4 * 100,
           "kept_oracle": kept_oracle, "exact_lookups": lookups}
print(f"A: {args.bits}-bit digests, n = {args.n:,}: collision in {measured:.0%} of {args.runs} runs (formula {formula:.0%})")
print(f"B: Bloom filter as the dedup: kept {kept_naive:,} of {args.keys:,} distinct keys (lost {summary['lost_pct']:.2f}%); filter {fill:.0%} full, final false-positive rate {summary['final_fp_pct']:.1f}%")
print(f"   Bloom filter in front of an exact set: kept {kept_oracle:,}; the exact set was consulted {lookups:,} times")

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