
from xnas.utils import get_logger

logger = get_logger(__name__)


def preprocess(image, label):
    import tensorflow as tf

    mean = [0.13066044]
    std = [0.3081079]

    
    image = tf.image.convert_image_dtype(image, dtype=tf.float32)
    image = (image - tf.reshape(mean, [1, 1, 1])) / tf.reshape(std, [1, 1, 1])

    label = tf.one_hot(label, depth=10, dtype=tf.int32)

    return image, label


def augment(image, label):
    import tensorflow as tf
    import tensorflow_addons as tfa

    pad = 4
    
    image = tf.image.resize_with_crop_or_pad(image,
                                             28 + pad * 2,
                                             28 + pad * 2)
    image = tf.image.random_crop(image, size=[28, 28, 1])
    
    image = tf.image.random_flip_left_right(image)
    
    image = tfa.image.random_cutout(tf.expand_dims(image, 0), (8, 8))
    image = tf.squeeze(image, axis=0)
    return image, label


def load_data():
    def load_train():
        import tensorflow as tf
        import tensorflow_datasets as tfds

        ds_train = tfds.load('mnist', as_supervised=True, split='train')
        ds_train = (
            ds_train
                .map(preprocess, num_parallel_calls=tf.data.AUTOTUNE)
                .cache()
                .map(augment, num_parallel_calls=tf.data.AUTOTUNE)
        )
        return ds_train

    def load_test():
        import tensorflow as tf
        import tensorflow_datasets as tfds

        ds_test = tfds.load('mnist', as_supervised=True, split='test')
        ds_test = (
            ds_test
                .map(preprocess, num_parallel_calls=tf.data.AUTOTUNE)
                .cache()
        )
        return ds_test

    train_size, valid_size = 60000, 10000

    return {
        'train_gen': load_train,
        'train_size': train_size,
        'valid_gen': load_test,
        'valid_size': valid_size,
        'types': ({'input_0': 'float32'}, 'int32'),
        'shapes': ({'input_0': (28, 28, 1)}, (10,)),
    }
