"""Dump the reference model's per-layer hidden state for one prompt.

Counterpart to DeepInfer's `DEEPINFER_DUMP_LAYERS`, so the two can be diffed layer by layer.
The reference residual stream is `hc_mult` parallel copies, so each layer is reduced with
`mean(dim=1)` -- the same convention `debug_dump.record` uses -- and only the final position
is kept, because that is the only row generation depends on.

    torchrun --nproc-per-node 8 ref_layer_dump.py --ids 0,128803,... --out /data/deepinfer-h20/ref
"""

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

import numpy as np
import torch
import torch.distributed as dist

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from collections.abc import Callable  # noqa: E402

from model import ModelArgs, Transformer  # noqa: E402
from safetensors.torch import load_model  # noqa: E402
from transformers import AutoTokenizer  # noqa: E402


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--ckpt-path", required=True)
    parser.add_argument("--config", required=True)
    parser.add_argument("--ids", required=True)
    parser.add_argument("--out", default="/data/deepinfer-h20/ref")
    parser.add_argument("--moe_layers", type=int, default=2, help="dump MoE internals for this many layers")
    args = parser.parse_args()

    ids = [int(x) for x in args.ids.replace(" ", "").split(",") if x]
    rank = int(os.getenv("RANK", "0"))
    local_rank = int(os.getenv("LOCAL_RANK", "0"))
    world_size = int(os.getenv("WORLD_SIZE", "1"))
    torch.cuda.set_device(local_rank)
    if world_size > 1:
        dist.init_process_group("nccl")

    torch.set_default_dtype(torch.bfloat16)
    torch.set_num_threads(8)
    with open(args.config) as f:
        model_args = ModelArgs(**json.load(f))
    model_args.temperature = 0.0
    model_args.max_batch_size = 1
    model_args.max_seq_len = 4096

    tokenizer = AutoTokenizer.from_pretrained(args.ckpt_path)
    with torch.device("cuda"):
        model = Transformer(model_args, tokenizer)
    load_model(model, os.path.join(args.ckpt_path, f"model{rank}-mp{world_size}.safetensors"))
    torch.set_default_device("cuda")

    captured: dict[str, np.ndarray] = {}

    def make_hook(layer_id: int) -> Callable[[torch.nn.Module, tuple, object], None]:
        def hook(_module: torch.nn.Module, _inputs: tuple[torch.Tensor, ...], output: object) -> None:
            # Block.forward returns (x, pre_mix); x is [b, s, hc_mult, dim].
            x = output[0] if isinstance(output, tuple) else output
            captured[f"layer{layer_id}"] = x.mean(dim=2)[0, -1].float().cpu().numpy()

        return hook

    def make_attn_hook(layer_id: int) -> Callable[[torch.nn.Module, tuple, object], None]:
        def hook(_module: torch.nn.Module, _inputs: tuple[torch.Tensor, ...], output: object) -> None:
            # The attention sub-layer output is already collapsed: [b, s, dim].
            x = output[0] if isinstance(output, tuple) else output
            captured[f"layer{layer_id}_attn"] = x[0, -1].float().cpu().numpy()

        return hook

    def embed_hook(_module: torch.nn.Module, _inputs: tuple[torch.Tensor, ...], output: object) -> None:
        captured["embed"] = output[0, -1].float().cpu().numpy()

    def make_gate_hook(layer_id: int) -> Callable[[torch.nn.Module, tuple, object], None]:
        def hook(_module: torch.nn.Module, _inputs: tuple[torch.Tensor, ...], output: object) -> None:
            # Gate.forward returns (weights [n, topk], indices [n, topk]).
            weights, indices = output
            captured[f"layer{layer_id}_topk_idx"] = indices[-1].to(torch.int32).cpu().numpy()
            captured[f"layer{layer_id}_topk_w"] = weights[-1].float().cpu().numpy()

        return hook

    def make_shared_hook(layer_id: int) -> Callable[[torch.nn.Module, tuple, object], None]:
        def hook(_module: torch.nn.Module, _inputs: tuple[torch.Tensor, ...], output: object) -> None:
            captured[f"layer{layer_id}_shared_out"] = output.view(-1, output.shape[-1])[-1].float().cpu().numpy()

        return hook

    def make_moe_hook(layer_id: int) -> Callable[[torch.nn.Module, tuple, object], None]:
        def hook(_module: torch.nn.Module, _inputs: tuple[torch.Tensor, ...], output: object) -> None:
            # The MoE output already includes the shared expert, matching our `_combined`.
            captured[f"layer{layer_id}_combined"] = output.view(-1, output.shape[-1])[-1].float().cpu().numpy()

        return hook

    model.embed.register_forward_hook(embed_hook)
    for layer_id, layer in enumerate(model.layers):
        layer.register_forward_hook(make_hook(layer_id))
        layer.attn.register_forward_hook(make_attn_hook(layer_id))
        if layer_id < args.moe_layers:
            layer.ffn.gate.register_forward_hook(make_gate_hook(layer_id))
            layer.ffn.shared_experts.register_forward_hook(make_shared_hook(layer_id))
            layer.ffn.register_forward_hook(make_moe_hook(layer_id))

    tokens = torch.tensor([ids], dtype=torch.long)
    with torch.inference_mode():
        output_ids, logits, _ = model.forward(tokens, 0)

    captured["final_row_logits"] = logits[0, -1].float().cpu().numpy()
    captured["next_token"] = np.array([int(output_ids[-1])], dtype=np.int32)

    if rank == 0:
        out = Path(args.out)
        out.mkdir(parents=True, exist_ok=True)
        path = out / "ref_layers.npz"
        np.savez(path, **captured)
        print(f"wrote {len(captured)} tensors to {path}", flush=True)
        print(json.dumps({"next_token": int(output_ids[-1]), "layers": len(model.layers)}), flush=True)

    if world_size > 1:
        dist.destroy_process_group()


if __name__ == "__main__":
    main()