#!/usr/bin/env python3
"""Offline, trace-driven simulation of per-layer slot tables with eviction policies and prefetch predictors.

Exact for hit/miss accounting (the router sequence is the recorded one). Stall time is *modelled* from measured constants:
t_read (seconds to fetch one expert from the tier) and t_layer (seconds of compute per layer), with a loader pool of `--loaders`
parallel readers: demand misses at a layer stall for ceil(misses/loaders) * t_read; a prefetch issued h layers ahead hides
h * t_layer of its latency. Outputs JSON + CSV under --out.
"""
import argparse, glob, json, math, os, sys, time
import numpy as np
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from titan.predictors import make as make_predictor
from titan.slots import EvictionPolicy, LFUPolicy

POLICIES = {"lru": EvictionPolicy, "lfu": LFUPolicy}


def load_traces(d):
    idx = json.load(open(os.path.join(d, "index.json")))
    out = []
    for p in idx["prompts"]:
        z = np.load(os.path.join(d, p["name"] + ".npz"))
        out.append(dict(p, topk=z["topk"].astype(np.int64), probs=z["probs"], phase=z["phase"]))
    return idx, out


class LayerSim:
    def __init__(self, cap, policy, E):
        self.cap = cap; self.pol = POLICIES[policy](); self.res = set(); self.E = E
        self.hits = self.misses = self.pf_loads = self.pf_used = self.demand_loads = self.evict = 0
        self.pf_pending = {}  # expert -> True if prefetched and not yet used
        self.late_pending = 0

    def _insert(self, e, protected):
        if e in self.res: return
        if len(self.res) >= self.cap:
            try: v = self.pol.victim(protected)
            except RuntimeError: return
            self.res.discard(v); self.pol.remove(v); self.evict += 1
            if self.pf_pending.pop(v, False): pass
        self.res.add(e)

    def prefetch(self, experts, protected):
        for e in experts:
            if e in self.res: continue
            if len(self.res) >= self.cap and all(x in protected for x in self.res): break
            self._insert(e, protected); self.pol.touch(e); self.pf_loads += 1; self.pf_pending[e] = True

    def access(self, needed):
        miss = 0
        for e in needed:
            if e in self.res:
                self.hits += 1
                if self.pf_pending.pop(e, False): self.pf_used += 1
            else:
                self.misses += 1; miss += 1; self.demand_loads += 1
                self._insert(e, set(needed))
            self.pol.touch(e)
        return miss


