"""Ground truth for the Neuron port: run the OFFICIAL DeepSeek-V4.1 reference model on the host CPU.

The downloaded reference (`inference/model.py`, `engram.py`, `convert.py`, ...) is used unmodified; only its
TileLang `kernel` module is swapped for the pure-torch `ref_kernels_torch` (same numerics, see that file).
DSpark/MTP (`dspark_block_size=0`) and vision (`vision_n_layers=0`) are disabled: DeepInfer runs text-only with
DEEPINFER_USE_NEXTN_DECODE=0, and neither affects the backbone logits.

Steps:
  1. convert the HF checkpoint with the official convert.py (MP=1, --expert-dtype fp4) unless already done;
  2. build reference ModelArgs from the HF config.json (same mapping as the official inference/config.json);
  3. per prompt (batch 1, fresh model): prefill the whole prompt, then N greedy decode steps (temperature 0,
     no EOS stop), saving logits (fp32), greedy tokens, per-layer block outputs, indexer top-k and engram hashes.

Usage (on the box, host CPU only; TORCH_DEVICE_BACKEND_AUTOLOAD=0 keeps torch_neuronx/NRT out of the process):
    TORCH_DEVICE_BACKEND_AUTOLOAD=0 ~/di/venv/bin/python -I examples/neuron/reference_golden.py \
        --hf-ckpt ~/di/tiny_ckpt --ref-dir ~/di/ref_dsv41/inference --converted ~/di/tiny_ckpt_ref_mp1 \
        --out ~/di/golden --summary-dir examples/neuron/golden
"""

from __future__ import annotations

import argparse
import dataclasses
import functools
import hashlib
import json
import os
import subprocess
import sys
import time

HERE = os.path.dirname(os.path.abspath(__file__))

PROMPTS = {
    # short: no window wrap, no compressed top-k selection
    "short": "The quick brown fox jumps over the lazy dog.",
    # medium (~200 tokens): wraps the 128-token sliding window and makes both indexer top-k levels select
    "medium": (
        "In the early morning the harbor was quiet. Fishing boats rocked gently against the wooden piers, "
        "their ropes creaking as the tide came in. An old man sat on an overturned crate, mending a net with "
        "slow, careful hands. He had worked these waters for more than fifty years, and he could read the "
        "weather in the color of the sky and the shape of the clouds. Today the sky was pale gray, and a thin "
        "mist hung over the water. Somewhere a gull cried, and the sound echoed off the stone walls of the "
        "warehouses. A young girl ran down the pier, her boots clattering on the boards. She stopped beside "
        "the old man and asked him whether the boats would go out today. He looked at the horizon for a long "
        "moment before answering. The sea, he said, decides that for us. We only listen."
    ),
}
# long (~700 tokens): more than candidate_topk_blocks * candidate_block_size = 512 compressed positions in the
# ratio-1 layers, so the candidate block pre-filter is active as well
PROMPTS["long"] = " ".join(
    [
        PROMPTS["medium"],
        "The history of computing is a story of abstraction. Early machines were programmed by rewiring "
        "panels and flipping switches, and every program was tied to the physical layout of the hardware. "
        "Assembly languages gave names to instructions, and compilers later translated whole expressions "
        "into sequences of those instructions. Operating systems hid the details of disks and terminals "
        "behind files and streams, while virtual memory let each process pretend it owned the entire machine.",
        "Networks added another layer. Packets travel through routers that know nothing about the "
        "applications sending them, and protocols stacked on top of one another turn unreliable links into "
        "dependable conversations. A web page that appears in a fraction of a second may involve dozens of "
        "servers, caches and databases spread across several continents, each one following rules that were "
        "agreed upon decades ago by people who could not have imagined how the system would grow.",
        "Mathematics offers a similar lesson. A proof about prime numbers does not care whether the numbers "
        "are written in decimal or binary, on paper or in silicon. The same theorem holds in every notation, "
        "and that independence is exactly what makes it useful. Engineers borrow this idea when they define "
        "interfaces: as long as both sides respect the contract, either side can be rewritten, optimized or "
        "replaced without the other noticing.",
        "Cooking, perhaps surprisingly, works the same way. A recipe names its ingredients and steps but "
        "leaves the choice of pan, stove and kitchen to the cook. Bread rises because yeast consumes sugar "
        "and releases gas, whether the dough is kneaded by hand in a village bakery or by a machine in a "
        "factory. Understanding the underlying chemistry lets a baker adapt when the flour is different, the "
        "air is humid, or the oven runs hot.",
        "Music, too, depends on structure. A melody can be transposed into another key, played on a different "
        "instrument, or slowed down to half its tempo, and listeners will still recognize it. Composers rely on "
        "patterns of tension and release, repetition and variation, to guide the ear through a piece. The notes "
        "on the page are only instructions; the performance brings them to life in a particular room on a "
        "particular evening, never exactly the same twice.",
        "Gardens change with every season. In spring the first green shoots push through the cold soil, and by "
        "summer the beds are crowded with tomatoes, beans and sunflowers that lean toward the light. Autumn "
        "brings the harvest and the slow work of clearing dead stems, while winter rests under frost and snow. "
        "A patient gardener learns that each year is different and that the soil remembers what was done to it.",
        "At the end of the day the harbor grew quiet again. The boats returned one by one, low in the water "
        "with their catch, and the old man folded his net and walked slowly home along the shore.",
    ]
)


