from typing import Optional, Union

import numpy as np
import torch
from PIL import Image
from torchvision.transforms import v2
from transformers.image_processing_utils import BaseImageProcessor, BatchFeature


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

    def __init__(
        self,
        image_size=512,
        keep_aspect_ratio=True,
        self_normalize=True,
        mean=None,
        std=None,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.image_size = image_size
        self.keep_aspect_ratio = keep_aspect_ratio
        self.self_normalize = self_normalize
        self.mean = mean or [123.675, 116.28, 103.53]
        self.std = std or [58.395, 57.12, 57.375]

    def preprocess(
        self,
        images: Union[Image.Image, np.ndarray, torch.Tensor, list],
        return_tensors: Optional[Union[str, torch.Tensor]] = None,
        **kwargs,
    ):
        if not isinstance(images, (list, tuple)):
            images = [images]
        pixel_values = [self._preprocess_image(image) for image in images]
        return BatchFeature(
            data={"pixel_values": torch.stack(pixel_values)},
            tensor_type=return_tensors,
        )

    def _preprocess_image(self, image):
        if isinstance(image, Image.Image):
            image = np.array(image.convert("RGB"), copy=True)
        if isinstance(image, np.ndarray):
            image = torch.from_numpy(image)
        if image.ndim != 3:
            raise ValueError("Expected an HWC or CHW RGB image.")
        if image.shape[-1] in (1, 3):
            image = image.permute(2, 0, 1)
        if image.shape[0] == 1:
            image = image.expand(3, -1, -1)
        if image.shape[0] != 3:
            raise ValueError("Cond-UNet requires one or three input channels.")
        height, width = image.shape[-2:]
        if self.keep_aspect_ratio:
            resize_factor = max(height, width) / self.image_size
            new_height = int(height / resize_factor)
            new_width = int(width / resize_factor)
            new_height += new_height % 2
            new_width += new_width % 2
            image = v2.functional.resize(image, [new_height, new_width])
            pad_left = (self.image_size - new_width) // 2
            pad_top = (self.image_size - new_height) // 2
            image = v2.functional.pad(image, fill=0, padding=[pad_left, pad_top])
        else:
            image = v2.functional.resize(image, [self.image_size, self.image_size])
        image = image.to(dtype=torch.float32)
        if image.max() <= 1:
            image = image * 255.0
        if self.self_normalize:
            mask = (image > 0).any(dim=0)
            if mask.any():
                valid_pixels = image[:, mask]
                mean = valid_pixels.mean()
                std = valid_pixels.std()
                if std > 1e-8:
                    return (image - mean) / std
                return image - mean
            return image.clone()

        mean = torch.tensor(self.mean, dtype=image.dtype).view(-1, 1, 1)
        std = torch.tensor(self.std, dtype=image.dtype).view(-1, 1, 1)
        return (image - mean) / std