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, 15, 16, 50), 'os.path.join', 'os.path.join', ({(16, 28, 16, 40): 'datasets_dir', (16, 41, 16, 49): '"""mnist/"""'}, {}), "(datasets_dir, 'mnist/')", False, 'import os\n'), ((18, 13, 18, 48), 'numpy.fromfile', 'np.fromfile', (), '', True, 'import numpy as np\n'), ((22, 13, 22, 48), 'numpy.fromfile', 'np.fromfile', (), '', True, 'import numpy as np\n'), ((26, 13, 26, 48), 'numpy.fromfile', 'np.fromfile', (), '', True, 'import numpy as np\n'), ((30, 13, 30, 48), 'numpy.fromfile', 'np.fromfile', (), '', True, 'import numpy as np\n'), ((36, 4, 36, 23), 'random.shuffle', 'shuffle', ({(36, 12, 36, 22): 'index_shuf'}, {}), '(index_shuf)', False, 'from random import shuffle\n'), ((41, 10, 41, 33), 'numpy.asarray', 'np.asarray', ({(41, 21, 41, 32): 'trX_shuffle'}, {}), '(trX_shuffle)', True, 'import numpy as np\n'), ((42, 10, 42, 33), 'numpy.asarray', 'np.asarray', ({(42, 21, 42, 32): 'trY_shuffle'}, {}), '(trY_shuffle)', True, 'import numpy as np\n'), ((8, 6, 8, 17), 'numpy.array', 'np.array', ({(8, 15, 8, 16): 'x'}, {}), '(x)', True, 'import numpy as np\n'), ((17, 14, 17, 62), 'os.path.join', 'os.path.join', ({(17, 27, 17, 35): 'data_dir', (17, 36, 17, 61): '"""train-images-idx3-ubyte"""'}, {}), "(data_dir, 'train-images-idx3-ubyte')", False, 'import os\n'), ((21, 14, 21, 62), 'os.path.join', 'os.path.join', ({(21, 27, 21, 35): 'data_dir', (21, 36, 21, 61): '"""train-labels-idx1-ubyte"""'}, {}), "(data_dir, 'train-labels-idx1-ubyte')", False, 'import os\n'), ((25, 14, 25, 61), 'os.path.join', 'os.path.join', ({(25, 27, 25, 35): 'data_dir', (25, 36, 25, 60): '"""t10k-images-idx3-ubyte"""'}, {}), "(data_dir, 't10k-images-idx3-ubyte')", False, 'import os\n'), ((29, 14, 29, 61), 'os.path.join', 'os.path.join', ({(29, 27, 29, 35): 'data_dir', (29, 36, 29, 60): '"""t10k-labels-idx1-ubyte"""'}, {}), "(data_dir, 't10k-labels-idx1-ubyte')", False, 'import os\n'), ((57, 14, 57, 29), 'numpy.asarray', 'np.asarray', ({(57, 25, 57, 28): 'trY'}, {}), '(trY)', True, 'import numpy as np\n'), ((58, 14, 58, 29), 'numpy.asarray', 'np.asarray', ({(58, 25, 58, 28): 'teY'}, {}), '(teY)', True, 'import numpy as np\n'), ((15, 11, 15, 27), 'numpy.array', 'np.array', ({(15, 20, 15, 26): 'ntrain'}, {}), '(ntrain)', True, 'import numpy as np\n')]