from __future__ import absolute_import, division

import keras
from keras.datasets import mnist
from keras.preprocessing.image import ImageDataGenerator
import numpy as np


def get_mnist_dataset():
    #(X_train, y_train), (X_test, y_test) = mnist.load_data()
    f = np.load('./mnist.npz')
    X_train = f['x_train']
    y_train = f['y_train']
    X_test  = f['x_test']
    y_test  = f['y_test']

    X_train = X_train.astype('float32')/255
    X_test  = X_test.astype('float32')/255
    X_train = X_train[..., None]
    X_test  = X_test[..., None]
    Y_train = keras.utils.to_categorical(y_train, 10)
    Y_test  = keras.utils.to_categorical(y_test, 10)

    return (X_train, Y_train), (X_test, Y_test)


def get_gen(set_name, batch_size, translate, scale, shuffle=True):
    if set_name == 'train':
       (X, Y),_ = get_mnist_dataset()
    elif set_name == 'test':
       _,(X, Y) = get_mnist_dataset()

    image_gen = ImageDataGenerator(
                 zoom_range=scale, \
                 width_shift_range=translate, \
                 height_shift_range=translate)
     
    gen = image_gen.flow(X,Y,batch_size=batch_size, shuffle=shuffle)
    return gen

