import random

import cv2

import numpy as np

def gaussian_noise(img, mean=0, sigma=0.03):

    img = img.copy()

    noise = np.random.normal(mean, sigma, img.shape)

    mask_overflow_upper = img+noise >= 1.0

    mask_overflow_lower = img+noise < 0

    noise[mask_overflow_upper] = 1.0

    noise[mask_overflow_lower] = 0

    img += noise

    return img

def brightness(img, low, high):

    value = random.uniform(low, high)

    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)

    hsv = np.array(hsv, dtype = np.float64)

    hsv[:,:,1] = hsv[:,:,1]*value

    hsv[:,:,1][hsv[:,:,1]>255]  = 255

    hsv[:,:,2] = hsv[:,:,2]*value 

    hsv[:,:,2][hsv[:,:,2]>255]  = 255

    hsv = np.array(hsv, dtype = np.uint8)

    img = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)

    return img

def fill(img, h, w):

    img = img.astype('float32')

    img = cv2.resize(img, (h, w), cv2.INTER_CUBIC)

    return img

def zoom(img, mask, value):

    if value > 1 or value < 0:

        print('Value for zoom should be less than 1 and greater than 0')

        return img, mask

    value = random.uniform(value, 1)

    h, w = img.shape[:2]

    h_taken = int(value*h)

    w_taken = int(value*w)

    h_start = random.randint(0, h-h_taken)

    w_start = random.randint(0, w-w_taken)

    img = img[h_start:h_start+h_taken, w_start:w_start+w_taken, :]

    img = fill(img, h, w)

    mask = mask[h_start:h_start+h_taken, w_start:w_start+w_taken, :]

    mask = fill(mask, h, w)

    return img, mask

def vertical_shift(img, mask, ratio=0.0):

    if ratio > 1 or ratio < 0:

        print('Value should be less than 1 and greater than 0')

        return img, mask

    ratio = random.uniform(-ratio, ratio)

    h, w = img.shape[:2]

    to_shift = h*ratio

    if ratio > 0:

        img = img[:, :int(w-to_shift), :]

        mask = mask[:, :int(w-to_shift), :]

    if ratio < 0:

        img = img[:, int(-1*to_shift):, :]

        mask = mask[:, int(-1*to_shift):, :]

    img = fill(img, h, w)

    mask = fill(mask, h, w)

    return img, mask

def horizontal_shift(img, mask, ratio=0.0):

    if ratio > 1 or ratio < 0:

        print('Value should be less than 1 and greater than 0')

        return img, mask

    ratio = random.uniform(-ratio, ratio)

    h, w = img.shape[:2]

    to_shift = w*ratio

    if ratio > 0:

        img = img[:, :int(w-to_shift), :]

        mask = mask[:, :int(w-to_shift), :]

    if ratio < 0:

        img = img[:, int(-1*to_shift):, :]

        mask = mask[:, int(-1*to_shift):, :]

    img = fill(img, h, w)

    mask = fill(mask, h, w)

    return img,mask