#!/usr/bin/env python3
"""Gate D profiling: where does a decode step spend its time with all experts resident? Times (synchronous, median of N) the
full step, one layer's attention, one layer's MoE block (3 gather_qmm + swiglu, top-8), and the three gather_qmm kernels alone,
and computes the effective bytes moved per second against the expert bytes touched (bandwidth-boundedness). Writes JSON.
"""
import argparse, json, os, platform, socket, statistics, subprocess, sys, time
import mlx.core as mx
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
from mlx_lm.models.activations import swiglu
from titan.registry import load_spec

def timeit(f, n=100, warm=10):
    for _ in range(warm): f()
    mx.synchronize(); ts = []
    for _ in range(n):
        t = time.perf_counter(); f(); mx.synchronize(); ts.append((time.perf_counter() - t) * 1e3)
    return {"median_ms": statistics.median(ts), "p10_ms": sorted(ts)[n // 10], "p90_ms": sorted(ts)[9 * n // 10]}

def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--model", required=True); ap.add_argument("--out", required=True); ap.add_argument("--n", type=int, default=100)
    a = ap.parse_args()
    spec = load_spec(a.model); model, cfg = load_model(Path(a.model)); tok = load_tokenizer(Path(a.model))
    L, E, K, D = spec.num_layers, spec.num_experts, spec.top_k, spec.hidden_size
    ids = mx.array(tok.encode("The capital of Canada is"))[None]
    cache = make_prompt_cache(model); y = mx.argmax(model(ids, cache=cache)[:, -1], -1); mx.eval(y)
    rec = {"host": socket.gethostname(), "chip": subprocess.run(["sysctl", "-n", "machdep.cpu.brand_string"], capture_output=True, text=True).stdout.strip(),
           "macos": platform.mac_ver()[0], "mlx": mx.__version__, "model": os.path.realpath(a.model), "L": L, "E": E, "K": K, "D": D,
           "bytes_per_expert": spec.bytes_per_expert, "expert_bytes_per_token": L * K * spec.bytes_per_expert, "n": a.n, "cmd": " ".join(sys.argv)}
    # 1. full decode step (sync)
    def step():
        out = model(y[None], cache=cache); mx.eval(out)
    rec["decode_step"] = timeit(step, a.n)
    # 2. one layer's blocks in isolation (layer 4), on a realistic hidden state
    layer = model.model.layers[4]
    h = mx.random.normal((1, 1, D)).astype(mx.float16); mx.eval(h)
    from mlx_lm.models.base import create_attention_mask
    c2 = make_prompt_cache(model); model(ids, cache=c2); mx.eval(c2[4].state)
    rec["attention_layer"] = timeit(lambda: mx.eval(layer.self_attn(layer.input_layernorm(h), None, c2[4])), a.n)
    rec["moe_block_layer"] = timeit(lambda: mx.eval(layer.mlp(layer.post_attention_layernorm(h))), a.n)
    sw = layer.mlp.switch_mlp
    xf = h.reshape(-1, D); idx = mx.array([[list(range(K))]], dtype=mx.uint32)
    xe = mx.expand_dims(xf, (-2, -3))
    def qmm3():
        up = sw.up_proj(xe, idx); g = sw.gate_proj(xe, idx); d = sw.down_proj(swiglu(g, up), idx); mx.eval(d)
    rec["gather_qmm_x3_layer"] = timeit(qmm3, a.n)
    rec["gather_qmm_up_only"] = timeit(lambda: mx.eval(sw.up_proj(xe, idx)), a.n)
    rec["router_layer"] = timeit(lambda: mx.eval(mx.argpartition(-mx.softmax(layer.mlp.gate(xf), axis=1, precise=True), kth=K - 1, axis=-1)[..., :K]), a.n)
    # 3. derived
    t_step = rec["decode_step"]["median_ms"]; t_moe = rec["moe_block_layer"]["median_ms"] * L; t_att = rec["attention_layer"]["median_ms"] * L
    rec["derived"] = {"moe_fraction_of_step_est": t_moe / t_step, "attention_fraction_est": t_att / t_step,
                      "expert_GBps_effective_in_moe_block": rec["expert_bytes_per_token"] / (t_moe / 1e3) / 1e9,
                      "expert_GBps_effective_full_step": rec["expert_bytes_per_token"] / (t_step / 1e3) / 1e9,
                      "tok_s_sync": 1e3 / t_step}
    os.makedirs(os.path.dirname(a.out), exist_ok=True); json.dump(rec, open(a.out, "w"), indent=2); print(json.dumps(rec, indent=1))

if __name__ == "__main__":
    main()