from __future__ import annotations

import numpy as np
from PIL import Image


def pad_image_to_square(image: Image.Image) -> Image.Image:
    """Pad an image with black pixels to create a square image.

    This keeps the localization coordinate space consistent with the MedGemma
    localization notebook while avoiding a second uint8 rescaling step.
    """
    rgb_image = image.convert("RGB")
    image_array = np.asarray(rgb_image, dtype=np.uint8)

    height, width = image_array.shape[:2]
    max_dim = max(height, width)

    pad_top = (max_dim - height) // 2
    pad_bottom = max_dim - height - pad_top
    pad_left = (max_dim - width) // 2
    pad_right = max_dim - width - pad_left

    padded = np.pad(
        image_array,
        ((pad_top, pad_bottom), (pad_left, pad_right), (0, 0)),
        mode="constant",
        constant_values=0,
    )
    return Image.fromarray(padded)