#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
#           This file was automatically generated from src/transformers/models/gemma4_unified/modular_gemma4_unified.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_gemma4_unified.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 math

import torch
from torchvision.transforms.v2 import functional as tvF

from ...image_processing_backends import TorchvisionBackend
from ...image_processing_utils import BatchFeature
from ...image_utils import ImageInput, PILImageResampling
from ...processing_utils import ImagesKwargs, Unpack
from ...utils import (
    TensorType,
    auto_docstring,
)


class Gemma4UnifiedImageProcessorKwargs(ImagesKwargs, total=False):
    """
    patch_size (`int`, *optional*):
        Size of each teacher image patch in pixels (before merging).
    max_soft_tokens (`int`, *optional*):
        Maximum number of soft (vision) tokens per image after patch merging.
        Must be one of {70, 140, 280, 560, 1120}.
    pooling_kernel_size (`int`, *optional*):
        Kernel size for merging teacher patches into model patches.
    """

    patch_size: int
    max_soft_tokens: int
    pooling_kernel_size: int


_SUPPORTED_SOFT_TOKENS = (70, 140, 280, 560, 1120)


def get_aspect_ratio_preserving_size(
    height: int,
    width: int,
    patch_size: int,
    max_patches: int,
    pooling_kernel_size: int,
) -> tuple[int, int]:
    """
    Image is resized to preserve aspect ratio so it fits within the patch budget.
    Target dimensions are the largest that:
    1) Produce at most `max_patches` patches when patchified with `patch_size`
    2) Have height and width divisible by `pooling_kernel_size * patch_size`
    """
    total_px = height * width
    target_px = max_patches * (patch_size**2)
    factor = math.sqrt(target_px / total_px)
    ideal_height = factor * height
    ideal_width = factor * width
    side_mult = pooling_kernel_size * patch_size

    # Round down to nearest multiple of side_mult
    target_height = int(math.floor(ideal_height / side_mult)) * side_mult
    target_width = int(math.floor(ideal_width / side_mult)) * side_mult

    # Handle edge cases where one or both dimensions round to 0
    if target_height == 0 and target_width == 0:
        raise ValueError(
            "Attempting to resize to a 0 x 0 image. Resized height should be divisble by "
            f"`pooling_kernel_size * patch_size`={pooling_kernel_size * patch_size}."
        )

    max_side_length = (max_patches // pooling_kernel_size**2) * side_mult
    if target_height == 0:
        target_height = side_mult
        target_width = min(
            int(math.floor(width / height)) * side_mult,
            max_side_length,
        )
    elif target_width == 0:
        target_width = side_mult
        target_height = min(
            int(math.floor(height / width)) * side_mult,
            max_side_length,
        )

    if target_height * target_width > target_px:
        raise ValueError(
            f"Resizing [{height}x{width}] to [{target_height}x{target_width}] "
            f"but this exceeds {max_patches} patches with patch_size {patch_size}"
        )

    return target_height, target_width


def convert_image_to_patches(image: "torch.Tensor", patch_size: int) -> "torch.Tensor":
    """
    Convert 3D tensor image of shape (num_channels, image_height, image_width) into 2D tensor of patches of shape
    (num_patches_height * num_patches_width, patch_size * patch_size * num_channels).
    """
    num_channels, image_height, image_width = image.shape
    num_patches_height = image_height // patch_size
    num_patches_width = image_width // patch_size
    patched_image = image.reshape(num_channels, num_patches_height, patch_size, num_patches_width, patch_size)
    patched_image = patched_image.permute(1, 3, 2, 4, 0)
    patched_image = patched_image.reshape(num_patches_height * num_patches_width, -1)
    return patched_image


# Adopted from Siglip2 (mask -> position ids)
def pad_along_first_dim(
    image: "torch.Tensor", positions: "torch.Tensor", target_length: int
) -> tuple["torch.Tensor", "torch.Tensor"]:
    """
    Pad the tensor along the first dimension.
    """
    current_length = image.shape[0]
    padding_length = target_length - current_length
    if padding_length > 0:
        padding = [0, 0] * (image.ndim - 1) + [0, padding_length]
        pos_padding = (0, 0, 0, padding_length)
        image = torch.nn.functional.pad(image, padding, mode="constant", value=0)
        positions = torch.nn.functional.pad(positions, pos_padding, mode="constant", value=-1)
    return image, positions


def patches_merge(
    patches: "torch.Tensor",
    positions_xy: "torch.Tensor",
    length: int,
) -> tuple["torch.Tensor", "torch.Tensor"]:
    """Merge k×k groups of small patches into larger patches.

    Given `L` input patches of dimension `D = patch_size² × 3`, merge groups of
    `k×k` spatially adjacent patches into `length` output patches of dimension
    `(k × patch_size)² × 3`. The spatial grouping is determined by integer-dividing
    the XY positions by `k`.

    Args:
        patches: (*, L, D) — input patches.
        positions_xy: (*, L, 2) — integer XY positions for each patch (-1 for padding).
        length: target number of output patches. Must satisfy L = length × k².

    Returns:
        merged_patches: (*, length, k²×D) — merged patch features.
        merged_positions: (*, length, 2) — new XY positions for merged patches.
    """
    patch_size = math.isqrt(patches.shape[-1] // 3)
    if patches.shape[-1] != patch_size * patch_size * 3:
        raise ValueError(f"Patch dimension {patches.shape[-1]} is not a valid `patch_size * patch_size * 3`")

    k = math.isqrt(patches.shape[-2] // length)
    if k * k * length != patches.shape[-2]:
        raise ValueError(f"Cannot merge {patches.shape} to {length}")

    # Compute target ordering for reordering patches into kernel-grouped order.
    # This ensures patches within each k×k kernel are contiguous.
    max_x = positions_xy[..., 0].max(dim=-1, keepdim=True)[0] + 1
    kernel_idxs = torch.div(positions_xy, k, rounding_mode="floor")
    num_patches_from_top_left = k * k * kernel_idxs[..., 0] + k * max_x * kernel_idxs[..., 1]

    position_within_kernel = torch.remainder(positions_xy, k)
    num_patches_from_top_left_of_kernel = position_within_kernel[..., 0] + position_within_kernel[..., 1] * k
    target_ordering = num_patches_from_top_left_of_kernel + num_patches_from_top_left

    # Reorder patches by computing the inverse permutation via argsort,
    # then gathering patches into kernel-grouped order.
    perm = target_ordering.long().argsort(dim=-1)  # inverse permutation
    # Expand perm indices to match patch feature dimension for gathering
    perm_expanded = perm.unsqueeze(-1).expand_as(patches)
    kernel_ordered_patches = patches.gather(-2, perm_expanded)

    batch_shape = patches.shape[:-2]

    # Reshape: (*, length*k*k, patch_size*patch_size*3) → (*, length, (k*patch_size)*(k*patch_size)*3)
    kernel_ordered_patches = kernel_ordered_patches.reshape(*batch_shape, length, k * k, patch_size, patch_size, 3)
    # Rearrange (l, a*b, p, q, c) → (l, a*p, b*q, c)
    kernel_ordered_patches = kernel_ordered_patches.reshape(*batch_shape, length, k, k, patch_size, patch_size, 3)
    kernel_ordered_patches = kernel_ordered_patches.permute(
        *range(len(batch_shape)), -6, -5, -3, -4, -2, -1
    )  # (..., l, k, p, k, q, c)
    merged_patches = kernel_ordered_patches.reshape(*batch_shape, length, k * patch_size * k * patch_size * 3)

    # Compute new positions for merged patches
    perm_pos = perm.unsqueeze(-1).expand_as(positions_xy)
    kernel_ordered_positions = positions_xy.float().gather(-2, perm_pos.long())

    # Handle padding: preserve -1 positions
    padding = (positions_xy == -1).all(dim=-1, keepdim=True)  # (..., L, 1)
    kernel_ordered_positions = kernel_ordered_positions * (~padding).float() + positions_xy.float() * padding.float()

    # Reshape positions and take min within each kernel to get the merged position
    kernel_ordered_positions = kernel_ordered_positions.reshape(*batch_shape, length, k * k, 2)
    new_positions = torch.div(kernel_ordered_positions, k, rounding_mode="floor")
    # For each merged patch, take the minimum position across the kernel
    new_positions = new_positions.min(dim=-2)[0].to(torch.long)

    return merged_patches, new_positions


@auto_docstring(custom_intro="Constructs a Gemma4 unified image processor.")
class Gemma4UnifiedImageProcessor(TorchvisionBackend):
    resample = PILImageResampling.BICUBIC
    image_mean = [0.0, 0.0, 0.0]
    image_std = [1.0, 1.0, 1.0]
    size = None
    default_to_square = True
    do_convert_rgb = True
    do_resize = True
    do_rescale = True
    do_normalize = False
    patch_size = 16
    max_soft_tokens = 280
    pooling_kernel_size = 3
    valid_kwargs = Gemma4UnifiedImageProcessorKwargs
    model_input_names = ["pixel_values", "image_position_ids", "num_soft_tokens_per_image"]

    def __init__(self, **kwargs: Unpack[Gemma4UnifiedImageProcessorKwargs]):
        super().__init__(**kwargs)

        if self.max_soft_tokens not in _SUPPORTED_SOFT_TOKENS:
            raise ValueError(f"`max_soft_tokens` must be one of {_SUPPORTED_SOFT_TOKENS}, got {self.max_soft_tokens}.")

    def _validate_preprocess_kwargs(self, **kwargs):
        # Gemma4Unified uses aspect_ratio_preserving_resize driven by patch_size,
        # max_soft_tokens, and pooling_kernel_size — not the standard `size`
        # parameter. Temporarily disable do_resize so the base validation
        # doesn't require `size` to be set.
        kwargs["do_resize"] = False
        super()._validate_preprocess_kwargs(**kwargs)

    def aspect_ratio_preserving_resize(
        self,
        image: torch.Tensor,
        patch_size: int,
        max_patches: int,
        pooling_kernel_size: int,
        resample: tvF.InterpolationMode,
    ) -> torch.Tensor:
        height, width = image.shape[-2], image.shape[-1]
        target_height, target_width = get_aspect_ratio_preserving_size(
            height=height,
            width=width,
            patch_size=patch_size,
            max_patches=max_patches,
            pooling_kernel_size=pooling_kernel_size,
        )

        if target_height == height and target_width == width:
            return image

        return tvF.resize(
            image,
            size=[target_height, target_width],
            interpolation=resample,
            antialias=True,
        )

    def preprocess(
        self,
        images: ImageInput,
        **kwargs: Unpack[Gemma4UnifiedImageProcessorKwargs],
    ) -> BatchFeature:
        return super().preprocess(images, **kwargs)

    def _preprocess(
        self,
        images: list["torch.Tensor"],
        do_resize: bool,
        resample: "PILImageResampling | tvF.InterpolationMode | int | None",
        do_rescale: bool,
        rescale_factor: float,
        do_normalize: bool,
        image_mean: float | list[float] | None,
        image_std: float | list[float] | None,
        return_tensors: str | TensorType | None,
        patch_size: int | None = None,
        max_soft_tokens: int | None = None,
        pooling_kernel_size: int | None = None,
        **kwargs,
    ) -> BatchFeature:
        if max_soft_tokens not in _SUPPORTED_SOFT_TOKENS:
            raise ValueError(f"`max_soft_tokens` must be one of {_SUPPORTED_SOFT_TOKENS}, got {max_soft_tokens}.")

        # Compute max_patches from max_soft_tokens and pooling_kernel_size
        max_patches = max_soft_tokens * pooling_kernel_size**2

        # Process each image individually: resize, rescale/normalize, patchify, pad.
        # Images have different aspect ratios and thus different resized dimensions,
        # so patchification and padding must happen per-image before stacking.
        pixel_values = []
        position_ids = []
        num_soft_tokens_per_image = []

        for image in images:
            # Step 1: Aspect-ratio-preserving resize
            if do_resize:
                image = self.aspect_ratio_preserving_resize(
                    image=image,
                    patch_size=patch_size,
                    max_patches=max_patches,
                    pooling_kernel_size=pooling_kernel_size,
                    resample=resample,
                )

            # Step 2: Rescale pixel values (typically to [0, 1]) and optionally identity normalize
            image = self.rescale_and_normalize(image, do_rescale, rescale_factor, do_normalize, image_mean, image_std)

            # Step 3: Patchify into teacher-size patches (16px)
            # (num_channels, height, width) → (num_teacher_patches, patch_size²*3)
            patch_height = image.shape[-2] // patch_size
            patch_width = image.shape[-1] // patch_size
            teacher_patches = convert_image_to_patches(image, patch_size)

            # Step 4: Compute teacher-level position IDs
            device = image.device
            patch_grid = torch.meshgrid(
                torch.arange(patch_width, device=device),
                torch.arange(patch_height, device=device),
                indexing="xy",
            )
            teacher_positions = torch.stack(patch_grid, dim=-1).reshape(teacher_patches.shape[0], 2)

            # Step 5: Merge k×k teacher patches into model patches via patches_merge
            # (num_teacher_patches, 768) → (num_model_patches, 6912)
            num_model_patches = teacher_patches.shape[0] // (pooling_kernel_size**2)
            merged_patches, merged_positions = patches_merge(
                teacher_patches.unsqueeze(0),  # add batch dim for patches_merge
                teacher_positions.unsqueeze(0),
                num_model_patches,
            )
            merged_patches = merged_patches.squeeze(0)  # remove batch dim
            merged_positions = merged_positions.squeeze(0)
            num_soft_tokens_per_image.append(merged_patches.shape[0])

            # Step 6: Pad merged patches and positions to max_soft_tokens
            merged_patches, merged_positions = pad_along_first_dim(merged_patches, merged_positions, max_soft_tokens)
            pixel_values.append(merged_patches)
            position_ids.append(merged_positions)

        # Stack into batch tensors
        pixel_values = torch.stack(pixel_values, dim=0)  # (batch, max_soft_tokens, model_patch_size²*3)
        position_ids = torch.stack(position_ids, dim=0)  # (batch, max_soft_tokens, 2)

        data = {
            "pixel_values": pixel_values,
            "image_position_ids": position_ids,
            "num_soft_tokens_per_image": num_soft_tokens_per_image,
        }
        return BatchFeature(data=data, tensor_type=return_tensors)


__all__ = ["Gemma4UnifiedImageProcessor"]
