"""Run the tiny V4.1 checkpoint through the real DeepInfer engine, teacher-forced on the reference golden.

Engine path: deepinfer.entrypoints.llm.LLM (scheduler, KV cache, model runner, sampler), EP=DP=CP=1, Neuron basic
mode defaults (eager, DEEPINFER_USE_NEXTN_DECODE=0). Teacher forcing uses the engine's own mechanism,
SamplingParams(golden_token_ids=...), so decode step i is fed the golden token i-1 exactly like the reference
(reference_golden.py --force-tokens-dir). Logits are captured by wrapping the sampler's `get_logits` (the fp32
[num_reqs, vocab] lm-head output before any sampling transform); per-layer block outputs by forward hooks on the
decoder layers (best effort, see --hidden).

Output: <out>/<prompt>.pt in the reference_golden.py format (prompt_ids, tokens = argmax per step, fed_tokens,
logits [steps+1, vocab] fp32, optional hidden), ready for compare_golden.py.

    cd ~/di/DeepInfer && ~/di/venv/bin/python examples/neuron/run_deepinfer_tiny.py \
        --hf-ckpt ~/di/tiny_ckpt --golden-dir ~/di/golden --out ~/di/deepinfer_out [--device cpu|neuron] \
        [--hidden] [--compare] [--threads 32]

Placeholder engine steps (is_placeholder=True, no real sequence) also run the model and the sampler; they are
filtered out of the recorded logits/hidden. `install_workarounds` patches, at runtime and only in this process,
NOTE-3 (fixed KV pool instead of the host-RAM-sized profile) -- the other bring-up bugs are fixed in-repo; the list
of workarounds that actually triggered is printed at the end. Set RUN_TINY_DEBUG=1 to print every sampler call.
"""

from __future__ import annotations

import argparse
import os
import sys
import time
import traceback
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parents[2]


def parse_args():
    ap = argparse.ArgumentParser()
    ap.add_argument("--hf-ckpt", default="~/di/tiny_ckpt")
    ap.add_argument("--golden-dir", default="~/di/golden")
    ap.add_argument("--out", default="~/di/deepinfer_out")
    ap.add_argument("--prompts", default="short,medium,long")
    ap.add_argument("--device", choices=["cpu", "neuron"], default="cpu", help="sets DEEPINFER_NEURON_DEVICE")
    ap.add_argument("--max-model-len", type=int, default=2048)
    ap.add_argument("--max-num-batched-tokens", type=int, default=1024)
    ap.add_argument("--threads", type=int, default=16)
    ap.add_argument("--num-kv-blocks", type=int, default=32, help="fixed KV block pool (128 tokens each); 0 = engine profiling")
    ap.add_argument("--hidden", action="store_true", help="also record per-layer block outputs (hc residual stream)")
    ap.add_argument("--compare", action="store_true", help="run compare_golden on every prompt afterwards")
    ap.add_argument("--attn-p-bf16", action="store_true", help="experiment: patch the flash_mla shim with the proposed bf16-P fix")
    ap.add_argument("--cp", type=int, default=1,
                    help="context-parallel group of one engine (decode replicated on every member, EP over the group): "
                         "CUDA_VISIBLE_DEVICES lists its cores (default 0..cp-1)")
    ap.add_argument("--ep", type=int, default=1,
                    help="expert parallel size (= data parallel size): one engine process per rank, every rank runs the "
                         "same teacher-forced prompts. Rank r uses entry r of CUDA_VISIBLE_DEVICES (default 0..ep-1), "
                         "i.e. one NeuronCore per process on --device neuron. Rank 0 writes <out>/<prompt>.pt, rank r>0 "
                         "<out>/rank<r>/<prompt>.pt")
    return ap.parse_args()


