#!/usr/bin/env python3
"""
MoE expert cache — REAL miss-rate under real eviction policies
==============================================================
moe_cache_sim.py measured a next-token prefetch *hit-rate* (a predictor-quality
proxy). This goes one level more concrete: an actual capacity-C expert cache with
real eviction, replaying the real token trace, counting real demand misses (cold
loads) — and, for the prefetch policy, also counting prefetch loads so its
bandwidth cost is not hidden.

Per layer, test stream of tokens; each token fires top-k of N experts which must be
resident. Cache holds C experts. Policies:
  - LRU            : demand-load on miss, evict least-recently-used.
  - LFU-static     : cache = top-C experts by train frequency, pinned, never changes.
  - CO-ACT prefetch: LRU, but after each token prefetch the B experts most co-activated
                     (learned on train) with the set that just fired, for the next token.
                     Counts BOTH demand misses and prefetch loads as bandwidth.

Honest metric: demand miss-rate = cold loads / expert-activations. Lower = better.
total-load-rate (prefetch policy only) exposes the bandwidth it spends to get there.
If CO-ACT cuts demand misses below LFU/LRU *without* blowing up total loads, the
co-activation structure pays for itself. If it only wins by loading everything, it
doesn't. train/test split, same trace as moe_cache_sim.py.

Usage: python moe_cache_miss.py [--model allenai/OLMoE-1B-7B-0924] [--caps 16,24,32]
"""
import argparse, json
import numpy as np
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from moe_cache_sim import CORPUS  # reuse the same contiguous real-text corpus


def lru_sim(test, cap, n_exp):
    """Demand-paged LRU. Returns (demand_misses, activations)."""
    cache = {}  # expert -> recency stamp
    clock = 0; misses = 0; acts = 0
    for d in test:
        for s in d:
            for e in s:
                acts += 1; clock += 1
                if e in cache:
                    cache[e] = clock
                else:
                    misses += 1
                    if len(cache) >= cap:
                        ev = min(cache, key=cache.get); del cache[ev]
                    cache[e] = clock
    return misses, acts


def lfu_static_sim(test, pinned, n_exp):
    misses = 0; acts = 0
    for d in test:
        for s in d:
            for e in s:
                acts += 1
                if e not in pinned:
                    misses += 1
    return misses, acts