def simulate(traces, L, E, K, cap, policy, pred_name, m, t_read, t_layer, loaders, fit_traces):
    pred = make_predictor(pred_name, L, E)
    depth = getattr(pred, "depth", 1)
    if fit_traces is not None: pred.fit([(t["topk"], t["probs"]) for t in fit_traces])
    tot = {"tokens": 0, "hits": 0, "misses": 0, "pf_loads": 0, "pf_used": 0, "demand_loads": 0, "stall_s": 0.0, "pred_hits": 0, "pred_total": 0, "evictions": 0}
    per_family = {}
    for tr in traces:
        layers = [LayerSim(cap, policy, E) for _ in range(L)]  # cold per prompt
        topk, probs = tr["topk"], tr["probs"]
        prev = None; prev_p = None
        fam = per_family.setdefault(tr["family"], {"tokens": 0, "hits": 0, "misses": 0, "stall_s": 0.0})
        for t in range(topk.shape[0]):
            cur = {("true", l): topk[t, l] for l in range(L)}  # read ONLY by the Oracle upper bound
            stall = 0.0
            if pred.horizon == "token":
                for l in range(L):
                    layers[l].prefetch(pred.predict(l, prev, prev_p, cur, m), set())
            for l in range(L):
                need = list(topk[t, l])
                if pred.horizon == "layer" and l > 0:
                    pass  # already issued at l-1
                miss = layers[l].access(need)
                if miss:
                    # demand misses: a layer-horizon prefetch for this layer was issued `depth` layers earlier and overlapped depth*t_layer;
                    # only misses that were *prefetched-but-late* benefit, so we model hidden = depth*t_layer for the late ones, 0 for unpredicted.
                    late = min(miss, layers[l].late_pending)
                    unpred = miss - late
                    stall += (math.ceil(late / loaders) * max(0.0, t_read - depth * t_layer) + math.ceil(unpred / loaders) * t_read) if late + unpred else 0.0
                    layers[l].late_pending = 0
                cur[l] = topk[t, l]; cur[("probs", l)] = probs[t, l]
                if pred.horizon == "layer" and l + depth < L:
                    p = pred.predict(l + depth, prev, prev_p, cur, m)
                    if pred_name != "none":
                        tot["pred_hits"] += len(set(p[:K]) & set(topk[t, l + depth])); tot["pred_total"] += K
                    layers[l + depth].prefetch(p, set())
                    layers[l + depth].late_pending = len(set(p) & set(topk[t, l + depth]) - layers[l + depth].res)
            prev, prev_p = topk[t], probs[t]
            tot["tokens"] += 1; fam["tokens"] += 1; tot["stall_s"] += stall; fam["stall_s"] += stall
            h = sum(x.hits for x in layers); mi = sum(x.misses for x in layers)
        tot["hits"] += sum(x.hits for x in layers); tot["misses"] += sum(x.misses for x in layers)
        tot["pf_loads"] += sum(x.pf_loads for x in layers); tot["pf_used"] += sum(x.pf_used for x in layers)
        tot["demand_loads"] += sum(x.demand_loads for x in layers); tot["evictions"] += sum(x.evict for x in layers)
        fam["hits"] += sum(x.hits for x in layers); fam["misses"] += sum(x.misses for x in layers)
    n = tot["hits"] + tot["misses"]
    res = {"capacity": cap, "policy": policy, "predictor": pred_name, "m": m, "hit_rate": tot["hits"] / n, "misses_per_token": tot["misses"] / tot["tokens"],
           "stall_ms_per_token": 1e3 * tot["stall_s"] / tot["tokens"], "prefetch_loads_per_token": tot["pf_loads"] / tot["tokens"],
           "prefetch_precision": (tot["pf_used"] / tot["pf_loads"]) if tot["pf_loads"] else None,
           "loads_per_token": (tot["pf_loads"] + tot["demand_loads"]) / tot["tokens"],
           "pred_recall_at_k": (tot["pred_hits"] / tot["pred_total"]) if tot["pred_total"] else None,
           "tokens": tot["tokens"], "families": {f: {"hit_rate": v["hits"] / (v["hits"] + v["misses"]), "stall_ms_per_token": 1e3 * v["stall_s"] / v["tokens"]} for f, v in per_family.items()}}
    return res


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--traces", required=True)
    ap.add_argument("--capacities", default="8,12,16,24,32,48")
    ap.add_argument("--policies", default="lru,lfu")
    ap.add_argument("--predictors", default="none,prev_rank,transition,transition2,transition3,probe,oracle")
    ap.add_argument("--m", type=int, default=8)
    ap.add_argument("--t-read-ms", type=float, default=0.43, help="measured per-expert fetch latency of the tier (laptop: 0.43 ms cold SSD, 0.27 ms page cache, 28 ms 1 GbE)")
    ap.add_argument("--t-layer-ms", type=float, default=0.128, help="per-layer decode compute (489 tok/s / 16 layers)")
    ap.add_argument("--loaders", type=int, default=1, help="parallel readers (current TITAN implementation: 1)")
    ap.add_argument("--out", required=True)
    a = ap.parse_args()
    idx, traces = load_traces(a.traces)
    L, E, K = idx["L"], idx["E"], idx["K"]
    fit = [t for t in traces if t["split"] == "fit"]; held = [t for t in traces if t["split"] == "heldout"]
    results = []
    for cap in [int(c) for c in a.capacities.split(",")]:
        for pol in a.policies.split(","):
            for pr in a.predictors.split(","):
                for split, trs in (("fit", fit), ("heldout", held)):
                    r = simulate(trs, L, E, K, cap, pol, pr, a.m, a.t_read_ms / 1e3, a.t_layer_ms / 1e3, a.loaders, fit)
                    r["split"] = split; results.append(r)
                    print(f"cap={cap:2d} {pol} {pr:10s} {split:7s} hit={r['hit_rate']:.3f} miss/tok={r['misses_per_token']:5.1f} stall={r['stall_ms_per_token']:6.2f}ms "
                          f"loads/tok={r['loads_per_token']:5.1f} pf_prec={r['prefetch_precision'] if r['prefetch_precision'] is None else round(r['prefetch_precision'],3)} recall@k={r['pred_recall_at_k'] if r['pred_recall_at_k'] is None else round(r['pred_recall_at_k'],3)}")
    os.makedirs(os.path.dirname(a.out), exist_ok=True)
    json.dump({"cmd": " ".join(sys.argv), "ts": time.strftime("%Y-%m-%dT%H:%M:%S"), "traces": a.traces, "L": L, "E": E, "K": K,
               "t_read_ms": a.t_read_ms, "t_layer_ms": a.t_layer_ms, "loaders": a.loaders, "results": results}, open(a.out, "w"), indent=1)


if __name__ == "__main__":
    main()