# Script for data augmentation functions

import numpy as np
from collections import deque
from PIL import Image
import cv2
from data.config import *

def imread_cv2(image_path):
    """
    Read image_path with cv2 format (H, W, C)
    if image is '.gif' outputs is a numpy array of {0,1}
    """
    image_format = image_path[-3:]
    if image_format == 'jpg':
        image = cv2.imread(image_path)
    else:
        image = np.array(Image.open(image_path))

    return image

def resize_cv2(image, heigh=1280, width=1918):
    return cv2.resize(image, (width, heigh), cv2.INTER_LINEAR)

def image_to_tensor(image, mean=0, std=1.):
    """Transform image (input is numpy array, read in by cv2) """
    if len(image.shape) == 2:
        image = image.reshape(image.shape[0], image.shape[1], 1)
    image = image.astype(np.float32)
    image = (image-mean)/std
    image = image.transpose((2,0,1))
    tensor = torch.from_numpy(image)

    return tensor

# --- Data Augmentation functions --- #
# A lot of functions can be found here:
# https://github.com/fchollet/keras/blob/master/keras/preprocessing/image.py#L223

# transform image and label
def randomHorizontalFlip(image, mask, p=0.5):
    """Do a random horizontal flip with probability p"""
    if np.random.random() < p:
        image = np.fliplr(image)
        mask = np.fliplr(mask)
    return image, mask

def randomVerticalFlip(image, mask, p=0.5):
    """Do a random vertical flip with probability p"""
    if np.random.random() < p:
        image = np.flipud(image)
        mask = np.flipud(mask)
    return image, mask

def randomHorizontalShift(image, mask, max_shift=0.05, p=0.5):
    """Do random horizontal shift with max proportion shift and with probability p
    Elements that roll beyond the last position are re-introduced at the first."""
    max_shift_pixels = int(max_shift*image.shape[1])
    shift = np.random.choice(np.arange(-max_shift_pixels, max_shift_pixels+1))
    if np.random.random() < p:
        image = np.roll(image, shift, axis=1)
        mask = np.roll(mask, shift, axis=1)
    return image, mask

def randomVerticalShift(image, mask, max_shift=0.05, p=0.5):
    """Do random vertical shift with max proportion shift and probability p
    Elements that roll beyond the last position are re-introduced at the first."""
    max_shift_pixels = int(max_shift*image.shape[0])
    shift = np.random.choice(np.arange(-max_shift_pixels, max_shift_pixels+1))
    if np.random.random() < p:
            image = np.roll(image, shift, axis=0)
            mask = np.roll(mask, shift, axis=0)
    return image, mask

def randomInvert(image, mask, p=0.5):
    """Randomly invert image with probability p"""
    if np.random.random() < p:
        image = 255 - image
        mask = mask
    return image, mask

def randomBrightness(image, mask, p=0.75):
    """With probability p, randomly increase or decrease brightness.
    See https://stackoverflow.com/questions/37822375/python-opencv-increasing-image-brightness-without-overflowing-uint8-array"""
    if np.random.random() < p:
        max_value = np.percentile(255-image, q=25) # avoid burning out white cars, so take image-specific maximum
        value = np.random.choice(np.arange(-max_value, max_value))
        if value > 0:
            image = np.where((255 - image) < value,255,image+value).astype(np.uint8)
        else:
            image = np.where(image < -value,0,image+value).astype(np.uint8)
    return image, mask

def randomHue(image, mask, p=0.25, max_value=75):
    """With probability p, randomly increase or decrease hue.
    See https://stackoverflow.com/questions/32609098/how-to-fast-change-image-brightness-with-python-opencv"""

    if np.random.random() < p:
        value = np.random.choice(np.arange(-max_value, max_value))
        hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
        hsv[:,:,0] = hsv[:,:,0] + value
        hsv = np.clip(hsv, a_min=0, a_max=255).astype(np.uint8)
        image = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)

    return image, mask