def hf_to_model_args(hf: dict) -> dict:
    """HF config.json -> reference ModelArgs kwargs (field mapping as in the official inference/config.json)."""
    t = hf["text_config"]
    v = hf["vision_config"]
    rs = t["rope_scaling"]
    return dict(
        vocab_size=t["vocab_size"],
        dim=t["hidden_size"],
        moe_inter_dim=t["moe_intermediate_size"],
        n_layers=t["num_hidden_layers"],
        n_mtp_layers=t["num_nextn_predict_layers"],
        dspark_block_size=t["dspark_block_size"],
        dspark_noise_token_id=t["dspark_noise_token_id"],
        dspark_target_layer_ids=tuple(t["dspark_target_layer_ids"]),
        dspark_markov_rank=t["dspark_markov_rank"],
        dspark_n_routed_experts=t["dspark_n_routed_experts"],
        dspark_n_activated_experts=t["dspark_num_experts_per_tok"],
        n_heads=t["num_attention_heads"],
        n_routed_experts=t["n_routed_experts"],
        n_shared_experts=t["n_shared_experts"],
        n_activated_experts=t["num_experts_per_tok"],
        score_func=t["scoring_func"],
        norm_topk_prob=t["norm_topk_prob"],
        route_scale=t["routed_scaling_factor"],
        swiglu_limit=t["swiglu_limit"],
        q_lora_rank=t["q_lora_rank"],
        head_dim=t["head_dim"],
        rope_head_dim=t["qk_rope_head_dim"],
        norm_eps=t["rms_norm_eps"],
        o_groups=t["o_groups"],
        o_lora_rank=t["o_lora_rank"],
        window_size=t["sliding_window"],
        kv_source_layers=tuple(t["kv_source_layer_ids"]),
        index_source_layers=tuple(t["index_source_layer_ids"]),
        original_seq_len=rs["original_max_position_embeddings"],
        rope_theta=t["rope_theta"],
        rope_factor=rs["factor"],
        beta_fast=rs["beta_fast"],
        beta_slow=rs["beta_slow"],
        index_n_heads=t["index_n_heads"],
        index_head_dim=t["index_head_dim"],
        index_topk=t["index_topk"],
        candidate_source_layer=t["candidate_source_layer_id"],
        candidate_topk_blocks=t["candidate_topk_blocks"],
        candidate_block_size=t["candidate_block_size"],
        hc_mult=t["hc_mult"],
        hc_sinkhorn_iters=t["hc_sinkhorn_iters"],
        hc_eps=t["hc_eps"],
        engram_layer_ids=tuple(t["engram_layer_ids"]),
        engram_vocab_size=t["engram_vocab_size"],
        engram_num_embeddings=tuple(t["engram_num_embeddings"]),
        engram_max_ngram_size=t["engram_max_ngram_size"],
        engram_pad_id=t["engram_pad_token_id"],
        engram_compressed_vocab_size=t["engram_compressed_vocab_size"],
        engram_n_heads=t["engram_n_heads"],
        engram_head_dim=t["engram_head_dim"],
        dtype="fp8" if hf["quantization_config"]["quant_method"] == "fp8" else "bf16",
        expert_dtype=hf["quantization_config"]["expert_dtype"],
        compress_rope_theta=t["compress_rope_theta"],
        compress_ratios=tuple(t["compress_ratios"]),
        vision_n_layers=v["num_hidden_layers"],
        vision_dim=v["hidden_size"],
        vision_n_heads=v["num_attention_heads"],
        vision_inter_dim=v["intermediate_size"],
        vision_patch_size=v["patch_size"],
        vision_rope_theta=v["rope_theta"],
        vision_downsample_ratio=v["downsample_ratio"],
        vision_max_n_token=v["max_image_tokens"],
        vision_min_pixels=v["min_pixels"],
        vision_max_wh_ratio=v["max_wh_ratio"],
        image_token_id=hf["image_token_id"],
    )


