"""Helper script to initalize the data needed for Task2.8"""

#Load MNIST Data
import os
import gzip
import numpy as np

#need to make sure urllib3 is installed via pip
from urllib.request import urlretrieve


def download(filename, source='http://yann.lecun.com/exdb/mnist/'):
    print(" Downloading %s" % filename )
    urlretrieve(source + filename, filename)

def load_mnist_images( filename ):
    if not os.path.exists( filename ):
         download(filename)
    with gzip.open(filename,'rb') as f:
        data = np.frombuffer(f.read() , np.uint8 , offset=16)
    data = data.reshape(-1, 28, 28)
    return data / np.float32(256)

def load_mnist_labels( filename ):
    if not os.path.exists( filename ):
        download( filename )
    with gzip.open(filename , 'rb') as f:
        data = np.frombuffer(f.read() , np.uint8 , offset=8)
    return data

def get_training_data():
    X_train = load_mnist_images('train-images-idx3-ubyte.gz')
    y_train = load_mnist_labels('train-labels-idx1-ubyte.gz')
    return [X_train, y_train]

def get_test_data():
    X_test = load_mnist_images('t10k-images-idx3-ubyte.gz')
    y_test = load_mnist_labels('t10k-labels-idx1-ubyte.gz')
    return [X_test, y_test]


# import data
X_train, y_train = get_training_data()
X_test, y_test = get_test_data()