import numpy as np
from PIL import Image


class background_mix(object):
    def __init__(self, patch=8, mix_patch=4):
        self.last_imgs=[]
        self.patch=patch
        self.mix_patch=mix_patch

    def __call__(self, img):
        img_ = np.array(img).copy()
        if len(self.last_imgs) == 0:
            self.last_imgs.append(img_)
            return img
        h, w, c = img_.shape
        grids =np.random.randint(self.patch, size=(self.mix_patch, 2))
        mask = np.zeros((h, w))
        for j in range(self.mix_patch):
            if not (grids[j][0] >=2 and grids[j][0]<=5 and grids[j][1] >=2 and grids[j][1]<=5):
                mask[grids[j][0]*h//self.patch:(grids[j][0]+1)*h//self.patch, grids[j][1]*w//self.patch:(grids[j][1]+1)*w//self.patch] = 1
        mask = (np.expand_dims(mask, axis=-1)).repeat(c, axis=-1)
        mixed_img_ = img_*(1-mask)+self.last_imgs[np.random.randint(len(self.last_imgs))]*(mask)
        self.last_imgs.append(img_)
        if len(self.last_imgs) > 10:
            self.last_imgs.pop(0)
        return Image.fromarray(mixed_img_.astype('uint8')).convert('RGB')