def coact_prefetch_sim(test, cap, co, freq, B, n_exp):
    """LRU + co-activation prefetch. Returns (demand_misses, prefetch_loads, acts)."""
    cache = {}; clock = 0; misses = 0; pf_loads = 0; acts = 0
    score_base = 1e-6 * freq
    for d in test:
        prev = None
        for s in d:
            # prefetch partners of the previous token's fired set
            if prev:
                score = co[list(prev)].sum(axis=0) + score_base
                for e in np.argsort(-score)[:B]:
                    e = int(e)
                    clock += 1
                    if e not in cache:
                        pf_loads += 1
                        if len(cache) >= cap:
                            ev = min(cache, key=cache.get); del cache[ev]
                        cache[e] = clock
                    else:
                        cache[e] = clock
            for e in s:
                acts += 1; clock += 1
                if e in cache:
                    cache[e] = clock
                else:
                    misses += 1
                    if len(cache) >= cap:
                        ev = min(cache, key=cache.get); del cache[ev]
                    cache[e] = clock
            prev = s
    return misses, pf_loads, acts


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default="allenai/OLMoE-1B-7B-0924")
    ap.add_argument("--caps", default="16,24,32")
    args = ap.parse_args()
    caps = [int(x) for x in args.caps.split(",")]

    dev = "mps" if torch.backends.mps.is_available() else "cpu"
    print(f"loading {args.model} on {dev} ...")
    tok = AutoTokenizer.from_pretrained(args.model)
    model = AutoModelForCausalLM.from_pretrained(
        args.model, dtype=torch.float16 if dev != "cpu" else torch.float32).to(dev).eval()
    cfg = model.config
    n_exp, topk, n_layers = cfg.num_experts, cfg.num_experts_per_tok, cfg.num_hidden_layers
    print(f"{n_exp} experts, top-{topk}, {n_layers} layers\n")

    seqs = [[] for _ in range(n_layers)]
    for text in CORPUS:
        ids = tok(text, return_tensors="pt").to(dev)
        with torch.no_grad():
            out = model(**ids, output_router_logits=True)
        for L in range(n_layers):
            lg = out.router_logits[L].float().cpu().numpy()
            if lg.ndim != 2 or lg.shape[1] != n_exp:
                seqs[L].append([]); continue
            top = np.argpartition(-lg, topk, axis=1)[:, :topk]
            seqs[L].append([set(r.tolist()) for r in top])

    layers_show = sorted(set([0, 1, n_layers//4, n_layers//2, 3*n_layers//4, n_layers-1]))
    print(f"{'='*86}\nREAL EXPERT CACHE MISS-RATE (demand misses / activation; lower=better)")
    print(f"top-{topk} of {n_exp} | train=1st half learns freq+coact, test=2nd half replayed")
    print(f"{'='*86}")
    hdr = f"{'layer':>5s} {'cap':>4s} {'LRU':>7s} {'LFU-st':>7s} {'COACT':>7s} {'COACT-tot':>9s}"
    print(hdr); print("-" * len(hdr))

    out_json = {}
    agg = {c: {'lru': [], 'lfu': [], 'coact': [], 'tot': []} for c in caps}
    for L in layers_show:
        docs = seqs[L]
        train = [d[:len(d)//2] for d in docs]
        test = [d[len(d)//2:] for d in docs]
        freq = np.zeros(n_exp); co = np.zeros((n_exp, n_exp))
        for d in train:
            for s in d:
                for e in s: freq[e] += 1
                sl = list(s)
                for i in range(len(sl)):
                    for j in range(len(sl)):
                        if i != j: co[sl[i], sl[j]] += 1
        rank = np.argsort(-freq)
        for c in caps:
            lm, la = lru_sim(test, c, n_exp)
            fm, fa = lfu_static_sim(test, set(rank[:c].tolist()), n_exp)
            cm, pf, ca = coact_prefetch_sim(test, c, co, freq, topk, n_exp)
            lru_r = lm/la; lfu_r = fm/fa; co_r = cm/ca; tot_r = (cm+pf)/ca
            print(f"{L:>5d} {c:>4d} {lru_r:>7.3f} {lfu_r:>7.3f} {co_r:>7.3f} {tot_r:>9.3f}")
            agg[c]['lru'].append(lru_r); agg[c]['lfu'].append(lfu_r)
            agg[c]['coact'].append(co_r); agg[c]['tot'].append(tot_r)
            out_json.setdefault(f"layer{L}", {})[f"cap{c}"] = dict(
                lru=lru_r, lfu_static=lfu_r, coact_demand=co_r, coact_total=tot_r)
    print("-" * len(hdr))
    for c in caps:
        m = {k: float(np.mean(v)) for k, v in agg[c].items()}
        print(f"{'mean':>5s} {c:>4d} {m['lru']:>7.3f} {m['lfu']:>7.3f} {m['coact']:>7.3f} {m['tot']:>9.3f}")
        out_json.setdefault("mean", {})[f"cap{c}"] = dict(
            lru=m['lru'], lfu_static=m['lfu'], coact_demand=m['coact'], coact_total=m['tot'])
    print("=" * len(hdr))
    print("LRU/LFU-st/COACT = demand miss-rate (cold loads / activation). COACT-tot adds")
    print("the prefetch loads (bandwidth COACT spends). COACT < LFU on demand AND COACT-tot")
    print("not far above it => co-activation prefetch pays for itself. Else it doesn't.")
    json.dump(out_json, open("/tmp/moe_cache_miss.json", "w"), indent=2)
    print("saved /tmp/moe_cache_miss.json")


if __name__ == "__main__":
    main()