# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0

"""Prompt construction, tokenization, static slot layout and position tables for GR00T N1.5 / N1.6 / N1.7 (CPU).

Sources:

* ``host-pipeline-parity.md`` §3 (tokenizer parity across transformers versions; the N1.5 bracketed
  ``repr([instruction])`` form is the only official behaviour, C4) and §4.2 (the three prompt templates).
* ``static-shape-plan.md`` §4.1 (token-count formulas), §4.2 (``S_pad`` / ``L_max`` table), §4.3 (text-pad policy:
  right padding for N1.5 / N1.7, text-slot padding for N1.6), §4.4 (position ids under padding), §4.5 (image slots).
* ``host-layout-helpers-validation.md`` §2.1-2.5 (N1.7 mRoPE position ids, interleaved ``[24, 20, 20]`` cos/sin with
  the HF ``1 / theta**x`` inverse frequencies, 48x48 pos-embed 4-tap interpolation in merge-window order, 2-D vision
  RoPE) -- the pure-torch functions below are ports of ``models/tt_dit/encoders/qwen3vl/{model,vision}_qwen3vl.py``
  and HF ``modeling_qwen3_vl.py`` that the report validated bit-for-bit against the goldens.

Tokenizers are the vendored Qwen2 BPE directories of the three reference checkouts, loaded with
``AutoTokenizer.from_pretrained(dir, trust_remote_code=False)`` (plain ``vocab.json`` + ``merges.txt``; identical ids
under transformers 4.51 / 4.57 / 5.12, host-pipeline-parity §3).  ``transformers`` is imported lazily.
"""

from __future__ import annotations

import os
import re
from dataclasses import dataclass
from functools import lru_cache
from pathlib import Path
from typing import Any, List, Optional, Sequence, Tuple

import torch

from models.experimental.gr00t.common.configs import GR00T_VERSIONS, GR00TConfig, LayoutConfig, get_config

DEFAULT_PROJECT_ROOT = Path("/home/deepgadget/experiments/gr00t")


def project_root() -> Path:
    """``$GR00T_PROJECT_ROOT`` or the project default (parent of the ``Isaac-GR00T*`` checkouts)."""
    return Path(os.environ.get("GR00T_PROJECT_ROOT", str(DEFAULT_PROJECT_ROOT)))


# host-pipeline-parity §3 / §10: the vendored tokenizer directories (no remote code needed)
_TOKENIZER_RELDIRS = {
    "n15": "Isaac-GR00T-n1.5/gr00t/model/backbone/eagle2_hg_model",
    "n16": "Isaac-GR00T-n1.6/gr00t/model/modules/nvidia/Eagle-Block2A-2B-v2",
    "n17": "Isaac-GR00T/tt_golden/backbone_shell/nvidia/Cosmos-Reason2-2B",
}


def tokenizer_dir(version: str) -> Path:
    """Directory of the vendored tokenizer (override with ``$GR00T_TOKENIZER_DIR_<N15|N16|N17>``)."""
    _check_version(version)
    env = os.environ.get(f"GR00T_TOKENIZER_DIR_{version.upper()}")
    path = Path(env) if env else project_root() / _TOKENIZER_RELDIRS[version]
    for required in ("vocab.json", "merges.txt", "tokenizer_config.json"):
        if not (path / required).is_file():
            raise FileNotFoundError(f"{version}: tokenizer dir {path} lacks {required}")
    return path


def _check_version(version: str) -> None:
    if version not in GR00T_VERSIONS:
        raise KeyError(f"unknown GR00T version {version!r}; expected one of {GR00T_VERSIONS}")


@lru_cache(maxsize=None)
def get_tokenizer(version: str):
    """Cached ``Qwen2Tokenizer`` of ``version`` (``trust_remote_code=False``; the v5 'incorrect regex pattern' warning
    is a false alarm for these tokenizers and must NOT be 'fixed', host-pipeline-parity §3)."""
    import transformers
    from transformers import AutoTokenizer

    prev = transformers.logging.get_verbosity()
    transformers.logging.set_verbosity_error()
    try:
        tok = AutoTokenizer.from_pretrained(str(tokenizer_dir(version)), trust_remote_code=False)
    finally:
        transformers.logging.set_verbosity(prev)
    cfg = get_config(version).prompt
    for name, tid in (("pad", cfg.pad_id), ("im_start", cfg.im_start_id), ("im_end", cfg.im_end_id)):
        got = tok.convert_tokens_to_ids(
            {"pad": "<|endoftext|>", "im_start": "<|im_start|>", "im_end": "<|im_end|>"}[name]
        )
        if got != tid:
            raise ValueError(f"{version}: tokenizer {name} id {got} != configs {tid}")
    return tok


# --------------------------------------------------------------------------------------------------------------------
# Language and prompt text
# --------------------------------------------------------------------------------------------------------------------

SYSTEM_USER_HEADER = "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n"
N15_GENERATION_SUFFIX = "<|im_end|>\n<|im_start|>assistant\n"
N16_SUFFIX = "<|im_end|>\n"
N17_HEADER = "<|im_start|>user\n"
N17_SUFFIX = "<|im_end|>\n"
IMG_CONTEXT = "<IMG_CONTEXT>"
IMG_START, IMG_END = "<img>", "</img>"
VISION_START, VISION_END, IMAGE_PAD = "<|vision_start|>", "<|vision_end|>", "<|image_pad|>"


