#!/usr/bin/env python3
"""
Offline MoE expert cache & MoCE closed-loop replayer for Qwen3.8-Flash-Next (qwen4exp).
Implements:
  - Open-loop and MoCE cache-prior reranking:
      z'_e = z_e + lambda_moce * delta_avg * 1[e in cache or e in original top-J]
  - Optional --keep-map pruning filter (filters top-64 trace candidates to kept experts)
  - Eviction policies:
      * lru: Least Recently Used
      * lfu: Least Frequently Used
      * ema: Popularity EMA on selected top-K experts (all n_experts decay every token)
      * near_pick_ema: Popularity EMA crediting unbiased router top-N (ds4-recipe doc 06)
      * inverse_router_ema: EMA of normalized router rank/logit utility (dawn's inverse-router idea)
      * belady_min: Optimal offline oracle (Belady's MIN algorithm, theoretical ceiling)
  - M5 Air latency prediction using measured hardware constants from bench/README.md.
"""
import argparse
import json
import math
import struct
from pathlib import Path

TRACE_MAGIC = b"Q4RT"  # Qwen4exp Routing Trace v1


def parse_trace_file(path: str, keep_map_path: str | None = None) -> dict:
    raw = Path(path).read_bytes()
    magic, ver, n_layers, n_experts, top_m, n_tokens = struct.unpack_from("<4sHHHH I", raw, 0)
    if magic != TRACE_MAGIC:
        raise ValueError(f"Bad trace magic: {magic}")
    rec_bytes = top_m * 2 + top_m * 2
    actual_tokens = (len(raw) - 16) // (n_layers * rec_bytes)
    if n_tokens == 0 or actual_tokens < n_tokens:
        n_tokens = actual_tokens

    keep_sets = None
    remap = None
    if keep_map_path:
        km = json.loads(Path(keep_map_path).read_text())
        keep_sets = []
        remap = []
        for l in range(n_layers):
            kept = km[str(l)]
            keep_sets.append(set(kept))
            remap.append({old_id: new_id for new_id, old_id in enumerate(kept)})
        n_experts = len(km["0"])

    layers_data = [[] for _ in range(n_layers)]
    off = 16
    ids_fmt = f"<{top_m}H"
    logits_fmt = f"<{top_m}e"
    id_bytes = top_m * 2
    for _t in range(n_tokens):
        for l in range(n_layers):
            ids = struct.unpack_from(ids_fmt, raw, off)
            logits = struct.unpack_from(logits_fmt, raw, off + id_bytes)
            off += rec_bytes
            if keep_sets is not None:
                ks = keep_sets[l]
                rm = remap[l]
                f_ids = []
                f_logits = []
                for eid, z in zip(ids, logits):
                    if eid in ks:
                        f_ids.append(rm[eid])
                        f_logits.append(z)
                ids = tuple(f_ids)
                logits = tuple(f_logits)
            layers_data[l].append((ids, logits))

    return {
        "version": ver,
        "n_layers": n_layers,
        "n_experts": n_experts,
        "top_m": top_m,
        "n_tokens": n_tokens,
        "layers_data": layers_data,
    }


def select_moce_top_k(
    ids: tuple[int, ...],
    logits: tuple[float, ...],
    cache_set: set[int],
    k_used: int = 10,
    lambda_moce: float = 0.0,
    protect_top_j: int = 3,
) -> list[int]:
    if not ids:
        return []
    if lambda_moce <= 0.0 or len(ids) <= k_used:
        return list(ids[:k_used])

    delta_avg = max(1e-4, logits[0] - logits[min(k_used - 1, len(logits) - 1)])
    bonus = lambda_moce * delta_avg
    scored = []
    for rank, (eid, z) in enumerate(zip(ids, logits)):
        if rank < protect_top_j or eid in cache_set:
            scored.append((z + bonus, -rank, eid))
        else:
            scored.append((z, -rank, eid))
    scored.sort(reverse=True)
    return [eid for _, _, eid in scored[:k_used]]


