import tensorlayer as T
from dragon.vm.tensorlayer.layers import Dense
from dragon.vm.tensorlayer.models import Model
import dragon.vm.tensorlayer as tl

import dragon as dg
import argparse
import numpy as np

X_train, y_train, X_val, y_val, X_test, y_test = T.files.load_mnist_dataset(shape=(-1, 784))


class MLP(Model):

    def __init__(self):
        super(MLP, self).__init__()
        self.dense1 = Dense(n_units=800, act=tl.act.relu, in_channels=784)
        self.dense2 = Dense(n_units=800, act=tl.act.relu, in_channels=800)
        self.dense3 = Dense(n_units=10, act=tl.act.relu, in_channels=800)

    def forward(self, x):
        z = self.dense1(x)
        z = self.dense2(z)
        out = self.dense3(z)
        return out


class Classifier(object):
    """The base classifier class."""

    # TensorSpec for graph execution
    image_spec = dg.Tensor([None, 3, 32, 32], 'float32')
    label_spec = dg.Tensor([None], 'int64')

    def __init__(self, optimizer):
        super(Classifier, self).__init__()
        self.net = MLP()
        self.optimizer = optimizer
        self.params = self.net.trainable_weights

    def step(self, image, label):
        with dg.GradientTape() as tape:
            logit = self.net(image)
            # logit = dg.cast(logit, 'float64')
            logit = dg.cast(dg.math.argmax(logit, -1), 'int32')
            #             label = dg.cast(label, 'float32')
            # print("logit :\n", logit, label)
            # loss = dg.losses.smooth_l1_loss([logit, label])
            #             loss = tl.losses.sparse_softmax_crossentropy(logit, label)
            loss = dg.math.sum(
                (logit - label) * (logit - label)
            )  # dg.losses.sparse_softmax_cross_entropy([logit, label])
        accuracy = dg.math.mean(dg.math.equal([logit, label]).astype('float32'))
        grads = tape.gradient(loss, self.params)
        self.optimizer.apply_gradients(zip(self.params, grads))
        return loss, accuracy, self.optimizer


if __name__ == '__main__':
    dg.autograph.set_execution('EAGER_MODE')
    # Define the model
    model = Classifier(dg.optimizers.SGD(base_lr=0.001, momentum=0.9, weight_decay=1e-4))
    # Main loop
    batch_size = 200
    for i in range(50):
        for X_batch, y_batch in T.iterate.minibatches(X_train, y_train, batch_size, shuffle=True):
            image = dg.EagerTensor(X_batch, copy=False)
            label = dg.EagerTensor(y_batch, copy=False, dtype='float32')
            loss, accuracy, _ = model.step(image, label)
            if i % 20 == 0:
                dg.logging.info(
                    'Iteration %d, lr = %s, loss = %.5f, accuracy = %.3f' %
                    (i, str(model.optimizer.base_lr), loss, accuracy)
                )