import numpy as np
import random
import torch
import numpy as np
from PIL import Image, ImageEnhance

class Compose():
    def __init__(self, transforms):
        self.transforms = transforms
    
    def __call__(self, img):
        for t in self.transforms:
            img = t(img)
        return img

class to_Tensor():
    def __call__(self,arr):
        if len(np.array(arr).shape) == 2:
            arr = np.array(arr)[:,:,None]
        arr = torch.from_numpy(np.array(arr).transpose(2,0,1))
        return arr

def imresize(im, size, interp='bilinear'):
    if interp == 'nearest':
        resample = Image.NEAREST
    elif interp == 'bilinear':
        resample = Image.BILINEAR
    elif interp == 'bicubic':
        resample = Image.BICUBIC
    else:
        raise Exception('resample method undefined!')
    return im.resize(size, resample)

class To_PIL_Image():
    def __call__(self, img):
        return to_pil_image(img)

class normalize():
    def __init__(self,mean,std):
        self.mean = torch.tensor(mean)
        self.std = torch.tensor(std)
    def __call__(self,img):
        self.mean = torch.as_tensor(self.mean,dtype=img.dtype,device=img.device)
        self.std = torch.as_tensor(self.std,dtype=img.dtype,device=img.device)
        return (img-self.mean)/self.std

class RandomVerticalFlip():
    def __init__(self, prob):
        self.prob = prob

    def __call__(self, img):
        if random.random() < self.prob:
            if isinstance(img, Image.Image):
                return img.transpose(Image.FLIP_TOP_BOTTOM)
            if isinstance(img, np.ndarray):
                return np.flip(img, axis=0)
        return img

class RandomHorizontallyFlip():
    def __init__(self, prob=0.5):
        self.prob = prob

    def __call__(self, img):
        if random.random() < self.prob:
            if isinstance(img, Image.Image):
                return img.transpose(Image.FLIP_LEFT_RIGHT)
            if isinstance(img, np.ndarray):
                return np.flip(img, axis=1)
        return img

class RandomRotate():
    def __init__(self, degree, prob=0.5):
        self.prob = prob
        self.degree = degree

    def __call__(self, img, interpolation=Image.BILINEAR):
        if random.random() < self.prob:
            rotate_detree = random.random() * 2 * self.degree - self.degree
            return img.rotate(rotate_detree, interpolation)
        return img

class RandomBrightness():
    def __init__(self, min_factor, max_factor, prob=0.5):
        """ :param min_factor: The value between 0.0 and max_factor
            that define the minimum adjustment of image brightness.
            The value  0.0 gives a black image,The value 1.0 gives the original image, value bigger than 1.0 gives more bright image.
            :param max_factor: A value should be bigger than min_factor.
            that define the maximum adjustment of image brightness.
            The value  0.0 gives a black image, value 1.0 gives the original image, value bigger than 1.0 gives more bright image.

        """
        self.prob = prob
        self.min_factor = min_factor
        self.max_factor = max_factor
    
    # def __brightness(self, img, factor):
    #     return img * (1.0 - factor) + img * factor

    # def __call__(self, img):
    #     if random.random() < self.prob:
    #         factor = np.random.uniform(self.min_factor, self.max_factor)
    #         return self.__brightness(img, factor)

    def __call__(self, img):
        if random.random() < self.prob:
            factor = np.random.uniform(self.min_factor, self.max_factor)
            enhancer_brightness = ImageEnhance.Brightness(img)
            return enhancer_brightness.enhance(factor)

        return img

class RandomContrast():
    def __init__(self, min_factor, max_factor, prob=0.5):
        """ :param min_factor: The value between 0.0 and max_factor
            that define the minimum adjustment of image contrast.
            The value  0.0 gives s solid grey image, value 1.0 gives the original image.
            :param max_factor: A value should be bigger than min_factor.
            that define the maximum adjustment of image contrast.
            The value  0.0 gives s solid grey image, value 1.0 gives the original image.
        """
        self.prob = prob
        self.min_factor = min_factor
        self.max_factor = max_factor

    def __call__(self, img):
        if random.random() < self.prob:
            factor = np.random.uniform(self.min_factor, self.max_factor)
            enhance_contrast = ImageEnhance.Contrast(img)
            return enhance_contrast.enhance(factor)
        return img

def to_pil_image(pic, mode=None):
    """Convert a tensor or an ndarray to PIL Image.

    See :class:`~torchvision.transforms.ToPIlImage` for more details.

    Args:
        pic (Tensor or numpy.ndarray): Image to be converted to PIL Image.
        mode (`PIL.Image mode`_): color space and pixel depth of input data (optional).

    .. _PIL.Image mode: http://pillow.readthedocs.io/en/3.4.x/handbook/concepts.html#modes

    Returns:
        PIL Image: Image converted to PIL Image.
    """
    # if not(_is_numpy_image(pic) or _is_tensor_image(pic)):
    #     raise TypeError('pic should be Tensor or ndarray. Got {}.'.format(type(pic)))

    npimg = pic
    if isinstance(pic, torch.FloatTensor):
        pic = pic.mul(255).byte()
    if torch.is_tensor(pic):
        npimg = np.transpose(pic.numpy(), (1, 2, 0))

    if not isinstance(npimg, np.ndarray):
        raise TypeError('Input pic must be a torch.Tensor or NumPy ndarray, ' +
                        'not {}'.format(type(npimg)))

    if npimg.shape[2] == 1:
        expected_mode = None
        npimg = npimg[:, :, 0]
        if npimg.dtype == np.uint8:
            expected_mode = 'L'
        elif npimg.dtype == np.int16:
            expected_mode = 'I;16'
        elif npimg.dtype == np.int32:
            expected_mode = 'I'
        elif npimg.dtype == np.float32:
            expected_mode = 'F'
        if mode is not None and mode != expected_mode:
            raise ValueError("Incorrect mode ({}) supplied for input type {}. Should be {}"
                             .format(mode, np.dtype, expected_mode))
        mode = expected_mode

    elif npimg.shape[2] == 4:
        permitted_4_channel_modes = ['RGBA', 'CMYK']
        if mode is not None and mode not in permitted_4_channel_modes:
            raise ValueError("Only modes {} are supported for 4D inputs".format(permitted_4_channel_modes))

        if mode is None and npimg.dtype == np.uint8:
            mode = 'RGBA'
    else:
        permitted_3_channel_modes = ['RGB', 'YCbCr', 'HSV']
        if mode is not None and mode not in permitted_3_channel_modes:
            raise ValueError("Only modes {} are supported for 3D inputs".format(permitted_3_channel_modes))
        if mode is None and npimg.dtype == np.uint8:
            mode = 'RGB'

    if mode is None:
        raise TypeError('Input type {} is not supported'.format(npimg.dtype))

    return Image.fromarray(npimg, mode=mode)