################### DataProcess.py ##################
import os, struct
import numpy as np
from array import array as pyarray 
from numpy import append, array, int8, uint8, zeros
from scipy import stats

def load_mnist(dataset="training", digits=None, path=None, asbytes=False, selection=None, 
    return_labels=True, return_indices=False, zscore=True, appendOne=True, isShuffle=True):
    """
    Loads MNIST files into a 3D numpy array.

    You have to download the data separately from [MNIST]_. It is recommended
    to set the environment variable ``MNIST`` to point to the folder where you
    put the data, so that you don't have to select path. On a Linux+bash setup,
    this is done by adding the following to your ``.bashrc``::

        export MNIST=/path/to/mnist

    Parameters
    ----------
    dataset : str 
        Either "training" or "testing", depending on which dataset you want to
        load. 
    digits : list 
        Integer list of digits to load. The entire database is loaded if set to
        ``None``. Default is ``None``.
    path : str 
        Path to your MNIST datafiles. The default is ``None``, which will try
        to take the path from your environment variable ``MNIST``. The data can
        be downloaded from http://yann.lecun.com/exdb/mnist/.
    asbytes : bool
        If True, returns data as ``numpy.uint8`` in [0, 255] as opposed to
        ``numpy.float64`` in [0.0, 1.0].
    selection : slice
        Using a `slice` object, specify what subset of the dataset to load. An
        example is ``slice(0, 20, 2)``, which would load every other digit
        until--but not including--the twentieth.
    return_labels : bool
        Specify whether or not labels should be returned. This is also a speed
        performance if digits are not specified, since then the labels file
        does not need to be read at all.
    return_indicies : bool
        Specify whether or not to return the MNIST indices that were fetched.
        This is valuable only if digits is specified, because in that case it
        can be valuable to know how far
        in the database it reached.
    zscore : boolean
        if True, do the zscore transformation for the raw image data. Default is true.
    appendOne : boolean
        if True, append one in the front of the image data representing the input to the bias node
    isShuffle : boolean
        if True, shuffle the data set.

    Returns
    -------
    images : ndarray
        Image data of shape ``(N, rows, cols)``, where ``N`` is the number of images. If neither labels nor inices are returned, then this is returned directly, and not inside a 1-sized tuple.
    labels : ndarray
        Array of size ``N`` describing the labels. Returned only if ``return_labels`` is `True`, which is default.
    indices : ndarray
        The indices in the database that were returned.

    Examples
    --------
    Assuming that you have downloaded the MNIST database and set the
    environment variable ``$MNIST`` point to the folder, this will load all
    images and labels from the training set:

    >>> images, labels = ag.io.load_mnist('training') # doctest: +SKIP

    Load 100 sevens from the testing set:    

    >>> sevens = ag.io.load_mnist('testing', digits=[7], selection=slice(0, 100), return_labels=False) # doctest: +SKIP

    """

    # The files are assumed to have these names and should be found in 'path'
    files = {
        'training': ('train-images-idx3-ubyte', 'train-labels-idx1-ubyte'),
        'testing': ('t10k-images-idx3-ubyte', 't10k-labels-idx1-ubyte'),
    }

    if path is None:
        try:
            path = os.environ['MNIST']
        except KeyError:
            raise ValueError("Unspecified path requires environment variable $MNIST to be set")

    try:
        images_fname = os.path.join(path, files[dataset][0])
        labels_fname = os.path.join(path, files[dataset][1])
    except KeyError:
        raise ValueError("Data set must be 'testing' or 'training'")

    # We can skip the labels file only if digits aren't specified and labels aren't asked for
    if return_labels or digits is not None:
        flbl = open(labels_fname, 'rb')
        magic_nr, size = struct.unpack(">II", flbl.read(8))
        labels_raw = pyarray("b", flbl.read())
        flbl.close()

    fimg = open(images_fname, 'rb')
    magic_nr, size, rows, cols = struct.unpack(">IIII", fimg.read(16))
    images_raw = pyarray("B", fimg.read())
    fimg.close()

    indices = range(size)
    N = len(indices)

    images = zeros((N, rows*cols), dtype=uint8)

    if return_labels:
        labels = zeros((N), dtype=int8)
    for i, index in enumerate(indices):
        images[i] = array(images_raw[ indices[i]*rows*cols : (indices[i]+1)*rows*cols ]).reshape(1,rows*cols)
        if return_labels:
            labels[i] = labels_raw[indices[i]]

    if not asbytes:
            images = images.astype(float)/255.0

    if zscore:
        images = zscoreTransform(images);
    #numpy.insert(arr, index, values, axis=None)
    if appendOne:
        images = np.insert(images, 0, 1, axis=1)

    #shuffle the images
    if isShuffle:
        np.random.shuffle(indices)
        images = images[indices]
        labels = labels[indices]

    
    #choose images from shuffled images
    indices = range(size)
    if digits:
        indices = [k for k in range(size) if labels[k] in digits]
    
    if selection:
        indices = indices[selection] 

    images = images[indices]
    labels = labels[indices]
    

    


    ret = (images,)
    if return_labels:
        ret += (labels,)
    if return_indices:
        ret += (indices,)
    if len(ret) == 1:
        return ret[0] # Don't return a tuple of one
    else:
        return ret

def zscoreTransform(data):
    """
    Do a z-score transformation for the data such that the mean is 0 and variance is 1.
    """
    zscoreData = stats.zscore(data);
    #convert NaN to zero. NaN mean this feacture is a constant. So this feature is zero after centering
    zscoreData[np.isnan(zscoreData)] = 0;
    return zscoreData

if __name__ == '__main__':
    myTrainData = load_mnist(dataset="training", path='../', digits = None, selection = None, zscore = True, isShuffle = True)
    import matplotlib.pyplot as plt

    images = myTrainData[0]
    for i in range(4):
        print i
        image = images[i]
        image = image[1:].reshape(28,28)
        plt.imshow(image, cmap=plt.get_cmap('gray'))
        plt.show()

