# MoonViT native-resolution image processor (port of the moonshotai/MoonViT-SO-400M and
# AI4Industry/MoonViT-V2 processors, which differ only in their limits).
#
# Output:
#   pixel_values     (total_patches, 3, patch_size, patch_size)
#   image_grid_thws  (num_images, 3) as (frames, height, width) in patches; frames is 1 for images
#
# `in_token_limit` counts ViT input tokens (14x14 patches, before the 2x2 merge), so the number of
# tokens the LLM sees is at most in_token_limit / 4.

import math

import numpy as np
import torch
from PIL import Image
from torchvision.transforms import functional as TF

from transformers.image_processing_utils import BaseImageProcessor, BatchFeature
from transformers.image_utils import ImageInput, make_flat_list_of_images


class MoonViTImageProcessor(BaseImageProcessor):
    model_input_names = ["pixel_values", "image_grid_thws"]

    def __init__(
        self,
        patch_size: int = 14,
        merge_kernel_size: int | list[int] | tuple[int, int] = (2, 2),
        in_token_limit: int = 4096,
        patch_limit_on_one_side: int | None = None,
        max_num_frames: int = 1,
        pad_input: bool = True,
        image_mean: list[float] | tuple[float, ...] = (0.5, 0.5, 0.5),
        image_std: list[float] | tuple[float, ...] = (0.5, 0.5, 0.5),
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.patch_size = patch_size
        if isinstance(merge_kernel_size, int):
            merge_kernel_size = (merge_kernel_size, merge_kernel_size)
        self.merge_kernel_size = list(merge_kernel_size)
        # Max ViT tokens (14x14 patches before the 2x2 merge) per image; the image is downscaled to fit.
        self.in_token_limit = in_token_limit
        # Max patches along one side (MoonViT-V2: 512, i.e. 7168px); None disables the per-side cap.
        self.patch_limit_on_one_side = patch_limit_on_one_side
        # Video frames covered by the temporal position embedding (MoonViT-V2: 4).
        self.max_num_frames = max_num_frames
        # True: pad right/bottom to a multiple of patch_size * merge_kernel_size; False: center-crop to it.
        self.pad_input = pad_input
        self.image_mean = list(image_mean)
        self.image_std = list(image_std)

    def resize_and_align(self, image: Image.Image) -> Image.Image:
        w, h = image.size
        p = self.patch_size
        scale = min(1.0, math.sqrt(self.in_token_limit / (max(1, w // p) * max(1, h // p))))
        if self.patch_limit_on_one_side is not None:
            side = self.patch_limit_on_one_side * p
            scale = min(scale, side / w, side / h)
        if scale < 1.0:
            image = image.resize((max(1, int(w * scale)), max(1, int(h * scale))), Image.Resampling.BICUBIC)

        w, h = image.size
        unit_h, unit_w = self.merge_kernel_size[0] * p, self.merge_kernel_size[1] * p
        if self.pad_input:
            image = TF.pad(image, [0, 0, (unit_w - w % unit_w) % unit_w, (unit_h - h % unit_h) % unit_h])
        else:
            image = TF.center_crop(image, [max(h - h % unit_h, unit_h), max(w - w % unit_w, unit_w)])
        return image

    def get_grid_thw(self, height: int, width: int, num_frames: int = 1) -> tuple[int, int, int]:
        """Grid (in patches) that `resize_and_align` produces for an image of this size."""
        w, h = self.resize_and_align(Image.new("RGB", (width, height))).size
        return num_frames, h // self.patch_size, w // self.patch_size

    def _to_pil(self, image) -> Image.Image:
        if isinstance(image, Image.Image):
            return image.convert("RGB")
        if isinstance(image, torch.Tensor):
            image = image.cpu().numpy()
        image = np.asarray(image)
        if image.ndim == 3 and image.shape[0] in (1, 3) and image.shape[-1] not in (1, 3):
            image = image.transpose(1, 2, 0)
        if image.dtype != np.uint8:
            image = (image * 255 if image.max() <= 1.0 else image).clip(0, 255).astype(np.uint8)
        return Image.fromarray(image).convert("RGB")

    def _patchify(self, frames: list[Image.Image]) -> tuple[torch.Tensor, list[int]]:
        p = self.patch_size
        x = torch.stack([TF.normalize(TF.to_tensor(f), self.image_mean, self.image_std) for f in frames])
        t, c, h, w = x.shape
        patches = x.reshape(t, c, h // p, p, w // p, p).permute(0, 2, 4, 1, 3, 5).reshape(-1, c, p, p)
        return patches, [t, h // p, w // p]

    def preprocess(self, images: ImageInput, return_tensors=None, **kwargs) -> BatchFeature:
        pixel_values, grid_thws = [], []
        for image in make_flat_list_of_images(images):
            patches, grid = self._patchify([self.resize_and_align(self._to_pil(image))])
            pixel_values.append(patches)
            grid_thws.append(grid)
        data = {"pixel_values": torch.cat(pixel_values), "image_grid_thws": torch.tensor(grid_thws, dtype=torch.long)}
        return BatchFeature(data=data, tensor_type=return_tensors)

    def preprocess_video(self, frames, return_tensors=None) -> BatchFeature:
        """One video = one sample with T frames; all frames share the first frame's resize/pad."""
        frames = [self._to_pil(f) for f in frames]
        if not 1 <= len(frames) <= self.max_num_frames:
            raise ValueError(f"Got {len(frames)} frames, but max_num_frames={self.max_num_frames}")
        w, h = frames[0].size
        target = self.resize_and_align(frames[0]).size
        frames = [self.resize_and_align(f) if f.size == (w, h) else f.resize(target) for f in frames]
        patches, grid = self._patchify(frames)
        data = {"pixel_values": patches, "image_grid_thws": torch.tensor([grid], dtype=torch.long)}
        return BatchFeature(data=data, tensor_type=return_tensors)


__all__ = ["MoonViTImageProcessor"]