"""
# Adapted from https://github.com/baaivision/EVA/tree/master/EVA-CLIP
"""

from torchvision import transforms
from torchvision.transforms.functional import InterpolationMode
from transformers.image_processing_utils import BatchFeature
from PIL import Image
from transformers.image_transforms import convert_to_rgb


class BaseProcessor:
    def __init__(self):
        self.transform = lambda x: x
        return

    def __call__(self, item):
        return self.transform(item)


class EvaClipImageBaseProcessor(BaseProcessor):
    def __init__(self, mean=None, std=None):
        self.mean = (0.48145466, 0.4578275, 0.40821073) if mean is None else mean
        self.std = (0.26862954, 0.26130258, 0.27577711) if std is None else std

        self.normalize = transforms.Normalize(self.mean, self.std)

    @property
    def image_mean(self):
        return self.mean


class EvaClipImageTrainProcessor(EvaClipImageBaseProcessor):
    def __init__(self, image_size=224, mean=None, std=None, min_scale=0.5, max_scale=1.0):
        super().__init__(mean=mean, std=std)

        self.transform = transforms.Compose(
            [
                convert_to_rgb,
                transforms.Resize(
                    image_size,
                    interpolation=InterpolationMode.BICUBIC,
                ),
                transforms.CenterCrop(image_size),
                transforms.ToTensor(),
                self.normalize,
            ]
        )

        self.image_size = image_size

    def preprocess(self, images, return_tensors):
        if isinstance(images, Image.Image):
            images = [images]
        else:
            assert isinstance(images, list)

        transformed_images = [self.transform(image).numpy() for image in images]
        data = {"pixel_values": transformed_images}

        return BatchFeature(data=data, tensor_type=return_tensors)

    def __call__(self, item):
        return self.transform(item)

    @property
    def crop_size(self):
        return {"height": self.image_size, "width": self.image_size}

    @property
    def size(self):
        return {"shortest_edge": self.image_size}