def setup_env(device: str, cp: int = 1) -> None:
    # Must happen before deepinfer is imported (see examples/_config.py).
    os.environ["DEEPINFER_ACCELERATOR"] = "neuron"
    os.environ["DEEPINFER_NEURON_DEVICE"] = device
    os.environ["DEEPINFER_USE_NEXTN_DECODE"] = "0"
    os.environ.setdefault("DEEPINFER_ENFORCE_EAGER", "1")
    os.environ.setdefault("DEEPINFER_ZERO_OVERHEAD_SCHEDULING", "0")
    os.environ.setdefault("CUDA_VISIBLE_DEVICES", ",".join(str(i) for i in range(cp)))  # mapped to the target
    os.environ.setdefault("RANK", "0")
    os.environ.setdefault("WORLD_SIZE", "1")
    os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
    os.environ.setdefault("DEEPINFER_USE_MULTIMODAL_POOL", "0")  # text only; avoids an 8 GB pinned host pool
    if str(REPO_ROOT) not in sys.path:
        sys.path.insert(0, str(REPO_ROOT))
    here = str(Path(__file__).resolve().parent)
    if here not in sys.path:
        sys.path.insert(0, here)


WORKAROUNDS: list[str] = []


def install_workarounds(num_kv_blocks: int, attn_p_bf16: bool = False) -> None:
    """Work around bugs in other modules without editing them; each one is reported (see the module docstring)."""
    # BUG-1 (engram gather_sf on gloo), BUG-2 (get_model is_cuda), BUG-4 (paged MQA logits page ids) and NOTE-5
    # (bind_cpu) were fixed in-repo on 2026-10-20; their workarounds are gone. Only NOTE-3 remains.

    # NOTE-3 (not a bug, a CPU-mode hazard): WorkerBase.profile_num_available_blocks runs a dummy prefill of
    # max_num_batched_tokens tokens (903 s on the host for 2048 tokens) and then sizes the KV cache from
    # mem_get_info(), which on cpu is host RAM -> 0.5 * 2 TB = 559 GiB of KV blocks. Use a fixed small pool instead.
    if num_kv_blocks > 0:
        os.environ["DEEPINFER_NEURON_NUM_KV_BLOCKS"] = str(num_kv_blocks)  # CP followers read it too
        import torch.distributed as _dist  # noqa: F401
        from deepinfer.worker.base import WorkerBase

        def profile_num_available_blocks(self):
            self.model_runner.set_block_size(self.cache_config.block_size)
            self.cache_config.num_gpu_blocks = num_kv_blocks
            WORKAROUNDS.append(f"NOTE-3: fixed {num_kv_blocks} KV blocks, no profile run")
            return num_kv_blocks

        WorkerBase.profile_num_available_blocks = profile_num_available_blocks

    # EXPERIMENT (proposed fix, see probe_layer0.py): FlashMLA-style P rounding to bf16 in the Neuron flash_mla shim
    if attn_p_bf16:
        import flash_mla.flash_mla_interface as fmi
        from probe_layer0 import sparse_attn_rows_p_bf16

        fmi._sparse_attn_rows = sparse_attn_rows_p_bf16
        WORKAROUNDS.append("EXPERIMENT: flash_mla _sparse_attn_rows with bf16 P (proposed fix)")


class Recorder:
    """Collects the lm-head logits of every sampler call (one request in flight -> one row per call)."""

    def __init__(self):
        self.step_kind = None  # (is_prompt, is_placeholder) of the engine step being executed
        self.n_placeholder = 0
        self.logits = []
        self.hidden_calls = []
        self._cur_hidden = {}

    def reset(self):
        self.n_placeholder = 0
        self.logits.clear()
        self.hidden_calls.clear()
        self._cur_hidden = {}


def install_logits_hook(rec: Recorder) -> None:
    import deepinfer.model_executor.layers.sampler as sampler_mod

    orig = sampler_mod.get_logits

    def get_logits(hidden_states, embedding, vocab_size):
        logits = orig(hidden_states, embedding, vocab_size)
        if rec.step_kind is not None and rec.step_kind[1]:
            rec.n_placeholder += 1  # placeholder step (no real sequence): not part of the request
            return logits
        rec.logits.append(logits.detach().float().cpu().clone())
        if os.environ.get("RUN_TINY_DEBUG"):
            import traceback as _tb

            st = " <- ".join(f"{f.name}:{f.lineno}" for f in _tb.extract_stack()[-6:-1][::-1])
            print(f"[run]   sampler call {len(rec.logits)}: logits {tuple(logits.shape)} argmax {int(logits[0].argmax())} "
                  f"hs {tuple(hidden_states.shape)} at {time.strftime('%H:%M:%S')} via {st}", flush=True)
        return logits

    sampler_mod.get_logits = get_logits

    from deepinfer.worker.worker import Worker

    orig_step = Worker.execute_model_step

    def execute_model_step(self, is_prompt, is_placeholder, *a, **k):
        rec.step_kind = (bool(is_prompt), bool(is_placeholder))
        try:
            return orig_step(self, is_prompt, is_placeholder, *a, **k)
        finally:
            rec.step_kind = None

    Worker.execute_model_step = execute_model_step


