# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

"""Standalone HF image processor for PanoVLM (FastViT, NCHW).

Reproduces the repo's ImageProcessorNCHW: RGB -> aspect-preserving pad-resize to
a fixed square -> rescale to [0,1] -> normalize -> NCHW pixel_values. Defaults
match training (``scripts_local/generate_panovlm.sh``): only ``image_size`` is
meant to be changed by users; ``image_resize_mode``/``image_mean``/``image_std``
are fixed.
"""
from __future__ import annotations

import numpy as np
from PIL import Image
from transformers.image_processing_utils import BaseImageProcessor, BatchFeature
from transformers.image_utils import ImageInput
from transformers.utils import TensorType


def _resize_pad(
    image: Image.Image, image_size: int, fill_mean: tuple[float, float, float]
) -> Image.Image:
    orig_w, orig_h = image.size
    scale = min(image_size / orig_w, image_size / orig_h)
    new_w, new_h = max(1, int(orig_w * scale)), max(1, int(orig_h * scale))
    image = image.resize((new_w, new_h))
    padded = Image.new(
        "RGB", (image_size, image_size), tuple(int(255 * x) for x in fill_mean)
    )
    padded.paste(image, ((image_size - new_w) // 2, (image_size - new_h) // 2))
    return padded


class PanoVLMImageProcessor(BaseImageProcessor):
    model_input_names = ["pixel_values"]

    def __init__(
        self,
        image_size: int = 1024,
        image_resize_mode: str = "pad",
        image_mean=(0.0, 0.0, 0.0),
        image_std=(1.0, 1.0, 1.0),
        **kwargs,
    ):
        if image_resize_mode not in ("pad", "square"):
            raise ValueError(
                f"Unknown image_resize_mode {image_resize_mode!r}; expected 'pad' or 'square'."
            )
        super().__init__(**kwargs)
        self.image_size = image_size
        self.image_resize_mode = image_resize_mode
        self.image_mean = tuple(image_mean)
        self.image_std = tuple(image_std)

    def _to_pil(self, image) -> Image.Image:
        if isinstance(image, Image.Image):
            img = image
        else:
            img = Image.fromarray(np.asarray(image))
        return img.convert("RGB") if img.mode != "RGB" else img

    def _process_one(self, image) -> np.ndarray:
        img = self._to_pil(image)
        if self.image_resize_mode == "pad":
            img = _resize_pad(img, self.image_size, self.image_mean)
        elif self.image_resize_mode == "square":
            img = img.resize((self.image_size, self.image_size))
        else:
            raise ValueError(
                f"Unknown image_resize_mode {self.image_resize_mode!r}; "
                "expected 'pad' or 'square'."
            )
        arr = np.asarray(img, dtype=np.float32) / 255.0
        arr = (arr - np.asarray(self.image_mean, np.float32)) / np.asarray(
            self.image_std, np.float32
        )
        return arr.transpose(2, 0, 1)  # HWC -> CHW

    def preprocess(
        self,
        images: ImageInput,
        return_tensors: str | TensorType | None = None,
        **kwargs,
    ) -> BatchFeature:
        # **kwargs accepted for HF API compat (do_rescale/do_normalize/etc.);
        # processing is fixed by the configured size/mean/std.
        if not isinstance(images, (list, tuple)):
            images = [images]
        pixel_values = np.stack([self._process_one(im) for im in images], axis=0)
        return BatchFeature(
            data={"pixel_values": pixel_values}, tensor_type=return_tensors
        )