#!/usr/bin/env python3

import numpy as np
import h5py
from theano import tensor
from blocks.bricks import Linear, Rectifier, Tanh, Logistic, Softmax, NDimensionalSoftmax, WithExtraDims, application, \
    Brick
from blocks.bricks.cost import CategoricalCrossEntropy, MisclassificationRate, Cost
from blocks.bricks import WEIGHT
from blocks.graph import ComputationGraph
from blocks.filter import VariableFilter
from blocks.initialization import IsotropicGaussian, Constant
from blocks.algorithms import GradientDescent, Scale, RMSProp, AdaGrad
from blocks.extensions.monitoring import DataStreamMonitoring
from blocks.main_loop import MainLoop
from blocks.extensions import FinishAfter, Printing, ProgressBar
from fuel.streams import DataStream
from fuel.schemes import SequentialScheme, ShuffledScheme
from fuel.transformers import Flatten
from fuel.datasets.hdf5 import H5PYDataset

"""
This code shows how to train a sequence to sequence translation model.
"""

# the text data is composed of questions and answers
text_data = [
    ('What color is the sky', 'It is blue.'),
    ('What color is the sun', 'The sun is yellow.'),
    ('What color is the moon', 'It is blue.'),
]


def get_words(sentence):
    for c in '?!.,':
        sentence = sentence.replace(c, ' ').replace('  ', ' ')

    return sentence.upper().split()


def normalise(text_data):
    td = []
    for q, a in text_data:
        q_words = get_words(q)
        a_words = get_words(a)

        td.append((q_words, a_words))

    return td

def get_word_index(word_list):
    return dict([(w, i) for i, w in enumerate(word_list)])

def get_word_list(text_data):
    words = set(['_OOV_', '_SOS_', '_EOS_'])

    for q, a in text_data:
        words.update(q)
        words.update(a)

    word_list = sorted(words)

    word_list_index = get_word_index(word_list)

    return word_list, word_list_index


def transform_sentence(sentence, word_list_index, max_length):
    s = [word_list_index['_SOS_']]
    for w in sentence:
        s.append(word_list_index.get(w, word_list_index['_OOV_']))

    for w in range(0, max_length - len(s)):
        s.append(word_list_index['_EOS_'])

    return s


def transform_data(text_data, word_list_index, max_length):
    id = []
    for q, a in text_data:
        qi = transform_sentence(q, word_list_index, max_length)
        ai = transform_sentence(a, word_list_index, max_length)

        id.append((qi, ai))

    return id


def split_q_and_a(index_data):
    qd = []
    ad = []

    for q, a in index_data:
        qd.append(q)
        ad.append(a)

    return qd, ad


def prepare_dataset(text_data):
    text_data = normalise(text_data)
    word_list, word_list_index = get_word_list(text_data)

    # print(word_list)
    # print(word_list_index)

    max_length = 0
    for q, a in text_data:
        max_length = max(max_length, len(q), len(a))

    # print(max_length)

    # add two for _SOS_ and _EOS_ word symbols
    index_data = transform_data(text_data, word_list_index, max_length + 2)

    print(index_data)

    train_features, train_targets = split_q_and_a(index_data)

    train_features = np.asarray(train_features, dtype=np.int32)
    train_targets = np.asarray(train_targets, dtype=np.int32)

    test_features = train_features
    test_targets = train_targets

    f = h5py.File('data/dataset.hdf5', mode='w')

    dt = h5py.special_dtype(vlen=np.unicode)
    words = f.create_dataset(
            "word_list",
            (
                len(word_list),
            ),
            dtype=dt)
    words.dims[0].label = 'words'
    for i, w in enumerate(word_list):
        words[i] = w

    features = f.create_dataset(
             'features',
             (
                 train_features.shape[0] + test_features.shape[0],
                 train_features.shape[1],
             ),
             dtype='int32'
     )
    features.dims[0].label = 'example'
    features.dims[1].label = 'index'

    features[...] = np.vstack(
             [train_features, test_features]
     )

    targets = f.create_dataset(
            'targets',
            (
                train_targets.shape[0] + test_targets.shape[0],
                train_targets.shape[1],
            ),
            dtype='uint8'
    )
    targets.dims[0].label = 'example'
    targets.dims[1].label = 'index'

    targets[...] = np.vstack(
            [train_targets, test_targets]
    )

    split_dict = {
        'word_list': {
            'word_list': (0, words.shape[0])
        },
        'train': {
            'features': (0, train_features.shape[0]),
            'targets':  (0, train_targets.shape[0])
        },
        'test':  {
            'features': (train_features.shape[0], train_features.shape[0] + test_features.shape[0]),
            'targets':  (train_targets.shape[0], train_targets.shape[0] + test_targets.shape[0])
        }
    }
    f.attrs['split'] = H5PYDataset.create_split_array(split_dict)

    f.flush()
    f.close()

    train_set = H5PYDataset('data/dataset.hdf5', which_sets=('train',), load_in_memory=True)
    test_set = H5PYDataset('data/dataset.hdf5', which_sets=('test',), load_in_memory=True)
    word_list = h5py.File('data/dataset.hdf5')['word_list']
    word_list_index = get_word_index(word_list)

    return train_set, test_set, word_list, word_list_index


