#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
#           This file was automatically generated from src/transformers/models/molmo2/modular_molmo2.py.
#               Do NOT edit this file manually as any edits will be overwritten by the generation of
#             the file from the modular. If any change should be done, please apply the change to the
#                          modular_molmo2.py file directly. One of our CI enforces this.
#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
# Copyright 2026 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import torch
from torch.nn import functional as F

from ...image_processing_backends import TorchvisionBackend
from ...image_processing_utils import BatchFeature
from ...image_utils import IMAGENET_STANDARD_MEAN, IMAGENET_STANDARD_STD, PILImageResampling, SizeDict
from ...processing_utils import Unpack, VideosKwargs
from ...utils import TensorType, auto_docstring, logging
from ...video_processing_utils import BaseVideoProcessor
from ...video_utils import VideoMetadata


logger = logging.get_logger(__name__)


class Molmo2VideoProcessorKwargs(VideosKwargs, total=False):
    """
    patch_size (`int`, *optional*):
        Side length in pixels of each ViT patch for video frames.
    pooling_size (`list[int]`, *optional*):
        `[pool_h, pool_w]` pooling window applied to video patch features.
    max_fps (`int`, *optional*):
        Maximum sampling rate in frames per second for short videos.
    """

    patch_size: int
    pooling_size: list[int]
    max_fps: int


def batch_pixels_to_patches(images: torch.Tensor, patch_size: int) -> torch.Tensor:
    """Reshape images of [num_images, h, w, channels] -> [num_images, n_patches, pixels_per_patch]"""
    num_crops, height, width, channels = images.shape
    num_patches_height = height // patch_size
    num_patches_width = width // patch_size
    images = images.reshape(num_crops, num_patches_height, patch_size, num_patches_width, patch_size, channels)
    images = images.transpose(2, 3)
    images = images.reshape(num_crops, num_patches_height * num_patches_width, patch_size * patch_size * channels)
    return images


