"""Reference next-token prediction from the official V4.1 inference code.

`generate.py` always wraps text in a chat template, so its outputs are not directly comparable
with a raw-token harness. This script takes explicit token ids and prints the greedy next token
(plus the top-5), which makes the comparison exact regardless of prompt formatting.

    torchrun --nproc-per-node 8 ref_next_token.py --ids 0,128803,2056,... --ckpt-path ...

Writes `<out>.json` with the ids, the chosen token and the top-5 so a harness can diff it.
"""

import argparse
import json
import os
import sys

import torch
import torch.distributed as dist

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

from model import ModelArgs, Transformer  # noqa: E402
from safetensors.torch import load_model  # 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, help="comma-separated prompt token ids")
    parser.add_argument("--out", default="/data/deepinfer-h20/ref/next_token.json")
    parser.add_argument("--topk", type=int, default=5)
    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

    # The engram hash table is built from the tokenizer's compressed token map, so it is not
    # optional even for a text-only forward.
    from transformers import AutoTokenizer

    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")

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

    row = logits[0, -1].float()
    top = torch.topk(row, args.topk)
    assert int(top.indices[0]) == int(output_ids[0, -1]), "argmax disagrees with the sampled id"

    payload = {
        "prompt_ids": ids,
        "next_token": int(top.indices[0]),
        "top_tokens": [int(x) for x in top.indices],
        "top_logits": [round(float(x), 4) for x in top.values],
    }
    if rank == 0:
        print(json.dumps(payload, ensure_ascii=False), flush=True)
        os.makedirs(os.path.dirname(args.out), exist_ok=True)
        with open(args.out, "w") as f:
            json.dump(payload, f, ensure_ascii=False, indent=2)
    if world_size > 1:
        dist.destroy_process_group()


if __name__ == "__main__":
    main()