#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
#           This file was automatically generated from src/transformers/models/llava_onevision/modular_llava_onevision.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_llava_onevision.py file directly. One of our CI enforces this.
#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
# Copyright 2024 the HuggingFace Inc. 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 numpy as np

from ...image_processing_backends import PilBackend
from ...image_processing_utils import BatchFeature, get_patch_output_size, select_best_resolution
from ...image_transforms import divide_to_patches
from ...image_utils import (
    OPENAI_CLIP_MEAN,
    OPENAI_CLIP_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    SizeDict,
)
from ...processing_utils import Unpack
from ...utils import TensorType, auto_docstring
from ...utils.import_utils import requires
from .image_processing_llava_onevision import LlavaOnevisionImageProcessorKwargs


@requires(backends=("vision", "torch", "torchvision"))
@auto_docstring
class LlavaOnevisionImageProcessorPil(PilBackend):
    model_input_names = ["pixel_values", "image_sizes", "batch_num_images"]
    valid_kwargs = LlavaOnevisionImageProcessorKwargs
    resample = PILImageResampling.BICUBIC
    image_mean = OPENAI_CLIP_MEAN
    image_std = OPENAI_CLIP_STD
    size = {"height": 384, "width": 384}
    default_to_square = False
    crop_size = None
    do_resize = True
    do_center_crop = None
    do_rescale = True
    do_normalize = True
    do_convert_rgb = True
    do_pad = True
    image_grid_pinpoints = [[384, 384], [384, 768], [384, 1152], [384, 1536], [384, 1920], [384, 2304], [768, 384], [768, 768], [768, 1152], [768, 1536], [768, 1920], [768, 2304], [1152, 384], [1152, 768], [1152, 1152], [1152, 1536], [1152, 1920], [1152, 2304], [1536, 384], [1536, 768], [1536, 1152], [1536, 1536], [1536, 1920], [1536, 2304], [1920, 384], [1920, 768], [1920, 1152], [1920, 1536], [1920, 1920], [1920, 2304], [2304, 384], [2304, 768], [2304, 1152], [2304, 1536], [2304, 1920], [2304, 2304]]  # fmt: skip

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

    @auto_docstring
    def preprocess(self, images: ImageInput, **kwargs: Unpack[LlavaOnevisionImageProcessorKwargs]) -> BatchFeature:
        if isinstance(images, (tuple, list)) and isinstance(images[0], (tuple, list)):
            # if the first element is a list, we assume that all elements are lists
            images = [x for x in images if x]  # handle text-only case
            batch_num_images = [len(x) for x in images]
        elif isinstance(images, (tuple, list)):
            # treat this as a single-image case for backward compatibility
            batch_num_images = [1] * len(images)
        else:
            batch_num_images = [1]
        return super().preprocess(images, batch_num_images, **kwargs)

    def _get_padding_size(self, original_resolution: tuple, target_resolution: tuple):
        """Get padding size for patching (returns tuple format for np.pad)."""
        original_height, original_width = original_resolution
        target_height, target_width = target_resolution
        paste_x, r_x = divmod(target_width - original_width, 2)
        paste_y, r_y = divmod(target_height - original_height, 2)
        return (paste_y, paste_y + r_y), (paste_x, paste_x + r_x)

    def _resize_for_patching(
        self,
        image: np.ndarray,
        target_resolution: tuple,
        resample: PILImageResampling,
    ) -> np.ndarray:
        """Resizes an image to a target resolution while maintaining aspect ratio."""
        new_height, new_width = get_patch_output_size(
            image, target_resolution, input_data_format=ChannelDimension.FIRST
        )
        resized_image = self.resize(image=image, size=SizeDict(height=new_height, width=new_width), resample=resample)

        return resized_image

    def _pad_for_patching(self, image: np.ndarray, target_resolution: tuple) -> np.ndarray:
        """Pad an image to a target resolution while maintaining aspect ratio."""
        new_resolution = get_patch_output_size(image, target_resolution, input_data_format=ChannelDimension.FIRST)
        padding_hw = self._get_padding_size(new_resolution, target_resolution)

        # For channels_first format (C, H, W), add (0, 0) for channel dimension
        # padding_hw is ((before_h, after_h), (before_w, after_w))
        # np.pad expects ((before_C, after_C), (before_H, after_H), (before_W, after_W))
        padding = ((0, 0), padding_hw[0], padding_hw[1])

        # Use np.pad directly for patching padding
        padded_image = np.pad(image, padding, mode="constant", constant_values=0)

        return padded_image

    def get_image_patches(
        self,
        image: np.ndarray,
        grid_pinpoints: list[list[int]],
        size: tuple,
        patch_size: int,
        resample: PILImageResampling,
    ) -> list[np.ndarray]:
        """Process an image with variable resolutions by dividing it into patches."""
        if not isinstance(grid_pinpoints, list):
            raise TypeError("grid_pinpoints must be a list of possible resolutions.")

        possible_resolutions = grid_pinpoints

        image_size = image.shape[-2:]
        best_resolution = select_best_resolution(image_size, possible_resolutions)
        resized_image = self._resize_for_patching(image, best_resolution, resample=resample)
        padded_image = self._pad_for_patching(resized_image, best_resolution)

        patches = divide_to_patches(padded_image, patch_size=patch_size)

        size_height, size_width = size
        resized_original_image = self.resize(
            image=image,
            size=SizeDict(height=size_height, width=size_width),
            resample=resample,
        )

        image_patches = [resized_original_image] + patches

        return image_patches

    def _pad_for_batching(
        self,
        pixel_values: list[np.ndarray],
    ) -> list[np.ndarray]:
        """Pads images on the `num_of_patches` dimension with zeros to form a batch of same number of patches."""
        max_patch = max(len(x) for x in pixel_values)
        # Use np.pad directly for patch dimension padding
        padded_values = []
        for image in pixel_values:
            # Padding format: ((before_dim0, after_dim0), (before_dim1, after_dim1), ...)
            padding = ((0, max_patch - image.shape[0]), (0, 0), (0, 0), (0, 0))
            padded_image = np.pad(image, padding, mode="constant", constant_values=0)
            padded_values.append(padded_image)

        return padded_values

    def _preprocess(
        self,
        images: list[np.ndarray],
        batch_num_images: list[int],
        do_resize: bool,
        size: SizeDict,
        image_grid_pinpoints: list[list[int]],
        resample: "PILImageResampling | int | None",
        do_center_crop: bool,
        crop_size: SizeDict,
        do_rescale: bool,
        rescale_factor: float,
        do_normalize: bool,
        image_mean: float | list[float] | None,
        image_std: float | list[float] | None,
        do_pad: bool | None,
        pad_size: SizeDict | None,
        disable_grouping: bool | None,
        return_tensors: str | TensorType | None,
        **kwargs,
    ) -> BatchFeature:
        """Custom preprocessing for LLaVA-NeXT with patch processing."""
        processed_images = []
        image_sizes = []
        # only single image patching is supported
        need_patching = [n == 1 for n in batch_num_images for _ in range(n)]

        # Backend's resize method handles resample conversion, so we can pass it directly
        # Determine the size tuple
        if size and size.height and size.width:
            size_tuple = (size.height, size.width)
        else:
            size_tuple = (size.shortest_edge, size.shortest_edge)

        # Determine the patch size
        if crop_size and crop_size.height:
            patch_size = crop_size.height
        elif size and size.height:
            patch_size = size.height
        else:
            patch_size = size.shortest_edge

        for i, image in enumerate(images):
            if need_patching[i]:
                # convert image into a list of patches
                # we intentionally use the same data format as the input data format
                image_patches = self.get_image_patches(
                    image,
                    image_grid_pinpoints,
                    size=size_tuple,
                    patch_size=patch_size,
                    resample=resample,
                )
            else:
                padded_image = self.pad_to_square(
                    image=image, background_color=tuple(int(x * 255) for x in self.image_mean)
                )
                image_patches = [padded_image]

            # preprocess patches
            pixel_values = []
            for patch in image_patches:
                if do_resize:
                    patch = self.resize(image=patch, size=size, resample=resample)

                if do_center_crop:
                    patch = self.center_crop(image=patch, size=crop_size)

                if do_rescale:
                    patch = self.rescale(image=patch, scale=rescale_factor)

                if do_normalize:
                    patch = self.normalize(image=patch, mean=image_mean, std=image_std)

                pixel_values.append(patch)

            pixel_values = np.array(pixel_values)
            processed_images.append(pixel_values)
            image_sizes.append(image.shape[-2:])

        if do_pad:
            processed_images = self._pad_for_batching(processed_images)

        return BatchFeature(
            data={
                "pixel_values": processed_images,
                "image_sizes": image_sizes,
                "batch_num_images": batch_num_images,
            },
            tensor_type=return_tensors,
        )

    def pad_to_square(
        self,
        image: np.ndarray,
        background_color: int | tuple[int, int, int] = 0,
    ) -> np.ndarray:
        """
        Pads an image to a square based on the longest edge.

        Args:
            image (`np.ndarray`):
                The image to pad. Shape: (num_channels, height, width) - always channels_first in backend.
            background_color (`int` or `tuple[int, int, int]`, *optional*, defaults to 0):
                The color to use for the padding.

        Returns:
            `np.ndarray`: The padded image.
        """
        # Backend always uses channels_first format: (num_channels, height, width)
        num_channels, height, width = image.shape

        if height == width:
            return image

        max_dim = max(height, width)

        # Ensure background_color is the correct shape
        if isinstance(background_color, int):
            background_color = [background_color]
        elif len(background_color) != num_channels:
            raise ValueError(
                f"background_color must have no more than {num_channels} elements to match the number of channels"
            )

        result = np.zeros((num_channels, max_dim, max_dim), dtype=image.dtype)
        for i, color in enumerate(background_color):
            result[i, :, :] = color
        if width > height:
            start = (max_dim - height) // 2
            result[:, start : start + height, :] = image
        else:
            start = (max_dim - width) // 2
            result[:, :, start : start + width] = image

        return result


__all__ = ["LlavaOnevisionImageProcessorPil"]