def GaussianBlur(image, mask, kernel=(1, 1),sigma=1,  p=0.5):
    """With probability p, apply Gaussian blur"""
    # TODO
    return image, mask

def randomRotate(image, mask, max_angle, p=0.5):
    """Perform random rotation with max_angle and probability p"""
    # TODO

    return(image, mask)

[(23, 'cv2.resize', 'cv2.resize', (['image', '(width, heigh)', 'cv2.INTER_LINEAR'], {}), False, 'import cv2\n'), (16, 'cv2.imread', 'cv2.imread', (['image_path'], {}), False, 'import cv2\n'), (43, 'numpy.random.random', 'np.random.random', ([], {}), True, 'import numpy as np\n'), (44, 'numpy.fliplr', 'np.fliplr', (['image'], {}), True, 'import numpy as np\n'), (45, 'numpy.fliplr', 'np.fliplr', (['mask'], {}), True, 'import numpy as np\n'), (50, 'numpy.random.random', 'np.random.random', ([], {}), True, 'import numpy as np\n'), (51, 'numpy.flipud', 'np.flipud', (['image'], {}), True, 'import numpy as np\n'), (52, 'numpy.flipud', 'np.flipud', (['mask'], {}), True, 'import numpy as np\n'), (59, 'numpy.arange', 'np.arange', (['(-max_shift_pixels)', '(max_shift_pixels + 1)'], {}), True, 'import numpy as np\n'), (60, 'numpy.random.random', 'np.random.random', ([], {}), True, 'import numpy as np\n'), (61, 'numpy.roll', 'np.roll', (['image', 'shift'], {'axis': '(1)'}), True, 'import numpy as np\n'), (62, 'numpy.roll', 'np.roll', (['mask', 'shift'], {'axis': '(1)'}), True, 'import numpy as np\n'), (69, 'numpy.arange', 'np.arange', (['(-max_shift_pixels)', '(max_shift_pixels + 1)'], {}), True, 'import numpy as np\n'), (70, 'numpy.random.random', 'np.random.random', ([], {}), True, 'import numpy as np\n'), (71, 'numpy.roll', 'np.roll', (['image', 'shift'], {'axis': '(0)'}), True, 'import numpy as np\n'), (72, 'numpy.roll', 'np.roll', (['mask', 'shift'], {'axis': '(0)'}), True, 'import numpy as np\n'), (77, 'numpy.random.random', 'np.random.random', ([], {}), True, 'import numpy as np\n'), (85, 'numpy.random.random', 'np.random.random', ([], {}), True, 'import numpy as np\n'), (86, 'numpy.percentile', 'np.percentile', (['(255 - image)'], {'q': '(25)'}), True, 'import numpy as np\n'), (98, 'numpy.random.random', 'np.random.random', ([], {}), True, 'import numpy as np\n'), (100, 'cv2.cvtColor', 'cv2.cvtColor', (['image', 'cv2.COLOR_BGR2HSV'], {}), False, 'import cv2\n'), (103, 'cv2.cvtColor', 'cv2.cvtColor', (['hsv', 'cv2.COLOR_HSV2BGR'], {}), False, 'import cv2\n'), (18, 'PIL.Image.open', 'Image.open', (['image_path'], {}), False, 'from PIL import Image\n'), (87, 'numpy.arange', 'np.arange', (['(-max_value)', 'max_value'], {}), True, 'import numpy as np\n'), (99, 'numpy.arange', 'np.arange', (['(-max_value)', 'max_value'], {}), True, 'import numpy as np\n'), (102, 'numpy.clip', 'np.clip', (['hsv'], {'a_min': '(0)', 'a_max': '(255)'}), True, 'import numpy as np\n'), (89, 'numpy.where', 'np.where', (['(255 - image < value)', '(255)', '(image + value)'], {}), True, 'import numpy as np\n'), (91, 'numpy.where', 'np.where', (['(image < -value)', '(0)', '(image + value)'], {}), True, 'import numpy as np\n')]