""" A bunch of util functions to build Seq2Seq models with Caffe2."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

from caffe2.python import core, recurrent
from caffe2.python.cnn import CNNModelHelper


class ModelHelper(CNNModelHelper):

    def __init__(self, init_params=True):
        super(ModelHelper, self).__init__(
            order='NCHW',  # this is only relevant for convolutional networks
            init_params=init_params,
        )
        self.non_trainable_params = []

    def AddParam(self, name, init=None, init_value=None, trainable=True):
        """Adds a parameter to the model's net and it's initializer if needed

        Args:
            init: a tuple (<initialization_op_name>, <initialization_op_kwargs>)
            init_value: int, float or str. Can be used instead of `init` as a
                simple constant initializer
            trainable: bool, whether to compute gradient for this param or not
        """
        if init_value is not None:
            assert init is None
            assert type(init_value) in [int, float, str]
            init = ('ConstantFill', dict(
                shape=[1],
                value=init_value,
            ))

        if self.init_params:
            param = self.param_init_net.__getattr__(init[0])(
                [],
                name,
                **init[1]
            )
        else:
            param = self.net.AddExternalInput(name)

        if trainable:
            self.params.append(param)
        else:
            self.non_trainable_params.append(param)

        return param


def rnn_unidirectional_encoder(
    model,
    embedded_inputs,
    input_lengths,
    initial_hidden_state,
    initial_cell_state,
    embedding_size,
    encoder_num_units,
    use_attention
):
    """ Unidirectional (forward pass) LSTM encoder."""

    with core.NameScope('', reset=True):
        outputs, final_hidden_state, _, final_cell_state = recurrent.LSTM(
            model=model,
            input_blob=embedded_inputs,
            seq_lengths=input_lengths,
            initial_states=(initial_hidden_state, initial_cell_state),
            dim_in=embedding_size,
            dim_out=encoder_num_units,
            scope='encoder',
            outputs_with_grads=([0] if use_attention else [1, 3]),
        )
    return outputs, final_hidden_state, final_cell_state


def rnn_bidirectional_encoder(
    model,
    embedded_inputs,
    input_lengths,
    initial_hidden_state,
    initial_cell_state,
    embedding_size,
    encoder_num_units,
    use_attention
):
    """ Bidirectional (forward pass and backward pass) LSTM encoder."""

    with core.NameScope('', reset=True):
        # Forward pass
        (
            outputs_fw,
            final_hidden_state_fw,
            _,
            final_cell_state_fw,
        ) = recurrent.LSTM(
            model=model,
            input_blob=embedded_inputs,
            seq_lengths=input_lengths,
            initial_states=(initial_hidden_state, initial_cell_state),
            dim_in=embedding_size,
            dim_out=encoder_num_units,
            scope='forward_encoder',
            outputs_with_grads=([0] if use_attention else [1, 3]),
        )

        # Backward pass
        reversed_embedded_inputs = model.net.ReversePackedSegs(
            [embedded_inputs, input_lengths],
            ['reversed_embedded_inputs'],
        )

        (
            outputs_bw,
            final_hidden_state_bw,
            _,
            final_cell_state_bw,
        ) = recurrent.LSTM(
            model=model,
            input_blob=reversed_embedded_inputs,
            seq_lengths=input_lengths,
            initial_states=(initial_hidden_state, initial_cell_state),
            dim_in=embedding_size,
            dim_out=encoder_num_units,
            scope='backward_encoder',
            outputs_with_grads=([0] if use_attention else [1, 3]),
        )

        outputs_bw = model.net.ReversePackedSegs(
            [outputs_bw, input_lengths],
            ['outputs_bw'],
        )

        # Concatenate forward and backward results
        outputs, _ = model.net.Concat(
            [outputs_fw, outputs_bw],
            ['outputs', 'outputs_dim'],
            axis=2,
        )

        final_hidden_state, _ = model.net.Concat(
            [final_hidden_state_fw, final_hidden_state_bw],
            ['final_hidden_state', 'final_hidden_state_dim'],
            axis=2,
        )

        final_cell_state, _ = model.net.Concat(
            [final_cell_state_fw, final_cell_state_bw],
            ['final_cell_state', 'final_cell_state_dim'],
            axis=2,
        )
    return outputs, final_hidden_state, final_cell_state