import random

import numpy as np

from PIL import ImageOps, ImageFilter, Image

import torchvision.transforms as torch_transforms

from torchvision.transforms import InterpolationMode

# import phasepack

##========================Natural Images========================================

class GaussianBlur(object):

    def __init__(self, p):

        self.p = p

    def __call__(self, img):

        if random.random() < self.p:

            sigma = random.random() * 1.9 + 0.1

            return img.filter(ImageFilter.GaussianBlur(sigma))

        else:

            return img

class Solarization(object):

    def __init__(self, p):

        self.p = p

    def __call__(self, img):

        if random.random() < self.p:

            return ImageOps.solarize(img)

        else:

            return img

###-----------------------------------------------------------------------------

## For SimCLR

class SimCLRTransform:

    """Transforms for SimCLR during training step of the pre-training stage.

    Transform::

        RandomResizedCrop(size=self.image_size)

        RandomHorizontalFlip()

        RandomApply([color_jitter], p=0.8)

        RandomGrayscale(p=0.2)

        RandomApply([GaussianBlur(kernel_size=int(0.1 * self.image_size))], p=0.5)

        transforms.ToTensor()

    Example::

        transform = SimCLRTrainDataTransform(image_size=32)

        x = sample()

        (xi, xj, xk) = transform(x) # xk is only for the online evaluator if used

    """

    def __init__(

        self, image_size: int = 256, jitter_strength: float = 1.0, normalize=None

    ) -> None:

        self.jitter_strength = jitter_strength

        self.image_size = image_size

        self.normalize = normalize

        kernel_size = int(0.1 * self.image_size)

        if kernel_size % 2 == 0:

            kernel_size += 1

        self.data_transform = torch_transforms.Compose([

            torch_transforms.RandomResizedCrop(size=self.image_size),

            torch_transforms.RandomHorizontalFlip(p=0.5),

            torch_transforms.RandomApply([

                    torch_transforms.ColorJitter(

                            0.8 * self.jitter_strength,

                            0.8 * self.jitter_strength,

                            0.8 * self.jitter_strength,

                            0.2 * self.jitter_strength)], p=0.8),

            torch_transforms.RandomGrayscale(p=0.2),

            torch_transforms.RandomApply(

                    [torch_transforms.GaussianBlur(kernel_size=kernel_size)],

                    p=0.5),

            torch_transforms.ToTensor(),

            torch_transforms.Normalize( mean=[0.485, 0.456, 0.406],

                                        std=[0.229, 0.224, 0.225])

        ])

    def __call__(self, sample):

        transform = self.data_transform

        xi = transform(sample)

        xj = transform(sample)

        return xi, xj

    def get_composition(self):

        return self.data_transform

## For BalowTwins and VICRegularization

class BarlowTwinsTransformOrig:

    def __init__(self, image_size = 256):

        self.transform = torch_transforms.Compose([

            torch_transforms.RandomResizedCrop(image_size,

                                    interpolation=InterpolationMode.BICUBIC),

            torch_transforms.RandomHorizontalFlip(p=0.5),

            torch_transforms.RandomApply(

                [torch_transforms.ColorJitter(brightness=0.4, contrast=0.4,

                                    saturation=0.2, hue=0.1)],

                p=0.8

            ),

            torch_transforms.RandomGrayscale(p=0.2),

            GaussianBlur(p=1.0),

            Solarization(p=0.0),

            torch_transforms.ToTensor(),

            torch_transforms.Normalize( mean=[0.485, 0.456, 0.406],

                                        std=[0.229, 0.224, 0.225])

        ])

        self.transform_prime = torch_transforms.Compose([

            torch_transforms.RandomResizedCrop(image_size,

                                    interpolation=InterpolationMode.BICUBIC),

            torch_transforms.RandomHorizontalFlip(p=0.5),

            torch_transforms.RandomApply(

                [torch_transforms.ColorJitter(brightness=0.4, contrast=0.4,

                                    saturation=0.2, hue=0.1)],

                p=0.8

            ),

            torch_transforms.RandomGrayscale(p=0.2),

            GaussianBlur(p=0.1),

            Solarization(p=0.2),

            torch_transforms.ToTensor(),

            torch_transforms.Normalize( mean=[0.485, 0.456, 0.406],

                                        std=[0.229, 0.224, 0.225])

        ])

    def __call__(self, x):

        y1 = self.transform(x)

        y2 = self.transform_prime(x)

        return y1, y2

    def get_composition(self):

        return (self.transform, self.transform_prime)

##======================== UltraSound Images ===================================

