"""
Implementation of "Attention is All You Need"
"""
import torch
import torch.nn as nn

from onmt.encoders.encoder import EncoderBase
from onmt.modules import MultiHeadedAttention
from onmt.modules.position_ffn import PositionwiseFeedForward
from onmt.utils.misc import sequence_mask


class TransformerEncoderLayer(nn.Module):
    """
    A single layer of the transformer encoder.

    Args:
        d_model (int): the dimension of keys/values/queries in
                   MultiHeadedAttention, also the input size of
                   the first-layer of the PositionwiseFeedForward.
        heads (int): the number of head for MultiHeadedAttention.
        d_ff (int): the second-layer of the PositionwiseFeedForward.
        dropout (float): dropout probability(0-1.0).
    """

    def __init__(self, d_model, heads, d_ff, dropout, attention_dropout,
                 gate, edges, max_relative_positions=0):
        super(TransformerEncoderLayer, self).__init__()

        self.self_attn = MultiHeadedAttention(
            heads, d_model, dropout=attention_dropout,
            max_relative_positions=max_relative_positions,
            edge_type=edges, gate=gate)
        self.feed_forward = PositionwiseFeedForward(d_model, d_ff, dropout)
        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)
        self.dropout = nn.Dropout(dropout)

    def forward(self, inputs, grh, mask):
        """
        Args:
            inputs (FloatTensor): ``(batch_size, src_len, model_dim)``
            mask (LongTensor): ``(batch_size, 1, src_len)``

        Returns:
            (FloatTensor):

            * outputs ``(batch_size, src_len, model_dim)``
        """
        input_norm = self.layer_norm(inputs)
        context, _ = self.self_attn(input_norm, input_norm, input_norm, grh,
                                    mask=mask, attn_type="self")
        out = self.dropout(context) + inputs
        return self.feed_forward(out)

    def update_dropout(self, dropout, attention_dropout):
        self.self_attn.update_dropout(attention_dropout)
        self.feed_forward.update_dropout(dropout)
        self.dropout.p = dropout


class TransformerEncoder(EncoderBase):
    """The Transformer encoder from "Attention is All You Need"
    :cite:`DBLP:journals/corr/VaswaniSPUJGKP17`

    .. mermaid::

       graph BT
          A[input]
          B[multi-head self-attn]
          C[feed forward]
          O[output]
          A --> B
          B --> C
          C --> O

    Args:
        num_layers (int): number of encoder layers
        d_model (int): size of the model
        heads (int): number of heads
        d_ff (int): size of the inner FF layer
        dropout (float): dropout parameters
        embeddings (onmt.modules.Embeddings):
          embeddings to use, should have positional encodings

    Returns:
        (torch.FloatTensor, torch.FloatTensor):

        * embeddings ``(src_len, batch_size, model_dim)``
        * memory_bank ``(src_len, batch_size, model_dim)``
    """

    def __init__(self, num_layers, d_model, heads, d_ff, dropout,
                 attention_dropout, embeddings, max_relative_positions,
                 aggregation, gate, edges):
        super(TransformerEncoder, self).__init__()

        self.embeddings = embeddings
        self.transformer = nn.ModuleList(
            [TransformerEncoderLayer(
                d_model, heads, d_ff, dropout, attention_dropout,
                gate, edges, max_relative_positions=max_relative_positions)
             for i in range(num_layers)])
        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)
        self.aggregation = aggregation
        if aggregation is not None:
            if aggregation == "dense":
                self.dense_linear = nn.ModuleList(
                    [nn.Linear(d_model+i*d_model, d_model) for i in range(num_layers)]
                )
                ## TODO: MAX_POOLING, LSTM
                self.aggregate_linear = nn.Linear((1+num_layers)*d_model, d_model)
                self.aggregate_layer_norm = nn.LayerNorm((1 + num_layers) * d_model, eps=1e-6)

            if aggregation == "jump":
                ## TODO: MAX_POOLING, LSTM
                self.aggregate_linear = nn.Linear((num_layers+1)*d_model,
                                                  d_model)
                self.aggregate_layer_norm = nn.LayerNorm((1 + num_layers) * d_model, eps=1e-6)

    @classmethod
    def from_opt(cls, opt, embeddings):
        """Alternate constructor."""
        return cls(
            opt.enc_layers,
            opt.enc_rnn_size,
            opt.heads,
            opt.transformer_ff,
            opt.dropout[0] if type(opt.dropout) is list else opt.dropout,
            opt.attention_dropout[0] if type(opt.attention_dropout)
            is list else opt.attention_dropout,
            embeddings,
            opt.max_relative_positions,
            opt.aggregation,
            opt.gate,
            opt.edges)

    def forward(self, src, grh, lengths=None):
        """See :func:`EncoderBase.forward()`"""
        self._check_args(src, lengths)
        assert src.size(0) == grh.size(-1), "srclen != grh_n"

        emb = self.embeddings(src) # batch, srclen, dim

        out = emb.transpose(0, 1).contiguous()
        mask = ~sequence_mask(lengths).unsqueeze(1)
        # Run the forward pass of every layer of the tranformer.
        out_list = []
        out_list.append(out)
        if self.aggregation == "dense":
            for i in range(len(self.transformer)):
                out = self.dense_linear[i](torch.cat(out_list, dim=-1))
                out = self.transformer[i](out, grh, mask)
                out_list.append(out)
            aggregate_out = torch.cat(out_list, dim=-1)
            aggregate_out = self.aggregate_layer_norm(aggregate_out)
            aggregate_out = self.aggregate_linear(aggregate_out)
            return emb, aggregate_out.transpose(0, 1).contiguous(), lengths
        else:
            for layer in self.transformer:
                out = layer(out, grh, mask)
                out_list.append(out)
            out = self.layer_norm(out)
            if self.aggregation == "jump":
                aggregate_out = torch.cat(out_list, dim=-1)
                aggregate_out = self.aggregate_layer_norm(aggregate_out)
                aggregate_out = self.aggregate_linear(aggregate_out)
                return emb, aggregate_out.transpose(0, 1).contiguous(), lengths
        return emb, out.transpose(0, 1).contiguous(), lengths

    def update_dropout(self, dropout, attention_dropout):
        self.embeddings.update_dropout(dropout)
        for layer in self.transformer:
            layer.update_dropout(dropout, attention_dropout)

