import numpy as np
import os
import typing

import tensorflow as tf

def load(batch_size: int) -> typing.Dict[str, typing.Any]:
    # MNIST is small, so for performance consideration load all to memory
    mnist_folderpath = './MNIST'
    train_images = np.load(os.path.join(mnist_folderpath, 'train_X.npy'))
    test_images = np.load(os.path.join(mnist_folderpath, 'test_X.npy'))
    # reshape for the channel dimension, change type to float32
    train_images = train_images.reshape(train_images.shape[0], 28, 28, 1).astype('float32')
    test_images = test_images.reshape(test_images.shape[0], 28, 28, 1).astype('float32')
    # get normalizer
    train_flat = train_images.reshape((-1, train_images.shape[-1]))
    test_flat = test_images.reshape((-1, test_images.shape[-1]))
    all_flat = np.concatenate([train_flat, test_flat], axis=0)
    mean = all_flat.mean(axis=0)
    std = all_flat.std(axis=0)
    # normalization
    train_images -= mean
    train_images /= std
    test_images -= mean
    test_images /= std
    # convert to tf.Dataset class
    train_dataset = tf.data.Dataset.from_tensor_slices(train_images)\
            .shuffle(train_images.shape[0])\
            .batch(batch_size)
    test_dataset = tf.data.Dataset.from_tensor_slices(test_images)\
            .shuffle(test_images.shape[0])\
            .batch(batch_size)
    data_dict = {
            'train_dataset':train_dataset,
            'test_dataset':test_dataset,
            'train_batch_count':np.ceil(train_images.shape[0]/batch_size).astype(int),
            'test_batch_count':np.ceil(test_images.shape[0]/batch_size).astype(int),
            'data_shape':train_images.shape[1:],
            }
    return data_dict

