# Copyright 2025 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.
"""Image processor class for Flava."""

import math
import random
from collections.abc import Iterable
from functools import lru_cache
from typing import Any

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

from ...image_processing_backends import TorchvisionBackend
from ...image_processing_utils import BatchFeature, get_size_dict
from ...image_transforms import group_images_by_shape, reorder_images
from ...image_utils import (
    OPENAI_CLIP_MEAN,
    OPENAI_CLIP_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    SizeDict,
)
from ...processing_utils import ImagesKwargs, Unpack
from ...utils import (
    TensorType,
    auto_docstring,
)


# These values are taken from CLIP
FLAVA_IMAGE_MEAN = OPENAI_CLIP_MEAN
FLAVA_IMAGE_STD = OPENAI_CLIP_STD
FLAVA_CODEBOOK_MEAN = [0.0, 0.0, 0.0]
FLAVA_CODEBOOK_STD = [1.0, 1.0, 1.0]
LOGIT_LAPLACE_EPS: float = 0.1


class FlavaImageProcessorKwargs(ImagesKwargs, total=False):
    """
    return_image_mask (`bool`, *optional*, defaults to `False`):
        Whether to return the image mask. Can be overridden by the `return_image_mask` parameter in `preprocess`.
    input_size_patches (`int`, *optional*, defaults to `14`):
        Number of patches in the image in height and width direction. 14x14 = 196 total patches. Can be overridden
        by the `input_size_patches` parameter in `preprocess`.
    total_mask_patches (`int`, *optional*, defaults to `75`):
        Total number of patches that should be masked. Can be overridden by the `total_mask_patches` parameter in
        `preprocess`.
    mask_group_min_patches (`int`, *optional*, defaults to `16`):
        Minimum number of patches that should be masked. Can be overridden by the `mask_group_min_patches`
        parameter in `preprocess`.
    mask_group_max_patches (`int`, *optional*):
        Maximum number of patches that should be masked. Can be overridden by the `mask_group_max_patches`
        parameter in `preprocess`.
    mask_group_min_aspect_ratio (`float`, *optional*, defaults to `0.3`):
        Minimum aspect ratio of the mask window. Can be overridden by the `mask_group_min_aspect_ratio` parameter
        in `preprocess`.
    mask_group_max_aspect_ratio (`float`, *optional*):
        Maximum aspect ratio of the mask window. Can be overridden by the `mask_group_max_aspect_ratio` parameter
        in `preprocess`.
    return_codebook_pixels (`bool`, *optional*, defaults to `False`):
        Whether to return the codebook pixel values.
    codebook_do_resize (`bool`, *optional*, defaults to `True`):
        Whether to resize the input for codebook to a certain. Can be overridden by the `codebook_do_resize`
        parameter in `preprocess`. `codebook_size`.
    codebook_size (`dict[str, int]`, *optional*, defaults to `{"height": 224, "width": 224}`):
        Resize the input for codebook to the given size. Can be overridden by the `codebook_size` parameter in
        `preprocess`.
    codebook_resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.LANCZOS`):
        Resampling filter to use if resizing the codebook image. With torchvision < 0.27, LANCZOS is not
        supported for torch Tensors and BICUBIC is used as the closest alternative. Can be overridden by the
        `codebook_resample` parameter in `preprocess`.
    codebook_do_center_crop (`bool`, *optional*, defaults to `True`):
        Whether to crop the input for codebook at the center. If the input size is smaller than
        `codebook_crop_size` along any edge, the image is padded with 0's and then center cropped. Can be
        overridden by the `codebook_do_center_crop` parameter in `preprocess`.
    codebook_crop_size (`dict[str, int]`, *optional*, defaults to `{"height": 224, "width": 224}`):
        Desired output size for codebook input when applying center-cropping. Can be overridden by the
        `codebook_crop_size` parameter in `preprocess`.
    codebook_do_rescale (`bool`, *optional*, defaults to `True`):
        Whether to rescale the input for codebook by the specified scale `codebook_rescale_factor`. Can be
        overridden by the `codebook_do_rescale` parameter in `preprocess`.
    codebook_rescale_factor (`int` or `float`, *optional*, defaults to `1/255`):
        Defines the scale factor to use if rescaling the codebook image. Can be overridden by the
        `codebook_rescale_factor` parameter in `preprocess`.
    codebook_do_map_pixels (`bool`, *optional*, defaults to `True`):
        Whether to map the pixel values of the codebook input to (1 - 2e)x + e. Can be overridden by the
        `codebook_do_map_pixels` parameter in `preprocess`.
    codebook_do_normalize (`bool`, *optional*, defaults to `True`):
        Whether or not to normalize the input for codebook with `codebook_image_mean` and `codebook_image_std`. Can
        be overridden by the `codebook_do_normalize` parameter in `preprocess`.
    codebook_image_mean (`Optional[Union[float, Iterable[float]]]`, *optional*, defaults to `[0, 0, 0]`):
        The sequence of means for each channel, to be used when normalizing images for codebook. Can be overridden
        by the `codebook_image_mean` parameter in `preprocess`.
    codebook_image_std (`Optional[Union[float, Iterable[float]]]`, *optional*, defaults to `[0.5, 0.5, 0.5]`):
        The sequence of standard deviations for each channel, to be used when normalizing images for codebook. Can
        be overridden by the `codebook_image_std` parameter in `preprocess`.
    """

    # Mask related params
    return_image_mask: bool
    input_size_patches: int
    total_mask_patches: int
    mask_group_min_patches: int
    mask_group_max_patches: int
    mask_group_min_aspect_ratio: float
    mask_group_max_aspect_ratio: float
    # Codebook related params
    return_codebook_pixels: bool
    codebook_do_resize: bool
    codebook_size: dict[str, int]
    codebook_resample: int
    codebook_do_center_crop: bool
    codebook_crop_size: dict[str, int]
    codebook_do_rescale: bool
    codebook_rescale_factor: int | float
    codebook_do_map_pixels: bool
    codebook_do_normalize: bool
    codebook_image_mean: float | Iterable[float]
    codebook_image_std: float | Iterable[float]


