import PIL
from PIL import ImageEnhance, Image, ImageFilter, ImageChops
import PIL.ImageOps
import numpy as np
randint = np.random.randint

def get_array_color_mode(x):
    """ Given a numpy array representing a single image, it returns the
        PIL color mode that will most likely work with it """
    x = x.squeeze()
    if x.ndim == 2:
        mode = "L"
    elif x.ndim == 3 and x.shape[2] == 1:
        mode = "L"
        x = x.squeeze()
    elif x.ndim == 3:
        mode = "RGB"
    else:
        assert False, "Incapable of interpreting array as an image"

    return mode

def pil2array(im):
    """ Given a PIL image it returns a numpy array representation """
    return np.asarray(im, dtype=np.uint8)

def array2pil(x):
    """ Given a numpy array containing image information returns a PIL image.
        Automatically handles mode, and even handles greyscale images with a
        channels axis
    """
    x = x.squeeze()
    return PIL.Image.fromarray(x, mode=get_array_color_mode(x))


# ==============================================================================
#                                                                  BATCH_PROCESS
# ==============================================================================
def batch_process(X, shape=None, mode=None):
    """ Given a batch of images as a numpy array, it does batch processing steps
        like resizing and color mode conversion.
        shape should be a 2-tuple (width, height) of the new desired
        dimensions.
        mode should be one of {"RGB", "L"}
    """
    assert X.dtype == np.uint8, "X should be 8 bit unsigned integers"
    if shape:
        width, height = shape
    else:
        width, height = X.shape[2], X.shape[1]

    n_samples = X.shape[0]

    # INITIALIZE NEW_BATCH ARRAY - By determining appropriate dimensions first
    if mode == "RGB" and X.ndim in {3,4}:
        n_channels = 3
        new_batch = np.zeros([n_samples, height, width, n_channels], dtype=np.uint8)
    elif mode == "L" and X.ndim in {3,4}:
        n_channels = 1
        new_batch = np.zeros([n_samples, height, width, n_channels], dtype=np.uint8)
    elif X.ndim == 4:
        n_channels = X.shape[3]
        new_batch = np.zeros([n_samples, height, width, n_channels], dtype=np.uint8)
    elif X.ndim == 3:
        n_channels = None
        new_batch = np.zeros([n_samples, height, width], dtype=np.uint8)
    else:
        assert False, "Cannot interpret X as a batch of images, check the dimensions"

    # PREPROCESS EACH IMAGE
    for i in range(n_samples):
        if shape:
            img = array2pil(X[i]).resize(shape, PIL.Image.BICUBIC)
        if mode in {"RGB", "L"}:
            img = img.convert(mode)

        if n_channels==1:
            img = np.asarray(img, dtype=np.uint8)
            new_batch[i] = np.expand_dims(img, axis=2)
        else:
            new_batch[i] = np.asarray(img, dtype=np.uint8)
    return new_batch


# ==============================================================================
#                                                                    RANDOM_CROP
# ==============================================================================
# def random_crop(im, min_scale=0.5, max_scale=1.0, preserve_size=False, resample=PIL.Image.NEAREST):
#     """
#     Args:
#         im:         PIL image
#         min_scale:   (float) minimum ratio along each dimension to crop from.
#         max_scale:   (float) maximum ratio along each dimension to crop from.
#         preserve_size: (bool) Should it resize back to original dims?
#         resample:       resampling method during rescale.
#     Returns:
#         PIL image of size crop_size, randomly cropped from `im`.
#     """
#     assert (min_scale < max_scale), "min_scale MUST be smaller than max_scale"
#     width, height = im.size
#     crop_width = np.random.randint(width*min_scale, width*max_scale)
#     crop_height = np.random.randint(height*min_scale, height*max_scale)
#     x_offset = np.random.randint(0, width - crop_width + 1)
#     y_offset = np.random.randint(0, height - crop_height + 1)
#     im2 = im.crop((x_offset, y_offset,
#                    x_offset + crop_width,
#                    y_offset + crop_height))
#     if preserve_size:
#         im2 = im2.resize(im.size, resample=resample)
#     return im2