def install_hidden_hooks(rec: Recorder, llm) -> str:
    """Hook the decoder layers of the driver worker's model. Returns a status string."""
    try:
        worker = llm.llm_engine.driver_worker
        runner = getattr(worker, "model_runner", None)
        model = getattr(runner, "model", None) if runner is not None else getattr(worker, "model", None)
        layers = model.model.layers
    except Exception as exc:  # noqa: BLE001
        return f"hidden hooks not installed: {type(exc).__name__}: {exc}"

    from deepinfer.model_executor.models.v4_modules.hc import HCLayerMixin

    n = len(layers)

    def hook(i):
        def fn(mod, inp, out):
            if rec.step_kind is None or rec.step_kind[1]:
                return None  # placeholder / non-request forward
            h, residual = out
            h = h.move() if hasattr(h, "move") else h
            try:
                # 2D output = before hc_post; the reference block output is the expanded hc stream
                x = HCLayerMixin.hc_post(h, residual) if h.dim() == 2 else h
            except Exception:  # noqa: BLE001
                x = h
            rec._cur_hidden[i] = x.detach().float().cpu().clone()
            if i == n - 1:
                rec.hidden_calls.append(dict(rec._cur_hidden))
                rec._cur_hidden = {}
            return None

        return fn

    for i, layer in enumerate(layers):
        layer.register_forward_hook(hook(i))
    return f"hidden hooks on {n} layers of {type(model).__name__}"


def main():
    args = parse_args()
    for k in ("hf_ckpt", "golden_dir", "out"):
        setattr(args, k, os.path.abspath(os.path.expanduser(getattr(args, k))))
    if args.ep > 1:
        launch_ep(args)
        return
    run(args)


def launch_ep(args) -> None:
    """EP=DP=args.ep: one spawned engine process per rank (what examples/_launch.py does for the V4 scripts)."""
    import multiprocessing as mp

    cores = os.environ.get("CUDA_VISIBLE_DEVICES") or ",".join(str(i) for i in range(args.ep))
    cores = cores.split(",")
    assert len(cores) >= args.ep, f"--ep {args.ep} needs {args.ep} entries in CUDA_VISIBLE_DEVICES, got {cores}"
    os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
    os.environ.setdefault("MASTER_PORT", "23333")
    ctx = mp.get_context("spawn")
    procs = [ctx.Process(target=_ep_rank_main, args=(args, rank, cores[rank]), daemon=False) for rank in range(args.ep)]
    for p in procs:
        p.start()
    # Fail fast: a dead rank leaves its peers blocked in a collective forever.
    while any(p.is_alive() for p in procs):
        if any(p.exitcode not in (None, 0) for p in procs):
            for p in procs:
                if p.is_alive():
                    p.terminate()
            break
        time.sleep(1)
    for p in procs:
        p.join()
    failed = [rank for rank, p in enumerate(procs) if p.exitcode != 0]
    if failed:
        raise RuntimeError(f"EP ranks failed: {failed}")


def _ep_rank_main(args, rank: int, core: str) -> None:
    # Before deepinfer / the Neuron runtime initialize in this process: one NeuronCore per rank.
    os.environ["CUDA_VISIBLE_DEVICES"] = core
    os.environ.pop("NEURON_RT_VISIBLE_CORES", None)  # PlatformNeuron derives it from CUDA_VISIBLE_DEVICES
    if rank > 0:
        args.out = os.path.join(args.out, f"rank{rank}")
    try:
        run(args, rank)
    except Exception:
        traceback.print_exc()
        sys.exit(1)