# Inspired from https://github.com/microsoft/unilm/blob/master/beit/masking_generator.py
class FlavaMaskingGenerator:
    def __init__(
        self,
        input_size: int | tuple[int, int] = 14,
        total_mask_patches: int = 75,
        mask_group_max_patches: int | None = None,
        mask_group_min_patches: int = 16,
        mask_group_min_aspect_ratio: float | None = 0.3,
        mask_group_max_aspect_ratio: float | None = None,
    ):
        if not isinstance(input_size, tuple):
            input_size = (input_size,) * 2
        self.height, self.width = input_size

        self.num_patches = self.height * self.width
        self.total_mask_patches = total_mask_patches

        self.mask_group_min_patches = mask_group_min_patches
        self.mask_group_max_patches = total_mask_patches if mask_group_max_patches is None else mask_group_max_patches

        mask_group_max_aspect_ratio = mask_group_max_aspect_ratio or 1 / mask_group_min_aspect_ratio
        self.log_aspect_ratio = (math.log(mask_group_min_aspect_ratio), math.log(mask_group_max_aspect_ratio))

    def __repr__(self):
        repr_str = "MaskingGenerator(%d, %d -> [%d ~ %d], max = %d, %.3f ~ %.3f)" % (
            self.height,
            self.width,
            self.mask_group_min_patches,
            self.mask_group_max_patches,
            self.total_mask_patches,
            self.log_aspect_ratio[0],
            self.log_aspect_ratio[1],
        )
        return repr_str

    def get_shape(self):
        return self.height, self.width

    def _mask(self, mask, max_mask_patches):
        delta = 0
        for _attempt in range(10):
            target_area = random.uniform(self.mask_group_min_patches, max_mask_patches)
            aspect_ratio = math.exp(random.uniform(*self.log_aspect_ratio))
            height = int(round(math.sqrt(target_area * aspect_ratio)))
            width = int(round(math.sqrt(target_area / aspect_ratio)))
            if width < self.width and height < self.height:
                top = random.randint(0, self.height - height)
                left = random.randint(0, self.width - width)

                num_masked = mask[top : top + height, left : left + width].sum()
                # Overlap
                if 0 < height * width - num_masked <= max_mask_patches:
                    zeros_pos = mask[top : top + height, left : left + width] == 0
                    mask[top : top + height, left : left + width][zeros_pos] = 1
                    delta += zeros_pos.sum()

                if delta > 0:
                    break
        return delta

    def __call__(self):
        mask = torch.zeros(self.get_shape(), dtype=torch.int)
        mask_count = 0
        while mask_count < self.total_mask_patches:
            max_mask_patches = self.total_mask_patches - mask_count
            max_mask_patches = min(max_mask_patches, self.mask_group_max_patches)

            delta = self._mask(mask, max_mask_patches)
            if delta == 0:
                break
            else:
                mask_count += delta

        return mask


