import logging
from typing import Optional, Union

from transformers.image_processing_utils import (
    BatchFeature,
)
from transformers.image_processing_utils_fast import (
    group_images_by_shape,
    reorder_images,
)
from transformers.image_utils import (
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    SizeDict,
    get_image_size,
    make_flat_list_of_images,
    valid_images,
)
from transformers.models.glm4v.image_processing_glm4v import smart_resize
from transformers.utils import (
    TensorType,
    is_torch_available,
    is_torchvision_available,
    is_torchvision_v2_available,
)
from transformers.video_utils import VideoInput

if is_torch_available():
    import torch


if is_torchvision_available():
    from transformers.image_utils import pil_torch_interpolation_mapping

    if is_torchvision_v2_available():
        from torchvision.transforms.v2 import functional as F
    else:
        from torchvision.transforms import functional as F

logger = logging.getLogger(__name__)


def new_patch__preprocess(
    self,
    images: list["torch.Tensor"],
    do_resize: bool,
    size: SizeDict,
    interpolation: Optional["F.InterpolationMode"],
    do_rescale: bool,
    rescale_factor: float,
    do_normalize: bool,
    image_mean: Optional[Union[float, list[float]]],
    image_std: Optional[Union[float, list[float]]],
    patch_size: int,
    temporal_patch_size: int,
    merge_size: int,
    do_convert_rgb: bool,
    input_data_format: Optional[Union[str, ChannelDimension]],
    device: Optional[Union[str, torch.device]],
    disable_grouping: Optional[bool],
):
    images = self._prepare_input_images(
        images=images,
        do_convert_rgb=do_convert_rgb,
        input_data_format=input_data_format,
        device=device,
    )

    height, width = get_image_size(images[0], channel_dim=ChannelDimension.FIRST)
    resized_height, resized_width = height, width

    # 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:
            resized_height, resized_width = smart_resize(
                num_frames=temporal_patch_size,
                height=height,
                width=width,
                temporal_factor=temporal_patch_size,
                factor=patch_size * merge_size,
            )
            stacked_images = F.resize(
                stacked_images,
                size=(resized_height, resized_width),
                interpolation=interpolation,
            )
        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():
        # Fused rescale and normalize
        stacked_images = self.rescale_and_normalize(
            stacked_images,
            do_rescale,
            rescale_factor,
            do_normalize,
            image_mean,
            image_std,
        )
        processed_images_grouped[shape] = stacked_images

    processed_images = reorder_images(processed_images_grouped, grouped_images_index)
    patches = torch.stack(processed_images, dim=0)
    if patches.shape[0] % temporal_patch_size != 0:
        repeats = patches[-1].unsqueeze(0).repeat(temporal_patch_size - 1, 1, 1, 1)
        patches = torch.cat([patches, repeats], dim=0)

    channel = patches.shape[1]
    grid_t = patches.shape[0] // temporal_patch_size
    grid_h, grid_w = resized_height // patch_size, resized_width // patch_size

    patches = patches.view(
        grid_t,
        temporal_patch_size,
        channel,
        grid_h // merge_size,
        merge_size,
        patch_size,
        grid_w // merge_size,
        merge_size,
        patch_size,
    )
    patches = patches.permute(0, 3, 6, 4, 7, 2, 1, 5, 8)
    flatten_patches = patches.reshape(
        grid_t * grid_h * grid_w,
        channel * temporal_patch_size * patch_size * patch_size,
    )

    return flatten_patches, (grid_t, grid_h, grid_w)


