# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""LRU, random and Clock on a loop; then a scan.

Part 1 replays a loop over n pages against a pool of c frames under LRU,
random replacement and Clock, for c from 1 to n. Part 2 references a hot set
of 50 pages, interrupts it with a one-time scan of 10,000 other pages, and
measures the hot set's hit rate afterwards under LRU, 2Q (with a ghost queue
of identifiers) and a "ring buffer for the scan" variant.
"""
import argparse, json, random
from collections import OrderedDict, deque

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--n", type=int, default=100, help="Pages in the loop.")
p.add_argument("--steps", type=int, default=100_000, help="References per pool size.")
p.add_argument("--hot", type=int, default=50, help="Part 2: hot pages.")
p.add_argument("--scan", type=int, default=10_000, help="Part 2: pages in the one-time scan.")
p.add_argument("--pool", type=int, default=100, help="Part 2: frames in the pool.")
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()
random.seed(args.seed)

def lru(n, c, steps):
    pool, hits = OrderedDict(), 0
    for t in range(steps):
        p = t % n
        if p in pool: hits += 1; pool.move_to_end(p)
        else:
            if len(pool) >= c: pool.popitem(last=False)                # evict least recent
            pool[p] = True
    return hits / steps

def rand(n, c, steps):
    pool, where, hits = [], {}, 0
    for t in range(steps):
        p = t % n
        if p in where: hits += 1
        elif len(pool) < c: where[p] = len(pool); pool.append(p)
        else:
            i = random.randrange(c); del where[pool[i]]; pool[i] = p; where[p] = i
    return hits / steps

def clock(n, c, steps):
    frames, ref, where, hand, hits = [None] * c, [False] * c, {}, 0, 0
    for t in range(steps):
        p = t % n
        if p in where: hits += 1; ref[where[p]] = True; continue
        while ref[hand]: ref[hand] = False; hand = (hand + 1) % c   # second chance
        if frames[hand] is not None: del where[frames[hand]]
        frames[hand] = p; where[p] = hand; ref[hand] = True; hand = (hand + 1) % c
    return hits / steps

loop = []
for c in sorted({1, *range(10, args.n + 1, 10), 90, 95, 99, args.n}):
    loop.append({"c": c, "LRU": lru(args.n, c, args.steps) * 100, "random": rand(args.n, c, args.steps) * 100,
                 "Clock": clock(args.n, c, args.steps) * 100})

# Part 2 ─ a hot set, a one-time scan, the hot set again
def trace():
    hot = [("h", i) for i in range(args.hot)]
    warm = [random.choice(hot) for _ in range(20 * args.hot)]
    scan = [("s", i) for i in range(args.scan)]
    after = [random.choice(hot) for _ in range(20 * args.hot)]
    return warm, scan, after

class LRU:
    def __init__(s, c): s.c, s.d = c, OrderedDict()
    def get(s, p, scan=False):
        if p in s.d: s.d.move_to_end(p); return True
        if len(s.d) >= s.c: s.d.popitem(last=False)
        s.d[p] = True; return False

class TwoQ:
    d = {}
    """2Q: new pages enter a FIFO (A1in); a page referenced again after leaving it (found in the ghost queue A1out) goes to the main LRU (Am)."""
    def __init__(s, c): s.kin, s.c = max(1, c // 4), c; s.a1, s.am, s.ghost = OrderedDict(), OrderedDict(), OrderedDict()
    def get(s, p, scan=False):
        if p in s.am: s.am.move_to_end(p); return True
        if p in s.a1: return True
        if p in s.ghost:
            del s.ghost[p]
            if len(s.am) >= s.c - s.kin: s.am.popitem(last=False)
            s.am[p] = True; return False
        if len(s.a1) >= s.kin:
            old, _ = s.a1.popitem(last=False); s.ghost[old] = True
            if len(s.ghost) > s.c: s.ghost.popitem(last=False)
        s.a1[p] = True; return False

class Ring(LRU):
    """Scan pages go through a small private ring of frames and never enter the shared pool."""
    def __init__(s, c): super().__init__(c); s.ring = deque(maxlen=8)
    def get(s, p, scan=False):
        if scan:
            hit = p in s.ring
            if not hit: s.ring.append(p)
            return hit
        return super().get(p)

scan_rows = []
warm, scan, after = trace()
for name, cls in (("LRU", LRU), ("2Q", TwoQ), ("ring buffer for the scan", Ring)):
    pool = cls(args.pool)
    for x in warm: pool.get(x)
    before = sum(pool.get(x) for x in warm) / len(warm)
    for x in scan: pool.get(x, scan=True)
    resident = sum(1 for i in range(args.hot) if ("h", i) in pool.d or ("h", i) in getattr(pool, "am", {}) or ("h", i) in getattr(pool, "a1", {}))
    scan_rows.append({"policy": name, "hot_hits_before_pct": before * 100, "hot_resident_after_pct": resident / args.hot * 100,
                      "hot_hits_after_pct": sum(pool.get(x) for x in after) / len(after) * 100})

for r in loop:
    if r["c"] in (50, 80, 90, 99, 100): print(f"c={r['c']:>3}  LRU {r['LRU']:5.1f}%  random {r['random']:5.1f}%  Clock {r['Clock']:5.1f}%")
for r in scan_rows:
    print(f"{r['policy']:<26} hot set: hits before the scan {r['hot_hits_before_pct']:5.1f}%, still resident after it {r['hot_resident_after_pct']:5.1f}%, hits over the next {20 * args.hot:,} references {r['hot_hits_after_pct']:5.1f}%")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"loop": loop, "scan": scan_rows}}, f)