import tensorflow as tf


class Augmentator(object):
    def __init__(self,type,size=1,mean=0,std=1):
        self.size = size
        if type=='scramble':
            self.augment = self.scramble
        elif type=='mix_scramble':
            self.augment = self.mix_scramble  
        elif type=='blur':
            self.augment = self.gaussian_blur
            self.pointwise_filter = tf.eye(3, batch_shape=[1, 1])

        elif type=='high_low_pass':
            self.augment = self.high_low_pass
            self.kernel = self.gaussian_kernel(size,mean,std)
            self.kernel = tf.tile(self.kernel[:, :, tf.newaxis, tf.newaxis], [1, 1, 3, 1])
            self.pointwise_filter = tf.eye(3, batch_shape=[1, 1])
            self.paddings = [[size,size],[size,size],[0,0]]
        elif type=='no_op':
            self.augment = self.no_op
