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', (['datasets_dir', '"""mnist/"""'], {}), False, 'import os\n'), (18, 'numpy.fromfile', 'np.fromfile', ([], {'file': 'fd', 'dtype': 'np.uint8'}), True, 'import numpy as np\n'), (22, 'numpy.fromfile', 'np.fromfile', ([], {'file': 'fd', 'dtype': 'np.uint8'}), True, 'import numpy as np\n'), (26, 'numpy.fromfile', 'np.fromfile', ([], {'file': 'fd', 'dtype': 'np.uint8'}), True, 'import numpy as np\n'), (30, 'numpy.fromfile', 'np.fromfile', ([], {'file': 'fd', 'dtype': 'np.uint8'}), True, 'import numpy as np\n'), (36, 'random.shuffle', 'shuffle', (['index_shuf'], {}), False, 'from random import shuffle\n'), (41, 'numpy.asarray', 'np.asarray', (['trX_shuffle'], {}), True, 'import numpy as np\n'), (42, 'numpy.asarray', 'np.asarray', (['trY_shuffle'], {}), True, 'import numpy as np\n'), (8, 'numpy.array', 'np.array', (['x'], {}), True, 'import numpy as np\n'), (17, 'os.path.join', 'os.path.join', (['data_dir', '"""train-images-idx3-ubyte"""'], {}), False, 'import os\n'), (21, 'os.path.join', 'os.path.join', (['data_dir', '"""train-labels-idx1-ubyte"""'], {}), False, 'import os\n'), (25, 'os.path.join', 'os.path.join', (['data_dir', '"""t10k-images-idx3-ubyte"""'], {}), False, 'import os\n'), (29, 'os.path.join', 'os.path.join', (['data_dir', '"""t10k-labels-idx1-ubyte"""'], {}), False, 'import os\n'), (57, 'numpy.asarray', 'np.asarray', (['trY'], {}), True, 'import numpy as np\n'), (58, 'numpy.asarray', 'np.asarray', (['teY'], {}), True, 'import numpy as np\n'), (15, 'numpy.array', 'np.array', (['ntrain'], {}), True, 'import numpy as np\n')]