import numpy as np
import pytest
from PIL import Image

from image_prediction.estimator.preprocessor.preprocessors.basic import BasicPreprocessor
from image_prediction.estimator.preprocessor.preprocessors.identity import IdentityPreprocessor
from image_prediction.estimator.preprocessor.utils import image_to_normalized_tensor, images_to_batch_tensor


def image_conversion_is_correct(image):
    tensor = image_to_normalized_tensor(image)
    image_re = Image.fromarray(np.uint8(tensor * 255), mode="RGB")
    return image == image_re and tensor.ndim == 3


def images_conversion_is_correct(images, tensor):
    if not (images or tensor.size > 0):
        return True
    return all([isinstance(tensor, np.ndarray), tensor.ndim == 4, tensor.shape[0] == len(images)])


def test_image_to_tensor(images):
    assert all(map(image_conversion_is_correct, images))


def test_images_to_batch_tensor(images):
    tensor = images_to_batch_tensor(images)
    assert images_conversion_is_correct(images, tensor)


def test_basic_preprocessor(images):
    tensor = BasicPreprocessor()(images)
    assert images_conversion_is_correct(images, tensor)


def test_identity_preprocessor(images):
    images_preprocessed = IdentityPreprocessor()(images)
    assert images_preprocessed == images