"""Actual Bloom-filter probes on synthetic keys. Python 3.10+, stdlib only.

Run: python experiment.py
No database, network, storage-latency or production benchmark is involved.
"""
from pathlib import Path
import csv
import hashlib
import json
import math
import platform
import statistics
import struct

COUNTS = [2000, 8000, 32000, 128000]
SEEDS = [11, 29, 47, 71, 101]
QUERIES = 50000
A = math.log(2) ** 2


def allocation(counts, budget_bits, weights=None):
    """Continuous optimum for sum(w_i exp(-a*b_i)); clamp b_i at zero."""
    weights = weights or [1.0] * len(counts)
    low, high = -100.0, 100.0  # logarithm of the Lagrange multiplier
    for _ in range(200):
        mid = (low + high) / 2
        bits = [max(0.0, (math.log(A * w / n) - mid) / A)
                for n, w in zip(counts, weights)]
        if sum(n * b for n, b in zip(counts, bits)) > budget_bits:
            low = mid
        else:
            high = mid
    return [max(0.0, (math.log(A * w / n) - high) / A)
            for n, w in zip(counts, weights)]


def byte_sizes(counts, bits, budget_bytes):
    """Largest-remainder rounding preserves the exact total byte budget."""
    raw = [n * b / 8 for n, b in zip(counts, bits)]
    result = [math.floor(x) for x in raw]
    remainder = budget_bytes - sum(result)
    assert 0 <= remainder <= len(counts)
    order = sorted(range(len(counts)), key=lambda i: raw[i] - result[i], reverse=True)
    for i in order[:remainder]:
        result[i] += 1
    assert sum(result) == budget_bytes
    return result


class Bloom:
    def __init__(self, n, size_bytes, seed, level):
        self.bits = bytearray(size_bytes)
        self.m = size_bytes * 8
        self.k = max(1, round(self.m / n * math.log(2))) if self.m else 0
        self.salt = f'bloom-v1:{seed}:{level}:'.encode()

    def positions(self, key):
        # SHAKE output supplies k pseudorandom 64-bit hash values. A fast DB
        # filter would use a different implementation and CPU/cache layout.
        raw = hashlib.shake_256(self.salt + key).digest(8 * self.k)
        return (h[0] % self.m for h in struct.iter_unpack('<Q', raw))

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

    def contains(self, key):
        if not self.m:
            return True  # no filter: every query must inspect the table
        return all(self.bits[bit >> 3] & (1 << (bit & 7)) for bit in self.positions(key))


def model_fpr(n, size_bytes, k):
    if size_bytes == 0:
        return 1.0
    m = size_bytes * 8
    return (1 - (1 - 1 / m) ** (k * n)) ** k


def write_csv(path, records):
    with path.open('w', newline='') as f:
        writer = csv.DictWriter(f, fieldnames=list(records[0]))
        writer.writeheader()
        writer.writerows(records)


def svg_bars(path, title, subtitle, rows, xmax, xtitle):
    # rows = (label, value, color). Numeric text makes the figure accessible.
    width, height = 900, 170 + len(rows) * 53
    s = [f'<svg xmlns="http://www.w3.org/2000/svg" width="{width}" height="{height}" viewBox="0 0 {width} {height}" role="img"><title>{title}</title>',
         '<rect width="100%" height="100%" fill="#faf9f6"/>',
         '<g font-family="Arial,sans-serif" fill="#222">',
         f'<text x="28" y="38" font-size="24" font-weight="600">{title}</text>',
         f'<text x="28" y="65" font-size="14" fill="#666">{subtitle}</text>']
    for i, (label, value, color) in enumerate(rows):
        y = 96 + i * 53
        s += [f'<text x="28" y="{y+21}" font-size="15">{label}</text>',
              f'<rect x="220" y="{y}" width="{value/xmax*560:.3f}" height="30" rx="2" fill="{color}"/>',
              f'<text x="{230+value/xmax*560:.3f}" y="{y+21}" font-size="15">{value:.2f}</text>']
    y = 106 + len(rows)*53
    s += [f'<text x="220" y="{y}" font-size="14" fill="#666">{xtitle}</text>', '</g></svg>']
    path.write_text('\n'.join(s))


