import numpy as np
import PIL.Image
import PIL.ImageOps
from transformers.image_processing_utils import BaseImageProcessor, BatchFeature
from transformers.image_transforms import to_pil_image

# MNIST training set statistics, as in mnistvit.preprocess.preprocess_mnist
MNIST_MEAN = 0.1307
MNIST_STD = 0.3081
NORMALIZATIONS = ("image", "mnist", None)


class MnistViTImageProcessor(BaseImageProcessor):
    """Image processor reproducing the preprocessing of mnistvit.

    Converts to grayscale, center crops to a square, resizes and normalizes.  By
    default, it reproduces `mnistvit.preprocess.read_digit_image`: each image is
    standardized with its own mean and standard deviation, and dark-on-bright images are
    inverted so that digits are bright on a dark background as in MNIST.

    Args:
        size (int, optional): Height and width of the output images.  Default: 28.
        normalization (str or None, optional): `'image'` standardizes each image with
            its own mean and standard deviation, which is robust to arbitrary images.
            `'mnist'` normalizes with the MNIST training set mean and standard deviation
            as in `mnistvit.preprocess.preprocess_mnist`, which reproduces the original
            test set evaluation but assumes MNIST-like images.  `None` keeps pixel
            values in [0, 1], as expected by the ONNX model, which standardizes and
            inverts internally.  Default: `'image'`.
        do_auto_invert (bool, optional): Whether to invert images with more bright than
            dark pixels.  Only applies if `normalization` is `'image'`.
            Default: `True`.
    """

    model_input_names = ["pixel_values"]
    # Class defaults let transformers v5 serialize an explicit `normalization=None`
    size = 28
    normalization = "image"
    do_auto_invert = True

    def __init__(
        self,
        size: int = 28,
        normalization: str | None = "image",
        do_auto_invert: bool = True,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)
        self.size = size
        self.normalization = normalization
        self.do_auto_invert = do_auto_invert

    def __call__(self, images, **kwargs) -> BatchFeature:
        return self.preprocess(images, **kwargs)

    def preprocess(
        self,
        images,
        size: int | None = None,
        normalization: str | None = "default",
        do_auto_invert: bool | None = None,
        return_tensors=None,
        **kwargs,
    ) -> BatchFeature:
        """Preprocess an image or a list of images.

        Args:
            images (PIL.Image.Image or numpy.ndarray or torch.Tensor or list): The image
                or images to preprocess.
            size (int, optional): Overrides `self.size`.
            normalization (str or None, optional): Overrides `self.normalization`.
            do_auto_invert (bool, optional): Overrides `self.do_auto_invert`.
            return_tensors (str or transformers.TensorType, optional): Type of the
                returned tensors, e.g. `'pt'` or `'np'`.

        Returns:
            transformers.BatchFeature: `pixel_values` of shape
                `(batch_size, 1, size, size)`.
        """
        size = self.size if size is None else size
        # None is a valid normalization, so a sentinel marks the default
        if normalization == "default":
            normalization = self.normalization
        if normalization not in NORMALIZATIONS:
            raise ValueError(
                f"normalization must be one of {NORMALIZATIONS}, got {normalization!r}"
            )
        do_auto_invert = (
            self.do_auto_invert if do_auto_invert is None else do_auto_invert
        )
        if not isinstance(images, (list, tuple)):
            images = [images]
        pixel_values = np.stack(
            [
                self._preprocess_single(image, size, normalization, do_auto_invert)
                for image in images
            ]
        )
        return BatchFeature(
            data={"pixel_values": pixel_values}, tensor_type=return_tensors
        )

    @staticmethod
    def _preprocess_single(
        image, size: int, normalization: str | None, do_auto_invert: bool
    ) -> np.ndarray:
        if isinstance(image, PIL.Image.Image):
            image = PIL.ImageOps.exif_transpose(image)
        else:
            image = to_pil_image(image)
        image = image.convert("L")
        width, height = image.size
        crop_size = min(width, height)
        left = int(round((width - crop_size) / 2.0))
        top = int(round((height - crop_size) / 2.0))
        image = image.crop((left, top, left + crop_size, top + crop_size))
        image = image.resize((size, size), PIL.Image.Resampling.BILINEAR)
        data = np.asarray(image, dtype=np.float32)[np.newaxis] / 255.0
        if normalization == "image":
            # For robustness, standardize with image statistics rather than MNIST
            # statistics
            data = (data - data.mean()) / max(data.std(ddof=1), 1e-7)
            if do_auto_invert and np.count_nonzero(data > 0) > data.size / 2:
                # More bright than dark pixels
                data = -data
        elif normalization == "mnist":
            data = (data - MNIST_MEAN) / MNIST_STD
        return data.astype(np.float32)