def new_patch_preprocess(
    self,
    images: ImageInput,
    videos: VideoInput = None,
    do_resize: Optional[bool] = None,
    size: Optional[dict[str, int]] = None,
    resample: Optional[Union["PILImageResampling", "F.InterpolationMode"]] = None,
    do_rescale: Optional[bool] = None,
    rescale_factor: Optional[float] = None,
    do_normalize: Optional[bool] = None,
    image_mean: Optional[Union[float, list[float]]] = None,
    image_std: Optional[Union[float, list[float]]] = None,
    patch_size: Optional[int] = None,
    temporal_patch_size: Optional[int] = None,
    merge_size: Optional[int] = None,
    do_convert_rgb: Optional[bool] = None,
    return_tensors: Optional[Union[str, TensorType]] = None,
    data_format: Optional[ChannelDimension] = ChannelDimension.FIRST,
    input_data_format: Optional[Union[str, ChannelDimension]] = None,
    device: Optional["torch.device"] = None,
    disable_grouping: Optional[bool] = False,
    **kwargs,
):
    do_resize = do_resize if do_resize is not None else self.do_resize
    size = size if size is not None else self.size
    resample = resample if resample is not None else self.resample
    do_rescale = do_rescale if do_rescale is not None else self.do_rescale
    rescale_factor = (
        rescale_factor if rescale_factor is not None else self.rescale_factor
    )
    do_normalize = do_normalize if do_normalize is not None else self.do_normalize
    image_mean = image_mean if image_mean is not None else self.image_mean
    image_std = image_std if image_std is not None else self.image_std
    patch_size = patch_size if patch_size is not None else self.patch_size
    temporal_patch_size = (
        temporal_patch_size
        if temporal_patch_size is not None
        else self.temporal_patch_size
    )
    merge_size = merge_size if merge_size is not None else self.merge_size
    do_convert_rgb = (
        do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb
    )

    # Make hashable for cache
    size = SizeDict(**size) if size is not None else None
    image_mean = tuple(image_mean) if image_mean is not None else None
    image_std = tuple(image_std) if image_std is not None else None

    self._validate_preprocess_kwargs(
        do_rescale=do_rescale,
        rescale_factor=rescale_factor,
        do_normalize=do_normalize,
        image_mean=image_mean,
        image_std=image_std,
        do_resize=do_resize,
        size=size,
        resample=resample,
        return_tensors=return_tensors,
        data_format=data_format,
    )
    interpolation = (
        pil_torch_interpolation_mapping[resample]
        if isinstance(resample, (PILImageResampling, int))
        else resample
    )

    if images is not None:
        images = make_flat_list_of_images(images)

    if images is not None and not valid_images(images):
        raise ValueError(
            "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
            "torch.Tensor, tf.Tensor or jax.ndarray."
        )

    data = {}
    if images is not None:
        pixel_values, vision_grid_thws = [], []
        for image in images:
            patches, image_grid_thw = self._preprocess(
                image,
                do_resize=do_resize,
                size=size,
                interpolation=interpolation,
                do_rescale=do_rescale,
                rescale_factor=rescale_factor,
                do_normalize=do_normalize,
                image_mean=image_mean,
                image_std=image_std,
                patch_size=patch_size,
                temporal_patch_size=temporal_patch_size,
                merge_size=merge_size,
                do_convert_rgb=do_convert_rgb,
                input_data_format=input_data_format,
                device=device,
                disable_grouping=disable_grouping,
            )
            pixel_values.extend(patches)
            vision_grid_thws.append(image_grid_thw)
        pixel_values = torch.stack(pixel_values)
        vision_grid_thws = torch.tensor(vision_grid_thws)
        data.update({"pixel_values": pixel_values, "image_grid_thw": vision_grid_thws})

    return BatchFeature(data=data, tensor_type=return_tensors)


def new_patch_get_number_of_image_patches(
    self, height: int, width: int, images_kwargs=None
):
    patch_size = images_kwargs.get("patch_size", None) or self.patch_size
    merge_size = images_kwargs.get("merge_size", None) or self.merge_size

    factor = patch_size * merge_size
    resized_height, resized_width = smart_resize(
        num_frames=self.temporal_patch_size,
        height=height,
        width=width,
        temporal_factor=self.temporal_patch_size,
        factor=factor,
    )
    grid_h, grid_w = resized_height // patch_size, resized_width // patch_size
    return grid_h * grid_w


def apply_patch_glm4v(model, processor):
    import types

    image_processor = processor.image_processor
    image_processor._preprocess = types.MethodType(
        new_patch__preprocess, image_processor
    )
    image_processor.preprocess = types.MethodType(new_patch_preprocess, image_processor)
    image_processor.get_number_of_image_patches = types.MethodType(
        new_patch_get_number_of_image_patches, image_processor
    )
