# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Scan resistance: LRU against an admission filter.

A Zipf request stream with a one-time scan in the middle, through plain LRU
and through LRU behind a count-min-sketch admission filter (admit a new key
only if it has been seen more often than the entry it would evict).
"""
import argparse, json
from collections import OrderedDict
import numpy as np

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--requests", type=int, default=2_000_000, help="Zipf requests in the stream.")
p.add_argument("--keys", type=int, default=100_000, help="Distinct keys in the Zipf stream.")
p.add_argument("--alpha", type=float, default=1.0, help="Zipf exponent.")
p.add_argument("--capacity", type=int, default=5_000, help="Cache size C, in entries.")
p.add_argument("--scan", type=int, default=200_000, help="Distinct one-time keys scanned in the middle of the stream.")
p.add_argument("--window", type=int, default=10_000, help="Requests per hit-rate window.")
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)

w = 1.0 / np.arange(1, args.keys + 1) ** args.alpha
zipf = rng.choice(args.keys, size=args.requests, p=w / w.sum()).tolist()
mid = args.requests // 2
trace = zipf[:mid] + list(range(args.keys, args.keys + args.scan)) + zipf[mid:]   # scan keys never repeat
scan_start, scan_end = mid, mid + args.scan


class CountMin:
    """4 rows of counters; every counter halves every `reset` increments (aging)."""
    def __init__(self, width, reset, rows=4):
        self.t = [[0] * width for _ in range(rows)]
        self.seeds = [0x9E3779B1, 0x85EBCA77, 0xC2B2AE3D, 0x27D4EB2F][:rows]
        self.w, self.reset, self.n = width, reset, 0
    def _cols(self, k):
        return [((k + 1) * s >> 7) % self.w for s in self.seeds]
    def add(self, k):
        for r, c in enumerate(self._cols(k)): self.t[r][c] += 1
        self.n += 1
        if self.n % self.reset == 0: self.t = [[x >> 1 for x in row] for row in self.t]
    def est(self, k):
        return min(self.t[r][c] for r, c in enumerate(self._cols(k)))


def run(admission):
    cache, hits, windows, sketch = OrderedDict(), [], [], CountMin(4 * args.capacity, 10 * args.capacity)
    for i, k in enumerate(trace):
        if admission: sketch.add(k)
        if k in cache:
            cache.move_to_end(k); hits.append(1)
        else:
            hits.append(0)
            if len(cache) < args.capacity:
                cache[k] = True
            else:
                victim = next(iter(cache))
                if not admission or sketch.est(k) > sketch.est(victim):
                    cache.popitem(last=False); cache[k] = True
        if (i + 1) % args.window == 0:
            windows.append(sum(hits[-args.window:]) / args.window)
    return windows


lru, filtered = run(False), run(True)
def zipf_hit(ws, a, b):                                  # mean hit rate over windows [a, b) of the Zipf part
    seg = ws[a // args.window: b // args.window]
    return sum(seg) / max(1, len(seg))
before = (zipf_hit(lru, mid - 200_000, mid), zipf_hit(filtered, mid - 200_000, mid))
after = (zipf_hit(lru, scan_end, scan_end + args.capacity * 4), zipf_hit(filtered, scan_end, scan_end + args.capacity * 4))
print(f"Zipf alpha {args.alpha}, {args.keys:,} keys, C = {args.capacity:,}, scan of {args.scan:,} keys at request {mid:,}")
print(f"hit rate before the scan:   LRU {before[0]:.1%}   with admission filter {before[1]:.1%}")
print(f"hit rate in the {args.capacity * 4:,} requests after the scan: LRU {after[0]:.1%}   with admission filter {after[1]:.1%}")

rows = [{"request": (i + 1) * args.window, "policy": name, "hit_pct": h * 100}
        for name, ws in (("LRU", lru), ("LRU + admission filter", filtered)) for i, h in enumerate(ws)]
rows = [r for r in rows if mid - 300_000 <= r["request"] <= scan_end + 300_000]          # around the scan
if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"windows": rows, "summary": [{
            "lru_before": before[0] * 100, "lru_after": after[0] * 100,
            "filter_before": before[1] * 100, "filter_after": after[1] * 100}]}}, f)