import torch
import torch.nn as nn
from torchvision import transforms
from torchvision.transforms import functional as TF

class BatchRandomGrayscale(nn.Module):
    def __init__(self, p=0.1):
        super(BatchRandomGrayscale, self).__init__()
        self.p = p

    def forward(self, images):
        '''
        It is assumed that all images have the same number of channels
        '''
        if torch.rand(1) < self.p:
            n_channels = TF.get_image_num_channels(images[0])
            return [TF.rgb_to_grayscale(image, num_output_channels=n_channels) for image in images]
        return images


class BatchColorJitter(nn.Module):
    def __init__(self, brightness=0, contrast=0, saturation=0, hue=0):
        super(BatchColorJitter, self).__init__()
        self.t = transforms.ColorJitter(brightness, contrast, saturation, hue)
        self.functions = [TF.adjust_brightness, TF.adjust_contrast, TF.adjust_saturation, TF.adjust_hue]

    def get_params(self):
        '''
        returns:
        `indices`: the order in which brightness, contrast, saturation and hue will be adjusted
        `brightness_factor`, `contrast_factor`, `saturation_factor` and `hue_factor`: float values or None
        '''
        return self.t.get_params(self.t.brightness, self.t.contrast, self.t.saturation, self.t.hue)

    def per_image_transform(self, image, indices, factors):
        for index in indices:
            if factors[index] is None:
                continue
            fn = self.functions[index]
            image = fn(image, factors[index])
        return image

    def forward(self, images):
        indices, *factors = self.get_params()
        return [self.per_image_transform(image, indices, factors) for image in images]


class BatchRandomHorizontalFlip(nn.Module):
    def __init__(self, p=0.5):
        super(BatchRandomHorizontalFlip, self).__init__()
        self.p = p

    def forward(self, images):
        if torch.rand(1) < self.p:
            return [TF.hflip(image) for image in images]
        return images


class BatchRandomVerticalFlip(nn.Module):
    def __init__(self, p=0.5):
        super(BatchRandomVerticalFlip, self).__init__()
        self.p = p

    def forward(self, images):
        if torch.rand(1) < self.p:
            return [TF.vflip(image) for image in images]
        return images


class BatchRandomAdjustSharpness(nn.Module):
    def __init__(self, sharpness_factor, p=0.5):
        super(BatchRandomAdjustSharpness, self).__init__()
        self.sharpness_factor = sharpness_factor
        self.p = p

    def forward(self, images):
        if torch.rand(1) < self.p:
            return [TF.adjust_sharpness(image, self.sharpness_factor) for image in images]
        return images


class BatchRandomPosterize(nn.Module):
    def __init__(self, bits, p=0.5):
        super(BatchRandomPosterize, self).__init__()
        self.bits = bits
        self.p = p

    def forward(self, images):
        if torch.rand(1) < self.p:
            return [TF.posterize(image, self.bits) for image in images]
        return images


class BatchRandomEqualize(nn.Module):
    def __init__(self, p=0.5):
        super(BatchRandomEqualize, self).__init__()
        self.p = p

    def forward(self, images):
        if torch.rand(1) < self.p:
            return [TF.equalize(image) for image in images]
        return images


class BatchRandomInvert(nn.Module):
    def __init__(self, p=0.5):
        super(BatchRandomInvert, self).__init__()
        self.p = p

    def forward(self, images):
        if torch.rand(1) < self.p:
            return [TF.invert(image) for image in images]
        return images


class BatchRandomSolarize(nn.Module):
    def __init__(self, threshold, p=0.5):
        super(BatchRandomSolarize, self).__init__()
        self.threshold = threshold
        self.p = p

    def forward(self, images):
        if torch.rand(1) < self.p:
            return [TF.solarize(image, self.threshold) for image in images]
        return images


class BatchResize(nn.Module):
    def __init__(self, size, interpolation=transforms.InterpolationMode.BILINEAR, max_size=None, antialias=None):
        super(BatchResize, self).__init__()
        self.size = size
        self.interpolation = interpolation
        self.max_size = max_size
        self.antialias = antialias

    def forward(self, images):
        return [TF.resize(image, self.size, self.interpolation, self.max_size, self.antialias) for image in images]


class BatchRandomApply(nn.Module):
    def __init__(self, batch_transforms, p=0.5):
        super(BatchRandomApply, self).__init__()
        self.batch_transforms = batch_transforms
        self.p = p

    def forward(self, images):
        if torch.rand(1) < self.p:
            for bt in self.batch_transforms:
                images = bt(images)
            return images
        return images


class BatchCenterCrop(nn.Module):
    def __init__(self, size):
        super(BatchCenterCrop, self).__init__()
        self.size = size

    def forward(self, images):
        return [TF.center_crop(image, self.size) for image in images]


class BatchPad(nn.Module):
    def __init__(self, padding, fill=0, padding_mode='constant'):
        super(BatchPad, self).__init__()
        self.padding, self.fill, self.padding_mode = padding, fill, padding_mode

    def forward(self, images):
        return [TF.pad(image, self.padding, self.fill, self.padding_mode) for image in images]


class BatchToTensor(nn.Module):
    def __init__(self):
        super(BatchToTensor, self).__init__()
        
    def forward(self, images):
        return [TF.to_tensor(image) for image in images]


class BatchCompose(nn.Module):
    def __init__(self, batch_transforms):
        super(BatchCompose, self).__init__()
        self.batch_transforms = batch_transforms

    def forward(self, images):
        for bt in self.batch_transforms:
            images = bt(images)
        return images


batch_transforms = nn.ModuleList([
    BatchColorJitter(0.3,0.3,0.3), 
    BatchRandomHorizontalFlip(),
    BatchRandomGrayscale(),
    BatchRandomAdjustSharpness(3),
    BatchRandomPosterize(bits=5, p=0.1),
    BatchRandomApply(nn.ModuleList([BatchPad([0,30])]), p=0.2),
    BatchResize([224,224]), 
    BatchToTensor()
])
batch_transform_train = BatchCompose(batch_transforms)

batch_transform_val = nn.ModuleList([
    BatchResize([224,224]), 
    BatchToTensor()
])
batch_transform_val = BatchCompose(batch_transform_val)