def formalize_language(text: str) -> str:
    """N1.6 / N1.7 ``formalize_language=True``: lower-case and strip every character that is not a word character or
    whitespace (``processing_gr00t_n1d6.py:369-372``, ``processing_gr00t_n1d7.py:512``)."""
    return re.sub(r"[^\w\s]", "", text.lower())


def n15_language(instruction: str) -> str:
    """The text the N1.5 chat template actually renders: ``str(np.array([instruction]))`` == ``repr([instruction])``
    (``['pick ...']``; Python switches to double quotes when the instruction contains ``'``).  Golden S = 296, not
    the plain-string 294 (host-pipeline-parity §3, golden-multisample §2)."""
    return repr([str(instruction)])


def language_text(version: str, instruction: str) -> str:
    """Instruction -> the exact text segment placed in the prompt."""
    _check_version(version)
    if version == "n15":
        return n15_language(instruction)
    return formalize_language(instruction)


def image_prefix_text(version: str, k: int) -> str:
    """Tokens preceding image ``k`` (1-based): N1.5 / N1.6 ``<image k><img>`` (Eagle ``replace_media_placeholder``),
    N1.7 ``<|vision_start|>``."""
    _check_version(version)
    if version == "n17":
        return VISION_START
    return f"<image {k}>{IMG_START}"


def build_prompt(version: str, instruction: str, tokens_per_image: Sequence[int]) -> str:
    """Render the full prompt string with the image runs already expanded (host-pipeline-parity §4.2 steps 5 / 3 / 2).

    * N1.5: ``SYS user\\n`` + ``<image k><img>`` + 256 x ``<IMG_CONTEXT>`` + ``</img>`` per image + ``repr([instr])``
      + ``<|im_end|>\\n<|im_start|>assistant\\n`` (``add_generation_prompt=True``).
    * N1.6: ``SYS user\\n`` + formalized text + ``<image k><img>`` + T_k x ``<IMG_CONTEXT>`` + ``</img>`` per image +
      ``<|im_end|>\\n`` (text **before** the images, no generation prompt).
    * N1.7: ``<|im_start|>user\\n`` + (``<|vision_start|>`` + T_k x ``<|image_pad|>`` + ``<|vision_end|>``) per image
      + formalized text + ``<|im_end|>\\n`` (Qwen3-VL template; ``<|image_pad|>`` expanded to ``t*h*w/4``).
    """
    _check_version(version)
    n_img = len(tokens_per_image)
    if n_img == 0:
        raise ValueError("build_prompt: at least one image is required (text-only prompts have no golden)")
    if version in ("n15", "n16") and n_img >= 10:
        raise NotImplementedError("Eagle prompts with >= 10 images render '<image 10>' as 7 tokens; not modelled")
    lang = language_text(version, instruction)
    if version == "n17":
        images = "".join(f"{VISION_START}{IMAGE_PAD * int(t)}{VISION_END}" for t in tokens_per_image)
        return N17_HEADER + images + lang + N17_SUFFIX
    images = "".join(
        f"{image_prefix_text(version, k + 1)}{IMG_CONTEXT * int(t)}{IMG_END}" for k, t in enumerate(tokens_per_image)
    )
    if version == "n15":
        return SYSTEM_USER_HEADER + images + lang + N15_GENERATION_SUFFIX
    return SYSTEM_USER_HEADER + lang + images + N16_SUFFIX


# --------------------------------------------------------------------------------------------------------------------
# Tokenization + structural verification
# --------------------------------------------------------------------------------------------------------------------


@dataclass(frozen=True)
class TokenizedPrompt:
    """One tokenized prompt at batch 1 with its verified structure.

    ``n_text`` counts the instruction BPE tokens **including** N1.5's two bracket tokens (``['`` 677, ``']`` 660) --
    the same convention as ``LayoutConfig.l_max``; ``text_span`` is the half-open index range of those tokens in the
    unpadded sequence; ``image_slots`` are the half-open runs of the image token id.
    """

    version: str
    text: str
    input_ids: torch.Tensor  # (1, S) int64
    attention_mask: torch.Tensor  # (1, S) int64, all ones
    seq_len: int
    n_images: int
    tokens_per_image: Tuple[int, ...]
    image_slots: Tuple[Tuple[int, int], ...]
    text_span: Tuple[int, int]
    n_text: int


def token_runs(ids: torch.Tensor, token_id: int) -> Tuple[Tuple[int, int], ...]:
    """Half-open ``[start, end)`` runs of ``token_id`` in a 1-D id tensor."""
    ids = ids.reshape(-1)
    hit = ids == token_id
    if not bool(hit.any()):
        return ()
    idx = torch.nonzero(hit).reshape(-1).tolist()
    runs: List[Tuple[int, int]] = []
    start = idx[0]
    prev = idx[0]
    for i in idx[1:]:
        if i != prev + 1:
            runs.append((start, prev + 1))
            start = i
        prev = i
    runs.append((start, prev + 1))
    return tuple(runs)


