data/image_processor.py

from __future__ import annotations

import math

import torch
from PIL import Image
from torchvision.transforms.v2 import functional as tvF
from transformers.image_utils import PILImageResampling
from models.config import VLMConfig


"""
Explanation:
    Assume patch_size = 16, pooling_kernel_size = 3, and max_teacher_patches = 1,500
    Suppose we have an input image with (height = 500, width = 600)

    total_px = height * width = 500 * 600 = 300,000
    target_px = max_teacher_patches * (patch_size**2) = 1,500 * (16**2) = 384,000
    factor = math.sqrt(target_px / total_px) = sqrt(384,000 / 300,000) = 1.1313 - factor by which each side will be multiplied
    ideal_height = factor * height = 500 * 1.1313 = 565.65
    ideal_width = factor * width = 600 * 1.1313 = 678.78
    side_mult = pooling_kernel_size * patch_size = 3 * 16 = 48 - each side must be a multiple of 48

    target_height = int(math.floor(ideal_height / side_mult)) * side_mult = int(math.floor(565.65 / 48)) * 48 = 528
    target_width = int(math.floor(ideal_width / side_mult)) * side_mult = int(math.floor(678.78 / 48)) * 48 = 672\
    
    max_side_length = (max_teacher_patches // pooling_kernel_size**2) * side_mult 
                    = (1,500 // 3**2) * 48 = 166 * 48 = 7,968

    [ (max_teacher_patches // pooling_kernel_size**2) counts the maximum number of model patches (48x48) and 
      * side_mult turns it into a number of pixels ]

"""
def get_aspect_ratio_preserving_size(
    height: int,
    width: int,
    teacher_patch_size: int,
    max_teacher_patches: int,
    pooling_kernel_size: int,
) -> tuple[int, int]:
    """
    Purpose:
        Determine target dimensions for image that is resized to preserve aspect ratio so it fits within the patch budget.
        Target dimensions are the largest that:
        1) Produce at most `max_teacher_patches` patches when patchified with `patch_size`
        2) Have height and width divisible by `pooling_kernel_size * patch_size` (i.e. size of model patch)

    Parameters:
     * height (int) : height of the image to resize

     * width (int) : width of the image to resize

     * teacher_patch_size (int) : length, in pixels, of one side of a patch that teacher uses

     * max_teacher_patches (int) : maximum number of patches (of size teacher_patch_size) that a resized image may contain

     * pooling_kernel_size (int) : length of one side of a pooling kernel that pools multiple teacher patches

    Returns:
        A tuple containing (target_height, target_width) - new dimensions the image should have after resizing
    """
    total_px = height * width
    target_px = max_teacher_patches * (teacher_patch_size**2)
    factor = math.sqrt(target_px / total_px)
    ideal_height = factor * height
    ideal_width = factor * width
    side_mult = pooling_kernel_size * teacher_patch_size

    # Round down to nearest multiple of side_mult
    target_height = int(math.floor(ideal_height / side_mult)) * side_mult
    target_width = int(math.floor(ideal_width / side_mult)) * side_mult

    # Handle edge cases where one or both dimensions round to 0
    if target_height == 0 and target_width == 0:
        raise ValueError(
            "Attempting to resize to a 0 x 0 image. Resized height should be divisble by "
            f"`pooling_kernel_size * patch_size`={pooling_kernel_size * teacher_patch_size}."
        )

    max_side_length = (max_teacher_patches // pooling_kernel_size**2) * side_mult
    if target_height == 0:
        target_height = side_mult
        target_width = min(
            int(math.floor(width / height)) * side_mult,
            max_side_length,
        )
    elif target_width == 0:
        target_width = side_mult
        target_height = min(
            int(math.floor(height / width)) * side_mult,
            max_side_length,
        )

    if target_height * target_width > target_px:
        raise ValueError(
            f"Resizing [{height}x{width}] to [{target_height}x{target_width}] "
            f"but this exceeds {max_teacher_patches} patches with patch_size {teacher_patch_size}"
        )

    return target_height, target_width


def convert_image_to_patches(image: torch.Tensor, patch_size: int) -> torch.Tensor:
    """
    Purpose:
        Convert 3D tensor image of shape (num_channels, image_height, image_width) into 2D tensor of patches of shape
        (num_patches_height * num_patches_width, patch_size * patch_size * num_channels).

    Parameters:
     * image (torch.Tensor) : tensor representing an image. Has dimensions (num_channels, image_height, image_width)

     * patch_size (int) : length, in pixels, of one side of a patch

    Returns:
        2D tensor of patches of shape 
        (num_patches, flat_patch_size) = (num_patches_height * num_patches_width, patch_size * patch_size * num_channels)
    """
    num_channels, image_height, image_width = image.shape
    num_patches_height = image_height // patch_size
    num_patches_width = image_width // patch_size
    patched_image = image.reshape(num_channels, num_patches_height, patch_size, num_patches_width, patch_size)
    patched_image = patched_image.permute(1, 3, 2, 4, 0)
    patched_image = patched_image.reshape(num_patches_height * num_patches_width, -1)
    return patched_image

"""
Explanation:
    flat_patches is a sequence of flattened patches. It's shape is (num_patches, flat_patch_size).
    positions is a sequence of [x, y] coordinate pairs. It's shape is (num_patches, 2).

    Assume we have flat_patches = [ [a, b, c, d], [e, f, g, h], [i, j, k, l] ] and positions = [ [0, 0], [0, 1], [1, 0] ].
    Assume target_length = 8

    current_length = flat_patches.shape[0] = 3
    padding_length = target_length - current_length = 8 - 3 = 5

    padding = [0, 0] * (image.ndim - 1) + [0, padding_length] = [0, 0] * (2 - 1)  + [0, 5] = 
            = [0, 0] + [0, 5] = [0, 0, 0, 5]

    pos_padding = (0, 0, 0, padding_length) = (0, 0, 0, 5) 

    image = torch.nn.functional.pad(image, padding, mode="constant", value=0)
          = [ [a, b, c, d], [e, f, g, h], [i, j, k, l], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0] ]

    positions = torch.nn.functional.pad(positions, pos_padding, mode="constant", value=-1)
              = [ [0, 0], [0, 1], [1, 0], [-1, -1], [-1, -1], [-1, -1], [-1, -1], [-1, -1]]
"""
def pad_along_first_dim(
    flat_patches: torch.Tensor, positions: torch.Tensor, target_length: int
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Purpose:
        Given a sequence of patch embeddings (flat_patches) and a sequence of flat patches' original 2D positions (positions),
        add padding to both so that they have specified length (target_length).

    Parameters:
     * flat_patches (torch.Tensor) : a 2D tensor of shape (num_patches, flat_patch_size) that contains a sequence of flattened patches

     * positions (torch.Tensor) : a 2D tensor of shape (num_patchs, flat_patch_size) s.t. positions[i] is the original 2D position of 
                                  flattened patch flat_patches[i] 

    Returns:
        A tuple (flat_patches, positions), where flat_patches has shape (target_length, flat_patch_size) and positions 
        has shape (target_length, 2)
    """
    current_length = flat_patches.shape[0]
    padding_length = target_length - current_length # the number of padding tokens to append
    
    if padding_length > 0:
        # padding and pos_padding are both tuples with 4 elements because the pad parameter of torch.nn.functional.pad()
        # is a tuple/list of even length giving padding amounts in pairs (left, right), ordered from the last dimension backward
        padding = [0, 0] * (flat_patches.ndim - 1) + [0, padding_length]
        pos_padding = (0, 0, 0, padding_length)

        # Specifying value=0 tells torch.nn.functional.pad() to pad flat_patches with vectors [0, ..., 0]
        # of size flat_patch_size
        flat_patches = torch.nn.functional.pad(flat_patches, padding, mode="constant", value=0)

        # Specifying value=0 tells torch.nn.functional.pad() to pad positions with vectors [-1, -1]
        positions = torch.nn.functional.pad(positions, pos_padding, mode="constant", value=-1)
    
    return flat_patches, positions


def patches_merge(
    patches: torch.Tensor,
    positions_xy: torch.Tensor,
    length: int,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Purpose:
        Merge k×k groups of small patches into larger patches.

        Given `L` input patches of dimension `D = patch_size² × 3`, merge groups of
        `k×k` spatially adjacent patches into `length` output patches of dimension
        `(k × patch_size)² × 3`. The spatial grouping is determined by integer-dividing
        the XY positions by `k`.

    Parameters:
        small_patches: (*, L, D) — input patches.
        positions_xy: (*, L, 2) — integer XY positions for each patch (-1 for padding).
        length: target number of output patches. Must satisfy L = length × k².

    Returns:
        merged_patches: (*, length, k²×D) — merged patch features.
        merged_positions: (*, length, 2) — new XY positions for merged patches.
    """
    patch_size = math.isqrt(patches.shape[-1] // 3)
    if patches.shape[-1] != patch_size * patch_size * 3:
        raise ValueError(f"Patch dimension {patches.shape[-1]} is not a valid `patch_size * patch_size * 3`")

    k = math.isqrt(patches.shape[-2] // length)
    if k * k * length != patches.shape[-2]:
        raise ValueError(f"Cannot merge {patches.shape} to {length}")

    # Compute target ordering for reordering patches into kernel-grouped order.
    # This ensures patches within each k×k kernel are contiguous.
    max_x = positions_xy[..., 0].max(dim=-1, keepdim=True)[0] + 1
    kernel_idxs = torch.div(positions_xy, k, rounding_mode="floor")
    num_patches_from_top_left = k * k * kernel_idxs[..., 0] + k * max_x * kernel_idxs[..., 1]

    position_within_kernel = torch.remainder(positions_xy, k)
    num_patches_from_top_left_of_kernel = position_within_kernel[..., 0] + position_within_kernel[..., 1] * k
    target_ordering = num_patches_from_top_left_of_kernel + num_patches_from_top_left

    # Reorder patches by computing the inverse permutation via argsort,
    # then gathering patches into kernel-grouped order.
    perm = target_ordering.long().argsort(dim=-1)  # inverse permutation
    # Expand perm indices to match patch feature dimension for gathering
    perm_expanded = perm.unsqueeze(-1).expand_as(patches)
    kernel_ordered_patches = patches.gather(-2, perm_expanded)

    batch_shape = patches.shape[:-2]

    # Reshape: (*, length*k*k, patch_size*patch_size*3) → (*, length, (k*patch_size)*(k*patch_size)*3)
    kernel_ordered_patches = kernel_ordered_patches.reshape(*batch_shape, length, k * k, patch_size, patch_size, 3)
    # Rearrange (l, a*b, p, q, c) → (l, a*p, b*q, c)
    kernel_ordered_patches = kernel_ordered_patches.reshape(*batch_shape, length, k, k, patch_size, patch_size, 3)
    kernel_ordered_patches = kernel_ordered_patches.permute(
        *range(len(batch_shape)), -6, -5, -3, -4, -2, -1
    )  # (..., l, k, p, k, q, c)
    merged_patches = kernel_ordered_patches.reshape(*batch_shape, length, k * patch_size * k * patch_size * 3)

    # Compute new positions for merged patches
    perm_pos = perm.unsqueeze(-1).expand_as(positions_xy)
    kernel_ordered_positions = positions_xy.float().gather(-2, perm_pos.long())

    # Handle padding: preserve -1 positions
    padding = (positions_xy == -1).all(dim=-1, keepdim=True)  # (..., L, 1)
    kernel_ordered_positions = kernel_ordered_positions * (~padding).float() + positions_xy.float() * padding.float()

    # Reshape positions and take min within each kernel to get the merged position
    kernel_ordered_positions = kernel_ordered_positions.reshape(*batch_shape, length, k * k, 2)
    new_positions = torch.div(kernel_ordered_positions, k, rounding_mode="floor")
    # For each merged patch, take the minimum position across the kernel
    new_positions = new_positions.min(dim=-2)[0].to(torch.long)

    return merged_patches, new_positions


class ImageProcessor:
    resample: PILImageResampling = PILImageResampling.BICUBIC
    do_resize: bool = True
    do_rescale: bool = True
    rescale_factor: float = 1 / 255  # uint8 [0, 255] -> float [0, 1]
    do_normalize: bool = False

    def __init__(self, cfg: VLMConfig) -> None:
        self.max_soft_tokens = cfg.max_soft_tokens
        self.pooling_kernel_size = cfg.pooling_kernel_size
        self.teacher_patch_size = cfg.teacher_patch_size
        self.max_teacher_patches =  cfg.max_teacher_patches # maximum number of teacher patches a resized image may consist of

    def aspect_ratio_preserving_resize(
        self,
        image: torch.Tensor,
    ) -> torch.Tensor:
        """
        Purpose:
            Resize image to preserve aspect ratio so it fits within the patch budget.
            Target dimensions are the largest that:
            1) Produce at most `max_teacher_patches` patches when patchified with `patch_size`
            2) Have height and width divisible by `pooling_kernel_size * patch_size` (i.e. size of model patch)

        Parameters:
         * image (torch.Tensor) : tensor representing an image. Has shape (C, H, W)

        Returns:
            Resized image - a tensor of shape (C, target_height, target_width)
        """
        height, width = image.shape[-2], image.shape[-1]
        
        target_height, target_width = get_aspect_ratio_preserving_size(
            height=height,
            width=width,
            teacher_patch_size=self.teacher_patch_size,
            max_teacher_patches=self.max_teacher_patches,
            pooling_kernel_size=self.pooling_kernel_size,
        )

        if target_height == height and target_width == width:
            return image

        return tvF.resize(
            image,
            size=[target_height, target_width],
            interpolation=self.resample,
            antialias=True,
        )

    def preprocess(
        self,
        image: torch.Tensor,
    ) -> tuple:
        """
        Purpose:
            Convert the image into a tensor of flattened model-size patches. 
            Also returns a sequence of flattened patches' (x, y) positions and the number
            of model-sized patches.

        Parameters:
         * self

         * image (torch.Tensor) : a tensor of shape (C, H, W)

        Returns:
            A tuple with three elements: 
            * merged_patches (torch.Tensor) - flattened model-sized patches. Has shape (num_image_tokens, model_patch_size²*3)
            * merged_positions (torch.Tensor) - flattened patches' 2D positions. Has shape  (num_image_tokens, 2)
            * num_image_tokens (int) - the number of soft tokens in the image
        """
        # Step 1: Aspect-ratio-preserving resize
        if self.do_resize:
            image = self.aspect_ratio_preserving_resize(image=image)

        # Step 2: Rescale pixel values (typically to [0, 1]) and optionally identity normalize
        image = image.to(torch.float32) * self.rescale_factor # Is this sufficient?

        # Step 3: Patchify into teacher-size patches (16px)
        # (num_channels, height, width) → (num_teacher_patches, patch_size²*3)
        patch_height = image.shape[-2] // self.teacher_patch_size
        patch_width = image.shape[-1] // self.teacher_patch_size
        teacher_patches = convert_image_to_patches(image, self.teacher_patch_size)

        # Step 4: Compute teacher-level position IDs
        device = image.device
        patch_grid = torch.meshgrid(
            torch.arange(patch_width, device=device),
            torch.arange(patch_height, device=device),
            indexing="xy",
        )
        teacher_positions = torch.stack(patch_grid, dim=-1).reshape(teacher_patches.shape[0], 2)

        # Step 5: Merge k×k teacher patches into model patches via patches_merge
        # (num_teacher_patches, 768) → (num_model_patches, 6912)
        num_model_patches = teacher_patches.shape[0] // (self.pooling_kernel_size**2)
        merged_patches, merged_positions = patches_merge(
            teacher_patches.unsqueeze(0),  # add batch dim for patches_merge
            teacher_positions.unsqueeze(0),
            num_model_patches,
        )

        merged_patches = merged_patches.squeeze(0)  # remove batch dim
        merged_positions = merged_positions.squeeze(0)
        num_image_tokens = merged_patches.shape[0]

        return merged_patches, merged_positions, num_image_tokens


    def __call__(self, image: Image.Image,) -> tuple:
        """
        Purpose:
            Convert the image into a tensor of flattened model-size patches. 
            Also returns a sequence of flattened patches' (x, y) positions and the number
            of model-sized patches.

        Parameters:
         * self

         * image (torch.Tensor) : a tensor of shape (C, H, W)

        Returns:
            A tuple with three elements: 
            * merged_patches (torch.Tensor) - flattened model-sized patches. Has shape (num_image_tokens, model_patch_size²*3)
            * merged_positions (torch.Tensor) - flattened patches' 2D positions. Has shape  (num_image_tokens, 2)
            * num_image_tokens (int) - the number of soft tokens in the image
        """
        # Convert each image in the list into a tensor
        image_as_tensor = tvF.pil_to_tensor(image) # uint8, shape (B, C, H, W), values [0, 255]

        return self.preprocess(image_as_tensor)