def replay_layer(
    layer_trace: list[tuple[tuple[int, ...], tuple[float, ...]]],
    n_experts: int,
    cache_cap: int,
    policy: str = "lru",
    k_used: int = 10,
    lambda_moce: float = 0.0,
    protect_top_j: int = 3,
    near_pick_n: int = 48,
    ema_decay: float = 0.99,
    warmup_tokens: int = 64,
) -> dict:
    if cache_cap < k_used:
        raise ValueError(f"cache_cap ({cache_cap}) must be >= k_used ({k_used})")
    if len(layer_trace) <= warmup_tokens:
        raise ValueError(f"trace length ({len(layer_trace)}) must exceed warmup_tokens ({warmup_tokens})")
    if cache_cap >= n_experts:
        eval_tokens = max(1, len(layer_trace) - warmup_tokens)
        return {"misses": 0, "eval_tokens": eval_tokens, "misses_per_tok": 0.0, "hit_rate": 1.0}

    next_use_lists: list[list[int]] = [[] for _ in range(n_experts)]
    next_use_ptr = [0] * n_experts
    if policy == "belady_min":
        for t_idx, (ids, _logits) in enumerate(layer_trace):
            for eid in ids[:k_used]:
                next_use_lists[eid].append(t_idx)

    cache_set: set[int] = set()
    last_used = [0] * n_experts
    freq = [0] * n_experts
    ema_score = [0.0] * n_experts
    ema_mult = 1.0
    inv_decay = 1.0 / ema_decay

    misses = 0
    total_picks = 0

    for t_idx, (ids, logits) in enumerate(layer_trace):
        chosen = select_moce_top_k(
            ids, logits, cache_set, k_used=k_used, lambda_moce=lambda_moce, protect_top_j=protect_top_j
        )
        if policy == "belady_min":
            for eid in chosen:
                while next_use_ptr[eid] < len(next_use_lists[eid]) and next_use_lists[eid][next_use_ptr[eid]] <= t_idx:
                    next_use_ptr[eid] += 1

        chosen_set = set(chosen)
        if t_idx >= warmup_tokens:
            missing_slots = max(0, k_used - len(chosen))
            total_picks += len(chosen) + missing_slots
            misses += missing_slots
        for eid in chosen:
            if eid not in cache_set:
                if t_idx >= warmup_tokens:
                    misses += 1
                if len(cache_set) >= cache_cap:
                    candidates = [c for c in cache_set if c not in chosen_set]
                    if not candidates:
                        candidates = list(cache_set)
                    if policy == "lru":
                        victim = min(candidates, key=lambda c: (last_used[c], c))
                    elif policy == "lfu":
                        victim = min(candidates, key=lambda c: (freq[c], last_used[c], c))
                    elif policy in ("ema", "near_pick_ema", "inverse_router_ema"):
                        victim = min(candidates, key=lambda c: (ema_score[c], last_used[c], c))
                    elif policy == "belady_min":
                        def next_time(c: int) -> tuple[int, int]:
                            ptr = next_use_ptr[c]
                            lst = next_use_lists[c]
                            return (lst[ptr] if ptr < len(lst) else 10**9, -c)
                        victim = max(candidates, key=next_time)
                    else:
                        raise ValueError(f"Unknown policy: {policy}")
                    cache_set.remove(victim)
                cache_set.add(eid)

            last_used[eid] = t_idx + 1
            freq[eid] += 1

        if policy in ("ema", "near_pick_ema", "inverse_router_ema"):
            # Lazy global decay: multiplying added credits by inv_decay^t is identical to
            # decaying all n_experts every step, and avoids O(n_experts) inner loops.
            ema_mult *= inv_decay
            if ema_mult > 1e100:
                scale = 1.0 / ema_mult
                for e in range(n_experts):
                    ema_score[e] *= scale
                ema_mult = 1.0

            if policy == "ema":
                for eid in chosen:
                    ema_score[eid] += ema_mult
            elif policy == "near_pick_ema":
                top_n = min(near_pick_n, len(ids))
                denom = max(1, top_n - k_used)
                for rank in range(top_n):
                    eid = ids[rank]
                    credit = 1.0 if rank < k_used else 0.35 * (1.0 - (rank - k_used) / denom)
                    ema_score[eid] += credit * ema_mult
            elif policy == "inverse_router_ema" and logits:
                z_max = logits[0]
                z_min = logits[-1]
                span = max(1e-4, z_max - z_min)
                for rank, (eid, z) in enumerate(zip(ids, logits)):
                    util = (z - z_min) / span
                    if rank < k_used:
                        util += 1.0
                    ema_score[eid] += util * ema_mult

    eval_tokens = max(1, len(layer_trace) - warmup_tokens)
    misses_per_tok = misses / eval_tokens
    hit_rate = 1.0 - (misses / max(1, total_picks))
    return {
        "misses": misses,
        "eval_tokens": eval_tokens,
        "misses_per_tok": misses_per_tok,
        "hit_rate": hit_rate,
    }