def _expected_prefix_ids(version: str, k: int) -> List[int]:
    if version == "n17":
        return [get_config(version).prompt.image_open_id]
    tok = get_tokenizer(version)
    return list(tok(image_prefix_text(version, k), add_special_tokens=False)["input_ids"])


def tokenize_prompt(version: str, text: str, tokens_per_image: Sequence[int]) -> TokenizedPrompt:
    """Tokenize a prompt from :func:`build_prompt` and verify its fixed structure token by token.

    Checks (raise ``ValueError`` on any drift -- static-shape-plan §10 item 2): header ids, one run of exactly ``T_k``
    image tokens per image wrapped by its prefix / close marker, the suffix ids, and that the instruction occupies the
    single remaining gap (N1.5 / N1.7 after the images, N1.6 before them).  ``n_text`` is derived from ``S`` minus the
    fixed tokens, never by tokenizing the instruction on its own (BPE merges across the text boundary would go
    unnoticed otherwise).
    """
    _check_version(version)
    cfg = get_config(version)
    pc = cfg.prompt
    tok = get_tokenizer(version)
    enc = tok([text], return_tensors="pt", padding=True)
    ids = enc["input_ids"].to(torch.int64)
    mask = enc["attention_mask"].to(torch.int64)
    if ids.shape[0] != 1 or int(mask.sum()) != ids.shape[1]:
        raise ValueError(f"{version}: expected an unpadded batch-1 encoding, got {tuple(ids.shape)}")
    s = ids.shape[1]
    flat = ids[0].tolist()
    n_img = len(tokens_per_image)
    runs = token_runs(ids[0], pc.image_token_id)
    if len(runs) != n_img:
        raise ValueError(f"{version}: {len(runs)} image-token runs for {n_img} images: {runs}")
    for k, ((a, b), t) in enumerate(zip(runs, tokens_per_image)):
        if b - a != int(t):
            raise ValueError(f"{version}: image {k + 1} run {a}:{b} has {b - a} tokens, expected {t}")
        prefix = _expected_prefix_ids(version, k + 1)
        if flat[a - len(prefix) : a] != prefix:
            raise ValueError(f"{version}: image {k + 1} prefix {flat[a - len(prefix):a]} != {prefix}")
        if flat[b] != pc.image_close_id:
            raise ValueError(f"{version}: image {k + 1} close token {flat[b]} != {pc.image_close_id}")
    header = list(pc.header_ids)
    if flat[: len(header)] != header:
        raise ValueError(f"{version}: header ids {flat[:len(header)]} != {header}")
    suffix = list(pc.suffix_ids)
    if flat[s - len(suffix) :] != suffix:
        raise ValueError(f"{version}: suffix ids {flat[s - len(suffix):]} != {suffix}")
    first_prefix_len = len(_expected_prefix_ids(version, 1))
    if version == "n16":
        # header | text | images | suffix
        text_start = len(header)
        text_end = runs[0][0] - first_prefix_len
        if runs[-1][1] + 1 != s - len(suffix):
            raise ValueError(f"{version}: tokens between the last image and the suffix")
    else:
        # header | images | text | suffix
        if runs[0][0] - first_prefix_len != len(header):
            raise ValueError(f"{version}: tokens between the header and the first image")
        text_start = runs[-1][1] + 1
        text_end = s - len(suffix)
    for k in range(1, n_img):  # images are back to back
        gap = runs[k][0] - len(_expected_prefix_ids(version, k + 1)) - (runs[k - 1][1] + 1)
        if gap != 0:
            raise ValueError(f"{version}: {gap} unexpected tokens between images {k} and {k + 1}")
    if text_end < text_start:
        raise ValueError(f"{version}: negative text span {text_start}:{text_end}")
    n_text = text_end - text_start
    expected_s = pc.seq_len(n_img, n_text - len(pc.text_extra_ids), int(tokens_per_image[0]))
    if len(set(int(t) for t in tokens_per_image)) == 1 and expected_s != s:
        raise ValueError(f"{version}: S={s} but the §4.1 formula gives {expected_s} (n_text={n_text})")
    return TokenizedPrompt(
        version=version,
        text=text,
        input_ids=ids,
        attention_mask=mask,
        seq_len=s,
        n_images=n_img,
        tokens_per_image=tuple(int(t) for t in tokens_per_image),
        image_slots=runs,
        text_span=(text_start, text_end),
        n_text=n_text,
    )


@lru_cache(maxsize=256)
def _encode_prompt_cached(version: str, instruction: str, tokens_per_image: Tuple[int, ...]) -> TokenizedPrompt:
    return tokenize_prompt(version, build_prompt(version, instruction, tokens_per_image), tokens_per_image)


def encode_prompt(version: str, instruction: str, tokens_per_image: Sequence[int]) -> TokenizedPrompt:
    """``build_prompt`` + ``tokenize_prompt``, cached per ``(version, instruction, tokens_per_image)`` (a pure
    function; the tokenizer call is ~0.5 ms and an instruction repeats for a whole episode).  The returned
    :class:`TokenizedPrompt` -- and its tensors -- are shared between calls: read-only."""
    return _encode_prompt_cached(str(version), str(instruction), tuple(int(t) for t in tokens_per_image))


# --------------------------------------------------------------------------------------------------------------------
# Static shape key and slot layout (static-shape-plan §4.2-4.5)
# --------------------------------------------------------------------------------------------------------------------


