import numpy as np
import os
from random import shuffle
datasets_dir = './../data/'

def one_hot(x,n):
	if type(x) == list:
		x = np.array(x)
	x = x.flatten()
	o_h = np.zeros((len(x),n))
	o_h[np.arange(len(x)),x] = 1
	return o_h

def mnist(ntrain=60000,ntest=10000,onehot=True):
    ntrain=np.array(ntrain).astype(int).squeeze()
    data_dir = os.path.join(datasets_dir,'mnist/')
    fd = open(os.path.join(data_dir,'train-images-idx3-ubyte'))
    loaded = np.fromfile(file=fd,dtype=np.uint8)
    trX = loaded[16:].reshape((60000,28*28)).astype(float)

    fd = open(os.path.join(data_dir,'train-labels-idx1-ubyte'))
    loaded = np.fromfile(file=fd,dtype=np.uint8)
    trY = loaded[8:].reshape((60000))

    fd = open(os.path.join(data_dir,'t10k-images-idx3-ubyte'))
    loaded = np.fromfile(file=fd,dtype=np.uint8)
    teX = loaded[16:].reshape((10000,28*28)).astype(float)

    fd = open(os.path.join(data_dir,'t10k-labels-idx1-ubyte'))
    loaded = np.fromfile(file=fd,dtype=np.uint8)
    teY = loaded[8:].reshape((10000))

    trY_shuffle = []
    trX_shuffle = []
    index_shuf = range(len(trY))
    shuffle(index_shuf)
    for i in index_shuf:
        trY_shuffle.append(trY[i])
        trX_shuffle.append(trX[i])

    trX = np.asarray(trX_shuffle)
    trY = np.asarray(trY_shuffle)
    
    trX = trX/255.
    teX = teX/255.

    trX = trX[:ntrain]
    trY = trY[:ntrain]

    teX = teX[:ntest]
    teY = teY[:ntest]

    if onehot:
        trY = one_hot(trY, 10)
        teY = one_hot(teY, 10)
    else:
        trY = np.asarray(trY)
        teY = np.asarray(teY)

    return trX,teX,trY,teY

[(16, 'os.path.join', 'os.path.join', False, 'import os\n'), (18, 'numpy.fromfile', 'np.fromfile', True, 'import numpy as np\n'), (22, 'numpy.fromfile', 'np.fromfile', True, 'import numpy as np\n'), (26, 'numpy.fromfile', 'np.fromfile', True, 'import numpy as np\n'), (30, 'numpy.fromfile', 'np.fromfile', True, 'import numpy as np\n'), (36, 'random.shuffle', 'shuffle', False, 'from random import shuffle\n'), (41, 'numpy.asarray', 'np.asarray', True, 'import numpy as np\n'), (42, 'numpy.asarray', 'np.asarray', True, 'import numpy as np\n'), (8, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (17, 'os.path.join', 'os.path.join', False, 'import os\n'), (21, 'os.path.join', 'os.path.join', False, 'import os\n'), (25, 'os.path.join', 'os.path.join', False, 'import os\n'), (29, 'os.path.join', 'os.path.join', False, 'import os\n'), (57, 'numpy.asarray', 'np.asarray', True, 'import numpy as np\n'), (58, 'numpy.asarray', 'np.asarray', True, 'import numpy as np\n'), (15, 'numpy.array', 'np.array', True, 'import numpy as np\n')]