import cPickle
import gzip
import os

import numpy as np

DATA_URL = 'http://deeplearning.net/data/mnist/'
FILE_NAME = 'mnist.pkl.gz'
DATA_DIRECTORY = '../data/'

def load_data(one_hot=True):
    """Return the MNIST data as a tuple containing the training data,
    the validation data, and the test data.

    The ``training_data`` is returned as a tuple with two entries.
    The first entry contains the actual training images.  This is a
    numpy ndarray with 50,000 entries.  Each entry is, in turn, a
    numpy ndarray with 784 values, representing the 28 * 28 = 784
    pixels in a single MNIST image.

    The second entry in the ``training_data`` tuple is a numpy ndarray
    containing 50,000 entries.  Those entries are the just the digit
    values (0...9) for the corresponding images contained in the first
    entry of the tuple, i.e. the labels.

    The ``validation_data`` and ``test_data`` are similar, except
    each contains only 10,000 images."""

    dataset_file = maybe_download()

    print '... loading data'
    f = gzip.open(dataset_file, 'rb')
    tr_d, va_d, te_d = cPickle.load(f)
    f.close()

    if one_hot:
        tr_d, va_d, te_d = [(x, to_one_hot(y)) for (x, y) in (tr_d, va_d, te_d)]

    return tr_d, va_d, te_d

def maybe_download():
    """Download the pickled MNIST data, unless it's already here."""
    filepath = os.path.join(DATA_DIRECTORY, FILE_NAME)
    if os.path.isfile(filepath):
        return filepath

    if not os.path.exists(DATA_DIRECTORY):
        os.mkdir(DATA_DIRECTORY)

    import urllib
    print 'Downloading data from %s' % DATA_URL
    filepath, _ = urllib.urlretrieve(DATA_URL + FILE_NAME, filepath)
    statinfo = os.stat(filepath)
    print 'Succesfully downloaded', filepath, statinfo.st_size, 'bytes.'
    return filepath

def to_one_hot(labels, num_classes=10):
    """"Convert class labels from scalars to one-hot vectors."""
    return np.eye(num_classes)[labels]
