"""
Common photometric transforms for data augmentation.
"""
import numpy as np
from PIL import Image
from torchvision import transforms as transforms
import cv2


# List all the available augmentations
available_augmentations = [
    "additive_gaussian_noise",
    "additive_speckle_noise",
    "random_brightness",
    "random_contrast",
    "additive_shade",
    "motion_blur",
]


class additive_gaussian_noise(object):
    """Additive gaussian noise."""

    def __init__(self, stddev_range=None):
        # If std is not given, use the default setting
        if stddev_range is None:
            self.stddev_range = [5, 95]
        else:
            self.stddev_range = stddev_range

    def __call__(self, input_image):
        # Get the noise stddev
        stddev = np.random.uniform(self.stddev_range[0], self.stddev_range[1])
        noise = np.random.normal(0.0, stddev, size=input_image.shape)
        noisy_image = (input_image + noise).clip(0.0, 255.0)

        return noisy_image


class additive_speckle_noise(object):
    """Additive speckle noise."""

    def __init__(self, prob_range=None):
        # If prob range is not given, use the default setting
        if prob_range is None:
            self.prob_range = [0.0, 0.005]
        else:
            self.prob_range = prob_range

    def __call__(self, input_image):
        # Sample
        prob = np.random.uniform(self.prob_range[0], self.prob_range[1])
        sample = np.random.uniform(0.0, 1.0, size=input_image.shape)

        # Get the mask
        mask0 = sample <= prob
        mask1 = sample >= (1 - prob)

        # Mask the image (here we assume the image ranges from 0~255
        noisy = input_image.copy()
        noisy[mask0] = 0.0
        noisy[mask1] = 255.0

        return noisy


class random_brightness(object):
    """Brightness change."""

    def __init__(self, brightness=None):
        # If the brightness is not given, use the default setting
        if brightness is None:
            self.brightness = 0.5
        else:
            self.brightness = brightness

        # Initialize the transformer
        self.transform = transforms.ColorJitter(brightness=self.brightness)

    def __call__(self, input_image):
        # Convert to PIL image
        if isinstance(input_image, np.ndarray):
            input_image = Image.fromarray(input_image.astype(np.uint8))

        return np.array(self.transform(input_image))


class random_contrast(object):
    """Additive contrast."""

    def __init__(self, contrast=None):
        # If the brightness is not given, use the default setting
        if contrast is None:
            self.contrast = 0.5
        else:
            self.contrast = contrast

        # Initialize the transformer
        self.transform = transforms.ColorJitter(contrast=self.contrast)

    def __call__(self, input_image):
        # Convert to PIL image
        if isinstance(input_image, np.ndarray):
            input_image = Image.fromarray(input_image.astype(np.uint8))

        return np.array(self.transform(input_image))


class additive_shade(object):
    """Additive shade."""

    def __init__(self, nb_ellipses=20, transparency_range=None, kernel_size_range=None):
        self.nb_ellipses = nb_ellipses
        if transparency_range is None:
            self.transparency_range = [-0.5, 0.8]
        else:
            self.transparency_range = transparency_range

        if kernel_size_range is None:
            self.kernel_size_range = [250, 350]
        else:
            self.kernel_size_range = kernel_size_range

    def __call__(self, input_image):
        # ToDo: if we should convert to numpy array first.
        min_dim = min(input_image.shape[:2]) / 4
        mask = np.zeros(input_image.shape[:2], np.uint8)
        for i in range(self.nb_ellipses):
            ax = int(max(np.random.rand() * min_dim, min_dim / 5))
            ay = int(max(np.random.rand() * min_dim, min_dim / 5))
            max_rad = max(ax, ay)
            x = np.random.randint(max_rad, input_image.shape[1] - max_rad)
            y = np.random.randint(max_rad, input_image.shape[0] - max_rad)
            angle = np.random.rand() * 90
            cv2.ellipse(mask, (x, y), (ax, ay), angle, 0, 360, 255, -1)

        transparency = np.random.uniform(*self.transparency_range)
        kernel_size = np.random.randint(*self.kernel_size_range)

        # kernel_size has to be odd
        if (kernel_size % 2) == 0:
            kernel_size += 1
        mask = cv2.GaussianBlur(mask.astype(np.float32), (kernel_size, kernel_size), 0)
        shaded = input_image[..., None] * (
            1 - transparency * mask[..., np.newaxis] / 255.0
        )
        shaded = np.clip(shaded, 0, 255)

        return np.reshape(shaded, input_image.shape)


class motion_blur(object):
    """Motion blur."""

    def __init__(self, max_kernel_size=10):
        self.max_kernel_size = max_kernel_size

    def __call__(self, input_image):
        # Either vertical, horizontal or diagonal blur
        mode = np.random.choice(["h", "v", "diag_down", "diag_up"])
        ksize = np.random.randint(0, int(round((self.max_kernel_size + 1) / 2))) * 2 + 1
        center = int((ksize - 1) / 2)
        kernel = np.zeros((ksize, ksize))
        if mode == "h":
            kernel[center, :] = 1.0
        elif mode == "v":
            kernel[:, center] = 1.0
        elif mode == "diag_down":
            kernel = np.eye(ksize)
        elif mode == "diag_up":
            kernel = np.flip(np.eye(ksize), 0)
        var = ksize * ksize / 16.0
        grid = np.repeat(np.arange(ksize)[:, np.newaxis], ksize, axis=-1)
        gaussian = np.exp(
            -(np.square(grid - center) + np.square(grid.T - center)) / (2.0 * var)
        )
        kernel *= gaussian
        kernel /= np.sum(kernel)
        blurred = cv2.filter2D(input_image, -1, kernel)

        return np.reshape(blurred, input_image.shape)


class normalize_image(object):
    """Image normalization to the range [0, 1]."""

    def __init__(self):
        self.normalize_value = 255

    def __call__(self, input_image):
        return (input_image / self.normalize_value).astype(np.float32)