# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import torch
import torch.nn as nn
import torch.nn.functional as F

from fairseq import utils
from fairseq.modules import (
    LayerNorm,
    MultiheadAttention,
    NormSelect,
)


class TransformerSentenceEncoderLayer(nn.Module):
    """
    Implements a Transformer Encoder Layer used in BERT/XLM style pre-trained
    models.
    """

    def __init__(
        self,
        embedding_dim: float = 768,
        ffn_embedding_dim: float = 3072,
        num_attention_heads: float = 8,
        dropout: float = 0.1,
        attention_dropout: float = 0.1,
        activation_dropout: float = 0.1,
        activation_fn: str = 'relu',
        add_bias_kv: bool = False,
        add_zero_attn: bool = False,
        export: bool = False,
        encoder_norm_self: str = 'layer',
        encoder_norm_ff: str = 'layer',
        encoder_normalize_before: bool = False,
    ) -> None:

        super().__init__()
        # Initialize parameters
        self.embedding_dim = embedding_dim
        self.dropout = dropout
        self.activation_dropout = activation_dropout

        # Initialize blocks
        self.activation_fn = utils.get_activation_fn(activation_fn)
        self.self_attn = MultiheadAttention(
            self.embedding_dim,
            num_attention_heads,
            dropout=attention_dropout,
            add_bias_kv=add_bias_kv,
            add_zero_attn=add_zero_attn,
            self_attention=True
        )

        # layer norm associated with the self attention layer
        # self.self_attn_layer_norm = LayerNorm(self.embedding_dim, export=export)
        self.self_attn_layer_norm = NormSelect(encoder_norm_self, self.embedding_dim, num_attention_heads)
        self.fc1 = nn.Linear(self.embedding_dim, ffn_embedding_dim)
        self.fc2 = nn.Linear(ffn_embedding_dim, self.embedding_dim)

        # layer norm associated with the position wise feed-forward NN
        # self.final_layer_norm = LayerNorm(self.embedding_dim, export=export)
        self.final_layer_norm = NormSelect(encoder_norm_ff, self.embedding_dim, num_attention_heads)
        self.encoder_normalize_before = encoder_normalize_before

    def forward(
        self,
        x: torch.Tensor,
        self_attn_mask: torch.Tensor = None,
        self_attn_padding_mask: torch.Tensor = None,
    ):
        """
        LayerNorm is applied either before or after the self-attention/ffn
        modules similar to the original Transformer imlementation.
        """
        residual = x
        x = self.maybe_layer_norm(self.self_attn_layer_norm, x, before=True, \
                                  pad_mask=self_attn_padding_mask)
        x, attn = self.self_attn(
            query=x,
            key=x,
            value=x,
            key_padding_mask=self_attn_padding_mask,
            need_weights=False,
            attn_mask=self_attn_mask,
        )
        x = F.dropout(x, p=self.dropout, training=self.training)
        x = residual + x
        #x = self.self_attn_layer_norm(x, pad_mask=self_attn_padding_mask, is_encoder=True)
        x = self.maybe_layer_norm(self.self_attn_layer_norm, x, after=True, \
                                  pad_mask=self_attn_padding_mask)
        residual = x
        
        x = self.maybe_layer_norm(self.final_layer_norm, x, before=True, \
                                  pad_mask=self_attn_padding_mask)
        x = self.activation_fn(self.fc1(x))
        x = F.dropout(x, p=self.activation_dropout, training=self.training)
        x = self.fc2(x)
        x = F.dropout(x, p=self.dropout, training=self.training)
        x = residual + x
        #x = self.final_layer_norm(x, pad_mask=self_attn_padding_mask, is_encoder=True)
        x = self.maybe_layer_norm(self.final_layer_norm, x, after=True, \
                                  pad_mask=self_attn_padding_mask)
        return x, attn

    def maybe_layer_norm(self, norm_fn, x, before=False, after=False, \
        pad_mask=None, is_encoder=True):
        assert before ^ after
        if after ^ self.encoder_normalize_before:
            return norm_fn(x)
        else:
            return x
