"""Op-by-op capture of the first decoder layers, in the official reference and in DeepInfer, for one prefill.

    # reference (pure CPU, torch kernels from ref_kernels_torch), twice with different thread counts = floor
    TORCH_DEVICE_BACKEND_AUTOLOAD=0 python -I examples/neuron/probe_layer0.py ref --out ~/di/probe/ref16.pt --threads 16
    TORCH_DEVICE_BACKEND_AUTOLOAD=0 python -I examples/neuron/probe_layer0.py ref --out ~/di/probe/ref7.pt --threads 7
    # DeepInfer engine (CPU mode), same prompt, prefill only
    python examples/neuron/probe_layer0.py deepinfer --out ~/di/probe/di.pt
    # table
    python -I examples/neuron/probe_layer0.py compare ~/di/probe/ref16.pt ~/di/probe/di.pt --floor ~/di/probe/ref7.pt

Captured names are "L{layer}.{op}"; every tensor is [seqlen, ...] (batch dim dropped), fp32 on CPU. The prompt is
the `short` golden prompt (BOS + 10 tokens); the table reports token 0 (BOS), token 1 and all tokens.
"""

from __future__ import annotations

import argparse
import functools
import json
import os
import sys
from pathlib import Path

HERE = Path(__file__).resolve().parent
LAYERS = (0, 1, 2)  # layer 2 only to catch layer 1's FFN hc_post (fused into layer 2's pre in DeepInfer)
COMPARE_LAYERS = (0, 1)


class Cap:
    def __init__(self):
        self.t = {}
        self.layer = None
        self.count = {}
        self.active = True

    def put(self, name, x, layer=None):
        import torch

        layer = self.layer if layer is None else layer
        if not self.active or layer is None or layer not in LAYERS:
            return
        key = f"L{layer}.{name}"
        n = self.count.get(key, 0)
        self.count[key] = n + 1
        if n:
            key = f"{key}#{n}"
        if isinstance(x, (tuple, list)):
            x = x[0]
        x = x.detach().float().cpu().clone()
        self.t[key] = x


