# Copyright 2026 WebBrain and the HuggingFace Inc. team. All rights reserved.
#
# Adapted from moonshotai/Kimi-VL-A3B-Instruct's image_processing_kimi_vl.py
# (Apache-2.0), which this repository's MoonViT-3d tower is itself a frozen,
# fingerprint-verified copy of (see VISION_ADAPTER_MANIFEST.json).
#
# 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 for Laguna XS 2.1 Vision."""

import math
from typing import Optional, Union

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_list_of_images, valid_images
from transformers.utils import TensorType

# MoonViT-3d was trained with plain (0.5, 0.5, 0.5) normalization, not the
# OpenAI CLIP statistics — matches moonshotai/Kimi-K2.6's own preprocessing.
LAGUNA_VISION_MEAN = (0.5, 0.5, 0.5)
LAGUNA_VISION_STD = (0.5, 0.5, 0.5)


class LagunaImageProcessor(BaseImageProcessor):
    model_type = "laguna"

    def __init__(
        self,
        patch_size: int = 14,
        pad_input: bool = True,
        image_mean: tuple[float, float, float] = LAGUNA_VISION_MEAN,
        image_std: tuple[float, float, float] = LAGUNA_VISION_STD,
        in_token_limit: int = 2048,
        merge_kernel_size: tuple[int, int] = (2, 2),
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.patch_size = patch_size
        self.pad_input = pad_input
        self.image_mean = image_mean
        self.image_std = image_std
        self.in_token_limit = in_token_limit
        self.merge_kernel_size = merge_kernel_size

    def rescale(self, image: Image.Image) -> Image.Image:
        w, h = image.size
        patch_size = self.patch_size
        merge_h, merge_w = self.merge_kernel_size

        if (w // patch_size) * (h // patch_size) > self.in_token_limit:
            scale = math.sqrt(self.in_token_limit / ((w // patch_size) * (h // patch_size)))
            new_w, new_h = max(1, int(w * scale)), max(1, int(h * scale))
            image = image.resize((new_w, new_h), Image.Resampling.BICUBIC)

        if self.pad_input:
            new_w, new_h = image.size
            pad_size_h = merge_h * patch_size
            pad_size_w = merge_w * patch_size
            pad_h = (pad_size_h - new_h % pad_size_h) % pad_size_h
            pad_w = (pad_size_w - new_w % pad_size_w) % pad_size_w
            image = TF.pad(image, (0, 0, pad_w, pad_h))
        else:
            new_w, new_h = image.size
            new_w = new_w - new_w % patch_size
            new_h = new_h - new_h % patch_size
            image = TF.center_crop(image, (new_h, new_w))

        w, h = image.size
        if w // patch_size >= 512 or h // patch_size >= 512:
            raise ValueError(
                f"Image too large after rescale: {w}x{h} exceeds the vision tower's "
                "512x512 patch-grid positional embedding range."
            )
        return image

    def to_tensor(self, image: Image.Image) -> torch.Tensor:
        return TF.to_tensor(image.convert("RGB"))

    def normalize(self, image: torch.Tensor) -> torch.Tensor:
        return TF.normalize(image, self.image_mean, self.image_std)

    def patchify(self, image: torch.Tensor) -> tuple[torch.Tensor, tuple[int, int]]:
        patch_size = self.patch_size
        c, h, w = image.shape
        patches = image.reshape(c, h // patch_size, patch_size, w // patch_size, patch_size)
        patches = patches.permute(1, 3, 0, 2, 4).contiguous().view(-1, c, patch_size, patch_size)
        return patches, (h // patch_size, w // patch_size)

    def _preprocess(self, image: ImageInput) -> tuple[torch.Tensor, tuple[int, int]]:
        image = self.rescale(image)
        image = self.to_tensor(image)
        image = self.normalize(image)
        return self.patchify(image)

    def preprocess(
        self,
        images: ImageInput,
        return_tensors: Optional[Union[str, TensorType]] = None,
        **kwargs,
    ) -> BatchFeature:
        images = make_list_of_images(images)
        if 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."
            )

        pixel_values, image_grid_hws = [], []
        for image in images:
            patches, grid_hw = self._preprocess(image)
            pixel_values.append(patches)
            image_grid_hws.append(grid_hw)

        pixel_values = torch.cat(pixel_values, dim=0)
        image_grid_hws = np.array(image_grid_hws)
        data = {"pixel_values": pixel_values, "image_grid_hws": image_grid_hws}
        return BatchFeature(data=data, tensor_type=return_tensors)


__all__ = ["LagunaImageProcessor"]