import tensorflow as tf
import copy


class ImagePool:
    def __init__(self, pool_size):
        self.pool_size = pool_size
        if self.pool_size > 0:
            self.num_imgs = 0
            self.images = []

    def query(self, images):
        if self.pool_size == 0:
            return images
        ret_imgs = []
        for i in range(images.shape[0]):
            image = tf.expand_dims(images[i], axis=0)
            if self.num_imgs < self.pool_size:
                self.num_imgs = self.num_imgs + 1
                self.images.append(image)
                ret_imgs.append(image)
            else:
                p = tf.random_uniform((1,), 0, 1).numpy()[0]
                if p > 0.5:
                    random_id = tf.random_uniform((1,), 0, self.pool_size - 1).numpy()[0].astype(int)
                    tmp = copy.copy(self.images[random_id])
                    self.images[random_id] = image
                    ret_imgs.append(tmp)
                else:
                    ret_imgs.append(image)
        ret_imgs = tf.concat(ret_imgs, 0)
        return ret_imgs