def convert(ref_dir: str, hf_ckpt: str, out: str) -> str:
    path = os.path.join(out, "model0-mp1.safetensors")
    if os.path.exists(path):
        print(f"[golden] using existing converted checkpoint {path}")
        return path
    cmd = [sys.executable, "-I", os.path.join(ref_dir, "convert.py"), "--hf-ckpt-path", hf_ckpt, "--save-path", out,
           "--model-parallel", "1", "--expert-dtype", "fp4"]
    print("[golden] $", " ".join(cmd))
    subprocess.run(cmd, check=True, cwd="/tmp")
    return path


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--hf-ckpt", default="~/di/tiny_ckpt")
    ap.add_argument("--ref-dir", default="~/di/ref_dsv41/inference", help="copy of the official inference/ dir")
    ap.add_argument("--converted", default="~/di/tiny_ckpt_ref_mp1")
    ap.add_argument("--out", default="~/di/golden")
    ap.add_argument("--summary-dir", default=os.path.join(HERE, "golden"))
    ap.add_argument("--steps", type=int, default=8, help="greedy decode steps after the prefill")
    ap.add_argument("--prompts", default=",".join(PROMPTS))
    ap.add_argument("--threads", type=int, default=16)
    ap.add_argument("--max-seq-len", type=int, default=2048)
    ap.add_argument(
        "--force-tokens-dir",
        default=None,
        help="teacher forcing: feed the `tokens` of <dir>/<prompt>.pt as decode inputs instead of this run's argmax "
        "(used to measure the reference's own numerical noise floor against an existing golden)",
    )
    args = ap.parse_args()
    for k in ("hf_ckpt", "ref_dir", "converted", "out", "summary_dir"):
        setattr(args, k, os.path.abspath(os.path.expanduser(getattr(args, k))))
    if args.force_tokens_dir:
        args.force_tokens_dir = os.path.abspath(os.path.expanduser(args.force_tokens_dir))
    os.makedirs(args.out, exist_ok=True)
    os.makedirs(args.summary_dir, exist_ok=True)

    import torch

    torch.set_num_threads(args.threads)
    assert "torch_neuronx" not in sys.modules, "run with TORCH_DEVICE_BACKEND_AUTOLOAD=0 (host CPU only)"

    # ---- reference code: real model.py / engram.py, torch kernels instead of TileLang ----
    sys.path.insert(0, HERE)
    import ref_kernels_torch

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

    assert ref_model.act_quant is ref_kernels_torch.act_quant
    assert os.path.dirname(os.path.abspath(ref_model.__file__)) == args.ref_dir
    # the compressed token map walks the whole vocab; build it once
    ref_engram.build_compressed_token_map = functools.lru_cache(maxsize=None)(ref_engram.build_compressed_token_map)

    ckpt_file = convert(args.ref_dir, args.hf_ckpt, args.converted)
    hf_cfg = json.load(open(os.path.join(args.hf_ckpt, "config.json")))
    margs = hf_to_model_args(hf_cfg)
    margs.update(
        max_batch_size=1,
        max_seq_len=args.max_seq_len,
        temperature=0.0,
        vision_n_layers=0,  # text only
        dspark_block_size=0,  # no DSpark / MTP modules
    )
    model_args = ref_model.ModelArgs(**margs)
    print("[golden]", model_args)
    tokenizer = AutoTokenizer.from_pretrained(args.converted)
    bos = hf_cfg["bos_token_id"]

    state = load_file(ckpt_file)
    dropped = sorted(k for k in state if k.startswith(("vision.", "aligner.", "mtp.", "image_")) or k.endswith(".bias_vl"))
    state = {k: v for k, v in state.items() if k not in set(dropped)}
    print(f"[golden] {len(state)} tensors for the backbone; dropped {len(dropped)} (vision/aligner/mtp/image_*/bias_vl)")

    torch.set_default_dtype(torch.bfloat16)
    for name in args.prompts.split(","):
        text = PROMPTS[name]
        ids = [bos] + tokenizer.encode(text, add_special_tokens=False)
        t0 = time.time()
        torch.manual_seed(0)
        with torch.device("cpu"):
            model = ref_model.Transformer(model_args, tokenizer)
        missing, unexpected = model.load_state_dict(state, strict=False)
        assert not missing and not unexpected, (missing, unexpected)
        model.eval()

        rec = {"hidden": [], "indexer": [], "engram": []}
        cur = {}

        def block_hook(i):
            def fn(mod, inp, out):
                cur.setdefault("hidden", {})[i] = out[0].detach().clone()
            return fn

        def idx_hook(i):
            def fn(mod, inp, out):
                cur.setdefault("indexer", {})[i] = out.detach().clone()
            return fn

        for i, layer in enumerate(model.layers):
            layer.register_forward_hook(block_hook(i))
            if layer.attn.indexer is not None:
                layer.attn.indexer.register_forward_hook(idx_hook(i))
        if model.engram_hash is not None:
            model.engram_hash.register_forward_hook(lambda m, i, o: cur.__setitem__("engram", o.detach().clone()))

        forced = None
        if args.force_tokens_dir:
            ref = torch.load(os.path.join(args.force_tokens_dir, f"{name}.pt"), weights_only=False)
            assert ref["prompt_ids"] == ids, "teacher-forcing golden was built from a different prompt"
            forced = ref["tokens"]
        tokens, fed, logits_all = [], [], []
        x = torch.tensor([ids], dtype=torch.long)
        pos = 0
        for step in range(args.steps + 1):
            cur.clear()
            out_ids, logits, _ = model(x, pos)
            pos += x.size(1)
            nxt = int(logits[0].argmax())
            assert nxt == int(out_ids[0])
            tokens.append(nxt)
            logits_all.append(logits[0].float().clone())
            rec["hidden"].append(torch.stack([cur["hidden"][i][0] for i in range(len(model.layers))]).bfloat16())
            rec["indexer"].append({k: v[0].int() for k, v in cur.get("indexer", {}).items()})
            rec["engram"].append(cur["engram"][0] if "engram" in cur else None)
            feed = forced[step] if forced is not None else nxt
            fed.append(feed)
            x = torch.tensor([[feed]], dtype=torch.long)
        logits_all = torch.stack(logits_all)
        dt = time.time() - t0

        top2 = logits_all.topk(2, dim=-1)
        margins = (top2.values[:, 0] - top2.values[:, 1]).tolist()
        payload = {
            "name": name,
            "prompt_text": text,
            "prompt_ids": ids,
            "tokens": tokens,  # argmax of each step: tokens[0] from the prefill, tokens[1:] from the decode steps
            "fed_tokens": fed[:-1],  # decode step i+1 was fed fed_tokens[i] (== tokens[i] unless teacher-forced)
            "teacher_forced_from": args.force_tokens_dir,
            "logits": logits_all,  # [steps+1, vocab] fp32; row i predicts tokens[i]
            "start_pos": [0] + [len(ids) + i for i in range(args.steps)],
            "hidden": rec["hidden"],  # per forward call: [n_layers, seqlen, hc_mult, dim] bf16 block outputs
            "indexer_topk": rec["indexer"],  # per forward call: {layer: [seqlen, topk] int32 (-1 = none, +window offset)}
            "engram_hashes": rec["engram"],  # per forward call: [seqlen, n_engram_layers, n_hash_cols] int64
            "model_args": dataclasses.asdict(model_args),
            "hf_config": hf_cfg,
        }
        pt = os.path.join(args.out, f"{name}.pt")
        torch.save(payload, pt)
        top5 = logits_all.topk(5, dim=-1)
        summary = {
            "name": name,
            "prompt_text": text,
            "prompt_ids": ids,
            "n_prompt_tokens": len(ids),
            "tokens": tokens,
            "decoded": tokenizer.decode(tokens),
            "top1_top2_margin": [round(m, 5) for m in margins],
            "top5_ids": top5.indices.tolist(),
            "top5_logits": [[round(v, 5) for v in row] for row in top5.values.tolist()],
            "logits_l2": [round(v, 4) for v in logits_all.norm(dim=-1).tolist()],
            "logits_mean": [round(v, 6) for v in logits_all.mean(dim=-1).tolist()],
            "full_golden": pt,
            "full_golden_sha256": hashlib.sha256(open(pt, "rb").read()).hexdigest(),
            "seconds": round(dt, 1),
        }
        with open(os.path.join(args.summary_dir, f"{name}.json"), "w") as fh:
            json.dump(summary, fh, indent=1)
        print(f"[golden] {name}: {len(ids)} prompt tokens -> {tokens} ({tokenizer.decode(tokens)!r}); "
              f"min top1-top2 margin {min(margins):.4f}; {dt:.1f}s; saved {pt}")


if __name__ == "__main__":
    main()