from __future__ import print_function
import plac

import tqdm
import numpy as np
from thinc.extra import datasets
from thinc.neural.base import Network
from thinc.neural.vec2vec import Affine, ReLu, Softmax
from thinc.neural.util import score_model
from thinc.neural.optimizers import linear_decay


class ReLuMLP(Network):
    Hidden = ReLu
    Output = Softmax
    width = 128
    depth = 3

    def setup(self, nr_out, nr_in, **kwargs):
        for i in range(self.depth):
            self.layers.append(self.Hidden(nr_out=self.width, nr_in=nr_in,
                name='hidden-%d' % i))
            nr_in = self.width
        self.layers.append(self.Output(nr_out=nr_out, nr_in=nr_in))
        self.set_weights(initialize=True)
        self.set_gradient()


def get_gradient(scores, labels):
    target = np.zeros(scores.shape)
    for i, label in enumerate(labels):
        target[i, int(label)] = 1.0
    return scores - target


def main(batch_size=128, nb_epoch=10, nb_classes=10):
    model = ReLuMLP(10, 784)
    train_data, check_data, test_data = datasets.mnist()
    
    with model.begin_training(train_data) as (trainer, optimizer):
        for examples, truth in trainer.iterate(model, train_data, check_data,
                                               nb_epoch=nb_epoch):
            examples = model.ops.asarray(examples)
            truth = model.ops.asarray(truth, dtype='i')
            guess, finish_update = model.begin_update(examples, dropout=0.3)
            gradient, loss = trainer.get_gradient(guess, truth)
            optimizer.set_loss(loss)
            finish_update(gradient, optimizer)
    print('Test: %.3f' % score_model(model, test_data))


if __name__ == '__main__':
    if 1:
        plac.call(main)
    else:
        import cProfile
        import pstats
        cProfile.runctx("plac.call(main)", globals(), locals(), "Profile.prof")
        s = pstats.Stats("Profile.prof")
        s.strip_dirs().sort_stats("time").print_stats()


from __future__ import print_function
import plac
from thinc.neural.vec2vec import ReLu, Softmax

from thinc.neural.toolz import pipe, clone
from thinc.loss import categorical_crossentropy
from thinc.optimizers import Adam
from thinc.extra import datasets
from thinc.util import score_model 


def main(depth=2, width=512, nb_epoch=10):
    # Input and output dimensions defined by data
    with Model.operators({'*': clone, '+' pipe}):
        model = ReLu(width) * depth + Softmax()
    
    train, dev = datasets.load_mnist()

    model.train(train, dev)

    (train_X, train_Y), (dev_X, dev_Y) = datasets.load_mnist()
    
    optimizer = Adam(0.001)
    with model.begin_training(train_X, train_Y) as trainer:
        for i in range(nb_epoch):
            for batch_X, batch_Y in trainer.iterate(train_X, train_Y):
                guess, finish_update = model.begin_update(examples, dropout=0.3)
                gradient, loss = categorical_crossentropy(guess, batch_Y)
                finish_update(gradient, optimizer)
            print(i, score_model(model, dev_X, dev_Y))
    with open('out.pickle', 'wb') as file_:
        pickle.dump(model, file_, -1)

Revise mnist example in anticipation of new API