"""Attention combination strategies.

This modules implements attention combination strategies for multi-encoder
scenario when we may want to combine the hidden states of the encoders in
more complicated fashion.

Currently there are two attention combination strategies flat and hierarchical
(see paper `Attention Combination Strategies for Multi-Source
Sequence-to-Sequence Learning <https://arxiv.org/pdf/1704.06567.pdf>`_).

The combination strategies may use the sentinel mechanism which allows the
decoder not to attend to the, and extract information on its own hidden state
(see paper `Knowing when to Look: Adaptive Attention via a Visual Sentinel for
Image Captioning  <https://arxiv.org/pdf/1612.01887.pdf>`_).
"""
from typing import Any, List, Union, Type, Tuple, NamedTuple
import tensorflow as tf

from neuralmonkey.dataset import Dataset
from neuralmonkey.decoding_function import (BaseAttention, AttentionLoopState,
                                            empty_attention_loop_state)
from neuralmonkey.model.model_part import ModelPart, FeedDict
from neuralmonkey.encoders.attentive import Attentive
from neuralmonkey.checking import assert_shape
from neuralmonkey.nn.projection import linear


class EncoderWrapper(ModelPart, Attentive):
    """Wrapper doing attention combination behaving as a single encoder.

    This class wraps encoders and performs the attention combination in such a
    way that for the decoder, it looks like a single encoder capable to
    generate a single context vector.
    """

    def __init__(self,
                 name: str,
                 encoders: List[Any],
                 attention_type: Type,
                 attention_state_size: int,
                 use_sentinels=False,
                 share_attn_projections=False) -> None:
        """Initializes the encoder wrapper.

        Args:
            name: Name of the encoder / its scope.
            encoders: List of encoders to be wrapped.
            attention_type: Type of the attention combination.
            attention_state_size: Dimension of the state projection of
                attention energy computation.
            use_sentinels: Flag whether the sentinel mechanism should be added
                to the attention combination.
            share_attn_projections: Flag whether the hidden state projection
                should be shared for the both the energies computation and
                context vector computation.
        """

        ModelPart.__init__(self, name, None, None)
        Attentive.__init__(self, attention_type)
        self.encoders = encoders
        self._attention_type = attention_type
        self._attention_state_size = attention_state_size
        self._use_sentinels = use_sentinels
        self._share_attn_projections = share_attn_projections

        self.encoded = tf.concat([e.encoded for e in encoders], 1)

    def create_attention_object(self):
        return self._attention_type(
            self.encoders,
            self._attention_state_size,
            "attention_{}".format(self.name),
            use_sentinels=self._use_sentinels,
            share_projections=self._share_attn_projections)

    def feed_dict(self, dataset: Dataset, train: bool) -> FeedDict:
        return {}

    @property
    def _attention_tensor(self):
        raise NotImplementedError("Encoder wrapper does not contain the"
                                  " attention tensor")

    @property
    def _attention_mask(self):
        raise NotImplementedError("Encoder wrapper does not contain the"
                                  " attention mask")


class MultiAttention(BaseAttention):
    """Base class for attention combination."""

    # pylint: disable=unused-argument
    def __init__(self,
                 encoders: List[Attentive],
                 attention_state_size: int,
                 scope: Union[tf.VariableScope, str],
                 share_projections: bool = False,
                 use_sentinels: bool = False) -> None:
        super().__init__(scope, None, attention_state_size)
        self._encoders = encoders
        self.attentions_in_time = []  # type: List[tf.Tensor]
        self._share_projections = share_projections
        self._use_sentinels = use_sentinels

        with tf.variable_scope(self.scope):
            self.attn_v = tf.get_variable(
                "attn_v", [1, 1, self.attention_state_size],
                initializer=tf.random_normal_initializer(stddev=.001))
    # pylint: enable=unused-argument

    def attention(self, decoder_state, decoder_prev_state,
                  decoder_input, _, step):
        """Get context vector for given decoder state."""
        raise NotImplementedError("Abstract method")

    @property
    def attn_size(self):
        return self.attention_state_size

    def _vector_logit(self,
                      projected_decoder_state: tf.Tensor,
                      vector_value: tf.Tensor,
                      scope: str) -> tf.Tensor:
        """Get logit for a single vector, e.g., sentinel vector."""
        assert_shape(projected_decoder_state, [-1, 1, -1])
        assert_shape(vector_value, [-1, -1])

        with tf.variable_scope("{}_logit".format(scope)):
            vector_bias = tf.get_variable(
                "vector_bias", [],
                initializer=tf.constant_initializer(0.0))

            proj_vector_for_logit = tf.expand_dims(
                linear(vector_value, self.attention_state_size,
                       scope="vector_projection"), 1)

            if self._share_projections:
                proj_vector_for_ctx = proj_vector_for_logit
            else:
                proj_vector_for_ctx = tf.expand_dims(
                    linear(vector_value, self.attention_state_size,
                           scope="vector_ctx_proj"), 1)

            vector_logit = tf.reduce_sum(
                self.attn_v *
                tf.tanh(projected_decoder_state + proj_vector_for_logit),
                [2]) + vector_bias
            assert_shape(vector_logit, [-1, 1])
            return proj_vector_for_ctx, vector_logit


class FlatMultiAttention(MultiAttention):
    """Flat attention combination strategy.

    Using this attention combination strategy, hidden states of the encoders
    are first projected to the same space (different projection for different
    encoders) and then we compute a joint distribution over all the hidden
    states. The context vector 
is then a weighted sum of another / then
    projection of the encoders hidden states. The sentinel vector can be added
    as an additional hidden state.

    See equations 8 to 10 in the Attention Combination Strategies paper.
    """
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        # pylint: disable=protected-access
    
    self._encoders_tensors = [e._attention_tensor for e in self._encoders]
        self._encoders_masks = [e._attention_mask for e in self._encoders]
        # pylint: enable=protected-access

        for e_m in self._encoders_masks:
            assert_shape(e_m, [-1, -1])

        for e_t in self._encoders_tensors:
            assert_shape(e_t, [-1, -1, -1])

        with tf.variable_scope(self.scope):
            self.encoder_projections_for_logits = \
                self.get_encoder_projections("logits_projections")

            self.encoder_attn_biases = [
                tf.get_variable(name="attn_bias_{}".format(i),
                                shape=[],
                                initializer=tf.constant_initializer(0.))
                for i in range(len(self._encoders_tensors))]

            if self._share_projections:
                self.encoder_projections_for_ctx = \
                    self.encoder_projections_for_logits
            else:
                self.encoder_projections_for_ctx = \
                    self.get_encoder_projections("context_projections")

            if self._use_sentinels:
                self._encoders_masks.append(
                    tf.ones([tf.shape(self._encoders_masks[0])[0], 1]))

            self.masks_concat = tf.concat(self._encoders_masks, 1)

    def initial_loop_