# Copyright 2023 The CROMA Authors and The HuggingFace Inc. team.
#
# 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 CROMA."""

from typing import Optional, Union

import numpy as np

from transformers.image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from transformers.image_transforms import 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, filter_out_non_signature_kwargs, logging


logger = logging.get_logger(__name__)


def _resize_multispectral(image: np.ndarray, size: dict[str, int], input_data_format: ChannelDimension) -> np.ndarray:
    target_height, target_width = size["height"], size["width"]

    if input_data_format == ChannelDimension.FIRST:
        image = np.transpose(image, (1, 2, 0))

    height, width, _ = image.shape
    if height == target_height and width == target_width:
        resized = image
    else:
        try:
            import cv2
        except ImportError as exc:
            raise ImportError(
                "Multispectral resize requires OpenCV (`opencv-python`) when input has more than 4 channels."
            ) from exc
        resized = cv2.resize(image, (target_width, target_height), interpolation=cv2.INTER_LINEAR)

    if input_data_format == ChannelDimension.FIRST:
        return np.transpose(resized, (2, 0, 1))
    return resized


def normalize_croma_channels(
    images: np.ndarray,
    input_data_format: ChannelDimension,
    use_8_bit: bool = False,
) -> np.ndarray:
    """
    Apply per-channel mean +/- 2*std normalization used by CROMA (SatMAE / SeCo style).

    Args:
        images (`np.ndarray`):
            Input image array in channels-first or channels-last format.
        input_data_format (`ChannelDimension`):
            Whether channels are first or last.
        use_8_bit (`bool`, *optional*, defaults to `False`):
            If `True`, scale to `[0, 255]` and cast to `uint8`. Otherwise scale to `[0, 1]`.
    """
    if input_data_format == ChannelDimension.LAST:
        images = np.transpose(images, (2, 0, 1))

    normalized_channels = []
    for channel in range(images.shape[0]):
        channel_data = images[channel].astype(np.float32)
        min_value = channel_data.mean() - 2 * channel_data.std()
        max_value = channel_data.mean() + 2 * channel_data.std()
        denominator = max_value - min_value
        if denominator == 0:
            denominator = 1.0

        if use_8_bit:
            channel_data = (channel_data - min_value) / denominator * 255.0
            channel_data = np.clip(channel_data, 0, 255).astype(np.uint8)
        else:
            channel_data = (channel_data - min_value) / denominator
            channel_data = np.clip(channel_data, 0, 1).astype(np.float32)

        normalized_channels.append(channel_data)

    normalized = np.stack(normalized_channels, axis=0)
    if input_data_format == ChannelDimension.LAST:
        return np.transpose(normalized, (1, 2, 0))
    return normalized


