#!/usr/bin/env python3
"""Bloom filters: a little space buys a low false-positive rate, DDIA 2e Ch. 4.

An LSM storage engine may have to check many SSTables to confirm a key does not
exist. A Bloom filter per SSTable gives a fast, approximate "is this key maybe
here?" so most of those disk reads can be skipped. The catch: it can report a
false positive (say "maybe" for a key that isn't there), never a false negative.

The surprise people get wrong is the *space*: you don't need bytes per key for a
low miss rate -- you need a handful of BITS. We build a Bloom filter in pure
Python, insert 1,000,000 keys, query 1,000,000 keys that are absent, and measure
the false-positive rate as we vary bits-per-key -- checking it against the
analytic (1 - e^(-k n / m))^k. The rate is independent of how big the keys are.

Pure standard library, deterministic hashing -> reproducible.
"""
import hashlib
import math

N = 1_000_000            # keys inserted
TEST = 1_000_000         # absent keys queried to measure false positives
BITS_PER_KEY = [4, 6, 8, 10, 12, 16]


class BloomFilter:
    def __init__(self, n, bits_per_key):
        self.m = n * bits_per_key
        # optimal number of hashes k = (m/n) ln 2
        self.k = max(1, round((self.m / n) * math.log(2)))
        self.bits = bytearray((self.m + 7) // 8)

    def _indexes(self, key):
        # Kirsch-Mitzenmacher double hashing from one SHA-256 digest -> k indexes.
        d = hashlib.sha256(key).digest()
        h1 = int.from_bytes(d[:8], "big")
        h2 = int.from_bytes(d[8:16], "big") | 1
        for i in range(self.k):
            yield (h1 + i * h2) % self.m

    def add(self, key):
        for idx in self._indexes(key):
            self.bits[idx >> 3] |= 1 << (idx & 7)

    def __contains__(self, key):
        return all(self.bits[idx >> 3] & (1 << (idx & 7)) for idx in self._indexes(key))


def main():
    present = [f"key-{i}".encode() for i in range(N)]
    absent = [f"absent-{i}".encode() for i in range(TEST)]

    print(f"Bloom filter: {N:,} keys inserted, {TEST:,} absent keys queried\n")
    print(f"{'bits/key':>8}  {'k':>2}  {'size':>8}  {'measured FP':>12}  {'analytic FP':>12}")
    for bpk in BITS_PER_KEY:
        bf = BloomFilter(N, bpk)
        for key in present:
            bf.add(key)
        false_pos = sum(1 for key in absent if key in bf)
        measured = false_pos / TEST
        analytic = (1 - math.exp(-bf.k * N / bf.m)) ** bf.k
        size_mb = len(bf.bits) / 1e6
        print(f"{bpk:>8}  {bf.k:>2}  {size_mb:6.2f} MB  {measured * 100:10.2f}%  {analytic * 100:10.2f}%")

    print("\nnote: the filter size depends only on the NUMBER of keys, not their length --")
    print("      1,000,000 keys at 10 bits/key is ~1.25 MB whether keys are 5 bytes or 5 KB.")


if __name__ == "__main__":
    main()