"""Aha demo: LFU beats LRU on a stationary popularity distribution, exactly
as everyone expects — and then the popular set shifts once, and LFU collapses
while LRU does not notice.

Three traces:
  1. stationary Zipf (s=1.0)  -> LFU wins by ~16 points
  2. hot set A, then hot set B -> LFU's hit rate falls off a cliff
  3. a sequential scan         -> the boundary condition, where LRU is the loser

All randomness comes from an explicitly seeded random.Random, so every number
printed here is byte-identical on every run.
"""

import bisect
import itertools
import random

from lru_cache import LRUCache, LFUCache

KEY_SPACE = 1000
CAPACITY = 20
PHASE_REQUESTS = 20_000
WINDOW = 2_000
HOT_A = list(range(0, 20))
HOT_B = list(range(500, 520))
SEED = 20260726


# ----------------------------------------------------------------- traces --


def zipf_trace(rng, n_requests, n_keys=KEY_SPACE, s=1.0):
    """Keys drawn from a Zipf distribution: key i has weight 1/i**s."""
    weights = [1.0 / (i ** s) for i in range(1, n_keys + 1)]
    total = sum(weights)
    cutoffs = list(itertools.accumulate(w / total for w in weights))
    return [min(bisect.bisect(cutoffs, rng.random()), n_keys - 1)
            for _ in range(n_requests)]


def hot_cold_trace(rng, hot, n_requests, hot_share=0.9, n_keys=KEY_SPACE):
    """90% of requests hit the 20-key hot set; 10% land anywhere at all."""
    out = []
    for _ in range(n_requests):
        if rng.random() < hot_share:
            out.append(rng.choice(hot))
        else:
            out.append(rng.randrange(n_keys))
    return out


def scan_trace(n_keys, n_requests):
    """Round and round a working set of n_keys, in order, forever."""
    return [i % n_keys for i in range(n_requests)]


# ------------------------------------------------------------------ replay --


def replay(cache, trace):
    """Read-through: a miss loads the value (always 1 — this toy is about
    which key survives, not what is stored) and returns the hit count."""
    before = cache.hits
    for key in trace:
        if cache.get(key) is None:
            cache.put(key, 1)
    return cache.hits - before


def tally(cache, trace):
    """Replay, and also count requests and misses per key."""
    requests, misses = {}, {}
    for key in trace:
        requests[key] = requests.get(key, 0) + 1
        if cache.get(key) is None:
            misses[key] = misses.get(key, 0) + 1
            cache.put(key, 1)
    return requests, misses


def rate(hits, total):
    return 100.0 * hits / total


class RandomCache:
    """Random replacement, for the boundary condition in part 3 only. It is
    deliberately not in lru_cache.py: it is a control, not a third policy the
    toy is teaching."""

    def __init__(self, capacity, seed):
        self.capacity = capacity
        self.values = {}
        self.rng = random.Random(seed)
        self.hits = 0
        self.misses = 0

    def get(self, key):
        if key in self.values:
            self.hits += 1
            return self.values[key]
        self.misses += 1
        return None

    def put(self, key, value):
        if key not in self.values and len(self.values) >= self.capacity:
            victim = self.rng.choice(sorted(self.values))
            del self.values[victim]
        self.values[key] = value


# -------------------------------------------------------------------- parts --


def part1_stationary():
    print("=" * 66)
    print("1. STATIONARY ZIPF (s=1.0, %d keys, %d requests, capacity %d)"
          % (KEY_SPACE, PHASE_REQUESTS, CAPACITY))
    print("=" * 66)
    trace = zipf_trace(random.Random(SEED), PHASE_REQUESTS)
    lru, lfu = LRUCache(CAPACITY), LFUCache(CAPACITY)
    lru_hits = replay(lru, trace)
    lfu_hits = replay(lfu, trace)
    print("  LRU hit rate  %6.2f%%   (%d hits / %d)"
          % (rate(lru_hits, len(trace)), lru_hits, len(trace)))
    print("  LFU hit rate  %6.2f%%   (%d hits / %d)"
          % (rate(lfu_hits, len(trace)), lfu_hits, len(trace)))
    print("  delta         %+6.2fpp  -- LFU wins, as expected"
          % (rate(lfu_hits, len(trace)) - rate(lru_hits, len(trace))))
    print()