def arange_for_pooling(
    index_grid: torch.Tensor,
    pool_h: int,
    pool_w: int,
) -> tuple[torch.Tensor, int, int]:
    """Group a [h, w] patch-index grid into pooling windows: `[num_rows * num_cols, pool_h * pool_w]`, padded
    with -1. Also returns the pooled grid's `(num_rows, num_cols)`."""
    height_padding = pool_h * ((index_grid.shape[0] + pool_h - 1) // pool_h) - index_grid.shape[0]
    width_padding = pool_w * ((index_grid.shape[1] + pool_w - 1) // pool_w) - index_grid.shape[1]
    index_grid = F.pad(
        index_grid,
        (width_padding // 2, (width_padding + 1) // 2, height_padding // 2, (height_padding + 1) // 2),
        mode="constant",
        value=-1,
    )
    num_rows, num_cols = index_grid.shape[0] // pool_h, index_grid.shape[1] // pool_w
    pooling_indices = (
        index_grid.reshape(num_rows, pool_h, num_cols, pool_w)
        .permute(0, 2, 1, 3)
        .reshape(num_rows * num_cols, pool_h * pool_w)
    )
    return pooling_indices, num_rows, num_cols


def resize_and_normalize_image(
    backend: "TorchvisionBackend",
    image: torch.Tensor,
    output_size: list[int],
    resample: PILImageResampling,
    do_rescale: bool,
    rescale_factor: float,
    do_normalize: bool,
    image_mean: list[float],
    image_std: list[float],
) -> torch.Tensor:
    # Under `torch.compile`, torchvision's uint8 bilinear resize can produce values slightly outside 0-255,
    # which overflow and wrap around (-1 becomes 255, so a near-black pixel becomes white). Resizing in
    # float, then rounding and clamping back, avoids that and gives the exact same bytes in eager mode.
    is_integer_input = not image.dtype.is_floating_point
    resized = backend.resize(
        image.float() if is_integer_input else image,
        size=SizeDict(height=output_size[0], width=output_size[1]),
        resample=resample,
        antialias=False,
    )
    if is_integer_input:
        resized = resized.round().clamp(0, 255).to(image.dtype)
    return backend.rescale_and_normalize(resized, do_rescale, rescale_factor, do_normalize, image_mean, image_std)


def build_resized_image(
    backend: "TorchvisionBackend",
    images: torch.Tensor,
    crop_size: int,
    resample: PILImageResampling,
    do_rescale: bool,
    rescale_factor: float,
    do_normalize: bool,
    image_mean: list[float],
    image_std: list[float],
    image_patch_size: int,
) -> tuple[torch.Tensor, torch.Tensor]:
    # `images`: a batch of same-shape images `[N, C, H, W]`; resize the whole batch at once.
    resized = resize_and_normalize_image(
        backend,
        images,
        [crop_size, crop_size],
        resample,
        do_rescale=do_rescale,
        rescale_factor=rescale_factor,
        do_normalize=do_normalize,
        image_mean=image_mean,
        image_std=image_std,
    )
    # [N, C, S, S] -> [N, 1, S, S, C]: one global (low-res) view per image.
    resized = resized.permute(0, 2, 3, 1).unsqueeze(1)
    # The per-patch index grid depends only on the (shared) shape, so it is built once.
    patches_per_side = crop_size // image_patch_size
    resized_index_grid = torch.arange(patches_per_side**2, dtype=torch.int32, device=images.device).reshape(
        patches_per_side, patches_per_side
    )
    return resized, resized_index_grid


@auto_docstring
class Molmo2VideoProcessor(BaseVideoProcessor):
    resample = PILImageResampling.BILINEAR
    size = {"height": 378, "width": 378}
    image_mean = IMAGENET_STANDARD_MEAN
    image_std = IMAGENET_STANDARD_STD
    patch_size = 14
    pooling_size = [3, 3]
    num_frames = 64
    do_sample_frames = True
    max_fps = 2
    valid_kwargs = Molmo2VideoProcessorKwargs
    model_input_names = ["pixel_values_videos", "video_token_pooling", "video_grids"]

    def __init__(self, **kwargs: Unpack[Molmo2VideoProcessorKwargs]):
        super().__init__(**kwargs)
        if self.size is not None and (self.size.get("height", None) is None or self.size.get("width", None) is None):
            raise ValueError("size must contain 'height' and 'width' keys.")

    def sample_frames(
        self,
        metadata: VideoMetadata,
        num_frames: int | None = None,
        fps: int | float | None = None,
        max_fps: int | float | None = None,
        **kwargs,
    ):
        if fps is not None and num_frames is not None:
            raise ValueError("`num_frames` and `fps` are mutually exclusive arguments, please use only one!")

        max_fps = max_fps if max_fps is not None else self.max_fps

        if metadata.fps is None:
            metadata.fps = fps if fps is not None else max_fps
            logger.warning_once(
                "Molmo2 inserts frame timestamps into video prompts, but the input video's `fps` was not provided "
                f"or could not be inferred. Defaulting to `fps={metadata.fps}`. Please provide `video_metadata` "
                "for more accurate timestamps."
            )
        if metadata.duration is None:
            metadata.duration = metadata.total_num_frames / metadata.fps

        if fps is not None:
            target_num_frames = int(metadata.duration * fps)
        else:
            target_num_frames = num_frames if num_frames is not None else self.num_frames
            if max_fps is not None and metadata.fps > max_fps:
                target_num_frames = min(target_num_frames, int(metadata.duration * max_fps))

        target_num_frames = max(min(target_num_frames, metadata.total_num_frames), 1)
        total = metadata.total_num_frames
        return torch.arange(0, total, total / target_num_frames).int()

    def _build_video_patches(
        self,
        video: torch.Tensor,
        crop_size: int,
        resample: PILImageResampling,
        do_rescale: bool,
        rescale_factor: float,
        do_normalize: bool,
        image_mean: list[float],
        image_std: list[float],
        image_patch_size: int,
        image_pooling_h: int,
        image_pooling_w: int,
    ) -> tuple[list[int], torch.Tensor, torch.Tensor]:
        # `build_resized_image` takes a whole `[N, C, H, W]` batch, so all frames of a video are resized in one call.
        resized_frames, resized_index_grid = build_resized_image(
            self,
            video,
            crop_size=crop_size,
            resample=resample,
            do_rescale=do_rescale,
            rescale_factor=rescale_factor,
            do_normalize=do_normalize,
            image_mean=image_mean,
            image_std=image_std,
            image_patch_size=image_patch_size,
        )
        # The pooling index grid is shape-dependent only, hence shared by every frame.
        pooling_indices, num_patch_rows, num_patch_cols = arange_for_pooling(
            resized_index_grid, image_pooling_h, image_pooling_w
        )
        # [T, 1, S, S, C] -> [T, n_patch, pixels_per_patch]
        patches = batch_pixels_to_patches(resized_frames.squeeze(1), image_patch_size)
        return [num_patch_rows, num_patch_cols], patches, pooling_indices

    def _preprocess(
        self,
        videos: list["torch.Tensor"],
        size: SizeDict,
        resample: PILImageResampling | None,
        do_rescale: bool,
        rescale_factor: float,
        do_normalize: bool,
        image_mean: float | list[float] | None,
        image_std: float | list[float] | None,
        patch_size: int,
        pooling_size: list[int],
        return_tensors: str | TensorType | None = None,
        **kwargs,
    ) -> BatchFeature:
        if size.height != size.width:
            raise ValueError(f"Molmo2 only supports a square `size`, got height={size.height}, width={size.width}.")
        crop_size = size.height
        image_pooling_h, image_pooling_w = pooling_size

        all_crops: list[torch.Tensor] = []
        all_pooled: list[torch.Tensor] = []
        all_grids: list[torch.Tensor] = []
        patch_offset = 0

        for video in videos:
            image_grid, patches, pooling_indices = self._build_video_patches(
                video,
                crop_size=crop_size,
                resample=resample,
                do_rescale=do_rescale,
                rescale_factor=rescale_factor,
                do_normalize=do_normalize,
                image_mean=image_mean,
                image_std=image_std,
                image_patch_size=patch_size,
                image_pooling_h=image_pooling_h,
                image_pooling_w=image_pooling_w,
            )
            num_frames, patches_per_frame = patches.shape[:2]
            frame_offsets = (
                patch_offset + torch.arange(num_frames, device=pooling_indices.device) * patches_per_frame
            ).view(-1, 1, 1)
            pooled_indices = torch.where(pooling_indices >= 0, pooling_indices + frame_offsets, pooling_indices)
            all_pooled.append(pooled_indices.reshape(-1, pooling_indices.shape[-1]))
            all_crops.append(patches)
            patch_offset += num_frames * patches_per_frame

            all_grids.append(torch.tensor([num_frames, image_grid[0], image_grid[1]], dtype=torch.int64))

        data = {
            "pixel_values_videos": torch.cat(all_crops, dim=0),
            "video_token_pooling": torch.cat(all_pooled, dim=0),
            "video_grids": torch.stack(all_grids, dim=0),
        }
        return BatchFeature(data, tensor_type=return_tensors)


__all__ = ["Molmo2VideoProcessor"]