# /// 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: measure the false-positive rate, and the cost of size.

1,000,000 keys at 5 to 20 bits per key with k = round(m/n ln 2); the
false-positive rate on 1,000,000 absent keys. Then 10^7 lookups in a 10 MB
and a 1 GB filter (NumPy), and in a blocked filter where all k bits of a key
sit in one 64-byte block.
"""
import argparse, json, time
import numpy as np

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--keys", type=int, default=1_000_000, help="Keys inserted.")
p.add_argument("--bits", type=lambda s: [float(v) for v in s.split(",")], default=[5, 8, 10, 14.4, 20], help="Bits per key.")
p.add_argument("--lookups", type=int, default=10_000_000, help="Lookups for the timing part.")
p.add_argument("--big-mb", type=int, default=1024, help="Size of the large filter, MB.")
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)
M1, M2 = np.uint64(0x9E3779B97F4A7C15), np.uint64(0xC2B2AE3D27D4EB4F)

def hashes(keys):
    h1 = (keys * M1) ^ (keys >> np.uint64(29)); h2 = (keys * M2) | np.uint64(1)
    return h1, h2

def positions(keys, k, m):
    h1, h2 = hashes(keys)
    return [((h1 + np.uint64(i) * h2) % np.uint64(m)).astype(np.int64) for i in range(k)]

ins = rng.integers(1, 2**62, args.keys, dtype=np.uint64); absent = rng.integers(2**62, 2**63, args.keys, dtype=np.uint64)
rows = []
for b in args.bits:
    m = int(b * args.keys); k = max(1, round(b * np.log(2)))
    bits = np.zeros(m, dtype=bool)
    for pos in positions(ins, k, m): bits[pos] = True
    present = np.ones(args.keys, dtype=bool)
    for pos in positions(absent, k, m): present &= bits[pos]
    fp = float(present.mean())
    rows.append({"bits_per_key": b, "k": k, "fp_pct": fp * 100, "formula_pct": (1 - np.exp(-k / b)) ** k * 100})
    print(f"{b:5.1f} bits/key, k = {k:2d}: false positives {fp:7.3%}   (formula {rows[-1]['formula_pct']:6.3f}%)")

timing = []
q = rng.integers(1, 2**62, args.lookups, dtype=np.uint64)
for label, mb, blocked in (("10 MB filter", 10, False), (f"{args.big_mb:,} MB filter", args.big_mb, False),
                           (f"{args.big_mb:,} MB blocked filter", args.big_mb, True)):
    m = mb * 8 * 2**20; bits = np.zeros(m // 8, dtype=np.uint8); k = 7
    h1, h2 = hashes(q)
    t = time.perf_counter()
    ok = np.ones(len(q), dtype=bool)
    if blocked:
        block = (h1 % np.uint64(m // 512)).astype(np.int64) * 512           # 64-byte block = 512 bits
        for i in range(k):
            pos = block + ((h2 >> np.uint64(9 * i)) % np.uint64(512)).astype(np.int64)
            ok &= (bits[pos >> 3] >> (pos & 7).astype(np.uint8)) & 1 == 1
    else:
        for i in range(k):
            pos = ((h1 + np.uint64(i) * h2) % np.uint64(m)).astype(np.int64)
            ok &= (bits[pos >> 3] >> (pos & 7).astype(np.uint8)) & 1 == 1
    dt = time.perf_counter() - t
    timing.append({"filter": label, "ns_per_lookup": dt / len(q) * 1e9}); print(f"{label:<26} {dt / len(q) * 1e9:6.1f} ns per lookup")
    del bits

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