import numpy as np
import chainer
from chainer import cuda, Function, gradient_check, Variable, optimizers, serializers, utils
from chainer import Link, Chain, ChainList
import chainer.functions as F
import chainer.links as L

from sklearn.datasets import fetch_mldata

mnist = fetch_mldata('MNIST original', data_home='../')

x_train = mnist.data[:60000]
y_train = mnist.target[:60000]

x_test  = mnist.data[60000:]
y_test  = mnist.target[60000:]


l2_unit = 1000

class MyChain(Chain):
    def __init__(self):
        super(MyChain, self).__init__(
            l1 = L.Linear(784, l2_unit),
            l2 = L.Linear(l2_unit, l2_unit),
            l3 = L.Linear(l2_unit, 10)
        )

    def __call__(self, x):
        h1 = F.relu(self.l1(x))
        h2 = F.relu(self.l2(h1))
        return self.l3(h2)

model = L.Classifier(MyChain())
optimizer = optimizers.Adam()
optimizer.setup(model)

train_size = 60000
test_size = 10000
batchsize = 1000

for epoch in range(1, 11):
    print('epoch:', epoch)
    indexes = np.random.permutation(train_size)
    sum_loss = 0
    sum_accuracy = 0
    sum_len = 0
    for i in range(0, train_size, batchsize):
        x = Variable(np.array(x_train[indexes[i : i + batchsize]], dtype=np.float32))
        t = Variable(np.array(y_train[indexes[i : i + batchsize]], dtype=np.int32))

        optimizer.update(model, x, t)

        sum_loss += float(model.loss.data) * len(t.data)
        sum_accuracy += float(model.accuracy.data) * len(t.data)

    print('loss     =', sum_loss / train_size)
    print('accuracy =', sum_accuracy / train_size)

    indexes = np.random.permutation(test_size)

    sum_loss = 0
    sum_accuracy = 0
    sum_len = 0
    for i in range(0, test_size, batchsize):
        x = Variable(np.array(x_test[indexes[i : i + batchsize]], dtype=np.float32))
        t = Variable(np.array(y_test[indexes[i : i + batchsize]], dtype=np.int32))

        sum_loss += float(model(x, t).data) * len(t.data)
        sum_accuracy += float(model.accuracy.data) * len(t.data)
        sum_len += len(t.data)

    print('test loss     =', sum_loss / test_size)
    print('test accuracy =', sum_accuracy / test_size)
    assert(sum_len == test_size)