"""Image, prompt, and batch processing for TrackingAnything."""

from __future__ import annotations

import math
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Iterable, List, Mapping, Optional, Sequence, Tuple

import numpy as np
from PIL import Image
import torch

from eaglevl.utils.locany.image_processing_locateanything import (
    LocateAnythingImageProcessor,
)


MEMORY_START_TOKEN = "<MEM_START>"
MEMORY_CONTEXT_TOKEN = "<MEM_CONTEXT>"
MEMORY_END_TOKEN = "<MEM_END>"
TRACK_QUERY_TOKEN = "<TRACK_QUERY>"
LOCATEANYTHING_NEGATIVE_ANSWER = "<box>none</box>"
LANGUAGE_PROMPT_MODE = "language"
MEMORY_ONLY_PROMPT_MODE = "memory_only"
MEMORY_TOKEN_COUNT = 48
TRACKING_SPECIAL_TOKENS = (
    MEMORY_START_TOKEN,
    MEMORY_CONTEXT_TOKEN,
    MEMORY_END_TOKEN,
    TRACK_QUERY_TOKEN,
)
# LocateAnything-3B already owns these three visual control embeddings.  They
# can serve as tokenizer-level aliases for the memory delimiters/placeholders
# because every context embedding is replaced before Qwen.  TRACK_QUERY is
# deliberately excluded: it always receives a new, independent vocabulary row.
BASE_TOKEN_ALIASES = {
    MEMORY_START_TOKEN: "<|vision_start|>",
    MEMORY_CONTEXT_TOKEN: "<|vision_pad|>",
    MEMORY_END_TOKEN: "<|vision_end|>",
}
IGNORE_INDEX = -100


def xywh_to_xyxy(box_xywh: Sequence[float]) -> np.ndarray:
    x, y, width, height = np.asarray(box_xywh, dtype=np.float32)
    return np.asarray([x, y, x + width, y + height], dtype=np.float32)


def normalize_xyxy(
    box_xyxy: Sequence[float],
    image_size: Tuple[int, int],
) -> np.ndarray:
    width, height = image_size
    scale = np.asarray([width, height, width, height], dtype=np.float32)
    return np.clip(np.asarray(box_xyxy, dtype=np.float32) / scale, 0.0, 1.0)


def square_memory_crop(
    image: Image.Image,
    box_xywh: Sequence[float],
    crop_scale: float = 1.25,
    output_size: int = 448,
) -> Tuple[Image.Image, np.ndarray, Dict[str, float]]:
    """Crop ``crop_scale * max(w,h)`` around a box using edge padding.

    Returns the resized square crop, target ``xyxy`` coordinates in that crop,
    and the exact affine metadata needed to audit/reproduce the transform.
    """

    x, y, width, height = [float(value) for value in box_xywh]
    if not all(math.isfinite(value) for value in (x, y, width, height)):
        raise ValueError("memory bbox contains a non-finite value")
    if width <= 0 or height <= 0:
        raise ValueError("memory bbox width and height must be positive")
    if crop_scale < 1.0:
        raise ValueError("crop_scale must be at least 1")

    image = image.convert("RGB")
    center_x = x + 0.5 * width
    center_y = y + 0.5 * height
    side = max(2.0, crop_scale * max(width, height))
    left = math.floor(center_x - 0.5 * side)
    top = math.floor(center_y - 0.5 * side)
    right = math.ceil(center_x + 0.5 * side)
    bottom = math.ceil(center_y + 0.5 * side)
    integer_side = max(right - left, bottom - top)
    # Rounding can differ by one pixel.  Keep a true square anchored at left/top.
    right = left + integer_side
    bottom = top + integer_side

    array = np.asarray(image)
    image_width, image_height = image.size
    pad_left = max(0, -left)
    pad_top = max(0, -top)
    pad_right = max(0, right - image_width)
    pad_bottom = max(0, bottom - image_height)
    if pad_left or pad_top or pad_right or pad_bottom:
        array = np.pad(
            array,
            ((pad_top, pad_bottom), (pad_left, pad_right), (0, 0)),
            mode="edge",
        )
    crop_left = left + pad_left
    crop_top = top + pad_top
    crop = Image.fromarray(array).crop(
        (
            crop_left,
            crop_top,
            crop_left + integer_side,
            crop_top + integer_side,
        )
    )
    crop = crop.resize((output_size, output_size), Image.Resampling.BICUBIC)

    resize_scale = output_size / integer_side
    box_in_crop = np.asarray(
        [
            (x - left) * resize_scale,
            (y - top) * resize_scale,
            (x + width - left) * resize_scale,
            (y + height - top) * resize_scale,
        ],
        dtype=np.float32,
    )
    box_in_crop = np.clip(box_in_crop, 0.0, float(output_size))
    metadata = {
        "crop_left": float(left),
        "crop_top": float(top),
        "crop_side": float(integer_side),
        "requested_crop_side": float(side),
        "resize_scale": float(resize_scale),
        "pad_left": float(pad_left),
        "pad_top": float(pad_top),
        "pad_right": float(pad_right),
        "pad_bottom": float(pad_bottom),
    }
    return crop, box_in_crop, metadata


