#!/usr/bin/env python
"""Batch-1 (and batched) end-to-end decode latency per token for budget-0, forced budgets, dense writes and the
learned controller, with warm-up and device synchronisation. Reports median / p95 per token and peak memory.

Usage: python -m evaluation.latency --config configs/experiments/routing_p4.yaml --prompt-len 64 --new 32 [--device mps]
"""
from __future__ import annotations

import argparse, json, sys, time
from pathlib import Path

import numpy as np
import torch

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from prisme.cache import RequestBatchCache  # noqa: E402
from prisme.generation import step_logits  # noqa: E402
from prisme.flops import action_cost_proxy, ops_per_token  # noqa: E402
from training.common import load_checkpoint, load_yaml, pick_device, sync, peak_memory_bytes, git_sha  # noqa: E402
from training.train_editor import build_model  # noqa: E402


@torch.no_grad()
def time_decode(model, cfg, device, B, prompt_len, n_new, budget, forced_k=None, cost=None, reps=3):
    times = []
    cache_bytes = 0
    for rep in range(reps + 1):
        cache = RequestBatchCache(layers=model.backbone.new_cache())
        prompt = torch.randint(0, cfg.vocab_size, (B, prompt_len), device=device)
        fa = None
        if forced_k is not None:
            a = 0 if forced_k == 1 else cfg.n_slots
            fa = torch.full((B, prompt_len, budget), a, device=device, dtype=torch.long)
        step_logits(model, prompt, cache, budget, cost=cost, forced_actions=fa)
        sync(device)
        tok = prompt[:, -1:]
        per_tok = []
        for i in range(n_new):
            fa1 = None if forced_k is None else torch.full((B, 1, budget), 0 if forced_k == 1 else cfg.n_slots, device=device, dtype=torch.long)
            t0 = time.perf_counter()
            lg, _, _ = step_logits(model, tok, cache, budget, cost=cost, forced_actions=fa1)
            sync(device)
            per_tok.append(time.perf_counter() - t0)
            tok = lg[:, -1].argmax(-1, keepdim=True)
        if rep > 0:
            times.extend(per_tok)
        cache_bytes = cache.nbytes()
    t = np.array(times) * 1000
    return {"median_ms": float(np.median(t)), "p95_ms": float(np.percentile(t, 95)), "mean_ms": float(t.mean()), "n": len(t), "cache_bytes": cache_bytes}


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--config", required=True)
    ap.add_argument("--device", default="auto")
    ap.add_argument("--prompt-len", type=int, default=64)
    ap.add_argument("--new", type=int, default=32)
    ap.add_argument("--batch-sizes", default="1,8")
    ap.add_argument("--out", default=None)
    args = ap.parse_args(argv)
    exp = load_yaml(args.config); device = pick_device(args.device)
    torch.manual_seed(0)
    cfg, model = build_model(exp); model = model.to(device).eval()
    meta = load_checkpoint(exp["checkpoint"], model, strict=False)
    T = args.prompt_len + args.new
    cost = torch.tensor(action_cost_proxy(cfg, T), device=device)
    recs = []
    for B in [int(b) for b in args.batch_sizes.split(",")]:
        def add(name, **kw):
            r = time_decode(model, cfg, device, B, args.prompt_len, args.new, **kw)
            rec = {"policy": name, "batch": B, **{k: v for k, v in kw.items() if k != "cost"}, **r}
            recs.append(rec); print(json.dumps(rec), flush=True)
        add("budget0", budget=0)
        for b in [1, 2, 4, 8]:
            add("forced_single", budget=b, forced_k=1)
            add("forced_pair", budget=b, forced_k=2)
        add("controller", budget=8, cost=cost)
        if not cfg.dense_writes:
            cfg.dense_writes = True
            for b in [1, 4, 8]:
                add("dense_forced", budget=b, forced_k=1)
            cfg.dense_writes = False
    analytic = {name: ops_per_token(cfg, T, n, k).total for name, n, k in [("budget0", 0, 1), ("4_k1", 4, 1), ("4_k2", 4, 2), ("8_k2", 8, 2), ("8_dense", 8, cfg.n_slots)]}
    out = {"config": args.config, "checkpoint_sha256": meta.get("weights_sha256"), "device": str(device), "torch": torch.__version__,
           "precision": "fp32", "prompt_len": args.prompt_len, "n_new": args.new, "code_commit": git_sha(), "peak_memory_bytes": peak_memory_bytes(device),
           "analytic_ops_per_token": analytic, "records": recs}
    path = Path(args.out or f"results/raw/latency_{Path(exp['checkpoint']).parent.name}_{device.type}.json")
    path.parent.mkdir(parents=True, exist_ok=True); path.write_text(json.dumps(out, indent=1)); print("wrote", path)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())