def run_self_test() -> None:
    import random
    rng = random.Random(42)
    n_experts = 64
    top_m = 24
    n_tokens = 400
    weights = [1.0 / ((i + 1) ** 0.85) for i in range(n_experts)]
    trace = []
    hot_burst = [0, 1, 2, 3, 4]
    for t in range(n_tokens):
        if t % 20 == 0:
            hot_burst = rng.sample(range(n_experts), 6)
        scored = []
        for e in range(n_experts):
            base = math.log(weights[e]) + rng.gauss(0, 0.8)
            if e in hot_burst:
                base += 1.5
            scored.append((base, e))
        scored.sort(reverse=True)
        ids = tuple(e for _, e in scored[:top_m])
        logits = tuple(float(z) for z, _ in scored[:top_m])
        trace.append((ids, logits))

    res = {}
    for pol in ["lru", "lfu", "ema", "near_pick_ema", "inverse_router_ema", "belady_min"]:
        r = replay_layer(trace, n_experts=n_experts, cache_cap=24, policy=pol, k_used=6, warmup_tokens=40)
        res[pol] = r["misses"]

    assert res["belady_min"] <= min(res["lru"], res["lfu"], res["ema"], res["near_pick_ema"], res["inverse_router_ema"]), res
    r_moce = replay_layer(trace, n_experts=n_experts, cache_cap=24, policy="near_pick_ema", k_used=6, lambda_moce=0.15, warmup_tokens=40)
    assert r_moce["misses"] < res["near_pick_ema"], (r_moce["misses"], res["near_pick_ema"])
    print("replay_moe_cache.py self-test: PASS", res, "moce:", r_moce["misses"])


def main():
    p = argparse.ArgumentParser(description="Replay MoE routing traces across cache policies and MoCE.")
    p.add_argument("--self-test", action="store_true")
    p.add_argument("--trace", help="Path to Q4RT binary trace file.")
    p.add_argument("--keep-map", default=None, help="Optional keep_map.json for pruned expert replay.")
    p.add_argument("--cache-cap", type=int, default=113, help="Resident experts per layer (113 = M5 Air 7.5 GB Q2_0).")
    p.add_argument("--k-used", type=int, default=10)
    p.add_argument("--lambda-moce", type=float, default=0.0)
    p.add_argument("--protect-top-j", type=int, default=3)
    p.add_argument("--ema-decay", type=float, default=0.99)
    p.add_argument("--near-pick-n", type=int, default=48)
    p.add_argument("--policies", default="lru,lfu,ema,near_pick_ema,inverse_router_ema,belady_min")
    args = p.parse_args()

    if args.self_test:
        run_self_test()
        return

    if not args.trace:
        p.error("--trace is required unless --self-test is set")

    tr = parse_trace_file(args.trace, keep_map_path=args.keep_map)
    policies = [x.strip() for x in args.policies.split(",") if x.strip()]
    summary = {
        "n_tokens": tr["n_tokens"],
        "n_experts": tr["n_experts"],
        "cache_cap": args.cache_cap,
        "lambda_moce": args.lambda_moce,
        "protect_top_j": args.protect_top_j,
        "policies": {},
    }
    for pol in policies:
        if pol == "belady_min" and args.lambda_moce > 0.0:
            continue
        total_misses_per_tok = 0.0
        hit_rates = []
        for l in range(tr["n_layers"]):
            r = replay_layer(
                tr["layers_data"][l],
                n_experts=tr["n_experts"],
                cache_cap=args.cache_cap,
                policy=pol,
                k_used=args.k_used,
                lambda_moce=args.lambda_moce,
                protect_top_j=args.protect_top_j,
                near_pick_n=args.near_pick_n,
                ema_decay=args.ema_decay,
            )
            total_misses_per_tok += r["misses_per_tok"]
            hit_rates.append(r["hit_rate"])
        mean_hit = sum(hit_rates) / len(hit_rates)
        miss_mb_per_tok = total_misses_per_tok * 1.3824
        summary["policies"][pol] = {
            "total_misses_per_token_48L": round(total_misses_per_tok, 2),
            "mean_hit_rate": round(mean_hit, 4),
            "ssd_mb_per_token_q2_0": round(miss_mb_per_tok, 2),
            "est_m5_air_decode_tok_s_1th_3.9GBs": round(1.0 / (miss_mb_per_tok / 3900.0 + 0.012), 2),
            "est_m5_air_decode_tok_s_4th_6.8GBs": round(1.0 / (miss_mb_per_tok / 6800.0 + 0.012), 2),
        }
    print(json.dumps(summary, indent=2))


if __name__ == "__main__":
    main()