def patch_box_coverage(
    box_xyxy: Sequence[float],
    crop_size: int = 448,
    patch_size: int = 14,
    merge_kernel_size: Tuple[int, int] = (2, 2),
) -> torch.Tensor:
    """Exact target coverage for every post-merger visual token.

    MoonViT concatenates each 2x2 group of 14x14 patches before ``mlp1``.
    Therefore one projected token corresponds to a 28x28 cell for the default
    merger.  Computing coverage on that merged grid aligns the embedding with
    the actual sequence entering the Resampler.
    """

    cell_height = patch_size * int(merge_kernel_size[0])
    cell_width = patch_size * int(merge_kernel_size[1])
    if crop_size % cell_height or crop_size % cell_width:
        raise ValueError("crop size must be divisible by the merged patch size")
    grid_height = crop_size // cell_height
    grid_width = crop_size // cell_width
    x1, y1, x2, y2 = [
        float(value) for value in np.asarray(box_xyxy, dtype=np.float32)
    ]

    cell_left = torch.arange(grid_width, dtype=torch.float32) * cell_width
    cell_top = torch.arange(grid_height, dtype=torch.float32) * cell_height
    cell_right = cell_left + cell_width
    cell_bottom = cell_top + cell_height
    intersection_width = (
        torch.minimum(cell_right[None, :], torch.tensor(x2))
        - torch.maximum(cell_left[None, :], torch.tensor(x1))
    ).clamp_min(0)
    intersection_height = (
        torch.minimum(cell_bottom[:, None], torch.tensor(y2))
        - torch.maximum(cell_top[:, None], torch.tensor(y1))
    ).clamp_min(0)
    coverage = intersection_height * intersection_width
    coverage = coverage / float(cell_height * cell_width)
    return coverage.flatten().clamp(0.0, 1.0)


@dataclass
class ProcessedMemory:
    pixel_values: torch.Tensor
    grid_hws: torch.Tensor
    slot_ids: torch.Tensor
    region_coverage: torch.Tensor
    relative_frame_gaps: torch.Tensor
    scores: torch.Tensor
    valid_mask: torch.Tensor
    crop_boxes_xyxy: torch.Tensor
    crop_metadata: Tuple[Optional[Dict[str, float]], ...]


