from typing import List, Optional, Union

import numpy as np
import torch
from PIL import Image
from transformers.feature_extraction_utils import BatchFeature, FeatureExtractionMixin
from transformers.image_utils import ImageFeatureExtractionMixin


class M2EncoderImageProcessor(FeatureExtractionMixin, ImageFeatureExtractionMixin):
    model_input_names = ["pixel_values"]

    def __init__(self, size: int = 224, resample: int = Image.BICUBIC, **kwargs):
        super().__init__(**kwargs)
        if isinstance(size, dict):
            size = int(size.get("height") or size.get("width"))
        self.size = size
        self.resample = resample

    def __call__(
        self,
        images,
        return_tensors: Optional[Union[str, torch.Tensor]] = None,
        **kwargs,
    ) -> BatchFeature:
        if not isinstance(images, (list, tuple)):
            images = [images]

        pixel_values: List[np.ndarray] = []
        for image in images:
            if not isinstance(image, Image.Image):
                image = Image.fromarray(np.asarray(image))
            image = image.convert("RGB")
            image = image.resize((self.size, self.size), resample=self.resample)
            array = np.asarray(image, dtype=np.float32) / 255.0
            array = np.transpose(array, (2, 0, 1))
            pixel_values.append(array)

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