#!/usr/bin/env python3
"""Collect routing traces: for every prompt, record per token and per layer the router probabilities (float16, all experts)
and the top-k expert ids actually selected, during greedy generation with all experts resident (mlx-lm model).
Traces feed the offline cache/predictor simulator (benchmarks/scripts/simulate_cache.py) and predictor fitting.
Output: one .npz per prompt under --out, plus an index JSON. Prefill and decode tokens are both recorded (phase array).
"""
import argparse, json, os, sys, time
import numpy as np
import mlx.core as mx
import yaml
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from mlx_lm.utils import load_model, load_tokenizer
from mlx_lm.models.cache import make_prompt_cache


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True)
    ap.add_argument("--workloads", default=str(Path(__file__).resolve().parents[1] / "configs" / "workloads.yaml"))
    ap.add_argument("--max-tokens", type=int, default=256)
    ap.add_argument("--out", required=True)
    a = ap.parse_args()
    model, cfg = load_model(Path(a.model)); tok = load_tokenizer(Path(a.model))
    wl = yaml.safe_load(open(a.workloads))
    layers = model.model.layers
    L = len(layers); E = cfg["num_experts"]; K = cfg["num_experts_per_tok"]
    probs_buf, ids_buf = [], []
    # hook the router: wrap each layer's mlp.__call__ to capture softmax probs and chosen ids
    for l, layer in enumerate(layers):
        mlp = layer.mlp
        def hooked(x, _mlp=mlp, _l=l):
            B, Ln, D = x.shape
            x_flat = x.reshape(-1, D)
            logits = _mlp.gate(x_flat)
            rw = mx.softmax(logits, axis=1, precise=True)
            idx = mx.argpartition(-rw, kth=_mlp.top_k - 1, axis=-1)[..., :_mlp.top_k]
            probs_buf.append((_l, rw.astype(mx.float16)))
            ids_buf.append((_l, idx))
            scores = mx.take_along_axis(rw, idx, axis=-1)
            if _mlp.norm_topk_prob:
                scores = scores / scores.sum(axis=-1, keepdims=True)
            y = _mlp.switch_mlp(x_flat, idx)
            y = (y * scores[..., None]).sum(axis=-2)
            return y.reshape(B, Ln, D)
        layer.mlp.__call__ = hooked  # instance-level override is not seen by nn.Module.__call__ -> patch the class attribute per layer instead
        layer.mlp = _Wrapped(hooked)
    os.makedirs(a.out, exist_ok=True)
    index = []
    eos = tok.eos_token_id
    for fam, splits in wl["families"].items():
        for split, prompts in splits.items():
            for pi, p in enumerate(prompts):
                ids = mx.array(tok.encode(p))[None]
                cache = make_prompt_cache(model)
                probs_buf.clear(); ids_buf.clear()
                t0 = time.perf_counter()
                logits = model(ids, cache=cache); y = mx.argmax(logits[:, -1], -1); mx.eval(y)
                n_prompt = ids.shape[1]
                P = [np.array(pb) for _, pb in probs_buf]; I = [np.array(ib) for _, ib in ids_buf]
                probs = [np.stack([P[l] for l in range(L)], 1)]  # [T, L, E]
                topk = [np.stack([I[l] for l in range(L)], 1)]   # [T, L, K]
                gen = [int(y.item())]
                for step in range(a.max_tokens - 1):
                    if gen[-1] == eos: break
                    probs_buf.clear(); ids_buf.clear()
                    logits = model(y[None], cache=cache); y = mx.argmax(logits[:, -1], -1); mx.eval(y)
                    P = [np.array(pb) for _, pb in probs_buf]; I = [np.array(ib) for _, ib in ids_buf]
                    probs.append(np.stack([P[l] for l in range(L)], 1)); topk.append(np.stack([I[l] for l in range(L)], 1))
                    gen.append(int(y.item()))
                probs = np.concatenate(probs, 0).astype(np.float16); topk = np.concatenate(topk, 0).astype(np.int16)
                phase = np.array([0] * n_prompt + [1] * (probs.shape[0] - n_prompt), dtype=np.int8)
                name = f"{fam}_{split}_{pi}"
                np.savez_compressed(os.path.join(a.out, name + ".npz"), probs=probs, topk=topk, phase=phase, tokens=np.array(ids.reshape(-1).tolist() + gen))
                index.append({"name": name, "family": fam, "split": split, "prompt": p[:80], "n_prompt": n_prompt, "n_gen": len(gen),
                              "seconds": time.perf_counter() - t0, "text": tok.decode(gen)[:120]})
                print(f"{name:24s} prompt={n_prompt:4d} gen={len(gen):4d} {index[-1]['seconds']:.1f}s")
    json.dump({"model": os.path.realpath(a.model), "L": L, "E": E, "K": K, "max_tokens": a.max_tokens, "ts": time.strftime("%Y-%m-%dT%H:%M:%S"),
               "cmd": " ".join(sys.argv), "prompts": index}, open(os.path.join(a.out, "index.json"), "w"), indent=2)


class _Wrapped:
    def __init__(self, fn): self.fn = fn
    def __call__(self, x): return self.fn(x)


if __name__ == "__main__":
    main()