class TrackingAnythingMemoryProcessor:
    """Preprocess M1-M6, or select M1-M5 from the unchanged six-slot cache."""

    def __init__(
        self,
        crop_size: int = 448,
        crop_scale: float = 1.25,
        patch_size: int = 14,
        merge_kernel_size: Tuple[int, int] = (2, 2),
        num_slots: int = 6,
    ) -> None:
        if num_slots not in (5, 6):
            raise ValueError("memory input must retain M1-M5 or M1-M6")
        self.num_slots = int(num_slots)
        self.crop_size = int(crop_size)
        self.crop_scale = float(crop_scale)
        self.patch_size = int(patch_size)
        self.merge_kernel_size = tuple(int(v) for v in merge_kernel_size)
        self.image_processor = LocateAnythingImageProcessor(
            patch_size=self.patch_size,
            merge_kernel_size=list(self.merge_kernel_size),
            in_token_limit=(self.crop_size // self.patch_size) ** 2,
        )
        merged_grid = self.crop_size // (
            self.patch_size * self.merge_kernel_size[0]
        )
        self.projected_token_count = merged_grid * merged_grid

    def __call__(
        self,
        image_paths: Sequence[Optional[str | Path]],
        boxes_xywh: Sequence[Sequence[float]],
        frame_indices: Sequence[int],
        current_frame_index: int,
        scores: Sequence[float],
        valid_mask: Optional[Sequence[bool]] = None,
        memory_images: Optional[Sequence[Optional[np.ndarray | Image.Image]]] = None,
    ) -> ProcessedMemory:
        # The immutable V6 cache/replay always retains six slots. Select the
        # model's prefix BEFORE crop loading; M6 is never encoded in 5-slot runs.
        if self.num_slots == 5:
            if any(len(v) not in (5, 6) for v in
                   (image_paths, boxes_xywh, frame_indices, scores)):
                raise ValueError("five-slot input requires a five/six-slot bank")
            image_paths = image_paths[:5]
            boxes_xywh = boxes_xywh[:5]
            frame_indices = frame_indices[:5]
            scores = scores[:5]
            valid_mask = None if valid_mask is None else valid_mask[:5]
            memory_images = None if memory_images is None else memory_images[:5]
        if not all(
            len(values) == self.num_slots
            for values in (
                image_paths,
                boxes_xywh,
                frame_indices,
                scores,
            )
        ):
            raise ValueError(f"TrackingAnything requires {self.num_slots} memory slots")
        if memory_images is None:
            memory_images = [None] * self.num_slots
        elif len(memory_images) != self.num_slots:
            raise ValueError(f"TrackingAnything requires {self.num_slots} memory images")
        inferred_valid = np.asarray(
            [
                (path is not None or memory_image is not None)
                and int(frame_index) >= 0
                and np.isfinite(np.asarray(box, dtype=np.float32)).all()
                and float(box[2]) > 0
                and float(box[3]) > 0
                for path, memory_image, box, frame_index in zip(
                    image_paths, memory_images, boxes_xywh, frame_indices
                )
            ],
            dtype=bool,
        )
        if valid_mask is not None:
            inferred_valid &= np.asarray(valid_mask, dtype=bool)

        crops: List[Image.Image] = []
        valid_slots: List[int] = []
        coverage = torch.zeros(self.num_slots, self.projected_token_count)
        crop_boxes = torch.full((self.num_slots, 4), float("nan"))
        crop_metadata: List[Optional[Dict[str, float]]] = [None] * self.num_slots
        gaps = torch.zeros(self.num_slots, dtype=torch.long)
        clean_scores = torch.zeros(self.num_slots, dtype=torch.float32)

        for slot in range(self.num_slots):
            if not inferred_valid[slot]:
                continue
            source_frame = int(frame_indices[slot])
            gap = int(current_frame_index) - source_frame
            if gap <= 0:
                raise ValueError(
                    f"non-causal memory in slot M{slot + 1}: "
                    f"source={source_frame}, current={current_frame_index}"
                )
            memory_image = memory_images[slot]
            if memory_image is None:
                with Image.open(image_paths[slot]) as handle:
                    image = handle.convert("RGB")
            elif isinstance(memory_image, Image.Image):
                image = memory_image.convert("RGB")
            else:
                array = np.asarray(memory_image, dtype=np.uint8)
                if array.ndim != 3 or array.shape[2] != 3:
                    raise ValueError(
                        f"memory image M{slot + 1} must be HxWx3 RGB, got {array.shape}"
                    )
                image = Image.fromarray(array, mode="RGB")
            crop, box_in_crop, metadata = square_memory_crop(
                image,
                boxes_xywh[slot],
                crop_scale=self.crop_scale,
                output_size=self.crop_size,
            )
            crops.append(crop)
            valid_slots.append(slot)
            coverage[slot] = patch_box_coverage(
                box_in_crop,
                crop_size=self.crop_size,
                patch_size=self.patch_size,
                merge_kernel_size=self.merge_kernel_size,
            )
            crop_boxes[slot] = torch.from_numpy(box_in_crop)
            crop_metadata[slot] = metadata
            gaps[slot] = gap
            score = float(scores[slot])
            clean_scores[slot] = score if math.isfinite(score) else 0.0

        if crops:
            processed = self.image_processor(
                images=crops, return_tensors="pt"
            )
            pixel_values = processed["pixel_values"]
            grid_hws = torch.as_tensor(
                processed["image_grid_hws"], dtype=torch.int32
            )
        else:
            pixel_values = torch.empty(
                0, 3, self.patch_size, self.patch_size, dtype=torch.float32
            )
            grid_hws = torch.empty(0, 2, dtype=torch.int32)

        return ProcessedMemory(
            pixel_values=pixel_values,
            grid_hws=grid_hws,
            slot_ids=torch.as_tensor(valid_slots, dtype=torch.long),
            region_coverage=coverage,
            relative_frame_gaps=gaps,
            scores=clean_scores,
            valid_mask=torch.from_numpy(inferred_valid),
            crop_boxes_xyxy=crop_boxes,
            crop_metadata=tuple(crop_metadata),
        )


