#!/usr/bin/env python3
"""
bloom.py — a working Bloom filter, its false-positive theory, and the gap
between them. Reference material for Project 4 (LSM storage engine).

Why an LSM engine needs this
----------------------------
An LSM tree stores a key's versions across many immutable SSTables. A read for
a key that does not exist must, in the worst case, touch every table on disk.
With 40 tables that is 40 random reads to answer "no". A Bloom filter per table
turns most of those into an in-memory rejection: it can say "definitely not
here" (never wrong) or "maybe here" (sometimes wrong). Read amplification for
absent keys collapses from O(#tables) to O(#tables * fpr) + 1.

The theory, derived
-------------------
m bits, n inserted keys, k independent hash functions.

  P(a specific bit is still 0 after all insertions)
      = (1 - 1/m)^(k*n)  ~  e^(-k*n/m)

  P(a specific bit is 1) = 1 - e^(-k*n/m)

  A false positive needs all k probed bits to be 1. Treating the bits as
  independent (they are not quite, which is the source of the small gap you
  will measure):

      fpr ~ (1 - e^(-k*n/m))^k                                      (1)

Minimise (1) over k: let x = e^(-kn/m). d/dk of k*ln(1-x) = 0 gives

      k_opt = (m/n) * ln 2   ~  0.693 * bits_per_key                (2)

Substituting (2) back into (1), each bit is 1 with probability exactly 1/2 --
the filter is at maximum entropy, which is the information-theoretic reason
this is the optimum -- and

      fpr_opt = 2^(-k_opt) = 2^(-0.693 * m/n)  ~  0.6185^(m/n)      (3)

Consequences worth memorising:
  * 10 bits/key -> ~1% fpr with k=7. This is why 10 is the near-universal
    default in RocksDB, LevelDB, Cassandra.
  * Each ADDITIONAL 10 bits/key divides the fpr by ~100 (0.6185^10 = 0.0082).
  * The fpr does not depend on key size or key distribution -- only on m/n and
    k. A Bloom filter over 20-byte keys and over 2-KB keys behaves identically.

    python3 bloom.py
"""

from __future__ import annotations

import hashlib
import math
import random


class BloomFilter:
    def __init__(self, n_expected: int, bits_per_key: float = 10.0):
        self.m = max(8, int(n_expected * bits_per_key))
        self.n_expected = n_expected
        # k must be an integer; round the real optimum and clamp to >= 1.
        self.k = max(1, round(bits_per_key * math.log(2)))
        self.bits = bytearray((self.m + 7) // 8)
        self.n = 0

    def _probes(self, key: bytes):
        """Kirsch-Mitzenmacher double hashing.

        Computing k independent hashes is wasteful. The Kirsch-Mitzenmacher
        result is that

            g_i(x) = h1(x) + i * h2(x)   (mod m)

        gives the same asymptotic false-positive rate as k independent hashes.
        So we take ONE 128-bit digest, split it into two 64-bit halves, and
        derive all k probes arithmetically. This is what production filters do,
        and it is why a Bloom lookup is one hash plus k array probes rather
        than k hashes.
        """
        d = hashlib.blake2b(key, digest_size=16).digest()
        h1 = int.from_bytes(d[:8], "little")
        h2 = int.from_bytes(d[8:], "little") | 1   # odd, so it generates the
        for i in range(self.k):                    # full residue ring mod 2^j
            yield (h1 + i * h2) % self.m

    def add(self, key: bytes) -> None:
        for p in self._probes(key):
            self.bits[p >> 3] |= 1 << (p & 7)
        self.n += 1

    def __contains__(self, key: bytes) -> bool:
        return all(self.bits[p >> 3] >> (p & 7) & 1 for p in self._probes(key))

    # -- analysis ----------------------------------------------------------
    def theoretical_fpr(self) -> float:
        return (1 - math.exp(-self.k * self.n / self.m)) ** self.k

    def fill_ratio(self) -> float:
        return sum(bin(b).count("1") for b in self.bits) / self.m

    @property
    def bytes_used(self) -> int:
        return len(self.bits)


def main() -> int:
    n = 100_000
    rng = random.Random(7)
    present = [f"key:{rng.getrandbits(64):016x}".encode() for _ in range(n)]
    absent = [f"nope:{rng.getrandbits(64):016x}".encode() for _ in range(200_000)]

    print(f"n = {n:,} inserted keys, probing {len(absent):,} absent keys\n")
    print(f"{'bits/key':>9} {'k':>3} {'fill':>7} {'theory':>10} {'measured':>10} "
          f"{'ratio':>7} {'RAM':>10}")
    print("-" * 64)
    for bpk in (4, 6, 8, 10, 12, 16, 20, 24):
        bf = BloomFilter(n, bpk)
        for key in present:
            bf.add(key)
        # Correctness invariant first: a Bloom filter must NEVER report a
        # false negative. If this fails, the filter is not a Bloom filter and
        # every performance number below is meaningless.
        assert all(key in bf for key in present), "FALSE NEGATIVE -- broken"
        fp = sum(1 for key in absent if key in bf)
        measured = fp / len(absent)
        theory = bf.theoretical_fpr()
        print(f"{bpk:>9} {bf.k:>3} {bf.fill_ratio():>7.4f} {theory:>10.5f} "
              f"{measured:>10.5f} {measured/theory if theory else 0:>6.2f}x "
              f"{bf.bytes_used/1e6:>8.2f} MB")

    print("\nfill ratio sits near 0.5 at every optimal k -- that is equation (2)")
    print("showing up in the data: at the optimum the filter is exactly half")
    print("full, carrying the maximum possible information per bit.")
    print()
    print("READ THE LAST TWO ROWS AS A MEASUREMENT FAILURE, NOT A RESULT.")
    print(f"At 24 bits/key the predicted fpr is ~1e-5, so {len(absent):,} absent")
    print("probes expect ~1.2 false positives. Observing 0 or 3 is pure Poisson")
    print("noise; the '0.00x ratio' means the experiment has no resolution here,")
    print("not that the filter beat theory. To measure a rate p you need on the")
    print("order of 100/p trials for a ~10% relative standard error -- so 1e-5")
    print("needs ~10 million probes. Always state the resolution of your")
    print("experiment before you report a ratio.\n")

    # The operational number: what a filter actually buys an LSM read path.
    print("Read amplification for an absent key, 40 SSTables on disk:")
    print(f"{'bits/key':>9} {'fpr':>9} {'disk reads':>12} {'vs no filter':>14}")
    print("-" * 48)
    for bpk, fpr in ((0, 1.0), (4, 0.147), (8, 0.0217), (10, 0.00819), (16, 0.000459)):
        reads = 40 * fpr
        label = "no filter" if bpk == 0 else f"{bpk}"
        print(f"{label:>9} {fpr:>9.5f} {reads:>12.3f} {40/max(reads,1e-9):>13.0f}x")
    print("\nAt 10 bits/key you pay 125 KB of RAM per 100k keys and turn 40 disk")
    print("reads into 0.33. That trade -- a little RAM for a lot of random I/O --")
    print("is the whole reason the LSM read path is viable.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())