# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy", "numba"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Draw a miss-ratio curve with SHARDS-style sampling.

One pass over a trace computes every request's reuse (stack) distance, which
gives the LRU miss ratio for every cache size at once. Checked against a real
LRU simulator, then repeated on a 1% sample of the keys.
"""
import argparse, json, time
from collections import OrderedDict
import numpy as np
from numba import njit

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--requests", type=int, default=2_000_000, help="Requests in the trace.")
p.add_argument("--keys", type=int, default=200_000, help="Distinct keys.")
p.add_argument("--alpha", type=float, default=0.8, help="Zipf exponent of key popularity.")
p.add_argument("--sizes", type=lambda s: [int(v) for v in s.split(",")], default=[100, 1_000, 10_000, 50_000],
               help="Cache sizes checked against the LRU simulator, comma-separated.")
p.add_argument("--rate", type=float, default=0.01, help="Sampling rate R for the SHARDS curve.")
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
trace = rng.choice(args.keys, size=args.requests, p=w / w.sum()).astype(np.int64)


@njit(cache=True)
def stack_distances(trace, n_keys):
    """Reuse distance of every request (distinct keys since its last use; -1 = first use).
    Fenwick tree over time: a 1 marks the latest access of each key."""
    n = len(trace)
    tree = np.zeros(n + 1, np.int64)
    last = np.full(n_keys, -1, np.int64)
    out = np.empty(n, np.int64)
    for t in range(n):
        k = trace[t]
        lt = last[k]
        if lt < 0:
            out[t] = -1
        else:
            s, i = 0, t                      # markers in (lt, t) = prefix(t) - prefix(lt + 1)
            while i > 0: s += tree[i]; i -= i & -i
            i = lt + 1
            while i > 0: s -= tree[i]; i -= i & -i
            out[t] = s
            i = lt + 1                       # remove the old marker
            while i <= n: tree[i] -= 1; i += i & -i
        i = t + 1                            # add the new marker
        while i <= n: tree[i] += 1; i += i & -i
        last[k] = t
    return out


def mrc(dist, sizes, scale=1.0):
    """Miss ratio at each size: P(distance >= C), first uses count as misses."""
    d = np.where(dist < 0, np.inf, dist * scale)
    return [float(np.mean(d >= c)) for c in sizes]


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


t = time.perf_counter(); exact = stack_distances(trace, args.keys); t_exact = time.perf_counter() - t
sim = [lru_miss(trace, c) for c in args.sizes]
exact_at = mrc(exact, args.sizes)

# SHARDS: keep keys whose (well-mixed) hash falls under R; scale their distances by 1/R
def mix(k):
    """splitmix64: a stable hash that spreads every key, including key 0, evenly."""
    z = (k.astype(np.uint64) + np.uint64(0x9E3779B97F4A7C15))
    z = (z ^ (z >> np.uint64(30))) * np.uint64(0xBF58476D1CE4E5B9)
    z = (z ^ (z >> np.uint64(27))) * np.uint64(0x94D049BB133111EB)
    return z ^ (z >> np.uint64(31))
keep = (mix(trace) % np.uint64(10_000)) < np.uint64(int(args.rate * 10_000))   # keep a key if its hash falls under R
t = time.perf_counter(); sampled = stack_distances(trace[keep], args.keys); t_sample = time.perf_counter() - t

print(f"trace: {args.requests:,} requests over {args.keys:,} keys, alpha {args.alpha}")
print(f"{'size':>8} {'exact MRC':>10} {'LRU sim':>9} {'sampled':>9}")
for c, e, s, m in zip(args.sizes, exact_at, sim, mrc(sampled, args.sizes, 1 / args.rate)):
    print(f"{c:>8,} {e:10.4f} {s:9.4f} {m:9.4f}")
print(f"sample: {keep.sum():,} of {args.requests:,} requests ({keep.mean():.2%}); "
      f"distance pass {t_sample:.2f} s against {t_exact:.2f} s exact (incl. compile on first call)")

# The whole curve on a log grid of sizes, exact and sampled
grid = np.unique(np.logspace(1, np.log10(args.keys), 40).astype(int)).tolist()
curve = [{"size": c, "curve": name, "miss": m}
         for name, dist, sc in (("exact", exact, 1.0), ("sampled", sampled, 1 / args.rate))
         for c, m in zip(grid, mrc(dist, grid, sc))]
checks = [{"size": c, "exact": e, "lru_sim": s} for c, e, s in zip(args.sizes, exact_at, sim)]

# C/N for h = 50% and 90%, from the exact curve, against h = (C/N)^(1 - alpha)
hit = dict(zip(grid, (1 - m for m in mrc(exact, grid))))
fit = []
for target in (0.5, 0.9):
    c = next((c for c in grid if hit[c] >= target), grid[-1])
    fit.append({"h": target, "measured_C_over_N": c / args.keys, "formula_C_over_N": target ** (1 / (1 - args.alpha))})
    print(f"h = {target:.0%}: measured C/N = {c / args.keys:.3f}   formula (C/N)^(1-alpha) gives {fit[-1]['formula_C_over_N']:.3f}")

err = max(abs(a - b) for a, b in zip(mrc(exact, grid), mrc(sampled, grid, 1 / args.rate)))
print(f"largest gap between the sampled and the exact curve: {err:.3f}")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"curve": curve, "checks": checks, "fit": fit, "summary": [{
            "max_sim_gap": max(abs(e - s) for e, s in zip(exact_at, sim)), "max_sample_gap": err,
            "sample_share": float(keep.mean())}]}}, f)