class TrackingAnythingProcessor:
    """Wrapper around the original LocateAnything processor."""

    def __init__(
        self,
        base_processor,
        crop_size: int = 448,
        crop_scale: float = 1.25,
        memory_token_count: int = MEMORY_TOKEN_COUNT,
        reuse_base_control_tokens: bool = True,
        memory_num_slots: int = 6,
    ) -> None:
        self.base_processor = base_processor
        self.tokenizer = base_processor.tokenizer
        self.memory_processor = TrackingAnythingMemoryProcessor(
            num_slots=memory_num_slots,
            crop_size=crop_size,
            crop_scale=crop_scale,
            patch_size=base_processor.image_processor.patch_size,
            merge_kernel_size=tuple(
                base_processor.image_processor.merge_kernel_size
            ),
        )
        self.memory_token_count = int(memory_token_count)
        if self.memory_token_count <= 0:
            raise ValueError("memory_token_count must be positive")
        self.reuse_base_control_tokens = bool(reuse_base_control_tokens)
        if not self.reuse_base_control_tokens:
            self.tokenizer.add_special_tokens(
                {
                    "additional_special_tokens": [
                        MEMORY_START_TOKEN,
                        MEMORY_CONTEXT_TOKEN,
                        MEMORY_END_TOKEN,
                    ]
                },
                replace_additional_special_tokens=False,
            )
        self.tokenizer.add_special_tokens(
            {"additional_special_tokens": [TRACK_QUERY_TOKEN]},
            replace_additional_special_tokens=False,
        )
        self.token_ids = self._resolve_token_ids()
        self._validate_token_ids()

    def _resolve_token_ids(self) -> Dict[str, int]:
        physical_tokens = {
            token: (
                BASE_TOKEN_ALIASES[token]
                if self.reuse_base_control_tokens
                and token in BASE_TOKEN_ALIASES
                else token
            )
            for token in TRACKING_SPECIAL_TOKENS
        }
        return {
            logical: int(self.tokenizer.convert_tokens_to_ids(physical))
            for logical, physical in physical_tokens.items()
        }

    def _ar_presence_token_ids(self) -> Dict[str, int]:
        none_ids = self.tokenizer.encode(
            "none",
            add_special_tokens=False,
        )
        coordinate_ids = [
            int(self.tokenizer.convert_tokens_to_ids(f"<{index}>"))
            for index in range(1001)
        ]
        if len(none_ids) != 1:
            raise RuntimeError("native negative 'none' must be one token")
        if any(
            right != left + 1
            for left, right in zip(coordinate_ids, coordinate_ids[1:])
        ):
            raise RuntimeError(
                "LocateAnything coordinate token IDs must be contiguous"
            )
        return {
            "box_start_token_id": int(
                self.tokenizer.convert_tokens_to_ids("<box>")
            ),
            "none_token_id": int(none_ids[0]),
            "coordinate_token_start_id": coordinate_ids[0],
            "coordinate_token_end_id": coordinate_ids[-1],
        }

    def add_special_tokens(self, model=None) -> Dict[str, int]:
        if not self.reuse_base_control_tokens:
            self.tokenizer.add_special_tokens(
                {
                    "additional_special_tokens": [
                        MEMORY_START_TOKEN,
                        MEMORY_CONTEXT_TOKEN,
                        MEMORY_END_TOKEN,
                    ]
                },
                replace_additional_special_tokens=False,
            )
        # Idempotent when __init__ already added it or when loading a prepared
        # TrackingAnything tokenizer.
        self.tokenizer.add_special_tokens(
            {"additional_special_tokens": [TRACK_QUERY_TOKEN]},
            replace_additional_special_tokens=False,
        )
        self.token_ids = self._resolve_token_ids()
        self._validate_token_ids(
            image_token_index=(
                getattr(model, "image_token_index", None)
                if model is not None
                else None
            )
        )
        if model is not None:
            current_vocab = model.get_input_embeddings().num_embeddings
            if current_vocab != len(self.tokenizer):
                model.resize_token_embeddings(
                    len(self.tokenizer), mean_resizing=False
                )
            model.config.memory_start_token_id = self.token_ids[
                MEMORY_START_TOKEN
            ]
            model.config.memory_context_token_id = self.token_ids[
                MEMORY_CONTEXT_TOKEN
            ]
            model.config.memory_end_token_id = self.token_ids[
                MEMORY_END_TOKEN
            ]
            model.config.track_query_token_id = self.token_ids[
                TRACK_QUERY_TOKEN
            ]
            for name, value in self._ar_presence_token_ids().items():
                setattr(model.config, name, value)
            if float(getattr(model.config, 'mtp_loss_weight', 0.0)) > 0:
                for name, token in {'mtp_mask_token_id':'<text_mask>',
                                    'mtp_null_token_id':'<null>',
                                    'mtp_box_end_token_id':'</box>',
                                    'mtp_eos_token_id':'<|im_end|>'}.items():
                    encoded = self.tokenizer.encode(token, add_special_tokens=False)
                    if len(encoded) != 1 or self.tokenizer.convert_ids_to_tokens(encoded[0]) != token:
                        raise ValueError(f'MTP requires native single token {token}')
                    setattr(model.config,name,int(encoded[0]))
            model.config.text_config.vocab_size = len(self.tokenizer)
            model.language_model.config.vocab_size = len(self.tokenizer)
            model.language_model.vocab_size = len(self.tokenizer)
            self.validate_model_alignment(model)
            model._tracking_tokenizer_aligned = True
        return dict(self.token_ids)

    def _validate_token_ids(
        self,
        image_token_index: Optional[int] = None,
    ) -> None:
        ids = {
            int(self.token_ids[token]) for token in TRACKING_SPECIAL_TOKENS
        }
        if len(ids) != len(TRACKING_SPECIAL_TOKENS):
            raise RuntimeError(
                "memory start/context/end and <TRACK_QUERY> must use four "
                "different token IDs"
            )
        track_id = int(self.token_ids[TRACK_QUERY_TOKEN])
        if not 0 <= track_id < len(self.tokenizer):
            raise RuntimeError("<TRACK_QUERY> token ID is outside the tokenizer")
        quad_id = int(self.tokenizer.convert_tokens_to_ids("<quad>"))
        if track_id == quad_id:
            raise RuntimeError("<TRACK_QUERY> must not reuse the <quad> token")
        if (
            image_token_index is not None
            and int(self.token_ids[MEMORY_CONTEXT_TOKEN])
            == int(image_token_index)
        ):
            raise RuntimeError(
                "memory context and current-image placeholders use the same "
                "token ID"
            )

    def validate_model_alignment(self, model) -> Dict[str, int]:
        """Fail fast if tokenizer, embeddings, heads, or control IDs drift."""

        vocab_size = len(self.tokenizer)
        track_id = int(self.token_ids[TRACK_QUERY_TOKEN])
        input_rows = int(model.get_input_embeddings().num_embeddings)
        lm_head_rows = int(model.language_model.lm_head.weight.shape[0])
        if input_rows != vocab_size:
            raise RuntimeError(
                f"input embedding rows ({input_rows}) != tokenizer size "
                f"({vocab_size})"
            )
        if lm_head_rows != vocab_size:
            raise RuntimeError(
                f"LM head rows ({lm_head_rows}) != tokenizer size "
                f"({vocab_size})"
            )
        if not 0 <= track_id < vocab_size:
            raise RuntimeError("<TRACK_QUERY> token ID is outside the vocabulary")

        expected_config_ids = {
            "memory_start_token_id": self.token_ids[MEMORY_START_TOKEN],
            "memory_context_token_id": self.token_ids[MEMORY_CONTEXT_TOKEN],
            "memory_end_token_id": self.token_ids[MEMORY_END_TOKEN],
            "track_query_token_id": track_id,
            **self._ar_presence_token_ids(),
        }
        for name, expected in expected_config_ids.items():
            actual = getattr(model.config, name, None)
            if actual != expected:
                raise RuntimeError(
                    f"model.config.{name}={actual!r}, expected {expected}"
                )
        image_token_index = getattr(
            model, "image_token_index",
            getattr(model.config, "image_token_index", None),
        )
        self._validate_token_ids(image_token_index=image_token_index)
        return {
            "vocab_size": vocab_size,
            "track_query_token_id": track_id,
            "image_token_index": int(image_token_index),
            **{
                name: int(value)
                for name, value in expected_config_ids.items()
            },
        }

    def _replace_logical_control_tokens(self, text: str) -> str:
        if not self.reuse_base_control_tokens:
            return text
        for logical, physical in BASE_TOKEN_ALIASES.items():
            text = text.replace(logical, physical)
        return text

    def memory_prefix(self) -> str:
        return (
            MEMORY_START_TOKEN
            + MEMORY_CONTEXT_TOKEN * self.memory_token_count
            + MEMORY_END_TOKEN
        )

    @staticmethod
    def tracking_instruction(language: str, prompt_mode: str) -> str:
        if prompt_mode == LANGUAGE_PROMPT_MODE:
            target_text = language.strip()
            if not target_text:
                raise ValueError("language prompt mode requires original language")
            return (
                "Locate and track the target initially described as\n"
                f"<ref>{target_text}</ref>\n"
                "in the current image.\n"
                f"{TRACK_QUERY_TOKEN}"
            )
        if prompt_mode == MEMORY_ONLY_PROMPT_MODE:
            return (
                "Locate and track the target represented by the visual memory\n"
                "in the current image.\n"
                f"{TRACK_QUERY_TOKEN}"
            )
        raise ValueError(f"unsupported tracking prompt mode: {prompt_mode!r}")

    def build_messages(
        self,
        current_image: Image.Image,
        language: str,
        answer: Optional[str] = None,
        prompt_mode: str = LANGUAGE_PROMPT_MODE,
    ) -> List[dict]:
        instruction = self.tracking_instruction(language, prompt_mode)
        messages: List[dict] = [
            {
                "role": "user",
                "content": [
                    {"text": self.memory_prefix()},
                    {"type": "image", "image": current_image},
                    {"text": instruction},
                ],
            }
        ]
        if answer is not None:
            messages.append({"role": "assistant", "content": answer})
        return messages

    def encode_conversation(
        self,
        current_image: Image.Image,
        language: str,
        answer: Optional[str],
        prompt_mode: str = LANGUAGE_PROMPT_MODE,
        mask_answer_ref_content: bool = False,
    ) -> Dict[str, torch.Tensor]:
        if answer is not None and mask_answer_ref_content:
            raise ValueError(
                "box-only tracking labels contain no assistant <ref> span"
            )
        full_messages = self.build_messages(
            current_image,
            language,
            answer,
            prompt_mode=prompt_mode,
        )
        full_text = self.base_processor.py_apply_chat_template(
            full_messages,
            tokenize=False,
            add_generation_prompt=answer is None,
        )
        full_text = self._replace_logical_control_tokens(full_text)
        encoded = self.base_processor(
            text=full_text,
            images=[current_image],
            return_tensors="pt",
            padding=False,
            truncation=True,
        )
        input_ids = encoded["input_ids"][0]
        attention_mask = encoded["attention_mask"][0]

        labels = torch.full_like(input_ids, IGNORE_INDEX)
        if answer is not None:
            prompt_messages = self.build_messages(
                current_image,
                language,
                answer=None,
                prompt_mode=prompt_mode,
            )
            prompt_text = self.base_processor.py_apply_chat_template(
                prompt_messages,
                tokenize=False,
                add_generation_prompt=True,
            )
            prompt_text = self._replace_logical_control_tokens(prompt_text)
            prompt_encoded = self.base_processor(
                text=prompt_text,
                images=[current_image],
                return_tensors="pt",
                padding=False,
                truncation=True,
            )
            prompt_ids = prompt_encoded["input_ids"][0]
            prefix_length = len(prompt_ids)
            if prefix_length >= len(input_ids):
                raise RuntimeError(
                    "assistant answer produced no supervised suffix"
                )
            if not torch.equal(input_ids[:prefix_length], prompt_ids):
                mismatch = torch.where(
                    input_ids[:prefix_length] != prompt_ids
                )[0]
                first_mismatch = (
                    int(mismatch[0]) if len(mismatch) else None
                )
                raise RuntimeError(
                    "prompt is not an exact prefix of the full conversation; "
                    f"first mismatch={first_mismatch}"
                )
            ref_start_id = self.tokenizer.convert_tokens_to_ids("<ref>")
            ref_end_id = self.tokenizer.convert_tokens_to_ids("</ref>")
            box_start_id = self.tokenizer.convert_tokens_to_ids("<box>")
            box_end_id = self.tokenizer.convert_tokens_to_ids("</box>")
            im_end_id = self.tokenizer.convert_tokens_to_ids("<|im_end|>")
            answer_suffix = input_ids[prefix_length:]
            end_offsets = torch.where(answer_suffix.eq(im_end_id))[0]
            if not len(end_offsets):
                raise RuntimeError(
                    "assistant suffix contains no <|im_end|> token"
                )
            supervision_end = prefix_length + int(end_offsets[0]) + 1
            labels[prefix_length:supervision_end] = input_ids[
                prefix_length:supervision_end
            ]
            supervised_ids = labels[labels.ne(IGNORE_INDEX)]
            if bool(supervised_ids.eq(ref_start_id).any()) or bool(
                supervised_ids.eq(ref_end_id).any()
            ):
                raise RuntimeError(
                    "assistant supervision must be box-only and contain no "
                    "<ref> tokens"
                )
            box_starts = torch.where(supervised_ids.eq(box_start_id))[0]
            box_ends = torch.where(supervised_ids.eq(box_end_id))[0]
            if (
                len(box_starts) != 1
                or len(box_ends) != 1
                or int(box_ends[0]) <= int(box_starts[0])
            ):
                raise RuntimeError(
                    "assistant supervision must contain exactly one ordered "
                    "<box>...</box> span"
                )
            if int(supervised_ids[0]) != int(box_start_id):
                raise RuntimeError(
                    "first assistant supervision token must be <box>"
                )
            if int(supervised_ids[-1]) != int(im_end_id):
                raise RuntimeError(
                    "last assistant supervision token must be <|im_end|>"
                )

        # TrackingAnything uses ordinary autoregressive generation. The
        # model's scoped mask adapter builds the 4-D SDPA mask and applies
        # empty-slot key masking.
        position_ids = torch.arange(len(input_ids), dtype=torch.long)

        memory_context_id = self.token_ids[MEMORY_CONTEXT_TOKEN]
        track_query_id = self.token_ids[TRACK_QUERY_TOKEN]
        ref_start_id = self.tokenizer.convert_tokens_to_ids("<ref>")
        ref_end_id = self.tokenizer.convert_tokens_to_ids("</ref>")
        ref_starts = torch.where(input_ids.eq(ref_start_id))[0]
        ref_ends = torch.where(input_ids.eq(ref_end_id))[0]
        language_token_mask = torch.zeros_like(input_ids, dtype=torch.bool)
        if prompt_mode == LANGUAGE_PROMPT_MODE:
            if (
                len(ref_starts) != 1
                or len(ref_ends) != 1
                or int(ref_ends[0]) <= int(ref_starts[0]) + 1
            ):
                raise RuntimeError(
                    "language prompt must contain exactly one non-empty "
                    "<ref>...</ref> span"
                )
            language_token_mask[
                int(ref_starts[0]) + 1 : int(ref_ends[0])
            ] = True
        elif len(ref_starts) or len(ref_ends):
            raise RuntimeError(
                "memory-only prompt unexpectedly contains a <ref> span"
            )
        return {
            "input_ids": input_ids,
            "labels": labels,
            "attention_mask": attention_mask,
            "position_ids": position_ids,
            "memory_token_mask": input_ids.eq(memory_context_id),
            "track_query_mask": input_ids.eq(track_query_id),
            "language_token_mask": language_token_mask,
            "pixel_values": encoded["pixel_values"],
            "image_grid_hws": torch.as_tensor(
                encoded["image_grid_hws"], dtype=torch.int32
            ),
        }


