"""Image processor for MolParser Mobile."""

from __future__ import annotations

from typing import List, Sequence, Union

import cv2
import numpy as np
import torch
from PIL import Image
from transformers import BaseImageProcessor
from transformers.feature_extraction_utils import BatchFeature


ImageInput = Union[str, Image.Image, np.ndarray, torch.Tensor]


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

    def __init__(
        self,
        image_size: int = 224,
        do_resize: bool = True,
        do_normalize: bool = True,
        image_mean: Sequence[float] = (0.485, 0.456, 0.406),
        image_std: Sequence[float] = (0.229, 0.224, 0.225),
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.image_size = int(image_size)
        self.do_resize = bool(do_resize)
        self.do_normalize = bool(do_normalize)
        self.image_mean = list(image_mean)
        self.image_std = list(image_std)

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

    def _to_pil(self, image: ImageInput) -> Image.Image:
        if isinstance(image, Image.Image):
            return image.convert("RGB")
        if isinstance(image, str):
            return Image.open(image).convert("RGB")
        if isinstance(image, torch.Tensor):
            tensor = image.detach().cpu()
            if tensor.ndim == 3 and tensor.shape[0] in {1, 3, 4}:
                tensor = tensor.permute(1, 2, 0)
            array = tensor.numpy()
        else:
            array = np.asarray(image)
        if array.dtype != np.uint8:
            if array.max() <= 1.0:
                array = array * 255.0
            array = np.clip(np.rint(array), 0, 255).astype(np.uint8)
        if array.ndim == 2:
            return Image.fromarray(array, mode="L").convert("RGB")
        if array.shape[-1] == 4:
            return Image.fromarray(array, mode="RGBA").convert("RGB")
        return Image.fromarray(array).convert("RGB")

    def _preprocess_one(self, image: ImageInput) -> np.ndarray:
        pil_image = self._to_pil(image)
        array = np.asarray(pil_image).astype(np.uint8)
        if self.do_resize:
            # Match deploy/transform.py: albumentations.Resize defaults to OpenCV INTER_LINEAR.
            array = cv2.resize(array, (self.image_size, self.image_size), interpolation=cv2.INTER_LINEAR)
        array = array.astype(np.float32) / 255.0
        if self.do_normalize:
            mean = np.asarray(self.image_mean, dtype=np.float32).reshape(1, 1, 3)
            std = np.asarray(self.image_std, dtype=np.float32).reshape(1, 1, 3)
            array = (array - mean) / std
        return np.transpose(array, (2, 0, 1))

    def preprocess(
        self,
        images: Union[ImageInput, Sequence[ImageInput]],
        return_tensors: str | None = None,
        **kwargs,
    ) -> BatchFeature:
        if not isinstance(images, (list, tuple)):
            images = [images]
        pixel_values: List[np.ndarray] = [self._preprocess_one(image) for image in images]
        data = {"pixel_values": np.stack(pixel_values, axis=0)}
        encoded = BatchFeature(data=data)
        if return_tensors is not None:
            encoded = encoded.convert_to_tensors(return_tensors)
        return encoded

    def __call__(self, images: Union[ImageInput, Sequence[ImageInput]], return_tensors: str | None = None, **kwargs):
        return self.preprocess(images=images, return_tensors=return_tensors, **kwargs)


__all__ = ["MolParserImageProcessor"]