def run(args, rank: int = 0):
    setup_env(args.device, args.cp)
    os.makedirs(args.out, exist_ok=True)

    import torch

    torch.set_num_threads(args.threads)
    import deepinfer  # noqa: F401
    from deepinfer import PLATFORM

    print(f"[run] platform device = {PLATFORM.device}", flush=True)
    if args.device == "cpu":
        assert str(PLATFORM.device) == "cpu", "refusing to run: target device is not the host CPU"

    install_workarounds(args.num_kv_blocks, args.attn_p_bf16)
    rec = Recorder()
    install_logits_hook(rec)

    from deepinfer.entrypoints.llm import LLM
    from deepinfer.sampling_params import SamplingParams

    t0 = time.time()
    llm = LLM(
        model=args.hf_ckpt,
        trust_remote_code=True,
        expert_parallel_size=args.ep * args.cp,
        context_parallel_size=args.cp,
        data_parallel_size=args.ep,
        engine_rank=rank,
        chunked_prefill=True,
        decode_sequence=False,
        max_model_len=args.max_model_len,
        max_num_batched_tokens=args.max_num_batched_tokens,
        max_num_seqs=1,
        gpu_memory_utilization=0.5,
    )
    print(f"[run] engine up in {time.time() - t0:.1f}s (ep={args.ep}, rank={rank})", flush=True)
    status = install_hidden_hooks(rec, llm) if args.hidden else "hidden not requested"
    print(f"[run] {status}", flush=True)

    for name in args.prompts.split(","):
        gold = torch.load(os.path.join(args.golden_dir, f"{name}.pt"), weights_only=False)
        ids, gtok = list(gold["prompt_ids"]), list(gold["tokens"])
        steps = len(gtok)
        rec.reset()
        sp = SamplingParams(temperature=0.0, max_tokens=steps, ignore_eos=True, golden_token_ids=gtok)
        t1 = time.time()
        outputs = llm.generate(prompt_token_ids=[ids], sampling_params=sp, use_tqdm=False)
        dt = time.time() - t1
        out_tokens = [o.token_id for o in outputs[0]]
        rows = [r for r in rec.logits]
        assert all(r.size(0) == 1 for r in rows), [tuple(r.shape) for r in rows]
        logits = torch.cat(rows)[-steps:] if len(rows) >= steps else torch.cat(rows)
        print(f"[run] {name}: {len(rows)} request sampler calls, {rec.n_placeholder} placeholder steps", flush=True)
        if len(rows) != steps:
            print(f"[run] WARNING {name}: {len(rows)} sampler calls for {steps} steps (chunked prefill?)", flush=True)
        payload = {
            "name": name,
            "prompt_ids": ids,
            "tokens": logits.argmax(-1).tolist(),
            "engine_output_tokens": out_tokens,
            "fed_tokens": gtok[: steps - 1],
            "logits": logits,
            "source": "deepinfer LLM engine, teacher-forced via SamplingParams.golden_token_ids",
            "device": args.device,
        }
        if rec.hidden_calls:
            hid = []
            for call in rec.hidden_calls[-steps:]:
                hid.append(torch.stack([call[i].reshape(call[i].shape[0], -1, gold["hidden"][0].shape[-1]) for i in sorted(call)]).bfloat16())
            payload["hidden"] = hid
        path = os.path.join(args.out, f"{name}.pt")
        torch.save(payload, path)
        match = sum(a == b for a, b in zip(payload["tokens"], gtok))
        print(f"[run] {name}: {len(ids)} prompt tokens, engine output {out_tokens} (golden-forced: {out_tokens == gtok}), "
              f"argmax==golden {match}/{steps}, {dt:.1f}s -> {path}", flush=True)
        if args.compare:
            import compare_golden

            res = compare_golden.compare(gold, payload)
            for s in res["steps"]:
                print(f"    step {s['step']}: match {s['match']} margin {s['margin_ref']:.3f} rel {s['rel']:.2e} cos {s['cos']:.6f} kl {s['kl']:.2e}")
            for L in res["layers"][:1]:
                print("    prefill hidden rel per layer:", L.get("error") or " ".join(f"{v:.1e}" for v in L["rel"]))
    print(f"[run] workarounds triggered: {WORKAROUNDS}", flush=True)
    llm.close()


if __name__ == "__main__":
    try:
        main()
    except Exception:
        traceback.print_exc()
        sys.exit(1)