def format_tracking_answer(
    language: str,
    normalized_xyxy: Optional[Sequence[float]] = None,
    target_present: bool = True,
) -> str:
    """Format a box-only positive or LocateAnything-native negative answer.

    ``language`` remains in the signature for call-site compatibility. It is
    used only in the input prompt and is never copied into the assistant label.
    """

    if not target_present:
        if normalized_xyxy is not None:
            raise ValueError("an absent target must not carry a bounding box")
        return LOCATEANYTHING_NEGATIVE_ANSWER
    if normalized_xyxy is None:
        raise ValueError("a present target requires a bounding box")
    box = np.asarray(normalized_xyxy, dtype=np.float64)
    if box.shape != (4,) or not np.isfinite(box).all():
        raise ValueError("normalized_xyxy must contain four finite values")
    if box[2] < box[0] or box[3] < box[1]:
        raise ValueError("normalized_xyxy has inverted corners")
    coordinates = np.rint(
        np.clip(box, 0.0, 1.0) * 1000.0
    ).astype(np.int64)
    coordinate_tokens = "".join(f"<{int(value)}>" for value in coordinates)
    return f"<box>{coordinate_tokens}</box>"


def format_bbox_answer(
    language: str,
    normalized_xyxy: Sequence[float],
) -> str:
    """Backward-compatible positive-only wrapper."""

    return format_tracking_answer(
        language=language,
        normalized_xyxy=normalized_xyxy,
        target_present=True,
    )