def train(train_set, test_set, l2_weight=1e-18):
    x = tensor.matrix('features')
    y = tensor.lmatrix('targets')

    n_classifiers = 3
    n_classes = 2
    l1 = Linear(
            name='l1',
            input_dim=2,
            output_dim=10,
            weights_init=IsotropicGaussian(0.1),
            biases_init=Constant(0)
    )
    l1.initialize()
    h1 = Logistic().apply(l1.apply(x))

    l2 = Linear(
            name='l1',
            input_dim=l1.output_dim,
            output_dim=n_classes * n_classifiers,
            weights_init=IsotropicGaussian(0.1),
            biases_init=Constant(0)
    )
    l2.initialize()
    l2 = l2.apply(h1)
    y_hat = MultiTargetSoftmax().apply(l2, n_classes, n_classifiers)

    cost = MultiTargetCategoricalCrossEntropy().apply(y, y_hat)
    error = MisclassificationRate().apply(y, y_hat)
    error.name = 'misclassification_rate'

    cg = ComputationGraph(cost)

    for w in VariableFilter(roles=[WEIGHT])(cg.variables):
        cost += l2_weight * (w ** 2).sum()
    cost.name = 'cost_with_regularization'

    # print('W1', W1.get_value())
    # print('W2', W2.get_value())

    algorithm = GradientDescent(
            cost=cost,
            parameters=cg.parameters,
            step_rule=RMSProp()
    )

    data_stream_train = Flatten(
            DataStream.default_stream(
                    train_set,
                    iteration_scheme=ShuffledScheme(train_set.num_examples, batch_size=80)
            )
    )

    data_stream_test = Flatten(
            DataStream.default_stream(
                    test_set,
                    iteration_scheme=SequentialScheme(test_set.num_examples, batch_size=1)
            )
    )

    monitor = DataStreamMonitoring(
            variables=[cost, error],
            data_stream=data_stream_test,
            prefix="test"
    )

    main_loop = MainLoop(
            data_stream=data_stream_train,
            algorithm=algorithm,
            extensions=[
                monitor,
                FinishAfter(after_n_epochs=100),
                Printing(),
                # ProgressBar()
            ]
    )

    main_loop.run()

    return x, y_hat


def test(x, y_hat):
    features_test = np.array(
            [
                [0, 0],
                [0, 1],
                [1, 0],
                [1, 1],
            ],
            dtype=np.float32
    )
    targets_test = np.array(
            [
                [0, 0, 0],
                [1, 0, 1],
                [1, 0, 1],
                [1, 1, 0],
            ],
            dtype=np.uint8
    )

    y_hat_eval = y_hat.eval({x: features_test})
    y_hat_eval_argmax = y_hat_eval.argmax(axis=1)

    print('Testing example')
    print('-' * 80)
    print('Features')
    print(features_test)
    print()
    print('Targets')
    print(targets_test)
    print()
    print('Probability predictions')
    print(y_hat_eval)
    print()
    print('ArgMax Predictions')
    print(y_hat_eval_argmax)


if __name__ == '__main__':
    train_set, test_set, word_list, word_list_index = prepare_dataset(text_data)

    # x, y_hat = train(train_set, test_set)
    #
    # test(x, y_hat)
