## @package crf
# Module caffe2.python.crf


import numpy as np
from caffe2.python import brew, core, model_helper, recurrent


"""
Due to a limitation in ReccurentNetworkOp, this layer only supports batch_size=1
In order to support batch_size > 1, we will have to implement the CRFUnit
and its gradient in C++ and handle the different batches there.
"""


class CRFWithLoss:
    def __init__(self, model, num_classes, transitions_blob=None):
        self.model = model
        self.num_classes = num_classes
        self.num_classes_padded = num_classes + 2  # After adding BOS and EOS
        if not transitions_blob:
            transitions_blob = self.model.param_init_net.UniformFill(
                [],
                [core.ScopedBlobReference("crf_transitions")],
                shape=[self.num_classes_padded, self.num_classes_padded],
                min=-1.0,
                max=1.0,
            )
        self.transitions = transitions_blob
        self.model.params.append(self.transitions)

    def crf_loss(self, predictions, labels, seq_lengths=None):
        # Since the transitions matrix is a shared parameter, need to
        # take a snapshot of it at the beginning since it can be updated
        # in between the operators that uses it when doing parallel updates
        transitions_snapshot = self.model.net.Copy(
            self.transitions, core.ScopedBlobReference("transitions_snapshot")
        )
        # Compute best path unary score from the logits
        path_unary_score = self._gather_entries_sum(
            predictions, labels, self.num_classes
        )
        # Append BOS and EOS entries to the predictions and labels
        predictions = CRFWithLoss.pad_predictions(
            predictions, self.model.param_init_net, self.model.net, self.num_classes
        )
        labels = CRFWithLoss.pad_labels(
            labels, self.model.param_init_net, self.model.net, self.num_classes
        )
        # Compute best path binary scores from the transitions matrix
        path_binary_score = self._path_binary_scores(
            labels, transitions_snapshot, seq_lengths
        )
        path_total_score = self.model.net.Add(
            [path_binary_score, path_unary_score],
            core.ScopedBlobReference("path_total"),
        )
        # Compute all paths score
        zero_index = self.model.param_init_net.ConstantFill([], shape=[1], value=0)
        initial_state = self.model.net.Gather(
            [predictions, zero_index],
            core.ScopedBlobReference("rnn_initial"),
            dense_gradient=True,
        )
        input_data, _ = self.model.net.RemovePadding(
            [predictions], padding_width=1, end_padding_width=0, outputs=2
        )
        input_data = self.model.net.ExpandDims(
            [input_data], core.ScopedBlobReference("rnn_input_data"), dims=[1]
        )
        # Due to a bug in RecurrentNetworkGradientOp, we need to copy the
        # transitions blob before sending it to the recurrent network
        transitions_copy = self.model.net.Copy(
            transitions_snapshot, core.ScopedBlobReference("transitions_copy")
        )
        all_paths_scores = self._crf_forward(
            input_data, initial_state, transitions_copy
        )
        loss = self.model.net.Sub(
            [all_paths_scores, path_total_score], core.ScopedBlobReference("crf_loss")
        )
        return loss

    def _path_binary_scores(self, labels, transitions, seq_lengths=None):
        column_ids, _ = self.model.net.RemovePadding(
            [labels], outputs=2, padding_width=1, end_padding_width=0
        )
        row_ids, _ = self.model.net.RemovePadding(
            [labels], outputs=2, padding_width=0, end_padding_width=1
        )
        # Since there is no multi-dimensional gather, I flatten the matrix to
        # a 1-d vector and transform the ids to (row_ids * num_columns +
        # column_ids) and do gather in 1-d
        num_columns_blob = self.model.net.ConstantFill(
            [row_ids], value=self.num_classes_padded
        )
        flattened_ids = self.model.net.Mul([row_ids, num_columns_blob])
        flattened_ids = self.model.net.Add([flattened_ids, column_ids])
        flattened_transitions = self.model.net.FlattenToVec([transitions])
        entries = self.model.net.Gather(
            [flattened_transitions, flattened_ids], dense_gradient=True
        )
        return self.model.ReduceFrontSum(entries)

    def _gather_entries_sum(self, in_data, indices, index_size):
        indices = self.model.net.Cast([indices], to="int64")
        index_size_blob = self.model.param_init_net.ConstantFill(
            [], shape=[1], value=index_size
        )
        query_one_hot = self.model.net.OneHot([indices, index_size_blob])
        flattend_query = self.model.net.FlattenToVec(query_one_hot)
        flattend_data = self.model.net.FlattenToVec(in_data)
        query_scores = self.model.net.DotProduct([flattend_query, flattend_data])
        final_sum = self.model.net.ReduceFrontSum([query_scores])
        return final_sum

    def _crf_forward(
        self, input_blob, initial_state, transitions_copy, seq_lengths=None
    ):
        # Build the RNN net and get the last timestep output
        out_last = self.build_crf_net(input_blob, initial_state, transitions_copy)
        out_last, _ = self.model.net.Reshape(
            [out_last], outputs=2, shape=(self.num_classes_padded,)
        )
        zero_segment_id = self.model.param_init_net.ConstantFill(
            [], value=0, shape=[self.num_classes_padded], dtype=core.DataType.INT32
        )

        # Compute the accumulated total score of all the paths
        accum_score = self.model.net.SortedSegmentRangeLogSumExp(
            [out_last, zero_segment_id]
        )
        accum_score, _ = self.model.net.Reshape(accum_score, outputs=2, shape=())
        return accum_score

    def build_crf_net(self, input_blob, initial_state, transitions):
        """
            Adds the crf_net recurrent operator to the model.

            model: model_helper.ModelHelper object new operators would be added
            to

            input_blob: the input sequence in a format T x N x D
            where T is sequence size, N - batch size and D - input dimension
            ##Only supports batch-size 1##

            seq_lengths: blob containing sequence lengths (unused)
            """

        scope = "crf_net"

        def s(name):
            ""
            # We have to manually scope due to our internal/external blob
            # relationships.
            return "{}/{}".format(str(scope), str(name))

        step_model = model_helper.ModelHelper(name="crf_step", param_model=self.model)
        input_t, cell_t_prev, _ = step_model.net.AddExternalInputs(
            core.ScopedBlobReference("input_t"),
            core.ScopedBlobReference("cell_t_prev"),
            transitions,
        )
        zero_segment_id = step_model.param_init_net.ConstantFill(
            [],
            [s("zero_segment_id")],
            value=0,
            shape=[self.num_classes_padded],
            dtype=core.DataType.INT32,
        )

        # A hack to bypass model cloning for test
        step_model.param_init_net.AddExternalOutput(zero_segment_id)
        """ the CRF step """
        # Do tile
        prev_transpose = brew.transpose(
            step_model, cell_t_prev, [s("prev_transpose")], axes=(0, 2, 1)
        )
        prev_tiled = step_model.net.Tile(
            prev_transpose, [s("prev_tiled")], tiles=self.num_classes_padded, axis=2
        )
        input_t_tiled = step_model.net.Tile(
            input_t, [s("input_t_tiled")], tiles=self.num_classes_padded, axis=1
        )
        input_with_prev = step_model.net.Add(
            [prev_tiled, input_t_tiled], [s("input_with_prev")]
        )
        all_with_transitions = step_model.net.Add(
            [input_with_prev, transitions],
            [s("prev_with_transitions")],
            broadcast=1,
            use_grad_hack=1,
        )
        all_with_transitions_reshaped, _ = step_model.net.Reshape(
            all_with_transitions,
            [s("all_with_transitions_reshaped"), s("all_with_transitions_orig")],
            shape=(self.num_classes_padded, self.num_classes_padded),
        )
        cell_t = step_model.net.SortedSegmentRangeLogSumExp(
            [all_with_transitions_reshaped, zero_segment_id], [s("cell_t")]
        )
        step_model.net.AddExternalOutputs(cell_t)
        """ recurrent network """
        cell_input_blob = initial_state
        out_all, out_last = recurrent.recurrent_net(
            net=self.model.net,
            cell_net=step_model.net,
            inputs=[(input_t, input_blob)],
            initial_cell_inputs=[(cell_t_prev, cell_input_blob)],
            links={cell_t_prev: cell_t},
            scope=scope,
            outputs_with_grads=(1,),
        )
        return out_last

    def update_predictions(self, classes):
        def crf_update_predictions_op(inputs, outputs):
            # This operator will compute the best path of classes by performing
            # Viterbi decoding and then updates the predictions to make the tag
            # On the best path has the highest score among the others
            predictions = inputs[0].data
            transitions = inputs[1].data
            predictions = inputs[0].data
            predictions_shape = inputs[0].shape
            outputs[0].reshape(predictions_shape)

            trellis = np.zeros(predictions_shape)
            backpointers = np.zeros(predictions_shape, dtype=np.int32)
            trellis[0] = predictions[0]

            for t in range(1, predictions_shape[0]):
                v = np.expand_dims(trellis[t - 1], 1) + transitions
                trellis[t] = predictions[t] + np.max(v, 0)
                backpointers[t] = np.argmax(v, 0)

            viterbi = [np.argmax(trellis[-1])]
            for bp in reversed(backpointers[1:]):
                viterbi.append(bp[viterbi[-1]])
            viterbi.reverse()

            new_predictions = np.zeros(predictions_shape)
            old_bests = []
            for i, w_predictions in enumerate(predictions):
                # Get the current tag with the maximum score
                new_predictions[i] = predictions[i]
                old_best = np.argmax(w_predictions)
                old_bests.append(old_best)
                # Swap the scores of the current best tag and the tag on the
                # Viterbi path
                w_predictions[viterbi[i]], w_predictions[old_best] = (
                    w_predictions[old_best],
                    w_predictions[viterbi[i]],
                )
                new_predictions[i] = w_predictions
            # Remove the BOS and EOS entries from the predictions matrix
            orig_predictions = new_predictions[1:-1, 0:-2]
            outputs[0].reshape(orig_predictions.shape)
            outputs[0].data[...] = orig_predictions

        padded_classes = CRFWithLoss.pad_predictions(
            classes, self.model.param_init_net, self.model.net, self.num_classes
        )
        new_classes = self.model.net.Python(crf_update_predictions_op)(
            [padded_classes, self.transitions],
            core.ScopedBlobReference("post_crf_classes"),
        )
        return new_classes

    @staticmethod
    def pad_labels(labels, init_net, net, num_classes):
        bos_i = num_classes
        eos_i = num_classes + 1
        bos_i_b = init_net.ConstantFill([], shape=[1], value=bos_i)
        eos_i_b = init_net.ConstantFill([], shape=[1], value=eos_i)
        labels = net.Cast([labels], to="int64")
        padded_labels, _ = net.Concat([bos_i_b, labels, eos_i_b], axis=0, outputs=2)
        return padded_labels

    @staticmethod
    def pad_predictions(predictions, init_net, net, num_classes):
        # This function will introduce two labels for beginning of sequence
        # And end of sequence, it will make the necessary udpates to the
        # the predictions blob

        low_score = -1000.0  # An arbitray very low number
        b_scores = np.array([[low_score] * num_classes + [0, low_score]]).astype(
            np.float32
        )

        e_scores = np.array([[low_score] * num_classes + [low_score, 0]]).astype(
            np.float32
        )

        b_scores = init_net.GivenTensorFill(
            [], "b_scores", shape=[1, num_classes + 2], values=b_scores
        )
        e_scores = init_net.GivenTensorFill(
            [], "e_scores", shape=[1, num_classes + 2], values=e_scores
        )

        zero_index = net.ConstantFill([], shape=[1], value=0)
        length = net.Gather([net.Shape([predictions]), zero_index])
        length = net.Cast(length, to="int32")
        t_range = net.LengthsRangeFill(length)
        padding = net.ConstantFill([t_range], value=low_score)
        padding = net.ExpandDims(padding, dims=[1])
        padded_predictions, _ = net.Concat(
            [predictions, padding, padding], outputs=2, axis=1
        )
        padded_predictions_concat, _ = net.Concat(
            [b_scores, padded_predictions, e_scores], outputs=2, axis=0
        )
        return padded_predictions_concat