import math

from transformers.image_utils import SizeDict
from transformers.models.idefics3.image_processing_idefics3 import (
    MAX_IMAGE_SIZE,
    Idefics3ImageProcessor,
    Idefics3ImageProcessorKwargs,
    _resize_output_size_rescale_to_max_len,
    _resize_output_size_scale_below_upper_bound,
)


def get_capped_longest_edge(height: int, width: int, longest_edge: int, tile_size: int, max_image_tiles: int) -> int:
    for edge in range(longest_edge, tile_size - 1, -tile_size):
        resized_height, resized_width = _resize_output_size_rescale_to_max_len(height, width, max_len=edge)
        resized_height, resized_width = _resize_output_size_scale_below_upper_bound(
            resized_height, resized_width, max_len=MAX_IMAGE_SIZE
        )
        aspect_ratio = resized_width / resized_height
        if resized_width >= resized_height:
            cols = math.ceil(resized_width / tile_size)
            rows = math.ceil(int(cols * tile_size / aspect_ratio) / tile_size)
        else:
            rows = math.ceil(resized_height / tile_size)
            cols = math.ceil(int(rows * tile_size * aspect_ratio) / tile_size)
        if rows * cols <= max_image_tiles:
            return edge
    return tile_size


class UltraVBertImageProcessorKwargs(Idefics3ImageProcessorKwargs, total=False):
    max_image_tiles: int | None


class UltraVBertImageProcessor(Idefics3ImageProcessor):
    max_image_tiles = None
    valid_kwargs = UltraVBertImageProcessorKwargs

    def _preprocess(self, images, *, do_resize, size, resample, do_image_splitting, max_image_size,
                    max_image_tiles=None, **kwargs):
        if do_resize and do_image_splitting and max_image_tiles is not None:
            tile_size = max_image_size["longest_edge"]
            images = [
                [
                    self.resize(
                        image,
                        SizeDict(
                            longest_edge=get_capped_longest_edge(
                                *image.shape[-2:], size.longest_edge, tile_size, max_image_tiles
                            )
                        ),
                        resample=resample,
                    )
                    for image in batch_images
                ]
                for batch_images in images
            ]
            do_resize = False
        return super()._preprocess(images, do_resize=do_resize, size=size, resample=resample,
                                   do_image_splitting=do_image_splitting, max_image_size=max_image_size, **kwargs)

    def get_number_of_image_patches(self, height: int, width: int, images_kwargs: dict | None = None):
        images_kwargs = dict(images_kwargs or {})
        max_image_tiles = images_kwargs.get("max_image_tiles", self.max_image_tiles)
        if max_image_tiles is not None:
            size = images_kwargs.get("size", self.size)
            tile_size = images_kwargs.get("max_image_size", self.max_image_size)["longest_edge"]
            longest_edge = get_capped_longest_edge(height, width, size["longest_edge"], tile_size, max_image_tiles)
            images_kwargs["size"] = {"longest_edge": longest_edge}
        return super().get_number_of_image_patches(height, width, images_kwargs)


__all__ = ["UltraVBertImageProcessor"]