@auto_docstring
class FlavaImageProcessor(TorchvisionBackend):
    valid_kwargs = FlavaImageProcessorKwargs
    resample = PILImageResampling.BICUBIC
    image_mean = FLAVA_IMAGE_MEAN
    image_std = FLAVA_IMAGE_STD
    size = {"height": 224, "width": 224}
    crop_size = {"height": 224, "width": 224}
    do_resize = True
    do_center_crop = True
    do_rescale = True
    do_normalize = True

    # Mask related params
    return_image_mask = False
    input_size_patches = 14
    total_mask_patches = 75
    mask_group_min_patches = 16
    mask_group_max_patches = None
    mask_group_min_aspect_ratio = 0.3
    mask_group_max_aspect_ratio = None
    # Codebook related params
    return_codebook_pixels = False
    codebook_do_resize = True
    codebook_size = {"height": 112, "width": 112}
    codebook_resample = PILImageResampling.LANCZOS
    codebook_do_center_crop = True
    codebook_crop_size = {"height": 112, "width": 112}
    codebook_do_rescale = True
    codebook_rescale_factor = 1 / 255
    codebook_do_map_pixels = True
    codebook_do_normalize = True
    codebook_image_mean = FLAVA_CODEBOOK_MEAN
    codebook_image_std = FLAVA_CODEBOOK_STD

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

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

    @classmethod
    def from_dict(cls, image_processor_dict: dict[str, Any], **kwargs):
        """
        Overrides the `from_dict` method from the base class to make sure parameters are updated if image processor is
        created using from_dict and kwargs e.g. `FlavaImageProcessor.from_pretrained(checkpoint, codebook_size=600)`
        """
        image_processor_dict = image_processor_dict.copy()
        if "codebook_size" in kwargs:
            image_processor_dict["codebook_size"] = kwargs.pop("codebook_size")
        if "codebook_crop_size" in kwargs:
            image_processor_dict["codebook_crop_size"] = kwargs.pop("codebook_crop_size")
        return super().from_dict(image_processor_dict, **kwargs)

    @lru_cache
    def masking_generator(
        self,
        input_size_patches,
        total_mask_patches,
        mask_group_min_patches,
        mask_group_max_patches,
        mask_group_min_aspect_ratio,
        mask_group_max_aspect_ratio,
    ) -> FlavaMaskingGenerator:
        return FlavaMaskingGenerator(
            input_size=input_size_patches,
            total_mask_patches=total_mask_patches,
            mask_group_min_patches=mask_group_min_patches,
            mask_group_max_patches=mask_group_max_patches,
            mask_group_min_aspect_ratio=mask_group_min_aspect_ratio,
            mask_group_max_aspect_ratio=mask_group_max_aspect_ratio,
        )

    def map_pixels(self, image: "torch.Tensor") -> "torch.Tensor":
        return (1 - 2 * LOGIT_LAPLACE_EPS) * image + LOGIT_LAPLACE_EPS

    def _standardize_kwargs(
        self,
        size: int | Iterable[int] | dict[str, int] | SizeDict | None = None,
        crop_size: int | Iterable[int] | dict[str, int] | SizeDict | None = None,
        default_to_square: bool | None = None,
        image_mean: float | list[float] | None = None,
        image_std: float | list[float] | None = None,
        codebook_size: int | Iterable[int] | dict[str, int] | SizeDict | None = None,
        codebook_crop_size: int | Iterable[int] | dict[str, int] | SizeDict | None = None,
        codebook_image_mean: float | list[float] | None = None,
        codebook_image_std: float | list[float] | None = None,
        codebook_resample: "PILImageResampling | tvF.InterpolationMode | int | None" = None,
        data_format: ChannelDimension | None = None,
        **kwargs,
    ) -> dict:
        """
        Update kwargs that need further processing before being validated
        Can be overridden by subclasses to customize the processing of kwargs.
        """
        kwargs = super()._standardize_kwargs(
            size=size,
            crop_size=crop_size,
            default_to_square=default_to_square,
            image_mean=image_mean,
            image_std=image_std,
            data_format=data_format,
            **kwargs,
        )
        if codebook_size is not None and not isinstance(codebook_size, SizeDict):
            codebook_size = SizeDict(**get_size_dict(size=codebook_size, default_to_square=default_to_square))
        if codebook_crop_size is not None and not isinstance(codebook_crop_size, SizeDict):
            codebook_crop_size = SizeDict(**get_size_dict(codebook_crop_size, param_name="codebook_crop_size"))
        if isinstance(codebook_image_mean, list):
            codebook_image_mean = tuple(codebook_image_mean)
        if isinstance(codebook_image_std, list):
            codebook_image_std = tuple(codebook_image_std)

        kwargs["codebook_size"] = codebook_size
        kwargs["codebook_crop_size"] = codebook_crop_size
        kwargs["codebook_image_mean"] = codebook_image_mean
        kwargs["codebook_image_std"] = codebook_image_std
        # Store codebook_resample as-is - resize() will handle conversion
        kwargs["codebook_resample"] = codebook_resample

        return kwargs

    def _preprocess_image(
        self,
        images: list["torch.Tensor"],
        do_resize: bool,
        size: SizeDict,
        resample: "PILImageResampling | tvF.InterpolationMode | int | None",
        do_center_crop: bool,
        crop_size: SizeDict,
        do_rescale: bool,
        rescale_factor: float,
        do_normalize: bool,
        do_map_pixels: bool,
        image_mean: float | list[float] | None,
        image_std: float | list[float] | None,
        disable_grouping: bool | None,
        return_tensors: str | TensorType | None,
    ) -> "torch.Tensor":
        # Group images by size for batched resizing
        grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping)
        resized_images_grouped = {}
        for shape, stacked_images in grouped_images.items():
            if do_resize:
                stacked_images = self.resize(image=stacked_images, size=size, resample=resample)
            resized_images_grouped[shape] = stacked_images
        resized_images = reorder_images(resized_images_grouped, grouped_images_index)

        # Group images by size for further processing
        # Needed in case do_resize is False, or resize returns images with different sizes
        grouped_images, grouped_images_index = group_images_by_shape(resized_images, disable_grouping=disable_grouping)
        processed_images_grouped = {}
        for shape, stacked_images in grouped_images.items():
            if do_center_crop:
                stacked_images = self.center_crop(stacked_images, crop_size)
            # Fused rescale and normalize
            stacked_images = self.rescale_and_normalize(
                stacked_images, do_rescale, rescale_factor, do_normalize, image_mean, image_std
            )
            if do_map_pixels:
                stacked_images = self.map_pixels(image=stacked_images)
            processed_images_grouped[shape] = stacked_images

        processed_images = reorder_images(processed_images_grouped, grouped_images_index)

        return processed_images

    def _preprocess(
        self,
        images: list["torch.Tensor"],
        do_resize: bool,
        size: SizeDict,
        resample: "PILImageResampling | tvF.InterpolationMode | 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,
        # Mask related params
        return_image_mask: bool | None,
        input_size_patches: int | None,
        total_mask_patches: int | None,
        mask_group_min_patches: int | None,
        mask_group_max_patches: int | None,
        mask_group_min_aspect_ratio: float | None,
        mask_group_max_aspect_ratio: float | None,
        # Codebook related params
        return_codebook_pixels: bool | None,
        codebook_do_resize: bool | None,
        codebook_size: SizeDict | None,
        codebook_resample: "PILImageResampling | tvF.InterpolationMode | int | None",
        codebook_do_center_crop: bool | None,
        codebook_crop_size: SizeDict | None,
        codebook_do_rescale: bool | None,
        codebook_rescale_factor: float | None,
        codebook_do_map_pixels: bool | None,
        codebook_do_normalize: bool | None,
        codebook_image_mean: float | list[float] | None,
        codebook_image_std: float | list[float] | None,
        disable_grouping: bool | None,
        return_tensors: str | TensorType | None,
        **kwargs,
    ) -> BatchFeature:
        processed_images = self._preprocess_image(
            images=images,
            do_resize=do_resize,
            size=size,
            resample=resample,
            do_center_crop=do_center_crop,
            crop_size=crop_size,
            do_rescale=do_rescale,
            rescale_factor=rescale_factor,
            do_normalize=do_normalize,
            do_map_pixels=False,
            image_mean=image_mean,
            image_std=image_std,
            disable_grouping=disable_grouping,
            return_tensors=return_tensors,
        )
        data = {
            "pixel_values": processed_images,
        }

        if return_codebook_pixels:
            codebook_processed_images = self._preprocess_image(
                images=images,
                do_resize=codebook_do_resize,
                size=codebook_size,
                resample=codebook_resample,  # resize() handles conversion from PILImageResampling to InterpolationMode
                do_center_crop=codebook_do_center_crop,
                crop_size=codebook_crop_size,
                do_rescale=codebook_do_rescale,
                rescale_factor=codebook_rescale_factor,
                do_normalize=codebook_do_normalize,
                do_map_pixels=codebook_do_map_pixels,
                image_mean=codebook_image_mean,
                image_std=codebook_image_std,
                disable_grouping=disable_grouping,
                return_tensors=return_tensors,
            )
            data["codebook_pixel_values"] = codebook_processed_images

        if return_image_mask:
            mask_generator = self.masking_generator(
                input_size_patches=input_size_patches,
                total_mask_patches=total_mask_patches,
                mask_group_min_patches=mask_group_min_patches,
                mask_group_max_patches=mask_group_max_patches,
                mask_group_min_aspect_ratio=mask_group_min_aspect_ratio,
                mask_group_max_aspect_ratio=mask_group_max_aspect_ratio,
            )
            masks = [mask_generator() for _ in range(len(images))]
            data["bool_masked_pos"] = masks

        return BatchFeature(data=data, tensor_type=return_tensors)


__all__ = ["FlavaImageProcessor"]