"""
Implementation of "Attention is All You Need"
"""
import torch
import torch.nn as nn

from onmt.encoders.encoder import EncoderBase
from onmt.modules import MultiHeadedAttention
from onmt.modules.position_ffn import PositionwiseFeedForward
from onmt.utils.misc import sequence_mask


class TransformerEncoderLayer(nn.Module):
    """
    A single layer of the transformer encoder.

    Args:
        d_model (int): the dimension of keys/values/queries in
                   MultiHeadedAttention, also the input size of
                   the first-layer of the PositionwiseFeedForward.
        heads (int): the number of head for MultiHeadedAttention.
        d_ff (int): the second-layer of the PositionwiseFeedForward.
        dropout (float): dropout probability(0-1.0).
    """

    def __init__(self, d_model, heads, d_ff, dropout, attention_dropout,
                 gate, edges, max_relative_positions=0):
        super(TransformerEncoderLayer, self).__init__()

        self.self_attn = MultiHeadedAttention(
            heads, d_model, dropout=attention_dropout,
            max_relative_positions=max_relative_positions,
            edge_type=edges, gate=gate)
        self.feed_forward = PositionwiseFeedForward(d_model, d_ff, dropout)
        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)
        self.dropout = nn.Dropout(dropout)

    def forward(self, inputs, grh, mask):
        """
        Args:
            inputs (FloatTensor): ``(batch_size, src_len, model_dim)``
            mask (LongTensor): ``(batch_size, 1, src_len)``

        Returns:
            (FloatTensor):

            * outputs ``(batch_size, src_len, model_dim)``
        """
        input_norm = self.layer_norm(inputs)
        context, _ = self.self_attn(input_norm, input_norm, input_norm, grh,
                                    mask=mask, attn_type="self")
        out = self.dropout(context) + inputs
        return self.feed_forward(out)

    def update_dropout(self, dropout, attention_dropout):
        self.self_attn.update_dropout(attention_dropout)
        self.feed_forward.update_dropout(dropout)
        self.dropout.p = dropout