@dataclass(frozen=True)
class ShapeKey:
    """The tuple every device shape is baked for (static-shape-plan D1 / §2): version, camera layout, text budget."""

    version: str
    n_images: int
    grid_thw: Optional[Tuple[int, int, int]]  # N1.7 ViT grid per image; None for SigLIP
    tokens_per_image: int
    s_pad: int  # S_pad_llm
    l_max: int
    layout_name: str


def select_layout(
    cfg: GR00TConfig,
    n_images: int,
    grid_thw: Optional[Tuple[int, int, int]],
    tokens_per_image: int,
    n_text: int,
    layout_name: Optional[str] = None,
) -> LayoutConfig:
    """Pick the ``LayoutConfig`` (static-shape-plan §4.2 rows) for an observation.

    With ``layout_name`` the named layout is used and checked for compatibility.  Otherwise the candidates are the
    layouts with the same camera layout (``n_images``, ``grid_thw``, ``tokens_per_image``) whose ``l_max`` admits
    ``n_text``; the one with the smallest ``s_pad_llm`` (then smallest ``l_max``) wins.  No candidate -> ``ValueError``
    (deployment asserts instead of silently recapturing, static-shape-plan D9 / §9).
    """
    if layout_name is not None:
        lay = cfg.layout(layout_name)
        if (lay.n_images, lay.grid_thw, lay.tokens_per_image) != (n_images, grid_thw, tokens_per_image):
            raise ValueError(
                f"{cfg.version}: layout {layout_name!r} is for n_images={lay.n_images}, grid={lay.grid_thw}, "
                f"T_img={lay.tokens_per_image}; observation has {n_images}, {grid_thw}, {tokens_per_image}"
            )
        if n_text > lay.l_max:
            raise ValueError(f"{cfg.version}: n_text={n_text} exceeds L_max={lay.l_max} of layout {layout_name!r}")
        return lay
    same_cams = [
        lay
        for lay in cfg.layouts
        if (lay.n_images, lay.grid_thw, lay.tokens_per_image) == (n_images, grid_thw, tokens_per_image)
    ]
    if not same_cams:
        raise ValueError(
            f"{cfg.version}: no layout for n_images={n_images}, grid={grid_thw}, T_img={tokens_per_image}; "
            f"have {[(l.name, l.n_images, l.grid_thw) for l in cfg.layouts]}"
        )
    fits = [lay for lay in same_cams if lay.l_max >= n_text]
    if not fits:
        raise ValueError(
            f"{cfg.version}: n_text={n_text} exceeds every text budget of the matching layouts "
            f"{[(l.name, l.l_max) for l in same_cams]}"
        )
    return min(fits, key=lambda l: (l.s_pad_llm, l.l_max))


@dataclass(frozen=True)
class SlotLayout:
    """Where every real token of one prompt sits in the fixed ``S_pad`` sequence (static-shape-plan §4.3-4.5)."""

    shape_key: ShapeKey
    text_padding: str  # "right" | "text_slot"
    seq_len: int  # S (real tokens)
    n_text: int
    input_ids_padded: torch.Tensor  # (1, S_pad) int64, pad id in the pad slots
    key_valid: torch.Tensor  # (1, S_pad) bool
    slot_of_token: torch.Tensor  # (S,) int64: slot index of real token i
    image_slots: Tuple[Tuple[int, int], ...]  # padded-sequence runs (== LayoutConfig.image_slots)
    text_slots: Tuple[int, int]  # real instruction tokens
    text_budget_slots: Tuple[int, int]  # the L_max region reserved for the instruction (incl. its pads)
    pad_slots: Tuple[Tuple[int, int], ...]
    position_ids: torch.Tensor  # (1, S_pad) int64: real token i -> i, k-th pad slot -> S + k


