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