# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Find the plateau and the knee.

A trace mixing three streams (a hot set, a cyclic scan, one-time keys), run
through LRU, FIFO and random eviction at a range of cache sizes.
"""
import argparse, json, random
from collections import OrderedDict, deque
import numpy as np

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--requests", type=int, default=1_000_000, help="Requests in the trace.")
p.add_argument("--hot", type=int, default=1_000, help="Keys in the hot set.")
p.add_argument("--scan", type=int, default=50_000, help="Keys in the cyclic scan.")
p.add_argument("--mix", type=lambda s: [float(v) for v in s.split(",")], default=[0.7, 0.2, 0.1],
               help="Shares of hot, scan and one-time requests, comma-separated.")
p.add_argument("--sizes", type=lambda s: [int(v) for v in s.split(",")],
               default=[500, 1_000, 2_000, 5_000, 20_000, 50_000, 60_000, 80_000, 100_000], help="Cache sizes, comma-separated.")
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)
random.seed(args.seed)

# The trace: hot keys 0..hot-1, scan keys cycle in order, one-time keys never repeat
stream = rng.choice(3, size=args.requests, p=args.mix)
trace, scan_pos, fresh = [], 0, 0
hot_draws = rng.integers(0, args.hot, args.requests)
for i, s in enumerate(stream.tolist()):
    if s == 0: trace.append(("h", int(hot_draws[i])))
    elif s == 1: trace.append(("s", scan_pos)); scan_pos = (scan_pos + 1) % args.scan
    else: trace.append(("o", fresh)); fresh += 1


def lru(c):
    cache, hits = OrderedDict(), 0
    for k in trace:
        if k in cache: hits += 1; cache.move_to_end(k)
        else:
            cache[k] = True
            if len(cache) > c: cache.popitem(last=False)
    return hits / len(trace)

def fifo(c):
    cache, order, hits = set(), deque(), 0
    for k in trace:
        if k in cache: hits += 1
        else:
            cache.add(k); order.append(k)
            if len(cache) > c: cache.discard(order.popleft())
    return hits / len(trace)

def rand(c):
    cache, slots, hits = {}, [], 0                      # key -> slot index, so a random victim is O(1)
    for k in trace:
        if k in cache: hits += 1
        elif len(slots) < c:
            cache[k] = len(slots); slots.append(k)
        else:
            i = random.randrange(c); del cache[slots[i]]; slots[i] = k; cache[k] = i
    return hits / len(trace)


rows = []
print(f"{'size':>8} {'LRU':>7} {'FIFO':>7} {'random':>7}")
for c in args.sizes:
    r = {"size": c, "LRU": lru(c), "FIFO": fifo(c), "random": rand(c)}
    rows.append(r)
    print(f"{c:>8,} {r['LRU']:7.1%} {r['FIFO']:7.1%} {r['random']:7.1%}")
ceiling = 1 - args.mix[2]
print(f"ceiling from one-time keys: {ceiling:.0%}")

long_rows = [{"size": r["size"], "policy": pol, "hit_pct": r[pol] * 100} for r in rows for pol in ("LRU", "FIFO", "random")]
if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"hits": long_rows, "summary": [{"ceiling_pct": ceiling * 100}]}}, f)