# Source: http://martin-thoma.com/classify-mnist-with-pybrain/
from struct import unpack
import gzip
from numpy import zeros, uint8

def load_test(imagefile='../Data/t10k-images-idx3-ubyte.gz', labelfile='../Data/t10k-labels-idx1-ubyte.gz', limit=-1):
    """Read input-vector (image) and target class (label, 0-9) and return
       it as list of tuples.
    """
    # Open the images with gzip in read binary mode
    images = gzip.open(imagefile, 'rb')
    labels = gzip.open(labelfile, 'rb')

    # Read the binary data

    # We have to get big endian unsigned int. So we need '>I'

    # Get metadata for images
    images.read(4)  # skip the magic_number
    number_of_images = images.read(4)
    number_of_images = unpack('>I', number_of_images)[0]
    rows = images.read(4)
    rows = unpack('>I', rows)[0]
    cols = images.read(4)
    cols = unpack('>I', cols)[0]

    # Get metadata for labels
    labels.read(4)  # skip the magic_number
    N = labels.read(4)
    N = unpack('>I', N)[0]

    if limit < 0:
        limit = N

    if number_of_images != N:
        raise Exception('number of labels did not match the number of images')

    # Get the data
    x = zeros((limit, rows * cols), dtype=uint8)  # Initialize numpy array
    y = zeros((limit), dtype=uint8)  # Initialize numpy array
    
    for i in range(N):

        if i % 1000 == 0:
            print("i: %i" % i)

        if i == limit:
            break

        for row in range(rows):
            for col in range(cols):
                tmp_pixel = images.read(1)  # Just a single byte
                tmp_pixel = unpack('>B', tmp_pixel)[0]
                x[i][cols * row + col] = tmp_pixel

        tmp_label = labels.read(1)
        y[i] = unpack('>B', tmp_label)[0]

    return (x, y)

def load_train(imagefile='../Data/train-images-idx3-ubyte.gz', labelfile='../Data/train-labels-idx1-ubyte.gz', limit=-1):
    """Read input-vector (image) and target class (label, 0-9) and return
       it as list of tuples.
    """
    # Open the images with gzip in read binary mode
    images = gzip.open(imagefile, 'rb')
    labels = gzip.open(labelfile, 'rb')

    # Read the binary data

    # We have to get big endian unsigned int. So we need '>I'

    # Get metadata for images
    images.read(4)  # skip the magic_number
    number_of_images = images.read(4)
    number_of_images = unpack('>I', number_of_images)[0]
    rows = images.read(4)
    rows = unpack('>I', rows)[0]
    cols = images.read(4)
    cols = unpack('>I', cols)[0]

    # Get metadata for labels
    labels.read(4)  # skip the magic_number
    N = labels.read(4)
    N = unpack('>I', N)[0]

    if limit < 0:
        limit = N

    if number_of_images != N:
        raise Exception('number of labels did not match the number of images')

    # Get the data
    x = zeros((limit, rows * cols), dtype=uint8)  # Initialize numpy array
    y = zeros((limit), dtype=uint8)  # Initialize numpy array
    
    for i in range(N):

        if i % 1000 == 0:
            print("i: %i" % i)

        if i == limit:
            break

        for row in range(rows):
            for col in range(cols):
                tmp_pixel = images.read(1)  # Just a single byte
                tmp_pixel = unpack('>B', tmp_pixel)[0]
                x[i][cols * row + col] = tmp_pixel

        tmp_label = labels.read(1)
        y[i] = unpack('>B', tmp_label)[0]

    return (x, y)