def run(out=None):
    out = Path(out or Path(__file__).parent)
    (out / 'data').mkdir(parents=True, exist_ok=True)
    (out / 'figures').mkdir(exist_ok=True)
    budget = sum(COUNTS) * 10 // 8
    policies = {'uniform': [10.0] * 4, 'optimized': allocation(COUNTS, budget*8)}
    rows, seeds, hist = [], [], []
    tests = []
    for policy, bits in policies.items():
        sizes = byte_sizes(COUNTS, bits, budget)
        for seed in SEEDS:
            filters = []
            for level, (n, size) in enumerate(zip(COUNTS, sizes)):
                bf = Bloom(n, size, seed, level)
                for j in range(n):
                    bf.add(f'present:{level}:{j}'.encode())
                missing = sum(not bf.contains(f'present:{level}:{j}'.encode()) for j in range(n))
                assert missing == 0
                tests.append({'test':'no_false_negatives','policy':policy,'seed':seed,'level':level,'keys':n,'passed':True})
                filters.append(bf)
            fp = [0] * 4
            histogram = [0] * 5
            for j in range(QUERIES):
                key = f'absent:{seed}:{j}'.encode()
                values = [int(bf.contains(key)) for bf in filters]
                for i, value in enumerate(values):
                    fp[i] += value
                histogram[sum(values)] += 1
            for level, (n, bf, false_positives) in enumerate(zip(COUNTS, filters, fp)):
                rows.append({'policy':policy, 'seed':seed, 'level':level+1, 'keys':n,
                             'filter_bytes':len(bf.bits), 'bits_per_key':bf.m/n, 'hashes':bf.k,
                             'queries':QUERIES, 'false_positives':false_positives,
                             'observed_fpr':false_positives/QUERIES,
                             'finite_occupancy_approx_fpr':model_fpr(n, len(bf.bits), bf.k)})
            seeds.append({'policy':policy,'seed':seed,'false_positive_probes':sum(fp),
                          'queries':QUERIES,'probes_per_query':sum(fp)/QUERIES,
                          'filter_bytes':sum(sizes)})
            for probes, count in enumerate(histogram):
                hist.append({'policy':policy,'seed':seed,'probes':probes,'queries':count})
    # Analytical checks include low budgets and equal-sized runs, not only the
    # particular geometry chosen for the main measurement.
    checks = 0
    for counts in [COUNTS, [1000]*4, [100, 10000, 100000]]:
        for b in [0, 0.1, 1, 2, 5, 10, 20]:
            optimum = allocation(counts, sum(counts)*b)
            assert abs(sum(n*x for n,x in zip(counts,optimum)) - sum(counts)*b) < 1e-6
            assert sum(math.exp(-A*x) for x in optimum) <= len(counts)*math.exp(-A*b)+1e-10
            checks += 1
    assert all(abs(x-10)<1e-10 for x in allocation([1000]*4,40000))
    assert Bloom(1,0,1,0).contains(b'anything')
    uniform = [x['probes_per_query'] for x in seeds if x['policy']=='uniform']
    optimized = [x['probes_per_query'] for x in seeds if x['policy']=='optimized']
    summary = {'counts':COUNTS,'seeds':SEEDS,'queries_per_seed':QUERIES,
               'total_absent_queries_per_policy':len(SEEDS)*QUERIES,'filter_bytes_per_policy':budget,
               'uniform_mean_probes_per_query':statistics.mean(uniform),
               'optimized_mean_probes_per_query':statistics.mean(optimized),
               'reduction_percent':100*(1-statistics.mean(optimized)/statistics.mean(uniform)),
               'uniform_seed_range':[min(uniform),max(uniform)],
               'optimized_seed_range':[min(optimized),max(optimized)],
               'no_false_negative_keys_checked':2*len(SEEDS)*sum(COUNTS),
               'allocation_cases_passed':checks,
               'continuous_uniform_expected_probes':4*math.exp(-A*10),
               'continuous_optimized_expected_probes':sum(math.exp(-A*x) for x in policies['optimized']),
               'measurement':'Actual Python Bloom-filter bit probes on synthetic keys; not database I/O or latency',
               'python':platform.python_version(),'hash':'SHAKE-256 pseudorandom 64-bit chunks per filter/key',
               'run_date':'2026-11-06'}
    write_csv(out/'data/filter-results.csv',rows)
    write_csv(out/'data/seed-results.csv',seeds)
    write_csv(out/'data/probe-histograms.csv',hist)
    (out/'data/summary.json').write_text(json.dumps(summary,indent=2)+'\n')
    (out/'data/test-results.json').write_text(json.dumps({'inserted_key_checks':tests,'allocation_cases_passed':checks},indent=2)+'\n')
    svg_bars(out/'figures/probes.svg','The same memory buys fewer false-positive probes',
             'Actual filter queries; 250,000 absent keys per policy, averaged across five seeds',
             [('Uniform',summary['uniform_mean_probes_per_query']*1000,'#999'),
              ('Optimized',summary['optimized_mean_probes_per_query']*1000,'#8b4345')],40,
             'False-positive table probes per 1,000 absent lookups (lower is better)')
    allocation_rows = []
    for i, n in enumerate(COUNTS):
        allocation_rows.extend([(f'{n:,} keys · uniform',10,'#aaa'),
                                (f'{n:,} keys · optimized',sizes[i]*8/n,'#8b4345')])
    svg_bars(out/'figures/allocation.svg','Small runs receive more bits per key',
             'Filter memory is fixed at 212,500 bytes in both policies', allocation_rows,20,'Bits per key')
    sweep=[]
    for b in range(0,21):
        xs = allocation(COUNTS,sum(COUNTS)*b)
        sweep.append({'bits_per_key_budget':b,'uniform_expected_probes':4*math.exp(-A*b),
                      'optimized_expected_probes':sum(math.exp(-A*x) for x in xs),
                      'optimized_bits_per_key':json.dumps(xs)})
    write_csv(out/'data/analytic-budget-sweep.csv',sweep)
    return summary


if __name__ == '__main__':
    print(json.dumps(run(),indent=2))