# ==============================================================================
#                                                         CROP_AND_PRESERVE_SIZE
# ==============================================================================
# def crop_and_preserve_size(im, crop_dims, offset, resample=PIL.Image.NEAREST):
#     """ Given a PIL image, the dimensions of the crop, and the offset of
#         the crop, it crops the image, and resizes it back to the original
#         dimensions.
#     Args:
#         im:         (PIL image)
#         crop_dims:  Dimensions of the crop region [width, height]
#         offset:     Position of the crop box from Top Left corner [x, y]
#         resample:   resamplimg method
#     """
#     crop_width, crop_height = crop_dims
#     x_offset, y_offset = offset
#     im2 = im.crop((x_offset, y_offset,
#                    x_offset + crop_width,
#                    y_offset + crop_height))
#     im2 = im2.resize(im.size, resample=resample)
#     return im2


# ==============================================================================
#                                                             RANDOM_90_ROTATION
# ==============================================================================
# def random_90_rotation(im):
#     """ Randomly rotates image in 90 degree increments
#         (90, -90, or 180 degrees) """
#     methods = [PIL.Image.ROTATE_90, PIL.Image.ROTATE_180, PIL.Image.ROTATE_270]
#     method = np.random.choice(methods)
#     return im.transpose(method=method)


# ==============================================================================
#                                                                 RANDOM_LR_FLIP
# ==============================================================================
def random_lr_flip(im):
    """ Randomly flips the image left-right with 0.5 probablility """
    if np.random.choice([0,1]) == 1:
        return im.transpose(method=PIL.Image.FLIP_LEFT_RIGHT)
    else:
        return im


# ==============================================================================
#                                                                 RANDOM_TB_FLIP
# ==============================================================================
def random_tb_flip(im):
    """ Randomly flips the image top-bottom with 0.5 probablility """
    if np.random.choice([0,1]) == 1:
        return im.transpose(method=PIL.Image.FLIP_TOP_BOTTOM)
    else:
        return im


# ==============================================================================
#                                                                  RANDOM_INVERT
# ==============================================================================
def random_invert(im):
    """ With a 0.5 probability, it inverts the colors
        NOTE: This does not work on RGBA images yet. """
    assert im.mode != "RGBA", "Does random_invert not support RGBA images"
    if np.random.choice([0,1]) == 1:
        return PIL.ImageOps.invert(im)
    else:
        return im


# ==============================================================================
#                                                                   RANDOM_SHIFT
# ==============================================================================
def random_shift(im, max=(5,5)):
    """ Randomly shifts an image.
    Args:
        im: (pil image)
        max: (tuple of two ints) max amount in each x y direction.
    """
    x_offset = np.random.randint(0, max[0])
    y_offset = np.random.randint(0, max[1])
    return ImageChops.offset(im, xoffset=x_offset, yoffset=y_offset)


# ==============================================================================
#                                                                    SHIFT_IMAGE
# ==============================================================================
def shift_image(im, shift):
    """ Returns a shifted copy of a PIL image.
    Args:
        im:     (PIL image)
        shift:  (tuple of two ints) How much to shift along each axis (x, y)
    """
    return ImageChops.offset(im, xoffset=shift[0], yoffset=shift[1])


# ==============================================================================
#                                                              RANDOM_BRIGHTNESS
# ==============================================================================
def random_brightness(im, sd=0.5, min=0, max=20):
    """Creates a new image which randomly adjusts the brightness of `im` by
       randomly sampling a brightness value centered at 1, with a standard
       deviation of `sd` from a normal distribution. Clips values to a
       desired min and max range.
    Args:
        im:   PIL image
        sd:   (float) Standard deviation used for sampling brightness value.
        min:  (int or float) Clip contrast value to be no lower than this.
        max:  (int or float) Clip contrast value to be no higher than this.
    Returns:
        PIL image with brightness randomly adjusted.
    """
    brightness = np.clip(np.random.normal(loc=1, scale=sd), min, max)
    enhancer = ImageEnhance.Brightness(im)
    return enhancer.enhance(brightness)


