# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Bloom filter size and false-positive rate.

Builds Bloom filters for 1,000,000 keys at 5, 10 and 15 bits per key (with
the best number of hash functions, k = round(bits × ln 2)), measures the
false-positive rate on 1,000,000 keys not in the set, and computes each
filter's size in MB.
"""
import argparse, json, math
import numpy as np

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--keys", type=int, default=1_000_000, help="Keys in the set.")
p.add_argument("--bits", default="5,10,15", help="Bits per key.")
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)
members = rng.integers(0, 2**62, args.keys, dtype=np.int64).astype(np.uint64)
others = rng.integers(2**62, 2**63 - 1, args.keys, dtype=np.int64).astype(np.uint64)  # disjoint from the members
def mix(x, seed):
    z = x + np.uint64(seed) * np.uint64(0x9E3779B97F4A7C15)
    z = (z ^ (z >> np.uint64(30))) * np.uint64(0xBF58476D1CE4E5B9); z = (z ^ (z >> np.uint64(27))) * np.uint64(0x94D049BB133111EB)
    return z ^ (z >> np.uint64(31))
rows = []
with np.errstate(over="ignore"):
    for b in map(float, args.bits.split(",")):
        m = int(b * args.keys); k = max(1, round(b * math.log(2))); bits = np.zeros(m, dtype=bool)
        h1, h2 = mix(members, 1), mix(members, 2) | np.uint64(1)
        for i in range(k): bits[((h1 + np.uint64(i) * h2) % np.uint64(m)).astype(np.int64)] = True     # double hashing
        o1, o2 = mix(others, 1), mix(others, 2) | np.uint64(1)
        hit = np.ones(args.keys, dtype=bool)
        for i in range(k): hit &= bits[((o1 + np.uint64(i) * o2) % np.uint64(m)).astype(np.int64)]
        rows.append({"bits_per_key": b, "k": k, "fp_pct": hit.mean() * 100, "formula_pct": (1 - math.exp(-k / b)) ** k * 100, "mb": m / 8 / 1e6})
for r in rows: print(f"{r['bits_per_key']:g} bits/key, k = {r['k']:>2}: false positives {r['fp_pct']:6.3f}% (formula {r['formula_pct']:6.3f}%), size {r['mb']:.3f} MB")

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