class TrackingAnythingCollator:
    """Pad text and concatenate variable-size current/memory vision batches."""

    def __init__(self, pad_token_id: int) -> None:
        self.pad_token_id = int(pad_token_id)

    def __call__(self, features: List[Mapping[str, object]]) -> Dict[str, object]:
        if not features:
            raise ValueError("cannot collate an empty batch")
        batch_size = len(features)
        max_length = max(len(item["input_ids"]) for item in features)

        def padded_stack(
            key: str,
            fill_value: int | bool,
            dtype: torch.dtype,
        ) -> torch.Tensor:
            output = torch.full(
                (batch_size, max_length), fill_value, dtype=dtype
            )
            for index, item in enumerate(features):
                value = torch.as_tensor(item[key], dtype=dtype)
                output[index, : len(value)] = value
            return output

        batch: Dict[str, object] = {
            "input_ids": padded_stack(
                "input_ids", self.pad_token_id, torch.long
            ),
            "labels": padded_stack("labels", IGNORE_INDEX, torch.long),
            "attention_mask": padded_stack(
                "attention_mask", 0, torch.long
            ),
            "position_ids": padded_stack("position_ids", 0, torch.long),
            "memory_token_mask": padded_stack(
                "memory_token_mask", False, torch.bool
            ),
            "track_query_mask": padded_stack(
                "track_query_mask", False, torch.bool
            ),
            "language_token_mask": padded_stack(
                "language_token_mask", False, torch.bool
            ),
            "pixel_values": torch.cat(
                [item["pixel_values"] for item in features], dim=0
            ),
            "image_grid_hws": torch.cat(
                [item["image_grid_hws"] for item in features], dim=0
            ),
        }

        memory_pixels = [
            item["memory_pixel_values"]
            for item in features
            if item["memory_pixel_values"].numel()
        ]
        batch["memory_pixel_values"] = (
            torch.cat(memory_pixels, dim=0)
            if memory_pixels
            else torch.empty(0, 3, 14, 14, dtype=torch.float32)
        )
        memory_grids = []
        assignments = []
        for batch_index, item in enumerate(features):
            grids = item["memory_grid_hws"]
            slots = item["memory_slot_ids"]
            if len(grids):
                memory_grids.append(grids)
                assignments.append(
                    torch.stack(
                        (
                            torch.full_like(slots, batch_index),
                            slots,
                        ),
                        dim=-1,
                    )
                )
        batch["memory_grid_hws"] = (
            torch.cat(memory_grids, dim=0)
            if memory_grids
            else torch.empty(0, 2, dtype=torch.int32)
        )
        batch["memory_batch_slot"] = (
            torch.cat(assignments, dim=0)
            if assignments
            else torch.empty(0, 2, dtype=torch.long)
        )

        for key in (
            "memory_region_coverage",
            "memory_relative_frame_gaps",
            "memory_scores",
            "memory_slot_valid_mask",
            "target_boxes_xyxy",
            "target_presence",
        ):
            batch[key] = torch.stack(
                [torch.as_tensor(item[key]) for item in features]
            )
        if "memory_gate_labels" in features[0]:
            if not all("memory_gate_labels" in item for item in features):
                raise ValueError(
                    "memory_gate_labels must be present for the whole batch"
                )
            batch["memory_gate_labels"] = torch.stack(
                [
                    torch.as_tensor(item["memory_gate_labels"])
                    for item in features
                ]
            )
        if "memory_reject_labels" in features[0]:
            if not all("memory_reject_labels" in item for item in features):
                raise ValueError(
                    "memory_reject_labels must be present for the whole batch"
                )
            batch["memory_reject_labels"] = torch.stack(
                [
                    torch.as_tensor(item["memory_reject_labels"])
                    for item in features
                ]
            )
        if "metadata" in features[0]:
            batch["metadata"] = [item["metadata"] for item in features]
        return batch


__all__ = [
    "IGNORE_INDEX",
    "BASE_TOKEN_ALIASES",
    "LANGUAGE_PROMPT_MODE",
    "LOCATEANYTHING_NEGATIVE_ANSWER",
    "MEMORY_CONTEXT_TOKEN",
    "MEMORY_END_TOKEN",
    "MEMORY_ONLY_PROMPT_MODE",
    "MEMORY_START_TOKEN",
    "MEMORY_TOKEN_COUNT",
    "TRACK_QUERY_TOKEN",
    "TRACKING_SPECIAL_TOKENS",
    "ProcessedMemory",
    "TrackingAnythingCollator",
    "TrackingAnythingMemoryProcessor",
    "TrackingAnythingProcessor",
    "format_bbox_answer",
    "format_tracking_answer",
    "normalize_xyxy",
    "patch_box_coverage",
    "square_memory_crop",
    "xywh_to_xyxy",
]