# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

import mxnet as mx
from mxnet import autograd as ag, ndarray as nd, gluon
from mxnet.gluon import Block, nn, rnn
import mxnet.optimizer as optim
import sys

# This example demonstrates how LSTM-CRF model can be implemented in Gluon to perform
# noun-phrase chunking as a sequence labeling task.

mx.random.seed(1)

# Helper functions to make the code more readable.
def to_scalar(x):
    return int(x.asscalar())

def argmax(vec):
    # return the argmax as a python int
    idx = nd.argmax(vec, axis=1)
    return to_scalar(idx)

def prepare_sequence(seq, word2idx):
    return nd.array([word2idx[w] for w in seq])

# Compute log sum exp is numerically more stable than multiplying probabilities
def log_sum_exp(vec):
    max_score = nd.max(vec).asscalar()
    return nd.log(nd.sum(nd.exp(vec - max_score))) + max_score

# Model
class BiLSTM_CRF(Block):
    def __init__(self, vocab_size, tag2idx, embedding_dim, hidden_dim):
        super(BiLSTM_CRF, self).__init__()
        with self.name_scope():
            self.embedding_dim = embedding_dim
            self.hidden_dim = hidden_dim
            self.vocab_size = vocab_size
            self.tag2idx = tag2idx
            self.tagset_size = len(tag2idx)

            self.word_embeds = nn.Embedding(vocab_size, embedding_dim)
            self.lstm = rnn.LSTM(hidden_dim // 2, num_layers=1, bidirectional=True)

            # Maps the output of the LSTM into tag space.
            self.hidden2tag = nn.Dense(self.tagset_size)

            # Matrix of transition parameters.  Entry i,j is the score of
            # transitioning *to* i *from* j.
            self.transitions = nd.random.normal(shape=(self.tagset_size, self.tagset_size))

            self.hidden = self.init_hidden()

    def init_hidden(self):
        return [nd.random.normal(shape=(2, 1, self.hidden_dim // 2)),
                nd.random.normal(shape=(2, 1, self.hidden_dim // 2))]

    def _forward_alg(self, feats):
        # Do the forward algorithm to compute the partition function
        alphas = [[-10000.] * self.tagset_size]
        alphas[0][self.tag2idx[START_TAG]] = 0.
        alphas = nd.array(alphas)

        # Iterate through the sentence
        for feat in feats:
            alphas_t = []  # The forward variables at this timestep
            for next_tag in range(self.tagset_size):
                # broadcast the emission score: it is the same regardless of
                # the previous tag
                emit_score = feat[next_tag].reshape((1, -1))
                # the ith entry of trans_score is the score of transitioning to
                # next_tag from i
                trans_score = self.transitions[next_tag].reshape((1, -1))
                # The ith entry of next_tag_var is the value for the
                # edge (i -> next_tag) before we do log-sum-exp
                next_tag_var = alphas + trans_score + emit_score
                # The forward variable for this tag is log-sum-exp of all the
                # scores.
                alphas_t.append(log_sum_exp(next_tag_var))
            alphas = nd.concat(*alphas_t, dim=0).reshape((1, -1))
        terminal_var = alphas + self.transitions[self.tag2idx[STOP_TAG]]
        alpha = log_sum_exp(terminal_var)
        return alpha

    def _get_lstm_features(self, sentence):
        self.hidden = self.init_hidden()
        length = sentence.shape[0]
        embeds = self.word_embeds(sentence).reshape((length, 1, -1))
        lstm_out, self.hidden = self.lstm(embeds, self.hidden)
        lstm_out = lstm_out.reshape((length, self.hidden_dim))
        lstm_feats = self.hidden2tag(lstm_out)
        return nd.split(lstm_feats, num_outputs=length, axis=0, squeeze_axis=True)

    def _score_sentence(self, feats, tags):
        # Gives the score of a provided tag sequence
        score = nd.array([0])
        tags = nd.concat(nd.array([self.tag2idx[START_TAG]]), *tags, dim=0)
        for i, feat in enumerate(feats):
            score = score + \
                self.transitions[to_scalar(tags[i+1]), to_scalar(tags[i])] + feat[to_scalar(tags[i+1])]
        score = score + self.transitions[self.tag2idx[STOP_TAG],
                                         to_scalar(tags[int(tags.shape[0]-1)])]
        return score

    def _viterbi_decode(self, feats):
        backpointers = []

        # Initialize the viterbi variables in log space
        vvars = nd.full((1, self.tagset_size), -10000.)
        vvars[0, self.tag2idx[START_TAG]] = 0

        for feat in feats:
            bptrs_t = []  # holds the backpointers for this step
            viterbivars_t = []  # holds the viterbi variables for this step

            for next_tag in range(self.tagset_size):
                # next_tag_var[i] holds the viterbi variable for tag i at the
                # previous step, plus the score of transitioning
                # from tag i to next_tag.
                # We don't include the emission scores here because the max
                # does not depend on them (we add them in below)
                next_tag_var = vvars + self.transitions[next_tag]
                best_tag_id = argmax(next_tag_var)
                bptrs_t.append(best_tag_id)
                viterbivars_t.append(next_tag_var[0, best_tag_id])
            # Now add in the emission scores, and assign vvars to the set
            # of viterbi variables we just computed
            vvars = (nd.concat(*viterbivars_t, dim=0) + feat).reshape((1, -1))
            backpointers.append(bptrs_t)

        # Transition to STOP_TAG
        terminal_var = vvars + self.transitions[self.tag2idx[STOP_TAG]]
        best_tag_id = argmax(terminal_var)
        path_score = terminal_var[0, best_tag_id]

        # Follow the back pointers to decode the best path.
        best_path = [best_tag_id]
        for bptrs_t in reversed(backpointers):
            best_tag_id = bptrs_t[best_tag_id]
            best_path.append(best_tag_id)
        # Pop off the start tag (we dont want to return that to the caller)
        start = best_path.pop()
        assert start == self.tag2idx[START_TAG]  # Sanity check
        best_path.reverse()
        return path_score, best_path

    def neg_log_likelihood(self, sentence, tags):
        feats = self._get_lstm_features(sentence)
        forward_score = self._forward_alg(feats)
        gold_score = self._score_sentence(feats, tags)
        return forward_score - gold_score

    def forward(self, sentence):  # dont confuse this with _forward_alg above.
        # Get the emission scores from the BiLSTM
        lstm_feats = self._get_lstm_features(sentence)

        # Find the best path, given the features.
        score, tag_seq = self._viterbi_decode(lstm_feats)
        return score, tag_seq

# Run training
START_TAG = "<START>"
STOP_TAG = "<STOP>"
EMBEDDING_DIM = 5
HIDDEN_DIM = 4

# Make up some training data
training_data = [(
    "the wall street journal reported today that apple corporation made money".split(),
    "B I I I O O O B I O O".split()
), (
    "georgia tech is a university in georgia".split(),
    "B I O O O O B".split()
)]

word2idx = {}
for sentence, tags in training_data:
    for word in sentence:
        if word not in word2idx:
            word2idx[word] = len(word2idx)

tag2idx = {"B": 0, "I": 1, "O": 2, START_TAG: 3, STOP_TAG: 4}

model = BiLSTM_CRF(len(word2idx), tag2idx, EMBEDDING_DIM, HIDDEN_DIM)
model.collect_params().initialize(mx.init.Xavier(magnitude=2.24), ctx=mx.cpu())
optimizer = gluon.Trainer(model.collect_params(), 'sgd', {'learning_rate': 0.01, 'wd': 1e-4})

# Check predictions before training
precheck_sent = prepare_sequence(training_data[0][0], word2idx)
precheck_tags = nd.array([tag2idx[t] for t in training_data[0][1]])
print(model(precheck_sent))

# Make sure prepare_sequence from earlier in the LSTM section is loaded
for epoch in range(300):  # again, normally you would NOT do 300 epochs, it is toy data
    for sentence, tags in training_data:
        # Step 1. Get our inputs ready for the network, that is,
        # turn them into Variables of word indices.
        # Remember to use autograd to record the calculation.
        with ag.record():
            sentence_in = prepare_sequence(sentence, word2idx)
            targets = nd.array([tag2idx[t] for t in tags])

            # Step 2. Run our forward pass.
            neg_log_likelihood = model.neg_log_likelihood(sentence_in, targets)

            # Step 3. Compute the loss, gradients, and update the parameters by
            # calling optimizer.step()
            neg_log_likelihood.backward()
        optimizer.step(1)

# Check predictions after training
precheck_sent = prepare_sequence(training_data[0][0], word2idx)
print(model(precheck_sent))

# Acknowledgement: this example is adopted from pytorch nlp tutorials.