class TransformerEncoder(EncoderBase):
    """The Transformer encoder from "Attention is All You Need"
    :cite:`DBLP:journals/corr/VaswaniSPUJGKP17`

    .. mermaid::

       graph BT
          A[input]
          B[multi-head self-attn]
          C[feed forward]
          O[output]
          A --> B
          B --> C
          C --> O

    Args:
        num_layers (int): number of encoder layers
        d_model (int): size of the model
        heads (int): number of heads
        d_ff (int): size of the inner FF layer
        dropout (float): dropout parameters
        embeddings (onmt.modules.Embeddings):
          embeddings to use, should have positional encodings

    Returns:
        (torch.FloatTensor, torch.FloatTensor):

        * embeddings ``(src_len, batch_size, model_dim)``
        * memory_bank ``(src_len, batch_size, model_dim)``
    """

    def __init__(self, num_layers, d_model, heads, d_ff, dropout,
                 attention_dropout, embeddings, max_relative_positions,
                 aggregation, gate, edges):
        super(TransformerEncoder, self).__init__()

        self.embeddings = embeddings
        self.transformer = nn.ModuleList(
            [TransformerEncoderLayer(
                d_model, heads, d_ff, dropout, attention_dropout,
                gate, edges, max_relative_positions=max_relative_positions)
             for i in range(num_layers)])
        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)
        self.aggregation = aggregation
        if aggregation is not None:
            if aggregation == "dense":
                self.dense_linear = nn.ModuleList(
                    [nn.Linear(d_model+i*d_model, d_model) for i in range(num_layers)]
                )
                ## TODO: MAX_POOLING, LSTM
                self.aggregate_linear = nn.Linear((1+num_layers)*d_model, d_model)
                self.aggregate_layer_norm = nn.LayerNorm((1 + num_layers) * d_model, eps=1e-6)

            if aggregation == "jump":
                ## TODO: MAX_POOLING, LSTM
                self.aggregate_linear = nn.Linear((num_layers+1)*d_model,
                                                  d_model)
                self.aggregate_layer_norm = nn.LayerNorm((1 + num_layers) * d_model, eps=1e-6)

    @classmethod
    def from_opt(cls, opt, embeddings):
        """Alternate constructor."""
        return cls(
            opt.enc_layers,
            opt.enc_rnn_size,
            opt.heads,
            opt.transformer_ff,
            opt.dropout[0] if type(opt.dropout) is list else opt.dropout,
            opt.attention_dropout[0] if type(opt.attention_dropout)
            is list else opt.attention_dropout,
            embeddings,
            opt.max_relative_positions,
            opt.aggregation,
            opt.gate,
            opt.edges)

    def forward(self, src, grh, lengths=None):
        """See :func:`EncoderBase.forward()`"""
        self._check_args(src, lengths)
        assert src.size(0) == grh.size(-1), "srclen != grh_n"

        emb = self.embeddings(src) # batch, srclen, dim

        out = emb.transpose(0, 1).contiguous()
        mask = ~sequence_mask(lengths).unsqueeze(1)
        # Run the forward pass of every layer of the tranformer.
        out_list = []
        out_list.append(out)
        if self.aggregation == "dense":
            for i in range(len(self.transformer)):
                out = self.dense_linear[i](torch.cat(out_list, dim=-1))
                out = self.transformer[i](out, grh, mask)
                out_list.append(out)
            aggregate_out = torch.cat(out_list, dim=-1)
            aggregate_out = self.aggregate_layer_norm(aggregate_out)
            aggregate_out = self.aggregate_linear(aggregate_out)
            return emb, aggregate_out.transpose(0, 1).contiguous(), lengths
        else:
            for layer in self.transformer:
                out = layer(out, grh, mask)
                out_list.append(out)
            out = self.layer_norm(out)
            if self.aggregation == "jump":
                aggregate_out = torch.cat(out_list, dim=-1)
                aggregate_out = self.aggregate_layer_norm(aggregate_out)
                aggregate_out = self.aggregate_linear(aggregate_out)
                return emb, aggregate_out.transpose(0, 1).contiguous(), lengths
        return emb, out.transpose(0, 1).contiguous(), lengths

    def update_dropout(self, dropout, attention_dropout):
        self.embeddings.update_dropout(dropout)
        for layer in self.transformer:
            layer.update_dropout(dropout, attention_dropout)

Implementation of "Attention is All You Need" A single layer of the transformer encoder.
 
     Args:
         d_model (int): the dimension of keys/values/queries in
                    MultiHeadedAttention, also the input size of
                    the first-layer of the PositionwiseFeedForward.
         heads (int): the number of head for MultiHeadedAttention.
         d_ff (int): the second-layer of the PositionwiseFeedForward.
         dropout (float): dropout probability(0-1.0). Args:
             inputs (FloatTensor): ``(batch_size, src_len, model_dim)``
             mask (LongTensor): ``(batch_size, 1, src_len)``
 
         Returns:
             (FloatTensor):
 
             * outputs ``(batch_size, src_len, model_dim)`` The Transformer encoder from "Attention is All You Need"
     :cite:`DBLP:journals/corr/VaswaniSPUJGKP17`
 
     .. mermaid::
 
        graph BT
           A[input]
           B[multi-head self-attn]
           C[feed forward]
           O[output]
           A --> B
           B --> C
           C --> O
 
     Args:
         num_layers (int): number of encoder layers
         d_model (int): size of the model
         heads (int): number of heads
         d_ff (int): size of the inner FF layer
         dropout (float): dropout parameters
         embeddings (onmt.modules.Embeddings):
           embeddings to use, should have positional encodings
 
     Returns:
         (torch.FloatTensor, torch.FloatTensor):
 
         * embeddings ``(src_len, batch_size, model_dim)``
         * memory_bank ``(src_len, batch_size, model_dim)`` ## TODO: MAX_POOLING, LSTM ## TODO: MAX_POOLING, LSTM Alternate constructor. See :func:`EncoderBase.forward()` # batch, srclen, dim # Run the forward pass of every layer of the tranformer.