#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
#           This file was automatically generated from src/transformers/models/embedding_gemma2/modular_embedding_gemma2.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_embedding_gemma2.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 numpy as np

from ...audio_utils import AudioInput, is_valid_audio
from ...image_utils import ImageInput, make_nested_list_of_images
from ...processing_utils import MultiModalData, ProcessingKwargs, ProcessorMixin, Unpack
from ...tokenization_utils_base import PreTokenizedInput, TextInput
from ...utils import (
    auto_docstring,
)
from ...utils.import_utils import requires
from ...video_utils import VideoInput, is_valid_video, make_batched_videos


class EmbeddingGemma2ProcessorKwargs(ProcessingKwargs, total=False):
    _defaults = {
        "text_kwargs": {
            "padding": True,
        },
        "images_kwargs": {
            "do_convert_rgb": True,
        },
        "audio_kwargs": {},
        "videos_kwargs": {"return_metadata": True},
    }


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 divisible 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


@auto_docstring
@requires(backends=("vision",))
class EmbeddingGemma2Processor(ProcessorMixin):
    valid_processor_kwargs = EmbeddingGemma2ProcessorKwargs

    def __init__(
        self,
        feature_extractor,
        image_processor,
        tokenizer,
        video_processor,
        chat_template=None,
        image_seq_length: int = 280,
        audio_seq_length: int = 750,
        audio_ms_per_token: int = 40,
        **kwargs,
    ):
        r"""
        image_seq_length (`int`, *optional*, defaults to 280):
            The number of soft tokens per image used for placeholder expansion.
        audio_seq_length (`int`, *optional*, defaults to 750):
            The maximum number of audio soft tokens per audio segment. Serves as an
            upper-bound cap when dynamic audio token counts are computed.
        audio_ms_per_token (`int`, *optional*, defaults to 40):
            Milliseconds of audio per output soft token. Used to dynamically compute
            the number of audio placeholder tokens as ``ceil(duration_ms / audio_ms_per_token)``.
            The default of 40 comes from the SSCP convolution's 4× time reduction on 10ms frames.
        """
        self.image_seq_length = image_seq_length
        self.image_token_id = tokenizer.image_token_id
        self.boi_token = tokenizer.boi_token
        self.eoi_token = tokenizer.eoi_token
        self.image_token = tokenizer.image_token

        # FIXME: add the token to config and ask Ryan to re-upload
        tokenizer.add_special_tokens({"additional_special_tokens": ["<|video|>"]})
        self.video_token = "<|video|>"
        self.video_token_id = tokenizer.convert_tokens_to_ids(self.video_token)

        # Audio token handling, mirroring the vision pattern.
        # audio_seq_length serves as the maximum cap on the number of audio soft tokens
        # any single audio segment can produce. With dynamic audio tokens, the actual
        # number of placeholders inserted per audio is computed from the audio duration.
        self.audio_seq_length = audio_seq_length
        # Milliseconds of audio per output soft token. The default of 40 comes from the
        # SSCP convolution's 4× time reduction applied to 10ms mel spectrogram frames.
        self.audio_ms_per_token = audio_ms_per_token
        self.audio_token_id = getattr(tokenizer, "audio_token_id", None)
        self.audio_token = getattr(tokenizer, "audio_token", None)
        self.boa_token = getattr(tokenizer, "boa_token", None)
        self.eoa_token = getattr(tokenizer, "eoa_token", None)

        super().__init__(
            feature_extractor=feature_extractor,
            image_processor=image_processor,
            tokenizer=tokenizer,
            video_processor=video_processor,
            chat_template=chat_template,
            **kwargs,
        )

    def prepare_inputs_layout(
        self,
        images: ImageInput | None = None,
        text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] = None,
        videos: VideoInput = None,
        audio: AudioInput = None,
        **kwargs,
    ):
        # When `text` is None, record per-sample counts for `audio` and `videos` before
        # `ProcessorMixin.prepare_inputs_layout` runs `make_list_of_audio` and `make_batched_videos`,
        # which flatten 2D nested per-sample lists into a 1D list of total items and discard sample
        # boundaries. `images` does not need this because `make_nested_list_of_images` preserves the
        # outer per-sample list structure.
        if not text:
            audio_per_sample = (
                [len(el) if isinstance(el, (list, tuple)) and not is_valid_audio(el) else 1 for el in audio]
                if isinstance(audio, (list, tuple)) and not is_valid_audio(audio)
                else None
            )
            videos_per_sample = (
                [
                    len(make_batched_videos(el)) if not (isinstance(el, (list, tuple)) and not el) else 0
                    for el in videos
                ]
                if isinstance(videos, (list, tuple)) and not is_valid_video(videos)
                else None
            )

        images, text, videos, audio = super().prepare_inputs_layout(
            images=images, text=text, videos=videos, audio=audio, **kwargs
        )

        # Model requires nested struct
        if images is not None:
            images = make_nested_list_of_images(images)

        # Normalize videos so len(videos) gives the number of videos, not frames
        if videos is not None:
            videos = make_batched_videos(videos)

        if not text:
            modality_counts = []
            if images is not None:
                modality_counts.append((self.image_token, [len(image_list) for image_list in images]))
            if videos is not None:
                modality_counts.append((self.video_token, videos_per_sample or [1] * len(videos)))
            if audio is not None:
                modality_counts.append((self.audio_token, audio_per_sample or [1] * len(audio)))

            if modality_counts:
                batch_sizes = {len(counts) for _, counts in modality_counts}
                if len(batch_sizes) > 1:
                    raise ValueError(
                        f"Received inconsistently sized modality batches when `text` is None: "
                        f"{[len(counts) for _, counts in modality_counts]}."
                    )
                batch_size = batch_sizes.pop()
                text = [
                    " ".join(token for token, counts in modality_counts for _ in range(counts[sample_idx]))
                    for sample_idx in range(batch_size)
                ]

        return images, text, videos, audio

    def validate_inputs(
        self,
        images: ImageInput | list[ImageInput] | None = None,
        text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] = None,
        videos: VideoInput = None,
        audio: AudioInput = None,
        **kwargs: Unpack[ProcessingKwargs],
    ):
        super().validate_inputs(images=images, text=text, **kwargs)

        # Unlike Gemma 4, any single modality on its own is a valid embedding input.
        if text is None and images is None and videos is None and audio is None:
            raise ValueError("You must provide at least one of `text`, `images`, `videos`, or `audio`.")

        if audio is not None and (self.audio_token is None or self.boa_token is None or self.eoa_token is None):
            raise ValueError("Audio inputs were provided, but the tokenizer does not have an `audio_token` defined.")

        if text is not None:
            n_images_in_text = [sample.count(self.image_token) for sample in text]
            if images is not None:
                if len(images) != len(text):
                    raise ValueError(
                        f"Received inconsistently sized batches of images ({len(images)}) and text ({len(text)})."
                    )

                n_images_in_images = [len(sublist) for sublist in images]
                if n_images_in_text != n_images_in_images:
                    raise ValueError(
                        f"The total number of {self.image_token} tokens in the prompts should be the same as the number of images passed."
                        f" Found {n_images_in_text} {self.image_token} tokens and {n_images_in_images} images per sample."
                    )
            elif any(n_images_in_text):
                raise ValueError(
                    f"Found {sum(n_images_in_text)} {self.image_token} tokens in the text but no images were passed."
                )

            n_videos_in_text = [sample.count(self.video_token) for sample in text]
            if videos is not None:
                if sum(n_videos_in_text) != len(videos):
                    raise ValueError(
                        f"The total number of {self.video_token} tokens in the prompts should be the same as the number of videos passed."
                        f" Found {sum(n_videos_in_text)} {self.video_token} tokens and {len(videos)} videos."
                    )
            elif any(n_videos_in_text):
                raise ValueError(
                    f"Found {sum(n_videos_in_text)} {self.video_token} tokens in the text but no videos were passed."
                )

            if self.audio_token is not None:
                n_audio_in_text = [sample.count(self.audio_token) for sample in text]
                if audio is not None:
                    n_audio_passed = len(audio) if isinstance(audio, (list, tuple)) else 1
                    if sum(n_audio_in_text) != n_audio_passed:
                        raise ValueError(
                            f"The total number of {self.audio_token} tokens in the prompts should be the same as the number of audio inputs passed."
                            f" Found {sum(n_audio_in_text)} {self.audio_token} tokens and {n_audio_passed} audio inputs."
                        )
                elif any(n_audio_in_text):
                    raise ValueError(
                        f"Found {sum(n_audio_in_text)} {self.audio_token} tokens in the text but no audio inputs were passed."
                    )

    def replace_image_token(self, image_inputs: dict, image_idx: int, **kwargs) -> str:
        num_soft_tokens = image_inputs["num_soft_tokens_per_image"][image_idx]
        return f"{self.boi_token}{self.image_token * num_soft_tokens}{self.eoi_token}"

    def replace_video_token(self, video_inputs: dict, video_idx: int, **kwargs) -> str:
        num_soft_tokens = video_inputs["num_soft_tokens_per_video"][video_idx]
        add_timestamps = kwargs.get("add_timestamps", self.video_processor.add_timestamps)

        # Visual-only mode: one block per frame, no timestamps
        if not add_timestamps:
            # `pixel_values_videos` is a flat frame sequence, so its leading axis indexes frames, not videos
            num_frames = int(video_inputs["num_frames_per_video"][video_idx])
            frame_str = f"{self.boi_token}{self.video_token * num_soft_tokens}{self.eoi_token}"
            return "".join([frame_str] * num_frames)

        metadata = video_inputs["video_metadata"][video_idx]

        if metadata.fps is None:
            raise ValueError(
                "Asked to build a prompt with frame timestamps, but no `fps` was provided in video metadata. "
                "The capture rate of already-decoded frames cannot be inferred. Please pass a `VideoMetadata` "
                "object with a valid `fps`, or set `add_timestamps=False`."
            )

        # mm:ss format for timestamps
        timestamp_str = [f"{int(seconds // 60):02d}:{int(seconds % 60):02d}" for seconds in metadata.timestamps]
        return " ".join(
            [f"{t} {self.boi_token}{self.video_token * num_soft_tokens}{self.eoi_token}" for t in timestamp_str]
        )

    def replace_audio_token(self, audio_inputs: dict, audio_idx: int, **kwargs) -> str:
        # TODO: Add tests for audio-only processor inputs.
        mask = audio_inputs["input_features_mask"][audio_idx]

        # Simulate two stride-2 conv blocks on the mask
        t = len(mask)
        for _ in range(2):
            t_out = (t + 2 - 3) // 2 + 1
            mask = mask[::2][:t_out]
            t = len(mask)

        return f"{self.boa_token}{self.audio_token * int(mask.sum())}{self.eoa_token}"

    def _get_num_multimodal_tokens(self, image_sizes=None, video_sizes=None, audio_lengths=None, **kwargs):
        """
        Computes the number of placeholder tokens needed for multimodal inputs with the given sizes.

        Args:
            image_sizes (`list[list[int]]`, *optional*):
                The input sizes formatted as (height, width) per each image.
            video_sizes (`list[list[int]]`, *optional*):
                The input sizes formatted as (num_frames, height, width) per each video.
            audio_lengths (`list[int]`, *optional*):
                The lengths of audio inputs in number of samples. Used to dynamically
                compute per-audio token counts.

        Returns:
            `MultiModalData`: A `MultiModalData` object holding number of tokens per each of the provided
            input modalities, along with other useful data.
        """

        images_kwargs = EmbeddingGemma2ProcessorKwargs._defaults.get("images_kwargs", {})
        images_kwargs.update(kwargs)
        patch_size = images_kwargs.get("patch_size", None) or self.image_processor.patch_size
        pooling_kernel_size = (
            images_kwargs.get("pooling_kernel_size", None) or self.image_processor.pooling_kernel_size
        )
        max_soft_tokens = images_kwargs.get("max_soft_tokens", None) or self.image_processor.max_soft_tokens

        max_patches = max_soft_tokens * pooling_kernel_size**2

        vision_data = {}
        if image_sizes is not None:
            num_image_tokens = []
            for image_size in image_sizes:
                target_h, target_w = get_aspect_ratio_preserving_size(
                    height=image_size[0],
                    width=image_size[1],
                    patch_size=patch_size,
                    max_patches=max_patches,
                    pooling_kernel_size=pooling_kernel_size,
                )
                patch_height = target_h // patch_size
                patch_width = target_w // patch_size
                num_image_tokens.append(patch_height * patch_width // pooling_kernel_size**2)

            num_image_patches = [1] * len(image_sizes)
            vision_data.update({"num_image_tokens": num_image_tokens, "num_image_patches": num_image_patches})

        if video_sizes is not None:
            videos_kwargs = EmbeddingGemma2ProcessorKwargs._defaults.get("videos_kwargs", {})
            videos_kwargs.update(kwargs)
            patch_size = videos_kwargs.get("patch_size", None) or self.video_processor.patch_size
            pooling_kernel_size = (
                videos_kwargs.get("pooling_kernel_size", None) or self.video_processor.pooling_kernel_size
            )
            max_soft_tokens = videos_kwargs.get("max_soft_tokens", None) or self.video_processor.max_soft_tokens

            max_patches = max_soft_tokens * pooling_kernel_size**2

            num_video_tokens = []
            for num_frames, height, width in video_sizes:
                target_h, target_w = get_aspect_ratio_preserving_size(
                    height=height,
                    width=width,
                    patch_size=patch_size,
                    max_patches=max_patches,
                    pooling_kernel_size=pooling_kernel_size,
                )
                patch_height = target_h // patch_size
                patch_width = target_w // patch_size
                num_video_tokens.append(num_frames * (patch_height * patch_width // pooling_kernel_size**2))
            vision_data.update({"num_video_tokens": num_video_tokens})

        if audio_lengths is not None:
            # Dynamically compute per-audio token counts from sample lengths.
            # audio_lengths are in number of samples; assume default sampling rate.
            sampling_rate = getattr(self.feature_extractor, "sampling_rate", 16_000)
            num_audio_tokens = [
                self._compute_audio_num_tokens(np.zeros(length), sampling_rate) for length in audio_lengths
            ]
            vision_data.update({"num_audio_tokens": num_audio_tokens})

        return MultiModalData(**vision_data)

    def _compute_audio_num_tokens(self, audio_waveform, sampling_rate: int) -> int:
        """Number of audio soft tokens, replicating the encoder's seq-length arithmetic.

        Mirrors EmbeddingGemma2AudioFeatureExtractor mel framing + the two stride-2 Conv2d
        subsampling layers in EmbeddingGemma2AudioSubSampleConvProjection, capped at
        ``audio_seq_length``. Must match ``audio_mask.sum()`` from the audio tower or
        vLLM's ``_merge_multimodal_embeddings`` will raise on a length mismatch.

        Args:
            audio_waveform: A 1-D array or list containing the raw audio samples.
            sampling_rate: The sampling rate of the audio waveform in Hz.

        Returns:
            The number of audio soft tokens to insert as placeholders.
        """
        num_samples = len(audio_waveform)

        # Step 1: mel frames (matches feature_extraction_embedding_gemma2.py)
        frame_size_for_unfold = self.feature_extractor.frame_length + 1
        pad_left = self.feature_extractor.frame_length // 2  # semicausal time padding
        num_mel_frames = (num_samples + pad_left - frame_size_for_unfold) // self.feature_extractor.hop_length + 1
        if num_mel_frames <= 0:
            return 0

        # Step 2: SSCP subsampling. Two stride-2 Conv2d layers (kernel=3, pad=1)
        # from EmbeddingGemma2AudioSubSampleConvProjection(Layer) in modeling_embedding_gemma2.py.
        # These are architectural constants, not exposed via EmbeddingGemma2AudioConfig
        # (note: config.conv_kernel_size=5 is the *conformer* depthwise conv, NOT
        # this one). The layer count tracks len(subsampling_conv_channels)==2.
        # Keep in sync if the audio encoder's subsampling stack changes.
        sscp_num_layers = 2
        sscp_kernel, sscp_stride, sscp_padding = 3, 2, 1
        t = num_mel_frames
        for _ in range(sscp_num_layers):
            # Conv1d/2d output-length along the time axis (dilation=1).
            t = (t + 2 * sscp_padding - sscp_kernel) // sscp_stride + 1

        # Cap at the configured maximum
        return min(t, self.audio_seq_length)

    @property
    def unused_input_names(self) -> list[str]:
        return ["num_soft_tokens_per_image", "num_soft_tokens_per_video"]


__all__ = ["EmbeddingGemma2Processor"]