class ClassifierTransform:

    def __init__(self, mode = "train", image_size = 256):

        data_mean = [0.485, 0.456, 0.406]

        data_std  = [0.229, 0.224, 0.225]

        train_transform = torch_transforms.Compose([

            torch_transforms.Resize(image_size,

                                    interpolation=InterpolationMode.BICUBIC),

            torch_transforms.RandAugment(num_ops=5, magnitude=5),

            torch_transforms.RandomHorizontalFlip(p=0.5),

            torch_transforms.RandomVerticalFlip(p=0.5),

            torch_transforms.ToTensor(),

            torch_transforms.Normalize(mean=data_mean, std=data_std)

        ])

        infer_transform = torch_transforms.Compose([

            torch_transforms.Resize(image_size,

                        interpolation=InterpolationMode.BICUBIC),

            torch_transforms.ToTensor(),

            torch_transforms.Normalize(mean=data_mean, std=data_std)

        ])

        if   mode == "train": self.transform = train_transform

        elif mode == "infer": self.transform = infer_transform

        else : raise ValueError("Unknown Mode set only `train` or `infer` allowed")

    def __call__(self, x):

        y = self.transform(x)

        return y

    def get_composition(self):

        return self.transform

##------------------------------------------------------------------------------

# def phasecongruence(image):

#     image = np.asarray(image)

#     [M, ori, ft, T] = phasepack.phasecongmono(image,

#                             nscale=5, minWaveLength=5)

#     out = ((M - M.min())/(M.max() - M.min()+1) *255.0).astype(np.uint8)

#     out = Image.fromarray(out).convert("RGB")

#     return out

class CustomInfoMaxTransform:

    def __init__(self, image_size = 256):

        self.transform = torch_transforms.Compose([

            torch_transforms.RandomResizedCrop(256,

                                    interpolation=InterpolationMode.BICUBIC),

            torch_transforms.RandomHorizontalFlip(p=0.5),

            # torch_transforms.Lambda(phasecongruence),

            torch_transforms.RandAugment(num_ops=5, magnitude=5),

            torch_transforms.ToTensor(),

            torch_transforms.Normalize( mean=[0.485, 0.456, 0.406],

                                        std=[0.229, 0.224, 0.225])

        ])

        self.transform_prime = torch_transforms.Compose([

            torch_transforms.RandomResizedCrop(image_size,

                                    interpolation=InterpolationMode.BICUBIC),

            torch_transforms.RandomHorizontalFlip(p=0.5),

            torch_transforms.RandomApply(

                [torch_transforms.ColorJitter(brightness=0.4, contrast=0.4,

                                    saturation=0.2, hue=0.1)],

                p=0.8

            ),

            torch_transforms.RandomGrayscale(p=0.2),

            GaussianBlur(p=0.1),

            Solarization(p=0.2),

            torch_transforms.ToTensor(),

            torch_transforms.Normalize( mean=[0.485, 0.456, 0.406],

                                        std=[0.229, 0.224, 0.225])

        ])

    def __call__(self, x):

        y1 = self.transform(x)

        y2 = self.transform_prime(x)

        return y1, y2

    def get_composition(self):

        return (self.transform, self.transform_prime)

##------------------------------------------------------------------------------

class AEncStandardTransform:

    def __init__(self, image_size = 256):

        data_mean = [0.485, 0.456, 0.406]

        data_std  = [0.229, 0.224, 0.225]

        train_transform = torch_transforms.Compose([

            torch_transforms.Resize(image_size,

                                    interpolation=InterpolationMode.BICUBIC),

            torch_transforms.RandAugment(num_ops=3, magnitude=3),

            torch_transforms.RandomHorizontalFlip(p=0.5),

            torch_transforms.RandomVerticalFlip(p=0.5),

            torch_transforms.ToTensor(),

            torch_transforms.Normalize(mean=data_mean, std=data_std)

        ])

        self.transform = train_transform

    def __call__(self, x):

        y = self.transform(x)

        return y, y

    def get_composition(self):

        return self.transform

class AEncInpaintTransform:

    def __init__(self, image_size = 256):

        data_mean = [0.485, 0.456, 0.406]

        data_std  = [0.229, 0.224, 0.225]

        self.transform = torch_transforms.Compose([

            torch_transforms.Resize(image_size,

                                    interpolation=InterpolationMode.BICUBIC),

            torch_transforms.RandAugment(num_ops=1, magnitude=1),

            torch_transforms.RandomHorizontalFlip(p=0.5),

            torch_transforms.RandomVerticalFlip(p=0.5),

            torch_transforms.ToTensor(),

            torch_transforms.Normalize(mean=data_mean, std=data_std)

        ])

        self.erase = torch_transforms.RandomErasing(p=1, value="random")

    def __call__(self, x):

        y2 = self.transform(x)

        y1 = self.erase(y2)

        return y1, y2

    def get_composition(self):

        return (self.transform, self.erase)