""" Code adapted from PixMix' paper. """

# fmt: off
from PIL import Image
from torch import nn

import numpy as np
from torch_uncertainty.transforms import Shear, Translate, augmentations


# fmt: on
def get_ab(beta):
    if np.random.random() < 0.5:
        a = np.float32(np.random.beta(beta, 1))
        b = np.float32(np.random.beta(1, beta))
    else:
        a = 1 + np.float32(np.random.beta(1, beta))
        b = -np.float32(np.random.beta(1, beta))
    return a, b


def add(img1: Image.Image, img2: Image.Image, beta: float) -> Image.Image:
    a, b = get_ab(beta)
    img1, img2 = img1 * 2 - 1, img2 * 2 - 1
    out = a * img1 + b * img2
    return (out + 1) / 2


def multiply(img1: Image.Image, img2: Image.Image, beta: float) -> Image.Image:
    a, b = get_ab(beta)
    img1, img2 = img1 * 2, img2 * 2
    out = (img1**a) * (img2.clip(1e-37) ** b)
    return out / 2


# Summarize mixing operations
mixings = [add, multiply]


class PixMix(nn.Module):
    """PixMix augmentation class.

    Args:
        mixing_set (MixingSet): Dataset to be mixed with.
        mixing_iterations (int): Number of mixing iterations.
        mixing_severity (float): Severity of mixing.
        all_ops (bool): Whether to use augmentations included in ImageNet-C.
            Defaults to True.

    Note:
        Default arguments are set to follow original guidelines.
    """

    def __init__(
        self,
        mixing_set,
        mixing_iterations: int = 4,
        augmentation_severity: float = 3,
        mixing_severity: float = 3,
        all_ops: bool = True,
    ):
        super().__init__()
        self.mixing_set = mixing_set
        self.num_mixing_images = len(mixing_set)
        self.mixing_iterations = mixing_iterations
        self.augmentation_severity = augmentation_severity
        self.mixing_severity = mixing_severity

        if not all_ops:
            allowed_augmentations = [
                aug for aug in augmentations if not aug.corruption_overlap
            ]
        else:
            allowed_augmentations = augmentations

        self.aug_instances = []
        for aug in allowed_augmentations:
            if aug == Shear or aug == Translate:
                self.aug_instances.append(aug(axis=0))
                self.aug_instances.append(aug(axis=1))
            else:
                self.aug_instances.append(aug())

    def __call__(self, img: Image.Image) -> np.ndarray:
        if np.random.random() < 0.5:
            mixed = self.augment_input(img)
        else:
            mixed = img

        for _ in range(np.random.randint(self.mixing_iterations + 1)):
            if np.random.random() < 0.5:
                aug_image_copy = self._augment(img)
            else:
                aug_image_copy = np.random.choice(self.num_mixing_images)

            mixed_op = np.random.choice(mixings)
            mixed = mixed_op(
                np.array(mixed), np.array(aug_image_copy), self.mixing_severity
            )
            mixed = np.clip(mixed, 0, 1)
        return mixed

    def _augment(self, image: Image.Image) -> np.ndarray:
        op = np.random.choice(self.aug_instances)
        if op.level_type == int:
            aug_level = self._sample_int(op.pixmix_max_level)
        else:
            aug_level = self._sample_float(op.pixmix_max_level)
        return op(image.copy(), aug_level)

    def _sample_level(self) -> float:
        return np.random.uniform(low=0.1, high=self.augmentation_severity)

    def _sample_int(self, maxval: int) -> int:
        """Helper method to scale `level` between 0 and maxval.

        Args:
            level: Level of the operation that will be between [0, maxval]
            maxval: Maximum value that the operation can have. This will be
            scaled to level/maxval.

        Returns:
            An int that results from scaling `maxval` according to `level`.
        """
        return int(self._sample_level() * maxval / 10)

    def _sample_float(self, maxval: float) -> float:
        """Helper function to scale `val` between 0 and maxval.

        Args:
            level: Level of the operation that will be in the range
                [0, `maxval`]
            maxval: Maximum value that the operation can have. This will be
                scaled to level/maxval.

        Returns:
            A float that results from scaling `maxval` according to `level`.
        """
        return float(self._sample_level()) * maxval / 10.0