def part2_shift():
    print("=" * 66)
    print("2. THE POPULAR SET SHIFTS ONCE (hot A -> hot B, same 90/10 shape)")
    print("=" * 66)
    rng = random.Random(SEED)
    phase1 = hot_cold_trace(rng, HOT_A, PHASE_REQUESTS)
    phase2 = hot_cold_trace(rng, HOT_B, PHASE_REQUESTS)

    caches = {"LRU": LRUCache(CAPACITY), "LFU": LFUCache(CAPACITY)}
    windows = {}
    summary = {}
    for name, cache in caches.items():
        h1 = replay(cache, phase1)
        evicted_before = cache.evictions
        per_window = []
        h2 = 0
        for start in range(0, len(phase2), WINDOW):
            got = replay(cache, phase2[start:start + WINDOW])
            per_window.append(rate(got, WINDOW))
            h2 += got
        windows[name] = per_window
        resident = set(cache.keys())
        summary[name] = (rate(h1, PHASE_REQUESTS), rate(h2, PHASE_REQUESTS),
                         len(resident & set(HOT_A)), len(resident & set(HOT_B)),
                         cache.evictions - evicted_before)

    print("  %-5s %10s %10s %15s %13s %13s"
          % ("", "phase 1", "phase 2", "holds dead-hot", "holds new-hot",
             "ph.2 evicts"))
    for name in ("LRU", "LFU"):
        p1, p2, dead, new, evicts = summary[name]
        print("  %-5s %9.2f%% %9.2f%% %12d/20 %10d/20 %13d"
              % (name, p1, p2, dead, new, evicts))
    print()
    print("  hit rate per %d-request window of phase 2:" % WINDOW)
    for name in ("LRU", "LFU"):
        print("    %-4s %s" % (name, " ".join("%5.1f" % w for w in windows[name])))
    print()

    probe = LFUCache(CAPACITY)
    requests, lfu_misses = tally(probe, phase1)
    _, lru_misses = tally(LRUCache(CAPACITY), phase1)
    worst = max(HOT_A, key=lambda k: lfu_misses.get(k, 0))
    rest = sorted(lfu_misses.get(k, 0) for k in HOT_A if k != worst)
    print("  phase 1 already shows the mechanism, per hot key:")
    print("    LFU: 19 of the 20 hot keys missed %d-%d times all phase;"
          % (rest[0], rest[-1]))
    print("         k%d missed %d of its %d requests, alone."
          % (worst, lfu_misses[worst], requests[worst]))
    print("    LRU: every hot key missed %d-%d times -- the same tax, spread out."
          % (min(lru_misses.get(k, 0) for k in HOT_A),
             max(lru_misses.get(k, 0) for k in HOT_A)))
    print()

    print("  what LFU holds the instant phase 2 begins (all %d residents):"
          % CAPACITY)
    counts = sorted(probe.counts.items(), key=lambda kv: (kv[1], kv[0]))
    for start in range(0, len(counts), 5):
        print(("    " + "  ".join("k%-3d=%-4d" % kv
                                  for kv in counts[start:start + 5])).rstrip())
    frozen = [c for _, c in counts if c > 1]
    print("    -> %d keys at counts %d..%d, and %d slot at count 1"
          % (len(frozen), min(frozen), max(frozen), len(counts) - len(frozen)))

    end = caches["LFU"].counts
    survivors = {k for k, c in counts if c > 1}
    still = [end[k] for k in survivors if k in end and end[k] > 1]
    print("    ...and after all %d requests of phase 2 and %d evictions:"
          % (PHASE_REQUESTS, summary["LFU"][4]))
    print("    %d of those %d keys are still resident, now at counts %d..%d."
          % (len(still), len(survivors), min(still), max(still)))
    print("    Nothing evicted them. Every one of those evictions hit one slot.")
    print()

    print("  the revolving door, step by step (same cache, hand-fed):")
    new_x, new_y = HOT_B[0], HOT_B[1]
    probe.get(new_x); probe.put(new_x, 1)
    print("    request k%d (new hot key)  -> miss, admitted at count %d;"
          " evicted k%d" % (new_x, probe.counts[new_x],
                            [k for k, _ in counts if k not in probe.counts][0]))
    probe.get(new_y); probe.put(new_y, 1)
    print("    request k%d (new hot key)  -> miss, admitted at count %d;"
          " evicted k%d" % (new_y, probe.counts[new_y],
                            new_x if new_x not in probe.counts else -1))
    print("    k%d needed %d more hits to outrank the weakest resident."
          % (new_x, min(frozen) - 1))
    print()


def part3_boundary():
    print("=" * 66)
    print("3. BOUNDARY: A SEQUENTIAL SCAN, WHERE LRU IS THE LOSER")
    print("=" * 66)
    for n_keys in (100, 101):
        trace = scan_trace(n_keys, 2000)
        lru = LRUCache(100)
        lru_hits = replay(lru, trace)
        line = "  %3d keys, capacity 100:  LRU %6.2f%%" % (
            n_keys, rate(lru_hits, len(trace)))
        for seed in (1, 2):
            rc = RandomCache(100, seed)
            hits = replay(rc, trace)
            line += "   random(seed=%d) %6.2f%%" % (seed, rate(hits, len(trace)))
        print(line)
    print()


if __name__ == "__main__":
    part1_stationary()
    part2_shift()
    part3_boundary()