#!/usr/bin/env python3
"""
W2 — An LSM read path in 80 lines, and the Bloom filter that saves it.  (~60 min)

Miniature of P04. Memtable + immutable sorted runs + per-run Bloom filter. Measure
read amplification for present and absent keys, with and without the filter, and
watch the 122x collapse.
"""
import hashlib, math, random

class Bloom:
    def __init__(self, n, bits_per_key=10):
        self.m = max(8, int(n * bits_per_key))
        self.k = max(1, round(bits_per_key * math.log(2)))
        self.bits = bytearray((self.m + 7) // 8)
    def _probes(self, key):
        d = hashlib.blake2b(key, digest_size=16).digest()
        h1 = int.from_bytes(d[:8], "little"); h2 = int.from_bytes(d[8:], "little") | 1
        for i in range(self.k):
            yield (h1 + i * h2) % self.m
    def add(self, key):
        for p in self._probes(key): self.bits[p >> 3] |= 1 << (p & 7)
    def __contains__(self, key):
        return all(self.bits[p >> 3] >> (p & 7) & 1 for p in self._probes(key))

class Run:
    """An immutable sorted run. block_reads counts what a real engine would fetch."""
    def __init__(self, items, bits_per_key=10):
        self.keys = sorted(items)
        self.d = dict(items)
        self.bloom = Bloom(len(items), bits_per_key) if bits_per_key else None
        for k, _ in items: 
            if self.bloom: self.bloom.add(k)
    def get(self, key, stats, use_bloom=True):
        if use_bloom and self.bloom is not None and key not in self.bloom:
            stats["bloom_rejects"] += 1
            return None                       # no disk read at all
        stats["block_reads"] += 1             # the expensive part
        return self.d.get(key)

class LSM:
    def __init__(self, n_runs, keys_per_run, bits_per_key=10, seed=0):
        rng = random.Random(seed)
        self.runs = []
        self.all_keys = []
        for r in range(n_runs):
            items = []
            for i in range(keys_per_run):
                k = f"key:{r}:{i}:{rng.getrandbits(32):08x}".encode()
                items.append((k, r * 1000 + i)); self.all_keys.append(k)
            self.runs.append(Run(items, bits_per_key))
    def get(self, key, use_bloom=True):
        stats = {"block_reads": 0, "bloom_rejects": 0}
        for run in reversed(self.runs):        # newest first
            v = run.get(key, stats, use_bloom)
            if v is not None:
                return v, stats
        return None, stats

N_RUNS, PER_RUN, PROBES = 40, 2000, 4000
rng = random.Random(1)
print(f"{N_RUNS} sorted runs x {PER_RUN} keys = {N_RUNS*PER_RUN:,} keys\n")
print(f"{'bits/key':>9} {'absent: reads':>14} {'present: reads':>15} {'bloom rejects':>14}")
print("-" * 56)
for bpk in (0, 4, 8, 10, 16):
    db = LSM(N_RUNS, PER_RUN, bpk)
    absent = [f"nope:{rng.getrandbits(48):012x}".encode() for _ in range(PROBES)]
    present = [rng.choice(db.all_keys) for _ in range(PROBES)]
    use = bpk > 0
    ra = sum(db.get(k, use)[1]["block_reads"] for k in absent) / PROBES
    rj = sum(db.get(k, use)[1]["bloom_rejects"] for k in absent) / PROBES
    rp = sum(db.get(k, use)[1]["block_reads"] for k in present) / PROBES
    label = "none" if bpk == 0 else str(bpk)
    print(f"{label:>9} {ra:>14.3f} {rp:>15.2f} {rj:>14.1f}")

print("\nRead the ABSENT column. Without a filter every miss touches all 40 runs.")
print("At 10 bits/key it touches 0.3 -- a ~120x reduction for ~125 KB of RAM per")
print("100k keys. That trade is why the LSM read path is viable at all.")
print("\nNow the PRESENT column, which contradicts the usual summary. The folklore is")
print("'Bloom filters help misses, not hits'. Measured: present-key reads fall from")
print("20.50 to 1.17, a 17x improvement. The reason is that a hit must still SKIP")
print("every newer run that lacks the key, and the filter skips those without a read.")
print("Without it a hit scans ~half the runs (20.5 of 40) before finding one.")
print("\nThe folklore is describing the ASYMPTOTE, not the common case: as the filter")
print("gets perfect, absent-key cost goes to 0 while present-key cost floors at 1 --")
print("the one unavoidable read. Both improve; only one can reach zero.")