# ------------------------------------------------------------------------------------------------ reference
def run_ref(args):
    import torch

    torch.set_num_threads(args.threads)
    sys.path.insert(0, str(HERE))
    import ref_kernels_torch

    sys.modules["kernel"] = ref_kernels_torch
    sys.path.insert(1, os.path.expanduser(args.ref_dir))
    import engram as ref_engram
    import model as M
    import reference_golden as G
    from safetensors.torch import load_file
    from transformers import AutoTokenizer

    ref_engram.build_compressed_token_map = functools.lru_cache(None)(ref_engram.build_compressed_token_map)
    hf = json.load(open(os.path.join(os.path.expanduser(args.hf_ckpt), "config.json")))
    ma = G.hf_to_model_args(hf)
    ma.update(max_batch_size=1, max_seq_len=2048, temperature=0.0, vision_n_layers=0, dspark_block_size=0)
    conv = os.path.expanduser(args.converted)
    tok = AutoTokenizer.from_pretrained(conv)
    sd = load_file(os.path.join(conv, "model0-mp1.safetensors"))
    sd = {k: v for k, v in sd.items() if not (k.startswith(("vision.", "aligner.", "mtp.", "image_")) or k.endswith(".bias_vl"))}
    torch.set_default_dtype(torch.bfloat16)
    model = M.Transformer(M.ModelArgs(**ma), tok)
    model.load_state_dict(sd, strict=False)
    cap = Cap()

    def put_hook(name):
        return lambda mod, i, o: cap.put(name, o)

    model.embed.register_forward_hook(lambda m, i, o: cap.put("embed", o[0], layer=0))
    for li, blk in enumerate(model.layers):
        blk.register_forward_pre_hook(lambda m, i, li=li: setattr(cap, "layer", li))
        a = blk.attn
        blk.attn_norm.register_forward_hook(put_hook("attn_norm"))
        blk.ffn_norm.register_forward_hook(put_hook("ffn_norm"))
        a.wq_a.register_forward_hook(put_hook("wq_a"))
        a.q_norm.register_forward_hook(put_hook("q_norm"))
        a.wq_b.register_forward_hook(put_hook("wq_b"))
        a.wkv.register_forward_hook(put_hook("wkv"))
        a.kv_norm.register_forward_hook(put_hook("kv_norm"))
        a.wo_b.register_forward_pre_hook(lambda m, i: cap.put("wo_b_in", i[0]))
        a.wo_b.register_forward_hook(put_hook("attn_out"))
        g = blk.ffn.gate
        g.register_forward_pre_hook(lambda m, i: cap.put("gate_logits", torch.nn.functional.linear(i[0].float(), m.weight.float())))
        g.register_forward_hook(lambda m, i, o: (cap.put("topk_weights", o[0]), cap.put("topk_idx", o[1].float()), None)[-1])
        se = blk.ffn.shared_experts
        se.register_forward_hook(put_hook("shared_out"))
        se.w1.register_forward_hook(put_hook("shared_w1"))
        se.w3.register_forward_hook(put_hook("shared_w3"))
        se.w2.register_forward_pre_hook(lambda m, i: cap.put("shared_act", i[0]))
        blk.ffn.register_forward_hook(put_hook("moe_out"))
        if blk.engram is not None:  # runs before the block (Transformer.forward), so pass the layer explicitly
            blk.engram.register_forward_hook(lambda m, i, o, li=li: cap.put("engram_out", o[0], layer=li))

    orig_split = M.hc_split_sinkhorn

    def hc_split_sinkhorn(mixes, *a, **k):
        pre, post, comb = orig_split(mixes, *a, **k)
        cap.put("mixes", mixes[0])
        cap.put("pre_mix", pre[0])
        cap.put("post_mix", post[0])
        cap.put("comb_mix", comb[0])
        return pre, post, comb

    M.hc_split_sinkhorn = hc_split_sinkhorn
    orig_pre, orig_post = M.Block.hc_pre, M.Block.hc_post
    M.Block.hc_pre = lambda self, x, m: (lambda y: (cap.put("hc_pre", y[0]), y)[1])(orig_pre(self, x, m))
    M.Block.hc_post = lambda self, x, r, p, c: (lambda y: (cap.put("hc_post", y[0]), y)[1])(orig_post(self, x, r, p, c))
    orig_sa = M.sparse_attn

    def sparse_attn(q, kv, sink, idx, scale):
        o = orig_sa(q, kv, sink, idx, scale)
        cap.put("q_rope", q[0])
        cap.put("kv_rope_quant", kv[0, : q.size(1)])  # prefill: kv rows == tokens (window part)
        cap.put("attn_core", o[0])
        return o

    M.sparse_attn = sparse_attn

    # MoE.forward (model.py:891-906), unchanged numerics, additionally recording the fp32 routed sum and the
    # DeepEP-combine-style result (routed partial rounded to bf16 before the shared expert is added)
    def moe_forward(self, x, image_mask=None):
        shape = x.size()
        x = x.view(-1, self.dim)
        weights, indices = self.gate(x, None if image_mask is None else image_mask.flatten())
        y = torch.zeros_like(x, dtype=torch.float32)
        counts = torch.bincount(indices.flatten(), minlength=self.n_routed_experts).tolist()
        for i in range(self.experts_start_idx, self.experts_end_idx):
            if counts[i] == 0:
                continue
            idx, top = torch.where(indices == i)
            y[idx] += self.experts[i](x[idx], weights[idx, top, None])
        cap.put("routed_fp32", y)
        shared = self.shared_experts(x)
        cap.put("moe_out_combine_style", (y.to(torch.bfloat16).float() + shared.float()).to(torch.bfloat16))
        y += shared
        return y.type_as(x).view(shape)

    M.MoE.forward = moe_forward

    gold = torch.load(os.path.join(os.path.expanduser(args.golden_dir), "short.pt"), weights_only=False)
    ids = gold["prompt_ids"]
    with torch.inference_mode():
        model(torch.tensor([ids]), 0)
    torch.save({"side": "ref", "ids": ids, "t": cap.t}, os.path.expanduser(args.out))
    print(f"[probe] ref: {len(cap.t)} tensors -> {args.out}")


