import random

import torch
import torchvision.transforms as T
import torchvision.transforms.functional as tf
import torch.nn.functional as F


class RandomResizedCrop(T.RandomResizedCrop):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def __call__(self, x, dtypes):
        """WARNING: torchvision v0.11. Wrapper to T.RandomResizedCrop.__call__"""
        i, j, h, w = self.get_params(x[0], self.scale, self.ratio)
        return [tf.resized_crop(img, i, j, h, w, self.size, self.interpolation) for img in x]

class ColorJitter(T.ColorJitter):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def forward(self, x, dtypes):
        fn_idx, brightness_factor, contrast_factor, saturation_factor, hue_factor = \
            self.get_params(self.brightness, self.contrast, self.saturation, self.hue)

        o = []
        for img, dtype in zip(x, dtypes):
            if dtype == 'albedo':
                for fn_id in fn_idx:
                    if fn_id == 0 and brightness_factor is not None:
                        img = tf.adjust_brightness(img, brightness_factor)
                    elif fn_id == 1 and contrast_factor is not None:
                        img = tf.adjust_contrast(img, contrast_factor)
                    elif fn_id == 2 and saturation_factor is not None:
                        img = tf.adjust_saturation(img, saturation_factor)
                    elif fn_id == 3 and hue_factor is not None:
                        img = tf.adjust_hue(img, hue_factor)
            o.append(img)
        return o

class RandomHorizontalFlip(T.RandomHorizontalFlip):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def flip_x(self, img, dtype):
        if dtype == 'normals':
            img[0] *= -1
        return tf.hflip(img)

    def forward(self, x, dtypes):
        if torch.rand(1) < self.p:
            return [self.flip_x(img, dtype) for img, dtype in zip(x, dtypes)]
        return x

class RandomVerticalFlip(T.RandomVerticalFlip):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def flip_y(self, img, dtype):
        if dtype == 'normals':
            img[1] *= -1
        return tf.vflip(img)

    def forward(self, x, dtypes):
        if torch.rand(1) < self.p:
            return [self.flip_y(img, dtype) for img, dtype in zip(x, dtypes)]
        return x

def deg0(x, y, z):
    return torch.stack([ x,  y,  z])
def deg90(x, y, z):
    return torch.stack([-y,  x,  z])
def deg180(x, y, z):
    return torch.stack([-x, -y,  z])
def deg270(x, y, z):
    return torch.stack([ y, -x,  z])

class RandomIncrementRotate:
    def __init__(self, p):
        self.p = p
        self.angles = [0, 90, 180, 270]

        # adjusts surface normals vector depending on rotation angle
        self.f = { 0: deg0, 90: deg90, 180: deg180, 270: deg270 }

    def rotate(self, img, theta, dtype):
        if dtype == 'normals':
            img = self.f[theta](*img)
        return tf.rotate(img, theta)

    def __call__(self, x, dtypes):
        if torch.rand(1) < self.p:
            theta = random.choice(self.angles)
            return [self.rotate(img, theta, dtype) for img, dtype in zip(x, dtypes)]
        return x

class NormalizeGeometry:
    def normalize(self, img, dtype):
        if dtype == 'normals':
            # perform [0, 1] -> [-1, 1] mapping
            img = 2*img - 1
            # normalize vector to unit sphere
            img = F.normalize(img, dim=0)
        return img

    def __call__(self, x, dtypes):
        return [self.normalize(img, dtype) for img, dtype in zip(x, dtypes)]

class RandomCrop(T.RandomCrop):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def forward(self, x, dtypes):
        img_size = tf.get_image_size(x[0])
        assert all(tf.get_image_size(y) == img_size  for y in x)
        i, j, h, w = self.get_params(x[0], self.size)
        return [tf.crop(img, i, j, h, w) for img in x]

class CenterCrop:
    def __init__(self, size):
        self.size = size

    def __call__(self, x, dtypes):
        return [tf.center_crop(img, self.size) for img in x]

class Resize(T.Resize):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def forward(self, x, dtypes):
        return [super(Resize, self).forward(img) for img in x]

class Identity():
    def __call__(self, x, dtypes):
        return x

class ToTensor:
    def __call__(self, x, dtypes):
        return [tf.to_tensor(img) for img in x]

class Pipeline:
    DATA_TYPES = ['input', 'normals', 'albedo']

    def __init__(self, *transforms, dtypes=None):
        assert all(d in Pipeline.DATA_TYPES for d in dtypes)
        self.dtypes = dtypes
        self.transforms = transforms

    def __call__(self, x):
        #print(self.dtypes)
        assert len(self.dtypes) == len(x)
        assert all(y.shape[1:] == x[0].shape[1:] for y in x)
        for f in self.transforms:
            x = f(x, self.dtypes)
        return x