import numpy as np

from imgaug import augmenters as iaa

import numpy as np

#def get_augumentation()

#    def eraser(input_img):

#        seq = iaa.Sequential([iaa.GaussianBlur((0,3.0)), iaa.Affine(translate_px={"x":(-40,40)})])

#        seq_det = seq.to_deterministic()

#        input_img = seq_det.augment_images(input_img)

#        return input_img

#    return eraser

def get_random_eraser(p=0.5, s_l=0.02, s_h=0.4, r_1=0.3, r_2=1/0.3, v_l=0, v_h=255, pixel_level=False):

    def eraser(input_img):

        img_h, img_w, img_c = input_img.shape

        p_1 = np.random.rand()

        if p_1 > p:

            return input_img

        while True:

            s = np.random.uniform(s_l, s_h) * img_h * img_w

            r = np.random.uniform(r_1, r_2)

            w = int(np.sqrt(s / r))

            h = int(np.sqrt(s * r))

            left = np.random.randint(0, img_w)

            top = np.random.randint(0, img_h)

            if left + w <= img_w and top + h <= img_h:

                break

        if pixel_level:

            c = np.random.uniform(v_l, v_h, (h, w, img_c))

        else:

            c = np.random.uniform(v_l, v_h)

        input_img[top:top + h, left:left + w, :] = c

        return input_img

    return eraser