# ==============================================================================
#                                                                RANDOM_CONTRAST
# ==============================================================================
def random_contrast(im, sd=0.5, min=0, max=10):
    """Creates a new image which randomly adjusts the contrast of `im` by
       randomly sampling a contrast value centered at 1, with a standard
       deviation of `sd` from a normal distribution. Clips values to a
       desired min and max range.
    Args:
        im:   PIL image
        sd:   (float) Standard deviation used for sampling contrast value.
        min:  (int or float) Clip contrast value to be no lower than this.
        max:  (int or float) Clip contrast value to be no higher than this.
    Returns:
        PIL image with contrast randomly adjusted.
    """
    contrast = np.clip(np.random.normal(loc=1, scale=sd), min, max)
    enhancer = ImageEnhance.Contrast(im)
    return enhancer.enhance(contrast)


# ==============================================================================
#                                                                    RANDOM_BLUR
# ==============================================================================
def random_blur(im, min=0, max=5):
    """ Creates a new image which applies a random amount of Gaussian Blur, with
        a blur radius that is randomly chosen to be in the range [min, max]
        inclusive.
    Args:
        im:   PIL image
        min:  (int) Min amount of blur desired.
        max:  (int) Max amount of blur desired.
    Returns:
        PIL image with random amount of blur applied.
    """
    blur_radius = randint(min, max+1)
    if blur_radius == 0:
        return im
    else:
        return im.filter(ImageFilter.GaussianBlur(radius=blur_radius))


# ==============================================================================
#                                                                   RANDOM_NOISE
# ==============================================================================
def random_noise(im, sd=5):
    """Creates a new image which has random noise.
       The intensity of the noise is determined by first randomly choosing the
       standard deviation of the noise as a value between 0 to `sd`.
       This value is then used as the standard deviation for randomly sampling
       individual pixel noise from a normal distribution.
       This random noise is added to the original image pixel values, and
       clipped to keep all values between 0-255.
    Args:
        im:   PIL image
        sd:   (int) Max Standard Deviation to select from.
    Returns:
        PIL image with random noise added.
    """
    mode = im.mode
    noise_sd = np.random.randint(0, sd)
    if noise_sd > 0:
        noise = np.random.normal(loc=0, scale=noise_sd, size=np.shape(im))
        im2 = np.asarray(im, dtype=np.float32) # prevent overflow
        im2 = np.clip(im2 + noise, 0, 255).astype(np.uint8)
        return array2pil(im2)
    else:
        return im