class CromaImageProcessor(BaseImageProcessor):
    """
    Image processor for CROMA models.

    CROMA accepts separate Sentinel-1 SAR (2 channels) and Sentinel-2 optical (12 channels) inputs. The processor
    applies per-channel mean +/- 2*std normalization and optional resizing.
    """

    model_input_names = ["sar_pixel_values", "optical_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.0,
        do_normalize: bool = True,
        use_8_bit: bool = False,
        do_convert_rgb: bool = False,
        **kwargs,
    ):
        super().__init__(**kwargs)
        size = size if size is not None else {"height": 120, "width": 120}
        self.do_resize = do_resize
        self.size = size
        self.resample = resample
        self.do_rescale = do_rescale
        self.rescale_factor = rescale_factor
        self.do_normalize = do_normalize
        self.use_8_bit = use_8_bit
        self.do_convert_rgb = do_convert_rgb

    def _preprocess_modality(
        self,
        images: ImageInput,
        expected_channels: Optional[int],
        do_resize: bool,
        size: dict[str, int],
        resample: PILImageResampling,
        do_rescale: bool,
        rescale_factor: float,
        do_normalize: bool,
        use_8_bit: bool,
        do_convert_rgb: bool,
        data_format: Union[str, ChannelDimension],
        input_data_format: Optional[Union[str, ChannelDimension]],
    ) -> list[np.ndarray]:
        images = make_flat_list_of_images(images)
        if not valid_images(images):
            raise ValueError("Invalid image type. Must be PIL, numpy, or torch tensor.")

        processed_images = []
        for image in images:
            image = to_numpy_array(image)
            if do_convert_rgb:
                image = self._convert_image_to_rgb(image)

            if input_data_format is None:
                if isinstance(image, np.ndarray) and image.ndim == 3:
                    # CROMA expects channels-first tensors; prefer FIRST when the leading axis looks like channels.
                    current_input_format = (
                        ChannelDimension.FIRST if image.shape[0] <= image.shape[-1] else ChannelDimension.LAST
                    )
                else:
                    try:
                        current_input_format = infer_channel_dimension_format(image)
                    except ValueError:
                        current_input_format = ChannelDimension.LAST
            else:
                current_input_format = input_data_format

            num_channels = image.shape[-1] if current_input_format == ChannelDimension.LAST else image.shape[0]
            if expected_channels is not None and num_channels != expected_channels:
                raise ValueError(
                    f"Expected {expected_channels} channels for this modality, but got {num_channels}."
                )

            if do_normalize:
                image = normalize_croma_channels(
                    image,
                    input_data_format=current_input_format,
                    use_8_bit=use_8_bit,
                )
                if use_8_bit and do_rescale:
                    image = image.astype(np.float32) / 255.0
            elif do_rescale:
                image = image * rescale_factor

            if do_resize:
                if num_channels != 3:
                    image = _resize_multispectral(image, size=size, input_data_format=current_input_format)
                else:
                    image = resize(image, size=size, resample=resample, input_data_format=current_input_format)

            image = to_channel_dimension_format(image, data_format, input_channel_dim=current_input_format)
            processed_images.append(image)

        return processed_images

    @filter_out_non_signature_kwargs()
    def preprocess(
        self,
        images: Optional[ImageInput] = None,
        sar_images: Optional[ImageInput] = None,
        optical_images: Optional[ImageInput] = None,
        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,
        use_8_bit: Optional[bool] = 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 = do_resize if do_resize is not None else self.do_resize
        size = size if size is not None else self.size
        size = get_size_dict(size, default_to_square=True)
        resample = resample if resample is not None else self.resample
        do_rescale = do_rescale if do_rescale is not None else self.do_rescale
        rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
        do_normalize = do_normalize if do_normalize is not None else self.do_normalize
        use_8_bit = use_8_bit if use_8_bit is not None else self.use_8_bit
        do_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb

        validate_preprocess_arguments(
            do_rescale=do_rescale,
            rescale_factor=rescale_factor,
            do_normalize=False,
            do_resize=do_resize,
            size=size,
            resample=resample,
        )

        if sar_images is None and optical_images is None:
            if images is None:
                raise ValueError("Provide `sar_images`, `optical_images`, or legacy `images` for preprocessing.")
            optical_images = images

        data = {}
        if sar_images is not None:
            data["sar_pixel_values"] = self._preprocess_modality(
                sar_images,
                expected_channels=2,
                do_resize=do_resize,
                size=size,
                resample=resample,
                do_rescale=do_rescale,
                rescale_factor=rescale_factor,
                do_normalize=do_normalize,
                use_8_bit=use_8_bit,
                do_convert_rgb=do_convert_rgb,
                data_format=data_format,
                input_data_format=input_data_format,
            )
        if optical_images is not None:
            data["optical_pixel_values"] = self._preprocess_modality(
                optical_images,
                expected_channels=12,
                do_resize=do_resize,
                size=size,
                resample=resample,
                do_rescale=do_rescale,
                rescale_factor=rescale_factor,
                do_normalize=do_normalize,
                use_8_bit=use_8_bit,
                do_convert_rgb=do_convert_rgb,
                data_format=data_format,
                input_data_format=input_data_format,
            )

        return BatchFeature(data=data, tensor_type=return_tensors)


__all__ = ["CromaImageProcessor", "normalize_croma_channels"]