def build_slot_layout(cfg: GR00TConfig, tp: TokenizedPrompt, layout: LayoutConfig) -> SlotLayout:
    """Place a tokenized prompt into the static layout.

    * ``"right"`` (N1.5 / N1.7): real tokens at ``[0, S)``, pad id 151643 at ``[S, S_pad)`` (§4.3 item 2).
    * ``"text_slot"`` (N1.6): header + text at ``[0, 14 + n_text)``, ``L_max - n_text`` pads, then the image blocks
      and the suffix shifted so that every image slot is fixed (§4.3 item 3); trailing pads if ``S + L_max - n_text
      < S_pad``.
    Position ids: real token ``i`` keeps its reference position ``i``; the ``k``-th pad slot gets ``S + k`` (finite,
    unique, never read after masking).  The resulting image slots must equal ``layout.image_slots`` -- a mismatch
    means tokenizer / template drift and raises.
    """
    s = tp.seq_len
    s_pad = layout.s_pad_llm
    l_max = layout.l_max
    if tp.n_text > l_max:
        raise ValueError(f"{cfg.version}: n_text={tp.n_text} > L_max={l_max} ({layout.name})")
    if layout.text_padding == "right":
        if s > s_pad:
            raise ValueError(f"{cfg.version}: S={s} > S_pad={s_pad} ({layout.name})")
        slot_of_token = torch.arange(s, dtype=torch.int64)
        pad_slots: Tuple[Tuple[int, int], ...] = ((s, s_pad),) if s < s_pad else ()
        t0, t1 = tp.text_span
        text_budget = (t0, t0 + l_max)
    elif layout.text_padding == "text_slot":
        t0, t1 = tp.text_span
        shift = l_max - tp.n_text
        total = s + shift
        if total > s_pad:
            raise ValueError(f"{cfg.version}: S + (L_max - n_text) = {total} > S_pad={s_pad} ({layout.name})")
        slot_of_token = torch.cat([torch.arange(t1, dtype=torch.int64), torch.arange(t1, s, dtype=torch.int64) + shift])
        pad_slots = tuple(r for r in ((t1, t1 + shift), (total, s_pad)) if r[1] > r[0])
        text_budget = (t0, t0 + l_max)
    else:
        raise ValueError(f"unknown text_padding policy {layout.text_padding!r}")

    ids_padded = torch.full((s_pad,), cfg.prompt.pad_id, dtype=torch.int64)
    ids_padded[slot_of_token] = tp.input_ids[0]
    key_valid = torch.zeros(s_pad, dtype=torch.bool)
    key_valid[slot_of_token] = True
    if int(key_valid.sum()) != s:
        raise AssertionError("slot placement is not injective")
    position_ids = torch.zeros(s_pad, dtype=torch.int64)
    position_ids[slot_of_token] = torch.arange(s, dtype=torch.int64)
    pad_idx = torch.nonzero(~key_valid).reshape(-1)
    position_ids[pad_idx] = s + torch.arange(pad_idx.numel(), dtype=torch.int64)

    image_slots = tuple((int(slot_of_token[a]), int(slot_of_token[b - 1]) + 1) for a, b in tp.image_slots)
    if image_slots != tuple(tuple(r) for r in layout.image_slots):
        raise ValueError(
            f"{cfg.version}: image slots {image_slots} != layout {layout.name!r} {layout.image_slots} "
            "(tokenizer / template drift?)"
        )
    text_slots = (
        (int(slot_of_token[t0]), int(slot_of_token[t1 - 1]) + 1) if t1 > t0 else (int(slot_of_token[t0 - 1]) + 1,) * 2
    )
    key = ShapeKey(
        version=cfg.version,
        n_images=layout.n_images,
        grid_thw=layout.grid_thw,
        tokens_per_image=layout.tokens_per_image,
        s_pad=s_pad,
        l_max=l_max,
        layout_name=layout.name,
    )
    return SlotLayout(
        shape_key=key,
        text_padding=layout.text_padding,
        seq_len=s,
        n_text=tp.n_text,
        input_ids_padded=ids_padded[None],
        key_valid=key_valid[None],
        slot_of_token=slot_of_token,
        image_slots=image_slots,
        text_slots=text_slots,
        text_budget_slots=text_budget,
        pad_slots=pad_slots,
        position_ids=position_ids[None],
    )


# --------------------------------------------------------------------------------------------------------------------
# LLM rotary tables (fp32, rotate-half layout: emb = cat(freqs, freqs))
# --------------------------------------------------------------------------------------------------------------------


# dtype of the `inv_freq` buffer the fp32 golden of each version was produced with (foundation-f2-host.md §4):
#   n16: Gr00tPolicy does `model.to(dtype=bfloat16)` and the dump then calls `model.float()`, so the non-persistent
#        `inv_freq` buffer is bf16-ROUNDED in the fp32 golden (and in deployment); the fp32-inv_freq table is 0.13 off
#        `llm_rotary_cos` at position 115, the bf16-rounded one 6e-8 -- the 4e-3 tap gate needs the rounding.
#   n17: the fp32 dump builds the model directly in fp32 -> fp32 `inv_freq` (golden matched to 6e-8).
#   n15: no rotary tap in the golden; `GR00T_N1_5.from_pretrained(torch_dtype=bf16)` creates the buffer with an
#        explicit fp32 dtype and nothing casts the module, so fp32 is the best-supported choice (UNVERIFIED).
LLM_ROPE_INV_FREQ_DTYPE = {"n15": torch.float32, "n16": torch.bfloat16, "n17": torch.float32}


def rope_inv_freq(head_dim: int, theta: float, inv_freq_dtype: torch.dtype = torch.float32) -> torch.Tensor:
    """HF ``compute_default_rope_parameters``: ``1 / theta ** (arange(0, d, 2) / d)`` in fp32, optionally rounded
    through ``inv_freq_dtype`` (bf16 when the reference model's buffer was cast, see ``LLM_ROPE_INV_FREQ_DTYPE``) and
    returned as fp32.  Use this form, not ``theta ** -x``: the two differ by one fp32 ulp and the bf16-rounded tables
    then disagree with the golden on 0.05-0.1 % of the entries (host-layout-helpers §2.2)."""
    inv = 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.int64).to(torch.float32) / head_dim))
    if inv_freq_dtype != torch.float32:
        inv = inv.to(inv_freq_dtype).to(torch.float32)
    return inv


