# Network
class HierarchicalConvnet(gluon.HybridBlock):
    def __init__(self, embedding_len, embedding_dim, num_channels, **kwargs):
        super(HierarchicalConvnet, self).__init__(**kwargs)

        with self.name_scope():

            self.embedder = gluon.nn.Embedding(embedding_len, embedding_dim)

            self.net_1 = gluon.nn.Conv1D(channels=num_channels, kernel_size=3, strides=1, padding=1,
                                         in_channels=embedding_dim, use_bias=True, activation='relu')

            self.net_2 = gluon.nn.Conv1D(channels=num_channels, kernel_size=3, strides=1, padding=1,
                                         in_channels=num_channels, use_bias=True, activation='relu')

            self.net_3 = gluon.nn.Conv1D(channels=num_channels, kernel_size=3, strides=1, padding=1,
                                         in_channels=num_channels, use_bias=True, activation='relu')

            self.net_4 = gluon.nn.Conv1D(channels=num_channels, kernel_size=3, strides=1, padding=1,
                                         in_channels=num_channels, use_bias=True, activation='relu')

    def embedding_initialize(self, glove_embedding):
        self.embedder.weight.set_data(glove_embedding.idx_to_vec)

    def hybrid_forward(self, F, sent):
        sent_embed = self.embedder(sent)
        sent = F.swapaxes(sent_embed, 1, 2)
        sent = self.net_1(sent)
        pooled_1 = F.max(sent, axis=2)  # gives dimension NC
        sent = self.net_2(sent)
        pooled_2 = F.max(sent, axis=2)
        sent = self.net_3(sent)
        pooled_3 = F.max(sent, axis=2)
        sent = self.net_4(sent)
        pooled_4 = F.max(sent, axis=2)
        return F.concat(pooled_1, pooled_2, pooled_3, pooled_4, dim=1)  # dimension is N X 4C


def build_mlp_classifier():
    net = gluon.nn.HybridSequential(prefix='mlp_')
    num_fc_1 = 512
    num_outputs = 3
    with net.name_scope():
        net.add(gluon.nn.Dense(num_fc_1, use_bias=True, activation='tanh'))
        net.add(gluon.nn.Dense(num_fc_1, use_bias=True, activation='tanh'))
        net.add(gluon.nn.Dense(num_outputs))
    return net


class SiameseNet(gluon.HybridBlock):
    def __init__(self, glove_embedding_len, glove_embedding_dim, num_channels, **kwargs):
        super(SiameseNet, self).__init__(**kwargs)
        with self.name_scope():
            self.encoder = HierarchicalConvnet(glove_embedding_len, glove_embedding_dim, num_channels, prefix='hcl_')
            self.classifier = build_mlp_classifier()

    def hybrid_forward(self, F, sent_1, sent_2):
        e_1 = self.encoder(sent_1)
        e_2 = self.encoder(sent_2)
        features = F.concat(e_1, e_2, F.abs(e_1 - e_2), (e_1 * e_2), dim=1)
        output = self.classifier(features)
        return output

    def encode(self, sent_1):
        emb = self.encoder(sent_1)
        return emb

#Training code Blob
ctx = [mx.gpu(i + params.gpu_offset) for i in range(params.num_gpu)]
# Other parameters from argparse......
kv = mx.kv.create('device')

data_iter = get_data_iter()  # dummy function, can provide actual implementation if needed

pretrained_embedding = text.embedding.create('glove', pretrained_file_name='glove.840B.300d.txt', vocabulary=src_vocab)

net = SiameseNet(len(pretrained_embedding), pretrained_embedding.vec_len, params.num_channels)
net.hybridize()
net.collect_params().initialize(mx.init.Xavier(), ctx=ctx)
net.encoder.embedding_initialize(pretrained_embedding)

trainer = gluon.Trainer(net.collect_params(), 'adam', {'learning_rate': learning_rate}, kvstore=kv)

epochs = params.epochs
softmax_loss = gluon.loss.SoftmaxCrossEntropyLoss()

for epoch in range(epochs):
        for i, batch in enumerate(train_iter):
            total_loss_sum = 0.0
            data_1 = gluon.utils.split_and_load(batch.data[0], ctx_list=ctx, batch_axis=0)
            data_2 = gluon.utils.split_and_load(batch.data[1], ctx_list=ctx, batch_axis=0)
            label = gluon.utils.split_and_load(batch.label[0], ctx_list=ctx, batch_axis=0)
            Ls = []
            outputs = []
            for data_x, data_y, data_z in zip(data_1, data_2, label):
                with autograd.record():
                    output = net(data_x, data_y)
                    loss = softmax_loss(output, data_z)
                    Ls.append(loss)
                    outputs.append(output)
                total_loss_sum += nd.sum(loss).asscalar()

            for L in Ls:
                L.backward()
            trainer.step(mini_batch_size)
            metric.update(label, outputs)

    name, val_acc = evaluate_accuracy(ctx, valid_iter, net)