# ------------------------------------------------------------------------------------------------ proposed fix
def sparse_attn_rows_p_bf16(q, gather_kv, topk, sm_scale, d_v, attn_sink, rms_norm_scale_factor, block=64):
    """Drop-in for neuron_ops/vendor/flash_mla/flash_mla_interface.py:_sparse_attn_rows with the numerics of the
    CUDA FlashMLA sparse kernels (e.g. csrc/cuda_kernels/sm100/prefill/sparse/fwd/head128/phase1.cuh:238-253,324):
    online softmax over blocks of `block` keys, P = exp(s - running_max) rounded to bf16 before the P@V GEMM,
    running sum and the final 1 / (l + exp(sink - m)) in fp32. Same return contract as the original."""
    import torch

    n, h_q, d_qk = q.shape
    out = torch.empty((n, h_q, d_v), dtype=torch.float32, device=q.device)
    max_logits = torch.empty((n, h_q), dtype=torch.float32, device=q.device)
    lse_out = torch.empty((n, h_q), dtype=torch.float32, device=q.device)
    rows = max(1, min(n, (1 << 26) // max(1, topk * d_qk)))
    for lo in range(0, n, rows):
        hi = min(lo + rows, n)
        kv, invalid = gather_kv(lo, hi)
        kv = kv.float()
        kv = torch.where(kv != kv, torch.zeros_like(kv), kv)
        logits = torch.bmm(q[lo:hi].float(), kv.transpose(1, 2))
        if rms_norm_scale_factor is not None:
            logits = logits * (rms_norm_scale_factor[lo:hi].unsqueeze(-1) * sm_scale)
        else:
            logits = logits * sm_scale
        logits = torch.where(invalid.unsqueeze(1), torch.full_like(logits, float("-inf")), logits)
        m = torch.full((hi - lo, h_q), -1e30, dtype=torch.float32, device=q.device)
        l = torch.zeros_like(m)
        acc = torch.zeros((hi - lo, h_q, d_v), dtype=torch.float32, device=q.device)
        for t0 in range(0, topk, block):
            sblk = logits[..., t0 : t0 + block]
            new = torch.maximum(m, sblk.amax(-1))
            sc = torch.exp(m - new)
            p = torch.exp(sblk - new.unsqueeze(-1))
            l = l * sc + p.sum(-1)
            acc = acc * sc.unsqueeze(-1) + torch.bmm(p.to(torch.bfloat16).float(), kv[:, t0 : t0 + block, :d_v])
            m = new
        denom = l + (torch.exp(attn_sink.float().view(1, h_q) - m) if attn_sink is not None else 0)
        lonely = l == 0
        o = acc / torch.where(denom == 0, torch.ones_like(denom), denom).unsqueeze(-1)
        out[lo:hi] = torch.where(lonely.unsqueeze(-1), torch.zeros_like(o), o)
        max_logits[lo:hi] = logits.max(dim=-1).values
        lse_out[lo:hi] = torch.where(lonely, torch.full_like(l, float("inf")), m + torch.log(l))
    return out, max_logits, lse_out


# ------------------------------------------------------------------------------------------------ DeepInfer
def run_deepinfer(args):
    sys.path.insert(0, str(HERE))
    import run_deepinfer_tiny as R

    R.setup_env(args.device)
    import torch

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

    assert str(PLATFORM.device) == "cpu" or args.device == "neuron"
    cap = Cap()
    cap.active = False
    rec = R.Recorder()
    R.install_logits_hook(rec)  # also tags placeholder steps

    import deepinfer.model_executor.models.v4_modules.attn as attn_mod
    import deepinfer.model_executor.models.v4_modules.hc as hc_mod
    import deepinfer.model_executor.models.v4_modules.mlp as mlp_mod

    def real_step():
        return rec.step_kind is not None and rec.step_kind[0] and not rec.step_kind[1]

    def wrap(mod, fname, fn_out):
        orig = getattr(mod, fname)

        def w(*a, **k):
            out = orig(*a, **k)
            if real_step():
                fn_out(out, a, k)
            return out

        setattr(mod, fname, w)

    wrap(hc_mod, "mhc_pre_apply_mix", lambda o, a, k: cap.put("hc_pre", o.reshape(-1, o.shape[-1])))
    wrap(hc_mod, "mhc_pre_norm_fn", lambda o, a, k: cap.put("mixes", o.reshape(-1, o.shape[-1])))

    def split_out(o, a, k):
        pre, post, comb = o
        cap.put("pre_mix", pre.reshape(-1, pre.shape[-2]))
        cap.put("post_mix", post.reshape(-1, post.shape[-2]))

    wrap(hc_mod, "mhc_pre_split_mixes", split_out)
    wrap(hc_mod, "sinkhorn_normalize", lambda o, a, k: cap.put("comb_mix", o.reshape(-1, *o.shape[-2:])))
    wrap(hc_mod, "mhc_post", lambda o, a, k: cap.put("hc_post_raw", o.reshape(-1, *o.shape[-2:])))
    wrap(hc_mod, "expand_to_mhc", lambda o, a, k: None)

    def fused_out(o, a, k):
        if o is None:
            return
        x = a[0]
        if x.dim() == 3 and x.shape[1] > 1:  # q [s, heads, head_dim]
            cap.put("q_rope", o)
        else:
            cap.put("kv_rope_quant", o.reshape(o.shape[0], -1))

    wrap(attn_mod, "fused_norm_rope_hadamard_cast_store_kvcache", fused_out)
    wrap(attn_mod, "prefill_attn_varlen_sparse", lambda o, a, k: cap.put("attn_core", o[0]))
    import tile_kernels.quant.swiglu_forward_kernel as swk

    orig_swiglu = mlp_mod.swiglu_forward_and_per_token_cast

    def swiglu(x, *a, **k):
        if real_step() and k.get("topk_weights") is None and (len(a) < 4 or a[3] is None):
            inter = x.shape[-1] // 2
            cap.put("shared_w1", x[:, :inter])
            cap.put("shared_w3", x[:, inter:])
            cap.put("shared_act", swk.swiglu_compute(x, None, None, None, k.get("clamp_value"), None))
        return orig_swiglu(x, *a, **k)

    mlp_mod.swiglu_forward_and_per_token_cast = swiglu
    if args.round_act_bf16:
        # experiment (not a proposed fix): round the SwiGLU output to bf16 before the FP8 cast, like the reference's
        # Expert.forward (`self.w2(x.to(dtype))`), to confirm that this rounding point explains the FFN deviation
        orig_compute = swk.swiglu_compute
        swk.swiglu_compute = lambda *a, **k: orig_compute(*a, **k).to(torch.bfloat16).float()
        print("[probe] SwiGLU output rounded to bf16 before the FP8 cast (experiment)")
    if args.fix_attn:
        import flash_mla.flash_mla_interface as fmi

        fmi._sparse_attn_rows = sparse_attn_rows_p_bf16
        print("[probe] using P-bf16 sparse attention (proposed fix)")
    wrap(mlp_mod, "moe_topk_gate", lambda o, a, k: (cap.put("topk_idx", o[0].float()), cap.put("topk_weights", o[1])))

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

    llm = LLM(model=os.path.expanduser(args.hf_ckpt), trust_remote_code=True, expert_parallel_size=1, context_parallel_size=1,
              data_parallel_size=1, engine_rank=0, chunked_prefill=True, decode_sequence=False, max_model_len=2048,
              max_num_batched_tokens=1024, max_num_seqs=1, gpu_memory_utilization=0.5)
    model = llm.llm_engine.driver_worker.model_runner.model
    q_lora = None

    def hook(name, fn=None):
        def h(m, i, o):
            if real_step():
                cap.put(name, fn(o) if fn else o)
        return h

    def bf16_of(o):
        # norms return ((fp8, sf), bf16) or (fp8, bf16) when return_bf16=True
        return o[1] if isinstance(o, tuple) else o

    model.model.embed_tokens.register_forward_hook(lambda m, i, o: cap.put("embed", o, layer=0) if real_step() else None)
    for li, layer in enumerate(model.model.layers):
        layer.register_forward_pre_hook(lambda m, i, li=li: setattr(cap, "layer", li) if real_step() else None)
        a = layer.self_attn
        q_lora = a.q_lora_rank
        layer.attn_layernorm.register_forward_hook(hook("attn_norm", bf16_of))
        layer.ffn_layernorm.register_forward_hook(hook("ffn_norm", bf16_of))
        a.q_a_kv_swa_proj.register_forward_hook(hook("q_a_kv"))
        a.q_a_norm.register_forward_hook(hook("q_norm", bf16_of))
        a.q_b_proj.register_forward_hook(hook("wq_b"))
        a.out_proj.register_forward_pre_hook(lambda m, i: cap.put("wo_b_in", i[0]) if real_step() else None)
        a.out_proj.register_forward_hook(hook("attn_out"))
        mlp = layer.mlp
        mlp.gate.register_forward_hook(hook("gate_logits"))
        orig_shared = mlp.shared_forward

        def shared_forward(*aa, _o=orig_shared, **kk):
            out = _o(*aa, **kk)
            if real_step() and out is not None:
                cap.put("shared_out", out)
            return out

        mlp.shared_forward = shared_forward
        mlp.register_forward_hook(hook("moe_out"))
        if layer.engram is not None:
            layer.engram.register_forward_hook(hook("engram_out"))

    gold = __import__("torch").load(os.path.join(os.path.expanduser(args.golden_dir), "short.pt"), weights_only=False)
    ids = gold["prompt_ids"]
    cap.active = True
    llm.generate(prompt_token_ids=[ids], sampling_params=SamplingParams(temperature=0.0, max_tokens=1, ignore_eos=True,
                                                                        golden_token_ids=gold["tokens"][:1]), use_tqdm=False)
    t = cap.t
    # split the fused q_a|kv projection like the reference's separate wq_a / wkv
    for li in LAYERS:
        k = f"L{li}.q_a_kv"
        if k in t:
            t[f"L{li}.wq_a"], t[f"L{li}.wkv"] = t[k][:, :q_lora], t[k][:, q_lora:]
        k = f"L{li}.wq_b"
        if k in t:
            H, hd = model.model.layers[li].self_attn.num_heads, model.model.layers[li].self_attn.qk_head_dim
            t[k] = t[k][:, : H * hd]
    # mhc_post order: inside layer L, call #0 is L's attention post; L's FFN post runs inside layer L+1 (fused
    # post+pre), i.e. "L{L+1}.hc_post_raw" call #0. Re-key to the reference's per-block hc_post / hc_post#1.
    raw = {k: v for k, v in t.items() if ".hc_post_raw" in k}
    for k in raw:
        t.pop(k)
    for li in LAYERS:
        attn_post = raw.get(f"L{li}.hc_post_raw" if li == 0 else f"L{li}.hc_post_raw#1")
        ffn_post = raw.get(f"L{li + 1}.hc_post_raw")
        if attn_post is not None:
            t[f"L{li}.hc_post"] = attn_post
        if ffn_post is not None:
            t[f"L{li}.hc_post#1"] = ffn_post
    torch_ = __import__("torch")
    torch_.save({"side": "deepinfer", "ids": ids, "t": t}, os.path.expanduser(args.out))
    print(f"[probe] deepinfer: {len(t)} tensors -> {args.out}; workarounds {R.WORKAROUNDS}")
    llm.close()


# ------------------------------------------------------------------------------------------------ compare
ORDER = ["embed", "mixes", "pre_mix", "post_mix", "comb_mix", "hc_pre", "attn_norm", "wq_a", "q_norm", "wq_b", "q_rope",
         "wkv", "kv_rope_quant", "attn_core", "wo_b_in", "attn_out", "hc_post", "mixes#1", "pre_mix#1", "post_mix#1",
         "comb_mix#1", "hc_pre#1", "ffn_norm", "gate_logits", "topk_idx", "topk_weights", "shared_w1", "shared_w3", "shared_act", "shared_out", "routed_fp32", "moe_out", "moe_out_combine_style", "hc_post#1",
         "engram_out"]


def run_compare(args):
    import torch

    A = torch.load(os.path.expanduser(args.a), weights_only=False)["t"]
    B = torch.load(os.path.expanduser(args.b), weights_only=False)["t"]
    F = torch.load(os.path.expanduser(args.floor), weights_only=False)["t"] if args.floor else None

    def rel(x, y):
        x, y = x.float(), y.float()
        if x.shape != y.shape:
            return f"shape {tuple(x.shape)}!={tuple(y.shape)}"
        return ((x - y).norm() / x.norm().clamp_min(1e-30)).item()

    def fmt(v):
        return f"{v:9.2e}" if isinstance(v, float) else f"{v:>9}"

    print(f"{'op':22s} {'shape':>18s} {'BOS':>9s} {'tok1':>9s} {'all':>9s} {'floor(all)':>10s}")
    def sort_topk(T):
        for k in list(T):
            if k.endswith("topk_idx"):
                idx = T[k]
                order = idx.argsort(-1)
                T[k] = idx.gather(-1, order)
                w = k.replace("topk_idx", "topk_weights")
                if w in T:
                    T[w] = T[w].gather(-1, order)

    for T in (A, B, F):
        if T is not None:
            sort_topk(T)
    for li in COMPARE_LAYERS:
        names = [n for n in ORDER if f"L{li}.{n}" in A]
        names += sorted(k.split(".", 1)[1] for k in A if k.startswith(f"L{li}.") and k.split(".", 1)[1] not in ORDER)
        for n in names:
            k = f"L{li}.{n}"
            if k not in B:
                print(f"{k:22s} {'(missing in candidate)':>18s}")
                continue
            a = A[k][0] if A[k].dim() >= 3 and A[k].shape[0] == 1 else A[k]
            b = B[k].reshape(a.shape) if B[k].numel() == a.numel() else B[k]
            f = F[k][0] if F is not None and k in F and F[k].dim() >= 3 and F[k].shape[0] == 1 else (F[k] if F is not None and k in F else None)
            cols = [rel(a[:1], b[:1]) if a.shape == b.shape else rel(a, b), rel(a[1:2], b[1:2]) if a.shape == b.shape else "", rel(a, b)]
            fl = rel(a, f) if f is not None else ""
            print(f"{k:22s} {str(tuple(a.shape)):>18s} " + " ".join(fmt(c) for c in cols) + f" {fmt(fl):>10s}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("side", choices=["ref", "deepinfer", "compare"])
    ap.add_argument("a", nargs="?")
    ap.add_argument("b", nargs="?")
    ap.add_argument("--floor", default=None)
    ap.add_argument("--out", default="~/di/probe/out.pt")
    ap.add_argument("--threads", type=int, default=16)
    ap.add_argument("--device", default="cpu")
    ap.add_argument("--fix-attn", action="store_true", help="deepinfer side: use the proposed P-bf16 attention")
    ap.add_argument("--round-act-bf16", action="store_true", help="deepinfer side: bf16-round SwiGLU output before FP8 cast")
    ap.add_argument("--hf-ckpt", default="~/di/tiny_ckpt")
    ap.add_argument("--converted", default="~/di/tiny_ckpt_ref_mp1")
    ap.add_argument("--ref-dir", default="~/di/ref_dsv41/inference")
    ap.add_argument("--golden-dir", default="~/di/golden")
    args = ap.parse_args()
    if args.side != "compare":
        os.makedirs(os.path.dirname(os.path.expanduser(args.out)), exist_ok=True)
    {"ref": run_ref, "deepinfer": run_deepinfer, "compare": run_compare}[args.side](args)


if __name__ == "__main__":
    main()