def rope_cos_sin_1d(
    position_ids: torch.Tensor, head_dim: int, theta: float, inv_freq_dtype: torch.dtype = torch.float32
) -> Tuple[torch.Tensor, torch.Tensor]:
    """1-D RoPE table for N1.5 / N1.6 (``Qwen3RotaryEmbedding.forward``: fp32 ``inv_freq @ pos``, ``cat(freqs,
    freqs)``, cos / sin): ``(S,) -> cos, sin (S, head_dim)`` fp32.  Pass ``inv_freq_dtype=torch.bfloat16`` to
    reproduce a model whose buffer was cast to bf16 (N1.6 golden and deployment)."""
    pos = position_ids.reshape(-1).to(torch.float32)
    inv = rope_inv_freq(head_dim, theta, inv_freq_dtype)
    freqs = pos[:, None] * inv[None, :]  # == (inv_freq[:, None] @ pos[None, :]).T (K=1 matmul is one product)
    emb = torch.cat((freqs, freqs), dim=-1)
    return emb.cos(), emb.sin()


def vision_llm_position_ids(start: int, grid_thw: Sequence[int], spatial_merge_size: int = 2) -> torch.Tensor:
    """M-RoPE ``(3, t * h/m * w/m)`` grid of one image block offset by ``start`` (HF ``get_vision_position_ids`` /
    tt_dit ``vision_position_ids``): ``T = start``, ``H = start + row``, ``W = start + col`` with ``W`` fastest."""
    t, h, w = (int(x) for x in grid_thw)
    gh, gw = h // spatial_merge_size, w // spatial_merge_size
    if gh * spatial_merge_size != h or gw * spatial_merge_size != w:
        raise ValueError(f"grid {grid_thw} is not divisible by the merge size {spatial_merge_size}")
    pos_t = (torch.arange(t) * 1).repeat_interleave(gh * gw) + start
    pos_h = (torch.arange(gh) + start).repeat_interleave(gw).repeat(t)
    pos_w = (torch.arange(gw) + start).repeat(gh * t)
    return torch.stack([pos_t, pos_h, pos_w], dim=0).to(torch.int64)


