# SPDX-License-Identifier: Apache-2.0
"""MiniMax H3 packed-sequence materialization for FL2VA and Ref2VA tasks.

Layout: [text L | imgvid_cond C | audio A(=t*2ch) | video_target V | pad P].
Builder rules:
- placeholder ids; block-derived pos infos/masks/cu_seqlens/document_id
- img_position_ids fp64 grid: text rows (row_idx,0,0); video/cond t counter
  continues text_len with temporal interp spans (frame_rescale 5/3 x
  frame_per_token (1,4,4,4,4)); each spatial sqrt_area axis uses evenly spaced
  coordinates excluding the right endpoint, then scales them by INTERP;
  audio channel-major blocks pinned to the w-grid extremes.
"""

from __future__ import annotations

from collections.abc import Mapping, Sequence

import numpy as np
import torch

MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES: tuple[tuple[int, ...], ...] = (
    (0,),
    (-1,),
    (0, -1),
)

MINIMAX_H3_TEXT_ID = -5
MINIMAX_H3_IMGVID_COND_ID = -11
MINIMAX_H3_AUDIO_REF_COND_ID = -17
MINIMAX_H3_AUDIO_FIRST_ID = -15
MINIMAX_H3_AUDIO_ID = -14
MINIMAX_H3_VIDEO_FIRST_ID = -3
MINIMAX_H3_VIDEO_ID = -2
MINIMAX_H3_VIDEO_LAST_ID = -4
MINIMAX_H3_PAD_ID = -1

_INTERP = 32
_T_GROUP = 5
_FRAME_PER_TOKEN = (1, 4, 4, 4, 4)
_FRAME_RESCALE = 5.0 / 3.0
_SEQ_ALIGN = 64
# MiniMax H3 packs video latents with a fixed [1, 2, 2] (t, h, w) patch.
_PATCH_T = 1
_PATCH_H = 2
_PATCH_W = 2


def _keyframe_cond_frame_indices(
    *,
    include_keyframe_cond: bool,
    keyframe_frame_indices: list[int] | tuple[int, ...] | None,
) -> list[int]:
    if not include_keyframe_cond:
        if keyframe_frame_indices is not None:
            raise ValueError("keyframe_frame_indices must be omitted when keyframe cond is not included")
        return []
    if keyframe_frame_indices is None:
        raise ValueError("strict fl2va packed layout requires keyframe_frame_indices")
    if any(isinstance(value, bool) or not isinstance(value, int) for value in keyframe_frame_indices):
        raise ValueError("strict fl2va packed layout requires integer keyframe_frame_indices")
    out = list(keyframe_frame_indices)
    if tuple(out) not in MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES:
        raise ValueError(
            "strict fl2va packed layout requires keyframe_frame_indices in "
            f"{MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES!r}, got {out!r}"
        )
    return out


def _resolve_keyframe_frame_indices(
    frame_indices: Sequence[int],
    *,
    frame_count: int | None,
) -> list[int]:
    if frame_indices and frame_count is None:
        raise ValueError("frame_count is required when keyframe_frame_indices are provided")
    if frame_count is None:
        return []
    if frame_count <= 0:
        raise ValueError("frame_count must be positive")
    seen: dict[int, int] = {}
    resolved: list[int] = []
    for block_index, semantic_index in enumerate(frame_indices):
        if semantic_index == -1:
            resolved_index = frame_count - 1
        elif 0 <= semantic_index < frame_count:
            resolved_index = semantic_index
        else:
            raise ValueError(f"keyframe frame index {semantic_index} must be -1 or in [0, {frame_count})")
        previous = seen.get(resolved_index)
        if previous is not None:
            raise ValueError(
                f"keyframe frame index at block {block_index} resolves to "
                f"{resolved_index}, already bound by block {previous}"
            )
        seen[resolved_index] = block_index
        resolved.append(resolved_index)
    return resolved


def _temporal_position_span(temporal_length: int) -> float:
    """Temporal position span for patch_t=1, in fp64.

    NOTE: intentionally NOT merged with ``_video_t_span``. This variant sums
    via numpy (pairwise summation), matching the fl2va anchor computation,
    while ``_video_t_span`` sums sequentially, matching the ref2va
    t-origin accumulation. The two orders diverge in the last ulp
    from n=16 onward, so each path must keep its own summation order.
    """
    spans = np.ones(int(temporal_length), dtype=np.float64) * _FRAME_RESCALE
    for token_index in range(_T_GROUP):
        spans[token_index::_T_GROUP] *= _FRAME_PER_TOKEN[token_index]
    return float(spans.sum())


