"""
Various routines to load/access MNIST data.
"""
from __future__ import absolute_import

import numpy

from ..amat import AMat

default_path = '/u/bergstrj/pub/data/mnist.amat'
"""the location of a file containing mnist data in .amat format"""


def head(n=10, path=None):
    """Load the first MNIST examples.

    Returns two matrices: x, y.  x has N rows of 784 columns.  Each row of x represents the
    28x28 grey-scale pixels in raster order.  y is a vector of N integers.  Each element y[i]
    is the label of the i'th row of x.
    
    """
    path = path if path is not None else default_path

    dat = AMat(path=path, head=n)

    return dat.input, numpy.asarray(dat.target, dtype='int64').reshape(dat.target.shape[0])

def all(path=None):
    return head(n=None, path=path)


def train_valid_test(path=None, ntrain=50000, nvalid=10000, ntest=10000):
    all_x, all_targ = all(path=path)

    train = all_x[0:ntrain], all_targ[0:ntrain]
    valid = all_x[ntrain:ntrain+nvalid], all_targ[ntrain:ntrain+nvalid]
    test = all_x[ntrain+nvalid:ntrain+nvalid+ntest], all_targ[ntrain+nvalid:ntrain+nvalid+ntest]

    return train, valid, test