def mrope_position_ids(
    input_ids: torch.Tensor, image_token_id: int, image_grid_thw: torch.Tensor, spatial_merge_size: int = 2
) -> torch.Tensor:
    """Pure-torch ``Qwen3VLModel.get_rope_index`` for one unpadded (or right-padded) sequence -> ``(3, S)`` int64.

    Text runs advance all three axes by their length; each image run consumes the next ``image_grid_thw`` row and
    advances the counter by ``max(h, w) // merge`` (``modeling_qwen3_vl.py`` ``get_rope_index``; identical to tt_dit
    ``mrope_position_ids`` and exact vs ``golden/n17/.../llm_position_ids_mrope.pt``, host-layout-helpers §2.1).
    Pads appended after the prompt are ordinary text tokens here, so they continue the counter.
    """
    ids = input_ids.reshape(-1)
    is_img = (ids == image_token_id).tolist()
    grids = [tuple(int(x) for x in row) for row in image_grid_thw.tolist()]
    segments: List[torch.Tensor] = []
    current = 0
    i = 0
    img_idx = 0
    n = len(is_img)
    while i < n:
        j = i
        while j < n and is_img[j] == is_img[i]:
            j += 1
        if not is_img[i]:
            seg = torch.arange(j - i, dtype=torch.int64).view(1, -1).expand(3, -1) + current
            current += j - i
        else:
            if img_idx >= len(grids):
                raise ValueError(f"more image runs than image_grid_thw rows ({len(grids)})")
            g = grids[img_idx]
            img_idx += 1
            n_tok = g[0] * (g[1] // spatial_merge_size) * (g[2] // spatial_merge_size)
            if n_tok != j - i:
                raise ValueError(f"image run of {j - i} tokens but grid {g} merges to {n_tok}")
            seg = vision_llm_position_ids(current, g, spatial_merge_size)
            current += max(g[1], g[2]) // spatial_merge_size
        segments.append(seg)
        i = j
    if img_idx != len(grids):
        raise ValueError(f"{len(grids)} image_grid_thw rows but only {img_idx} image runs in the ids")
    return torch.cat(segments, dim=1)


def apply_interleaved_mrope(freqs: torch.Tensor, mrope_section: Sequence[int]) -> torch.Tensor:
    """``(3, S, d/2)`` per-axis freqs -> ``(S, d/2)``: slots ``0,3,6,..`` T, ``1,4,..`` H, ``2,5,..`` W up to
    ``3 * section[axis]``, the tail stays T (``Qwen3VLTextRotaryEmbedding.apply_interleaved_mrope``; for [24,20,20]:
    ``[T,H,W] x 20`` then ``T,T,T,T``)."""
    out = freqs[0].clone()
    for axis, offset in ((1, 1), (2, 2)):
        length = int(mrope_section[axis]) * 3
        out[..., offset:length:3] = freqs[axis, ..., offset:length:3]
    return out


def mrope_cos_sin(
    position_ids: torch.Tensor,
    head_dim: int,
    theta: float,
    mrope_section: Sequence[int],
    inv_freq_dtype: torch.dtype = torch.float32,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Interleaved M-RoPE table for N1.7: ``(3, S) -> cos, sin (S, head_dim)`` fp32, rotate-half layout
    (``cos[:, :64] == cos[:, 64:]``).  HF inv_freq form -> within 6e-8 of ``llm_rotary_cos/sin`` and 100 % identical
    after bf16 rounding (host-layout-helpers §1.1 #4, §2.2)."""
    if position_ids.ndim != 2 or position_ids.shape[0] != 3:
        raise ValueError(f"mrope position ids must be (3, S), got {tuple(position_ids.shape)}")
    if sum(int(x) for x in mrope_section) != head_dim // 2:
        raise ValueError(f"mrope_section {tuple(mrope_section)} must sum to head_dim/2 = {head_dim // 2}")
    inv = rope_inv_freq(head_dim, theta, inv_freq_dtype)
    freqs = position_ids.to(torch.float32)[:, :, None] * inv[None, None, :]  # (3, S, d/2)
    freqs = apply_interleaved_mrope(freqs, mrope_section)
    emb = torch.cat((freqs, freqs), dim=-1)
    return emb.cos(), emb.sin()


# --------------------------------------------------------------------------------------------------------------------
# N1.7 vision host tables (per grid; merge-window patch order, host-layout-helpers §2.3-2.5)
# --------------------------------------------------------------------------------------------------------------------


def vision_bilinear_indices_and_weights(
    grid_thw: torch.Tensor, num_grid_per_side: int = 48, spatial_merge_size: int = 2
) -> Tuple[torch.Tensor, torch.Tensor]:
    """4-tap bilinear taps of the learned ``48x48`` position table for every patch, in merge-window order:
    ``indices (4, P) int64``, ``weights (4, P) fp32`` (HF ``fast_pos_embed_interpolate`` /
    ``vision_utils.get_vision_bilinear_indices_and_weights``; bitwise equal to tt_dit ``vision_pos_embeds``)."""
    side = num_grid_per_side
    m = spatial_merge_size
    idx_parts: List[List[torch.Tensor]] = [[] for _ in range(4)]
    w_parts: List[List[torch.Tensor]] = [[] for _ in range(4)]
    for t, h, w in grid_thw.tolist():
        t, h, w = int(t), int(h), int(w)
        if h % m or w % m:
            raise ValueError(f"grid ({t},{h},{w}) not divisible by merge size {m}")
        h_grid = torch.linspace(0, side - 1, h)
        w_grid = torch.linspace(0, side - 1, w)
        h_floor = h_grid.int()
        w_floor = w_grid.int()
        h_ceil = (h_floor + 1).clamp(max=side - 1)
        w_ceil = (w_floor + 1).clamp(max=side - 1)
        h_frac = h_grid - h_floor
        w_frac = w_grid - w_floor
        h_floor_off = h_floor * side
        h_ceil_off = h_ceil * side
        corner_idx = [
            (h_floor_off[:, None] + w_floor[None, :]).flatten(),
            (h_floor_off[:, None] + w_ceil[None, :]).flatten(),
            (h_ceil_off[:, None] + w_floor[None, :]).flatten(),
            (h_ceil_off[:, None] + w_ceil[None, :]).flatten(),
        ]
        corner_w = [
            ((1 - h_frac)[:, None] * (1 - w_frac)[None, :]).flatten(),
            ((1 - h_frac)[:, None] * w_frac[None, :]).flatten(),
            (h_frac[:, None] * (1 - w_frac)[None, :]).flatten(),
            (h_frac[:, None] * w_frac[None, :]).flatten(),
        ]
        h_idx = torch.arange(h).view(h // m, m)
        w_idx = torch.arange(w).view(w // m, m)
        reorder = (h_idx[:, :, None, None] * w + w_idx[None, None, :, :]).transpose(1, 2).flatten().repeat(t)
        for i in range(4):
            idx_parts[i].append(corner_idx[i][reorder].to(torch.int64))
            w_parts[i].append(corner_w[i][reorder].to(torch.float32))
    return torch.stack([torch.cat(p) for p in idx_parts]), torch.stack([torch.cat(p) for p in w_parts])


def vision_pos_embed(pos_embed_weight: torch.Tensor, indices: torch.Tensor, weights: torch.Tensor) -> torch.Tensor:
    """``(P, hidden)`` position embedding added right after the N1.7 patch linear: ``sum_i w_i * table[idx_i]``."""
    if pos_embed_weight.ndim != 2 or indices.shape != weights.shape or indices.shape[0] != 4:
        raise ValueError("vision_pos_embed: expected table (N, C), indices/weights (4, P)")
    return (pos_embed_weight[indices].to(torch.float32) * weights[:, :, None]).sum(0)


def vision_rope_position_ids(grid_thw: torch.Tensor, spatial_merge_size: int = 2) -> torch.Tensor:
    """``(P, 2)`` ``(row, col)`` of every patch in merge-window order (HF ``rot_pos_emb`` /
    ``vision_utils.get_vision_position_ids``)."""
    m = spatial_merge_size
    out = []
    for t, h, w in grid_thw.tolist():
        t, h, w = int(t), int(h), int(w)
        hpos = torch.arange(h).unsqueeze(1).expand(-1, w).reshape(h // m, m, w // m, m).transpose(1, 2).flatten()
        wpos = torch.arange(w).unsqueeze(0).expand(h, -1).reshape(h // m, m, w // m, m).transpose(1, 2).flatten()
        out.append(torch.stack([hpos, wpos], dim=-1).repeat(t, 1))
    return torch.cat(out, dim=0).to(torch.int64)


def vision_rope_cos_sin(
    grid_thw: torch.Tensor, head_dim: int = 64, spatial_merge_size: int = 2, rope_theta: float = 10000.0
) -> Tuple[torch.Tensor, torch.Tensor]:
    """2-D vision RoPE ``cos, sin (P, head_dim)`` fp32: 16 frequencies per axis, ``emb = cat(freqs, freqs)`` with
    ``freqs = [row * inv, col * inv]`` (``Qwen3VLVisionRotaryEmbedding(head_dim // 2)`` + ``rot_pos_emb``; exact vs
    HF, host-layout-helpers §2.5)."""
    pos = vision_rope_position_ids(grid_thw, spatial_merge_size)
    dim = head_dim // 2
    inv = 1.0 / (rope_theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
    freqs = (pos.to(torch.float32).unsqueeze(-1) * inv).flatten(1)
    emb = torch.cat((freqs, freqs), dim=-1)
    return emb.cos(), emb.sin()


def vision_cu_seqlens(grid_thw: torch.Tensor) -> Tuple[int, ...]:
    """Attention block boundaries: one block of ``h * w`` patches per frame (``[0, 448, 896, 1344, 1792]`` for DROID)."""
    bounds = [0]
    for t, h, w in grid_thw.tolist():
        for _ in range(int(t)):
            bounds.append(bounds[-1] + int(h) * int(w))
    return tuple(bounds)


# --------------------------------------------------------------------------------------------------------------------
# Per-call caches (pure functions of hashable keys; the tensors they return are shared -> read-only)
# --------------------------------------------------------------------------------------------------------------------

GridKey = Tuple[Tuple[int, int, int], ...]


def grid_key(grid_thw: torch.Tensor) -> GridKey:
    """Hashable ``((t, h, w), ...)`` of an ``image_grid_thw`` tensor (the cache key of the per-grid tables)."""
    g = torch.as_tensor(grid_thw)
    if g.ndim != 2 or g.shape[1] != 3:
        raise ValueError(f"image_grid_thw must be (n_img, 3), got {tuple(g.shape)}")
    return tuple((int(t), int(h), int(w)) for t, h, w in g.tolist())


@lru_cache(maxsize=32)
def vision_tables_cached(
    grid: GridKey, num_grid_per_side: int, spatial_merge_size: int, head_dim: int, rope_theta: float
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, Tuple[int, ...]]:
    """``(pos_embed_indices (4, P), pos_embed_weights (4, P), rope_cos (P, d), rope_sin (P, d), cu_seqlens)`` of one
    grid -- the N1.7 vision host tables, constant per static layout (:func:`vision_bilinear_indices_and_weights`,
    :func:`vision_rope_cos_sin`, :func:`vision_cu_seqlens`; 0.7 ms per call uncached)."""
    g = torch.tensor(grid, dtype=torch.int64)
    idx, wts = vision_bilinear_indices_and_weights(g, num_grid_per_side, spatial_merge_size)
    cos, sin = vision_rope_cos_sin(g, head_dim, spatial_merge_size, rope_theta)
    return idx, wts, cos, sin, vision_cu_seqlens(g)


@lru_cache(maxsize=256)
def mrope_tables_cached(
    input_ids_padded: Tuple[int, ...],
    image_token_id: int,
    grid: GridKey,
    spatial_merge_size: int,
    head_dim: int,
    theta: float,
    mrope_section: Tuple[int, ...],
    inv_freq_dtype: torch.dtype,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """``(position_ids (3, S_pad), cos (S_pad, d), sin (S_pad, d))`` of one padded id sequence and grid
    (:func:`mrope_position_ids` + :func:`mrope_cos_sin`, 0.3 ms per call uncached; the ids repeat for a whole
    episode of one instruction)."""
    ids = torch.tensor(input_ids_padded, dtype=torch.int64)
    g = torch.tensor(grid, dtype=torch.int64)
    pos3 = mrope_position_ids(ids, image_token_id, g, spatial_merge_size)
    cos, sin = mrope_cos_sin(pos3, head_dim, theta, mrope_section, inv_freq_dtype)
    return pos3, cos, sin


__all__ = [
    "IMG_CONTEXT",
    "LLM_ROPE_INV_FREQ_DTYPE",
    "SYSTEM_USER_HEADER",
    "ShapeKey",
    "SlotLayout",
    "TokenizedPrompt",
    "apply_interleaved_mrope",
    "build_prompt",
    "build_slot_layout",
    "encode_prompt",
    "formalize_language",
    "get_tokenizer",
    "grid_key",
    "image_prefix_text",
    "language_text",
    "mrope_cos_sin",
    "mrope_position_ids",
    "mrope_tables_cached",
    "n15_language",
    "rope_cos_sin_1d",
    "rope_inv_freq",
    "select_layout",
    "token_runs",
    "tokenize_prompt",
    "tokenizer_dir",
    "vision_bilinear_indices_and_weights",
    "vision_cu_seqlens",
    "vision_llm_position_ids",
    "vision_pos_embed",
    "vision_rope_cos_sin",
    "vision_rope_position_ids",
    "vision_tables_cached",
]