import theano
import numpy
from theano import tensor
from blocks.model import Model
from blocks.bricks import Linear, Tanh
from ctc_cost import CTC
from blocks.initialization import IsotropicGaussian, Constant
from fuel.datasets import IterableDataset
from fuel.streams import DataStream
from blocks.algorithms import (GradientDescent, Scale,
                               StepClipping, CompositeRule)
from blocks.extensions.monitoring import TrainingDataMonitoring
from blocks.main_loop import MainLoop
from blocks.extensions import FinishAfter, Printing
from blocks.bricks.recurrent import SimpleRecurrent
from blocks.graph import ComputationGraph
try:
    import cPickle as pickle
except:
    import pickle

floatX = theano.config.floatX


@theano.compile.ops.as_op(itypes=[tensor.lvector],
                          otypes=[tensor.lvector])
def print_pred(y_hat):
    blank_symbol = 4
    res = []
    for i, s in enumerate(y_hat):
        if (s != blank_symbol) and (i == 0 or s != y_hat[i - 1]):
            res += [s]
    return numpy.asarray(res)

n_epochs = 200
x_dim = 4
h_dim = 9
num_classes = 4

with open("ctc_test_data.pkl", "rb") as pkl_file:
    try:
        data = pickle.load(pkl_file)
        inputs = data['inputs']
        labels = data['labels']
        # from S x T x B x D to S x T x B
        inputs_mask = numpy.max(data['mask_inputs'], axis=-1)
        labels_mask = data['mask_labels']
    except:
        data = pickle.load(pkl_file, encoding='bytes')
        inputs = data[b'inputs']
        labels = data[b'labels']
        # from S x T x B x D to S x T x B
        inputs_mask = numpy.max(data[b'mask_inputs'], axis=-1)
        labels_mask = data[b'mask_labels']



print('Building model ...')
# T x B x F
x = tensor.tensor3('x', dtype=floatX)
# T x B
x_mask = tensor.matrix('x_mask', dtype=floatX)
# L x B
y = tensor.matrix('y', dtype=floatX)
# L x B
y_mask = tensor.matrix('y_mask', dtype=floatX)

x_to_h = Linear(name='x_to_h',
                input_dim=x_dim,
                output_dim=h_dim)
x_transform = x_to_h.apply(x)
rnn = SimpleRecurrent(activation=Tanh(),
                      dim=h_dim, name="rnn")
h = rnn.apply(x_transform)
h_to_o = Linear(name='h_to_o',
                input_dim=h_dim,
                output_dim=num_classes + 1)
h_transform = h_to_o.apply(h)
# T x B x C+1
y_hat = tensor.nnet.softmax(
    h_transform.reshape((-1, num_classes + 1))
).reshape((h.shape[0], h.shape[1], -1))
y_hat.name = 'y_hat'

y_hat_mask = x_mask
cost = CTC().apply(y, y_hat, y_mask, y_hat_mask, 'normal_scale')
cost.name = 'CTC'
# Initialization
for brick in (rnn, x_to_h, h_to_o):
    brick.weights_init = IsotropicGaussian(0.01)
    brick.biases_init = Constant(0)
    brick.initialize()

print('Bulding DataStream ...')
dataset = IterableDataset({'x': inputs,
                           'x_mask': inputs_mask,
                           'y': labels,
                           'y_mask': labels_mask})
stream = DataStream(dataset)

print('Bulding training process...')
algorithm = GradientDescent(cost=cost,
                            parameters=ComputationGraph(cost).parameters,
                            step_rule=CompositeRule([StepClipping(10.0),
                                                     Scale(0.02)]))
monitor_cost = TrainingDataMonitoring([cost],
                                      prefix="train",
                                      after_epoch=True)

# sample number to monitor
sample = 8

y_hat_max_path = print_pred(tensor.argmax(y_hat[:, sample, :], axis=1))
y_hat_max_path.name = 'Viterbi'
monitor_output = TrainingDataMonitoring([y_hat_max_path],
                                        prefix="y_hat",
                                        every_n_epochs=1)

length = tensor.sum(y_mask[:, sample]).astype('int32')
tar = y[:length, sample].astype('int32')
tar.name = '_Target_Seq'
monitor_target = TrainingDataMonitoring([tar],
                                        prefix="y",
                                        every_n_epochs=1)

model = Model(cost)
main_loop = MainLoop(data_stream=stream, algorithm=algorithm,
                     extensions=[monitor_cost, monitor_output,
                                 monitor_target,
                                 FinishAfter(after_n_epochs=n_epochs),
                                 Printing()],
                     model=model)

print('Starting training ...')
main_loop.run()