def minimax_h3_packed_sequence(
    *,
    text_len: int,
    latent_t: int,
    latent_h: int,
    latent_w: int,
    audio_t: int,
    audio_channel: int = 2,
    include_keyframe_cond: bool,
    keyframe_frame_indices: list[int] | tuple[int, ...] | None = None,
    frame_count: int | None = None,
) -> dict[str, torch.Tensor]:
    """Build the packed-sequence structural fields for one CFG branch.

    The used length is padded up to a multiple of 64.
    """
    ph, pw = latent_h // _PATCH_H, latent_w // _PATCH_W
    frame_rows = ph * pw
    cond_frame_indices = _keyframe_cond_frame_indices(
        include_keyframe_cond=include_keyframe_cond,
        keyframe_frame_indices=keyframe_frame_indices,
    )
    resolved_cond_frame_indices = _resolve_keyframe_frame_indices(
        cond_frame_indices,
        frame_count=frame_count,
    )
    cond_rows = len(cond_frame_indices) * frame_rows
    video_rows = latent_t * frame_rows
    audio_rows = audio_t * audio_channel
    used = text_len + cond_rows + audio_rows + video_rows
    seq_len = ((used + _SEQ_ALIGN - 1) // _SEQ_ALIGN) * _SEQ_ALIGN

    text_sl = slice(0, text_len)
    cond_sl = slice(text_len, text_len + cond_rows)
    audio_sl = slice(cond_sl.stop, cond_sl.stop + audio_rows)
    video_sl = slice(audio_sl.stop, audio_sl.stop + video_rows)
    pad_sl = slice(video_sl.stop, seq_len)

    input_ids = torch.full((seq_len,), MINIMAX_H3_PAD_ID, dtype=torch.int64)
    input_ids[text_sl] = MINIMAX_H3_TEXT_ID
    if cond_rows:
        input_ids[cond_sl] = MINIMAX_H3_IMGVID_COND_ID
    input_ids[audio_sl] = MINIMAX_H3_AUDIO_ID
    input_ids[audio_sl.start] = MINIMAX_H3_AUDIO_FIRST_ID
    input_ids[video_sl] = MINIMAX_H3_VIDEO_ID
    input_ids[video_sl.start] = MINIMAX_H3_VIDEO_FIRST_ID
    input_ids[video_sl.stop - 1] = MINIMAX_H3_VIDEO_LAST_ID

    image_mask = torch.zeros(seq_len, dtype=torch.bool)
    image_mask[cond_sl] = True
    image_mask[video_sl] = True
    audio_mask = torch.zeros(seq_len, dtype=torch.bool)
    audio_mask[audio_sl] = True

    img_pos = torch.cat(
        [
            torch.arange(cond_sl.start, cond_sl.stop),
            torch.arange(video_sl.start, video_sl.stop),
        ]
    )
    update_mask = torch.zeros(img_pos.shape[0], dtype=torch.bool)
    update_mask[cond_rows:] = True
    audio_pos = torch.arange(audio_sl.start, audio_sl.stop)
    text_pos = torch.arange(0, text_len)

    g = torch.zeros(seq_len, 3, dtype=torch.float64)
    g[text_sl, 0] = torch.arange(text_len, dtype=torch.float64)

    t_grid = _video_t_grid(latent_t, float(text_len))
    sqrt_area = np.sqrt(latent_h * latent_w)
    h_grid = _axis_from_sqrt_area(latent_h, _PATCH_H, sqrt_area)
    w_grid = _axis_from_sqrt_area(latent_w, _PATCH_W, sqrt_area)
    hh, ww = torch.meshgrid(h_grid, w_grid, indexing="ij")
    frame = torch.stack([hh.reshape(-1), ww.reshape(-1)], dim=-1)
    video_g = torch.empty(latent_t, frame_rows, 3, dtype=torch.float64)
    video_g[:, :, 0] = t_grid[:, None]
    video_g[:, :, 1:] = frame[None]
    for block_index, pixel_index in enumerate(resolved_cond_frame_indices):
        sl = slice(
            cond_sl.start + block_index * frame_rows,
            cond_sl.start + (block_index + 1) * frame_rows,
        )
        if pixel_index == 0:
            cond_t = float(text_len)
        elif frame_count is not None and pixel_index == frame_count - 1:
            cond_t = float(text_len) + _temporal_position_span(latent_t) - _FRAME_RESCALE
        else:
            raise ValueError(
                f"fl2va packed layout only supports first/last keyframe anchors, got resolved frame index {pixel_index}"
            )
        g[sl, 0] = cond_t
        g[sl, 1:] = frame
    g[video_sl] = video_g.reshape(-1, 3)
    audio_t_grid = float(text_len) + torch.arange(audio_t, dtype=torch.float64)
    g[audio_sl, 0] = audio_t_grid.repeat(audio_channel)
    g[audio_sl, 2] = torch.cat(
        [
            torch.full((audio_t,), float(w_grid[0]), dtype=torch.float64),
            torch.full((audio_rows - audio_t,), float(w_grid[-1]), dtype=torch.float64),
        ]
    )

    token_tags = torch.full((seq_len,), -1, dtype=torch.long)  # PADDING
    token_tags[text_sl] = 1  # TEXT (fl2va image-segment override happens upstream)
    token_tags[audio_sl] = 2  # AUDIO
    token_tags[img_pos] = 0  # VIDEO

    cu = torch.tensor([0, used, seq_len], dtype=torch.int32)
    doc = torch.zeros(seq_len, dtype=torch.int32)
    doc[pad_sl] = 1
    return {
        "seq_len": torch.tensor(seq_len),
        "input_ids": input_ids,
        "image_mask": image_mask,
        "audio_mask": audio_mask,
        "img_pos": img_pos,
        "audio_pos": audio_pos,
        "text_pos": text_pos,
        "update_mask": update_mask,
        "img_position_ids": g,
        "token_tags": token_tags,
        "cu_seqlens": cu,
        "document_id": doc,
        # Extent of the video segment within the packed sequence.
        "latent_grid": torch.tensor([latent_t, ph, pw], dtype=torch.int64),
        "video_row_start": torch.tensor(video_sl.start, dtype=torch.int64),
    }


def _positive_int(
    block: Mapping[str, object],
    key: str,
    path: str,
    *,
    allow_zero: bool = False,
) -> int:
    value = block.get(key)
    if isinstance(value, bool) or not isinstance(value, int):
        raise ValueError(f"{path}.{key} must be an integer")
    if value < 0 or (value == 0 and not allow_zero):
        predicate = "non-negative" if allow_zero else "positive"
        raise ValueError(f"{path}.{key} must be {predicate}")
    return int(value)


def _axis_from_sqrt_area(dim: int, patch: int, sqrt_area: float) -> torch.Tensor:
    ratio = dim / sqrt_area
    left = (1.0 - ratio) * 1.0 / 2.0
    right = left + ratio * 1.0
    grid = np.linspace(left, right, dim // patch, endpoint=False) * _INTERP
    return torch.from_numpy(grid).to(torch.float64)


def _video_t_grid(n: int, origin: float) -> torch.Tensor:
    spans = torch.tensor(
        [_FRAME_RESCALE * _FRAME_PER_TOKEN[k % _T_GROUP] for k in range(n)],
        dtype=torch.float64,
    )
    return origin + torch.cat([torch.zeros(1, dtype=torch.float64), spans[:-1].cumsum(0)])


def _video_t_span(n: int) -> float:
    # Sequential fp64 summation on purpose — see _temporal_position_span for
    # why the two span implementations must not be unified.
    return sum(_FRAME_RESCALE * _FRAME_PER_TOKEN[k % _T_GROUP] for k in range(n))


def _range_for_slice(sl: slice) -> torch.Tensor:
    return torch.arange(sl.start, sl.stop, dtype=torch.long)


def _cat_ranges(parts: list[torch.Tensor]) -> torch.Tensor:
    if parts:
        return torch.cat(parts)
    return torch.empty(0, dtype=torch.long)


def minimax_h3_packed_sequence_ref2va_blocks(
    *,
    text_len: int,
    latent_t: int,
    latent_h: int,
    latent_w: int,
    audio_t: int,
    ref_blocks: Sequence[Mapping[str, object]],
    audio_channel: int = 2,
    seq_len: int | None = None,
) -> dict[str, torch.Tensor]:
    """General ref2va-family packed layout.

    ``ref_blocks`` are consumed in request/plan order:
    - ``{"kind": "image", "latent_h": H, "latent_w": W}``
    - ``{"kind": "audio", "ref_audio_t": T}``
    - ``{"kind": "video"|"video_audio", "ref_audio_t": T,
       "latent_t": RT, "latent_h": RH, "latent_w": RW}``

    Video-bearing blocks pack their audio rows immediately before their video
    rows; both share the same temporal origin and advance by the longer of the
    audio and video spans. Standalone audio advances the target origin by its
    own T, and image blocks advance it by one integer slot.
    """
    if not isinstance(ref_blocks, Sequence) or isinstance(ref_blocks, (str, bytes)):
        raise ValueError("ref_blocks must be a sequence")

    parsed: list[dict[str, object]] = []
    ref_visual_rows = 0
    ref_audio_rows = 0
    for index, raw in enumerate(ref_blocks):
        path = f"ref_blocks[{index}]"
        if not isinstance(raw, Mapping):
            raise ValueError(f"{path} must be an object")
        kind = raw.get("kind", raw.get("type"))
        if not isinstance(kind, str) or not kind:
            raise ValueError(f"{path}.kind must be a non-empty string")
        if kind == "image":
            rh = _positive_int(raw, "latent_h", path)
            rw = _positive_int(raw, "latent_w", path)
            rows = (rh // _PATCH_H) * (rw // _PATCH_W)
            item = {"kind": kind, "latent_h": rh, "latent_w": rw, "rows": rows}
            ref_visual_rows += rows
        elif kind == "audio":
            rt = _positive_int(raw, "ref_audio_t", path, allow_zero=True)
            rows = rt * audio_channel
            item = {"kind": kind, "ref_audio_t": rt, "audio_rows": rows}
            ref_audio_rows += rows
        elif kind in ("video", "video_audio"):
            rt = _positive_int(raw, "ref_audio_t", path, allow_zero=True)
            vt = _positive_int(raw, "latent_t", path)
            vh = _positive_int(raw, "latent_h", path)
            vw = _positive_int(raw, "latent_w", path)
            frame_rows = (vh // _PATCH_H) * (vw // _PATCH_W)
            audio_rows = rt * audio_channel
            video_rows = vt * frame_rows
            item = {
                "kind": kind,
                "ref_audio_t": rt,
                "latent_t": vt,
                "latent_h": vh,
                "latent_w": vw,
                "frame_rows": frame_rows,
                "audio_rows": audio_rows,
                "video_rows": video_rows,
            }
            ref_audio_rows += audio_rows
            ref_visual_rows += video_rows
        else:
            raise ValueError(f"{path}.kind unsupported for ref2va: {kind!r}")
        parsed.append(item)

    ph, pw = latent_h // _PATCH_H, latent_w // _PATCH_W
    frame_rows = ph * pw
    video_rows = latent_t * frame_rows
    audio_rows = audio_t * audio_channel
    ref_rows = ref_visual_rows + ref_audio_rows
    used = text_len + ref_rows + audio_rows + video_rows
    if seq_len is None:
        seq_len = ((used + _SEQ_ALIGN - 1) // _SEQ_ALIGN) * _SEQ_ALIGN
    if seq_len < used:
        raise ValueError(f"seq_len {seq_len} < used rows {used}")

    text_sl = slice(0, text_len)
    cursor = text_len
    block_slices: list[dict[str, object]] = []
    for item in parsed:
        kind = str(item["kind"])
        if kind == "image":
            rows = int(item["rows"])
            visual_sl = slice(cursor, cursor + rows)
            cursor = visual_sl.stop
            block_slices.append({**item, "visual_sl": visual_sl})
        elif kind == "audio":
            rows = int(item["audio_rows"])
            audio_sl = slice(cursor, cursor + rows)
            cursor = audio_sl.stop
            block_slices.append({**item, "audio_sl": audio_sl})
        else:
            a_rows = int(item["audio_rows"])
            v_rows = int(item["video_rows"])
            audio_sl = slice(cursor, cursor + a_rows)
            visual_sl = slice(audio_sl.stop, audio_sl.stop + v_rows)
            cursor = visual_sl.stop
            block_slices.append({**item, "audio_sl": audio_sl, "visual_sl": visual_sl})

    audio_sl = slice(cursor, cursor + audio_rows)
    video_sl = slice(audio_sl.stop, audio_sl.stop + video_rows)
    pad_sl = slice(video_sl.stop, seq_len)

    input_ids = torch.full((seq_len,), MINIMAX_H3_PAD_ID, dtype=torch.int64)
    input_ids[text_sl] = MINIMAX_H3_TEXT_ID
    image_mask = torch.zeros(seq_len, dtype=torch.bool)
    audio_mask = torch.zeros(seq_len, dtype=torch.bool)

    ref_img_pos_parts: list[torch.Tensor] = []
    ref_audio_pos_parts: list[torch.Tensor] = []
    g = torch.zeros(seq_len, 3, dtype=torch.float64)
    g[text_sl, 0] = torch.arange(text_len, dtype=torch.float64)

    target_area = np.sqrt(latent_h * latent_w)
    h_grid = _axis_from_sqrt_area(latent_h, _PATCH_H, target_area)
    w_grid = _axis_from_sqrt_area(latent_w, _PATCH_W, target_area)
    hh, ww = torch.meshgrid(h_grid, w_grid, indexing="ij")
    target_frame = torch.stack([hh.reshape(-1), ww.reshape(-1)], dim=-1)

    t_cursor = float(text_len)
    for item in block_slices:
        kind = str(item["kind"])
        if kind == "image":
            visual_sl = item["visual_sl"]
            assert isinstance(visual_sl, slice)
            input_ids[visual_sl] = MINIMAX_H3_IMGVID_COND_ID
            image_mask[visual_sl] = True
            ref_img_pos_parts.append(_range_for_slice(visual_sl))
            rh = int(item["latent_h"])
            rw = int(item["latent_w"])
            area = np.sqrt(rh * rw)
            ref_hh, ref_ww = torch.meshgrid(
                _axis_from_sqrt_area(rh, _PATCH_H, area),
                _axis_from_sqrt_area(rw, _PATCH_W, area),
                indexing="ij",
            )
            g[visual_sl, 0] = t_cursor
            g[visual_sl, 1] = ref_hh.reshape(-1)
            g[visual_sl, 2] = ref_ww.reshape(-1)
            t_cursor += 1.0
        elif kind == "audio":
            audio_ref_sl = item["audio_sl"]
            assert isinstance(audio_ref_sl, slice)
            ref_t = int(item["ref_audio_t"])
            input_ids[audio_ref_sl] = MINIMAX_H3_AUDIO_REF_COND_ID
            audio_mask[audio_ref_sl] = True
            ref_audio_pos_parts.append(_range_for_slice(audio_ref_sl))
            ref_t_grid = t_cursor + torch.arange(ref_t, dtype=torch.float64)
            g[audio_ref_sl, 0] = ref_t_grid.repeat(audio_channel)
            if ref_t:
                g[audio_ref_sl, 2] = torch.cat(
                    [
                        torch.full((ref_t,), float(w_grid[0]), dtype=torch.float64),
                        torch.full(
                            (int(item["audio_rows"]) - ref_t,),
                            float(w_grid[-1]),
                            dtype=torch.float64,
                        ),
                    ]
                )
            t_cursor += float(ref_t)
        else:
            audio_ref_sl = item["audio_sl"]
            visual_sl = item["visual_sl"]
            assert isinstance(audio_ref_sl, slice)
            assert isinstance(visual_sl, slice)
            ref_t = int(item["ref_audio_t"])
            vt = int(item["latent_t"])
            vh = int(item["latent_h"])
            vw = int(item["latent_w"])
            input_ids[audio_ref_sl] = MINIMAX_H3_AUDIO_REF_COND_ID
            input_ids[visual_sl] = MINIMAX_H3_IMGVID_COND_ID
            audio_mask[audio_ref_sl] = True
            image_mask[visual_sl] = True
            ref_audio_pos_parts.append(_range_for_slice(audio_ref_sl))
            ref_img_pos_parts.append(_range_for_slice(visual_sl))

            ref_area = np.sqrt(vh * vw)
            rv_h_grid = _axis_from_sqrt_area(vh, _PATCH_H, ref_area)
            rv_w_grid = _axis_from_sqrt_area(vw, _PATCH_W, ref_area)
            rv_hh, rv_ww = torch.meshgrid(rv_h_grid, rv_w_grid, indexing="ij")

            ref_t_grid = t_cursor + torch.arange(ref_t, dtype=torch.float64)
            g[audio_ref_sl, 0] = ref_t_grid.repeat(audio_channel)
            if ref_t:
                g[audio_ref_sl, 2] = torch.cat(
                    [
                        torch.full((ref_t,), float(rv_w_grid[0]), dtype=torch.float64),
                        torch.full(
                            (int(item["audio_rows"]) - ref_t,),
                            float(rv_w_grid[-1]),
                            dtype=torch.float64,
                        ),
                    ]
                )

            rv_frame = torch.stack([rv_hh.reshape(-1), rv_ww.reshape(-1)], dim=-1)
            rv_g = torch.empty(vt, int(item["frame_rows"]), 3, dtype=torch.float64)
            rv_g[:, :, 0] = _video_t_grid(vt, t_cursor)[:, None]
            rv_g[:, :, 1:] = rv_frame[None]
            g[visual_sl] = rv_g.reshape(-1, 3)
            t_cursor += max(float(ref_t), _video_t_span(vt))

    input_ids[audio_sl] = MINIMAX_H3_AUDIO_ID
    input_ids[audio_sl.start] = MINIMAX_H3_AUDIO_FIRST_ID
    input_ids[video_sl] = MINIMAX_H3_VIDEO_ID
    input_ids[video_sl.start] = MINIMAX_H3_VIDEO_FIRST_ID
    input_ids[video_sl.stop - 1] = MINIMAX_H3_VIDEO_LAST_ID
    audio_mask[audio_sl] = True
    image_mask[video_sl] = True

    audio_t_grid = t_cursor + torch.arange(audio_t, dtype=torch.float64)
    g[audio_sl, 0] = audio_t_grid.repeat(audio_channel)
    g[audio_sl, 2] = torch.cat(
        [
            torch.full((audio_t,), float(w_grid[0]), dtype=torch.float64),
            torch.full((audio_rows - audio_t,), float(w_grid[-1]), dtype=torch.float64),
        ]
    )

    video_g = torch.empty(latent_t, frame_rows, 3, dtype=torch.float64)
    video_g[:, :, 0] = _video_t_grid(latent_t, t_cursor)[:, None]
    video_g[:, :, 1:] = target_frame[None]
    g[video_sl] = video_g.reshape(-1, 3)

    target_img_pos = _range_for_slice(video_sl)
    target_audio_pos = _range_for_slice(audio_sl)
    img_pos = _cat_ranges(ref_img_pos_parts + [target_img_pos])
    audio_pos = _cat_ranges(ref_audio_pos_parts + [target_audio_pos])

    update_mask = torch.zeros(img_pos.shape[0], dtype=torch.bool)
    update_mask[ref_visual_rows:] = True
    audio_update_mask = torch.zeros(audio_pos.shape[0], dtype=torch.bool)
    audio_update_mask[ref_audio_rows:] = True
    text_pos = torch.arange(0, text_len)

    token_tags = torch.full((seq_len,), -1, dtype=torch.long)  # PADDING
    token_tags[text_sl] = 1  # TEXT
    if ref_audio_rows:
        token_tags[_cat_ranges(ref_audio_pos_parts)] = 2  # AUDIO (reference)
    token_tags[audio_sl] = 2  # AUDIO (target)
    token_tags[img_pos] = 0  # VIDEO (refs + target)

    cu = torch.tensor([0, used, seq_len], dtype=torch.int32)
    doc = torch.zeros(seq_len, dtype=torch.int32)
    doc[pad_sl] = 1
    return {
        "seq_len": torch.tensor(seq_len),
        "input_ids": input_ids,
        "image_mask": image_mask,
        "audio_mask": audio_mask,
        "img_pos": img_pos,
        "audio_pos": audio_pos,
        "text_pos": text_pos,
        "update_mask": update_mask,
        "audio_update_mask": audio_update_mask,
        "img_position_ids": g,
        "token_tags": token_tags,
        "cu_seqlens": cu,
        "document_id": doc,
        # Extent of the video segment within the packed sequence.
        "latent_grid": torch.tensor([latent_t, ph, pw], dtype=torch.int64),
        "video_row_start": torch.tensor(video_sl.start, dtype=torch.int64),
    }


__all__ = [
    "MINIMAX_H3_AUDIO_FIRST_ID",
    "MINIMAX_H3_AUDIO_ID",
    "MINIMAX_H3_AUDIO_REF_COND_ID",
    "MINIMAX_H3_IMGVID_COND_ID",
    "MINIMAX_H3_PAD_ID",
    "MINIMAX_H3_TEXT_ID",
    "MINIMAX_H3_VIDEO_FIRST_ID",
    "MINIMAX_H3_VIDEO_ID",
    "MINIMAX_H3_VIDEO_LAST_ID",
    "minimax_h3_packed_sequence",
    "minimax_h3_packed_sequence_ref2va_blocks",
]