# ==============================================================================
#                                                                  RANDOM_SHADOW
# ==============================================================================
# def random_shadow(im, shadow, intensity=(0.0, 0.7), crop_range=(0.02, 0.25)):
#     """ Given an image of the scene, and an image of a shadow pattern,
#         It will take random crops from the shadow pattern, and perform
#         random rotations and flips of that crop, before overlaying the
#         shadow on the scene image.
#         The intensity of the shadow is randomly chosen from zero
#         intensity to a max of `max_intensity`.
#         NOTE: This was designed to make use of shadow images being
#         black and white in color (but same colorspace mode as scene image).
#     Args:
#         im:             (PIL image) Image of scene
#         shadow:         (PIL image)
#                         Image of shadow pattern to take a crop from
#         intensity:      (tuple of two floats)(default = (0.0, 0.7))
#                         Min and max values (between 0 to 1) specifying how
#                         strong to make the shadows.
#         crop_range:     (tuple of two floats)(default=(0.02, 0.25))
#                         Min and Max scale for random crop sizes from
#                         the shadow image.
#     Examples:
#         shadow = PIL.Image.open("shadow_aug.png")
#         image = PIL.Image.open("scene.jpg")
#         random_shadow(image, shadow=shadow, max_intensity=0.7, crop_range=(0.02, 0.4))
#     """
#     width, height = im.size
#     mode = im.mode
#     assert im.mode == shadow.mode, "Scene image and shadow image must be same colorspace mode"
#
#     # Take random crop from shadow image
#     min_crop_scale, max_crop_scale = crop_range
#     shadow = random_crop(shadow, min_scale=min_crop_scale, max_scale=max_crop_scale, preserve_size=False)
#     shadow = shadow.resize((width, height), resample=PIL.Image.BILINEAR)
#
#     # random flips, rotations, and color inversion
#     shadow = random_tb_flip(random_lr_flip(random_90_rotation(shadow)))
#     shadow = random_invert(shadow)
#     # Ensure same shape as scene image after flips and rotations
#     shadow = shadow.resize((width, height), resample=PIL.Image.BILINEAR)
#
#     # Scale the shadow into proportional intensities (0-1)
#     intensity_value = np.random.rand(1)
#     min, max = intensity
#     intensity_value = (intensity_value*(max - min))+min # remapped to min,max range
#     shadow = np.divide(shadow, 255)
#     shadow = np.multiply(intensity_value, shadow)
#
#     # Overlay the shadow
#     overlay = (np.multiply(im, 1-shadow)).astype(np.uint8)
#     return PIL.Image.fromarray(overlay, mode=mode)
#
#
# # ==============================================================================
# #                                                                RANDOM_ROTATION
# # ==============================================================================
# def random_rotation(im, max=10, include_corners=True, resample=PIL.Image.NEAREST):
#     """ Creates a new image which is rotated by a random amount between
#         [-max, +max] inclusive.
#     Args:
#         im:              (PIL image)
#         max:             (int) Max angle (in degrees in either direction).
#         include_corners: (bool)
#                 If True, then the image canvas is expanded at first to
#                 fit the rotated corners, and then rescaled back to
#                 original image size.
#                 If False, then the original image canvas remains intact,
#                 and the corners of the rotated image that fall outside
#                 this box are clipped off.
#     Returns:
#         PIL image with random rotation applied.
#     """
#     original_dims = im.size
#     angle = randint(-max, max+1)
#     if angle == 0:
#         return im
#     else:
#         im2 = im.rotate(angle, resample=resample, expand=include_corners)
#         if include_corners:
#             im2 = im2.resize(original_dims, resample=resample)
#         return im2


# ==============================================================================
#                                                         RANDOM_TRANSFORMATIONS
# ==============================================================================
def random_transformations(
    X,
    # shadow=(0.6, 0.9),
    # shadow_file="shadow_pattern.jpg",
    # shadow_crop_range=(0.02, 0.5),
    # rotate=180,
    # crop=0.5,
    lr_flip=True,
    tb_flip=True,
    brightness=(0.5, 0.4, 4),
    contrast=(0.5, 0.3, 5),
    blur=3,
    noise=10,
    target_shape=None,
    ):

    img_shape = X.shape
    image = array2pil(X)
    # original_dims = image.size

    if lr_flip and np.random.choice([True, False]):
        image = image.transpose(method=PIL.Image.FLIP_LEFT_RIGHT)

    if tb_flip and np.random.choice([True, False]):
        image = image.transpose(method=PIL.Image.FLIP_TOP_BOTTOM)



    # if brightness is not None:
    #     image = random_brightness(image, sd=brightness[0], min=brightness[1], max=brightness[2])
    # if contrast is not None:
    #     image = random_contrast(image, sd=contrast[0], min=contrast[1], max=contrast[2])
    # if blur is not None:
    #     image = random_blur(image, 0, blur)
    # if noise:
    #     image = random_noise(image, sd=noise)




    print('the original shape is {}, the target shape is {}'.format(img_shape, target_shape))
    if img_shape != target_shape:
        print('do resize BICUBIC')
        image = image.resize([target_shape[1], target_shape[0]], resample=PIL.Image.BICUBIC)

    image = np.asarray(image, dtype=np.uint8)
    return image
