# Copyright 2025 S5 authors and Hugging Face converters.
"""Image processor for S5 remote-sensing models (ImageNet mean/std)."""

from __future__ import annotations

from typing import Optional, Union

from transformers.image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from transformers.image_transforms import convert_to_rgb, resize, to_channel_dimension_format
from transformers.image_utils import (
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    infer_channel_dimension_format,
    make_flat_list_of_images,
    to_numpy_array,
    valid_images,
    validate_preprocess_arguments,
)
from transformers.utils import TensorType, logging

try:
    from transformers.utils import filter_out_non_signature_kwargs
except ImportError:  # older transformers
    def filter_out_non_signature_kwargs():
        def decorator(fn):
            return fn

        return decorator


try:
    from configuration_s5 import IMAGENET_MEAN, IMAGENET_STD
except ImportError:
    from .configuration_s5 import IMAGENET_MEAN, IMAGENET_STD


logger = logging.get_logger(__name__)


class S5ImageProcessor(BaseImageProcessor):
    """Resize / rescale / ImageNet-normalize RGB remote-sensing images for S5."""

    model_input_names = ["pixel_values"]

    def __init__(
        self,
        do_resize: bool = False,
        size: Optional[dict[str, int]] = None,
        resample: PILImageResampling = PILImageResampling.BILINEAR,
        do_rescale: bool = True,
        rescale_factor: float = 1 / 255,
        do_normalize: bool = True,
        image_mean: Optional[Union[float, list[float]]] = None,
        image_std: Optional[Union[float, list[float]]] = None,
        do_convert_rgb: bool = True,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.do_resize = do_resize
        self.size = size if size is not None else {"height": 512, "width": 512}
        self.resample = resample
        self.do_rescale = do_rescale
        self.rescale_factor = rescale_factor
        self.do_normalize = do_normalize
        self.image_mean = list(IMAGENET_MEAN) if image_mean is None else image_mean
        self.image_std = list(IMAGENET_STD) if image_std is None else image_std
        self.do_convert_rgb = do_convert_rgb

    @filter_out_non_signature_kwargs()
    def preprocess(
        self,
        images: ImageInput,
        do_resize: Optional[bool] = None,
        size: Optional[dict[str, int]] = None,
        resample: Optional[PILImageResampling] = None,
        do_rescale: Optional[bool] = None,
        rescale_factor: Optional[float] = None,
        do_normalize: Optional[bool] = None,
        image_mean: Optional[Union[float, list[float]]] = None,
        image_std: Optional[Union[float, list[float]]] = None,
        return_tensors: Optional[Union[str, TensorType]] = None,
        data_format: Union[str, ChannelDimension] = ChannelDimension.FIRST,
        input_data_format: Optional[Union[str, ChannelDimension]] = None,
        do_convert_rgb: Optional[bool] = None,
    ):
        do_resize = self.do_resize if do_resize is None else do_resize
        size = get_size_dict(self.size if size is None else size, default_to_square=True)
        resample = self.resample if resample is None else resample
        do_rescale = self.do_rescale if do_rescale is None else do_rescale
        rescale_factor = self.rescale_factor if rescale_factor is None else rescale_factor
        do_normalize = self.do_normalize if do_normalize is None else do_normalize
        image_mean = self.image_mean if image_mean is None else image_mean
        image_std = self.image_std if image_std is None else image_std
        do_convert_rgb = self.do_convert_rgb if do_convert_rgb is None else do_convert_rgb

        images = make_flat_list_of_images(images)
        if not valid_images(images):
            raise ValueError("Invalid image type. Must be PIL, numpy, or torch tensor.")

        validate_preprocess_arguments(
            do_rescale=do_rescale,
            rescale_factor=rescale_factor,
            do_normalize=do_normalize,
            image_mean=image_mean,
            image_std=image_std,
            do_resize=do_resize,
            size=size,
            resample=resample,
        )

        processed = []
        for image in images:
            if do_convert_rgb:
                try:
                    image = convert_to_rgb(image)
                except Exception:
                    image = to_numpy_array(image)
            image = to_numpy_array(image)
            if input_data_format is None:
                try:
                    inferred = infer_channel_dimension_format(image)
                except ValueError:
                    inferred = ChannelDimension.LAST
            else:
                inferred = input_data_format
            if do_resize:
                image = resize(
                    image,
                    size=(size["height"], size["width"]),
                    resample=resample,
                    input_data_format=inferred,
                )
            if do_rescale:
                image = image * rescale_factor
            if do_normalize:
                image = self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=inferred)
            processed.append(to_channel_dimension_format(image, data_format, input_channel_dim=inferred))

        return BatchFeature(data={"pixel_values": processed}, tensor_type=return_tensors)


__all__ = ["S5ImageProcessor"]