"""Transformer module classes."""

import abc
import collections
import math
from typing import Optional, Tuple

import torch
from torch import nn

from ... import data, defaults, special
from . import base, position


class Error(Exception):
    pass


class AttentionOutput(collections.UserList):
    """Tracks an attention output during the forward pass.

    This object can be passed into a hook to the Transformer forward pass
    in order to modify its behavior such that certain attention weights are
    returned, and then stored in self.outputs.
    """

    def __call__(
        self,
        module: nn.Module,
        module_in: Tuple[torch.Tensor, torch.Tensor, torch.Tensor],
        module_out: Tuple[torch.Tensor, torch.Tensor],
    ) -> None:
        """Stores the second return argument of `module`.

        This is intended to be called on a multiehaded attention, which returns
        both the contextualized representation, and the attention weights.

        Args:
            module (nn.Module): ignored.
            module_in (Tuple[torch.Tensor, torch.Tensor, torch.Tensor]):
                ignored.
            module_out (Tuple[torch.Tensor, torch.Tensor]): Output from
                the module. The second tensor is the attention weights.
        """
        _, attention = module_out
        self.append(attention)


class TransformerModule(base.BaseModule):
    """Abstract base module for transformers.

    Args:
        *args: passed to superclass.
        attention_heads (int, optional): number of attention heads.
        hidden_size (int, optional): size of the hidden layer.
        layers (int, optional): number of layers.
        max_length (int, optional): max length for input.
        **kwargs: passed to superclass.
    """

    attention_heads: int
    esq: float
    hidden_size: int
    layers: int
    module: nn.TransformerEncoder
    positional_encoding: position.PositionalEncoding

    def __init__(
        self,
        *args,
        attention_heads: int = defaults.ATTENTION_HEADS,
        hidden_size: int = defaults.HIDDEN_SIZE,
        layers: int = defaults.LAYERS,
        max_length: int = defaults.MAX_LENGTH,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        self.attention_heads = attention_heads
        self.esq = math.sqrt(self.embedding_size)
        self.hidden_size = hidden_size
        self.layers = layers
        self.module = self.get_module()
        self.positional_encoding = position.PositionalEncoding(
            self.embedding_size, max_length
        )

    def embed(
        self, symbols: torch.Tensor, embeddings: nn.Embedding
    ) -> torch.Tensor:
        """Embeds the symbols and adds positional encoding."""
        embedded = self.esq * embeddings(symbols)
        return self.dropout_layer(embedded + self.positional_encoding(symbols))

    @abc.abstractmethod
    def get_module(self) -> base.BaseModule: ...


class TransformerEncoder(TransformerModule):
    """Transformer encoder.

    Our implementation uses "pre-norm", i.e., it applies layer normalization
    before attention. Because of this, we are not currently able to make use
    of PyTorch's nested tensor feature.

    After:
        Xiong, R., Yang, Y., He, D., Zheng, K., Zheng, S., Xing, C., ..., and
        Liu, T.-Y. 2020. In Proceedings of the 37th International Conference
        on Machine Learning, pages 10524-10533.
    """

    def forward(
        self,
        symbols: data.PaddedTensor,
        embeddings: nn.Embedding,
        *args,
        **kwargs,
    ) -> torch.Tensor:
        """Encodes the symbols with the TransformerEncoder.

        Args:
            symbols (data.PaddedTensor).
            embeddings (nn.Embedding).
            *args: ignored.
            **kwargs: ignored.

        Returns:
            torch.Tensor: sequence of encoded symbols.
        """
        embedded = self.embed(symbols.tensor, embeddings)
        return self.module(embedded, src_key_padding_mask=symbols.mask)

    def get_module(self) -> nn.TransformerEncoder:
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=self.embedding_size,
            dim_feedforward=self.hidden_size,
            nhead=self.attention_heads,
            dropout=self.dropout,
            activation="relu",
            norm_first=True,
            batch_first=True,
        )
        return nn.TransformerEncoder(
            encoder_layer=encoder_layer,
            num_layers=self.layers,
            norm=nn.LayerNorm(self.embedding_size),
            # This silences a warning about the use of nested tensors,
            # currently an experimental feature.
            # TODO(#225): Re-enable if nested tensors are generalized to work
            # with `norm_first`.
            enable_nested_tensor=False,
        )

    @property
    def name(self) -> str:
        return "transformer"

    @property
    def output_size(self) -> int:
        return self.embedding_size


class FeatureInvariantTransformerEncoder(TransformerEncoder):
    """Transformer encoder with feature invariance.

    This is only sensibly used in a configuration where there is a shared
    source and features encoder, as in the following YAML snippet:

        source_encoder:
          class_path: yoyodyne.models.modules.TransformerEncoder
            init_args:
              ...
        features_encoder: true

    There is already space in the embeddings for the type embeddings; the
    caller just has to indicate whether the symbols are source or target
    symbols.

    After:
        Wu, S., Cotterell, R., and Hulden, M. 2021. Applying the transformer to
        character-level transductions. In Proceedings of the 16th Conference of
        the European Chapter of the Association for Computational Linguistics:
        Main Volume, pages 1901-1907.
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def embed(
        self,
        symbols: data.PaddedTensor,
        embeddings: nn.Embedding,
        is_source: bool,
    ) -> torch.Tensor:
        """Embeds the symbols and adds type and positional encodings."""
        embedded = self.esq * embeddings(symbols.tensor)
        type_embedded = self.esq * embeddings(
            torch.where(
                symbols.mask,
                special.SOURCE_IDX if is_source else special.FEATURES_IDX,
                special.PAD_IDX,
            )
        )
        return self.dropout_layer(
            embedded + type_embedded + self.positional_encoding(symbols.tensor)
        )

    def forward(
        self,
        symbols: data.PaddedTensor,
        embeddings: nn.Embedding,
        is_source: bool,
        *args,
        **kwargs,
    ) -> torch.Tensor:
        """Encodes the symbols with the TransformerEncoder.

        Args:
            symbols (data.PaddedTensor).
            embeddings (nn.Embedding).
            is_source (bool): is this being used to encode source or features?
            *args: ignored.
            **kwargs: ignored.

        Returns:
            torch.Tensor: sequence of encoded symbols.
        """
        embedded = self.embed(symbols, embeddings, is_source)
        return self.module(embedded, src_key_padding_mask=symbols.mask)

    @property
    def name(self) -> str:
        return "feature-invariant transformer"


class WrappedTransformerDecoder(nn.TransformerDecoder):
    """Wraps TransformerDecoder API for better variable naming."""

    def forward(
        self,
        source_encoded: torch.Tensor,
        source_mask: torch.Tensor,
        target: torch.Tensor,
        target_mask: Optional[torch.Tensor],
        causal_mask: torch.Tensor,
    ) -> torch.Tensor:
        return super().forward(
            memory=source_encoded,
            memory_key_padding_mask=source_mask,
            tgt=target,
            tgt_key_padding_mask=target_mask,
            tgt_is_causal=True,
            tgt_mask=causal_mask,
        )


class TransformerDecoder(TransformerModule):
    """Transformer decoder.

    Args:
        *args: passed to superclass.
        decoder_input_size (int).
        **kwargs: passed to superclass.
    """

    decoder_input_size: int
    # Constructed inside __init__.
    module: WrappedTransformerDecoder

    def __init__(
        self,
        *args,
        decoder_input_size: int = defaults.EMBEDDING_SIZE,
        **kwargs,
    ):
        self.decoder_input_size = decoder_input_size
        super().__init__(*args, **kwargs)

    def forward(
        self,
        source_encoded: torch.Tensor,
        source_mask: torch.Tensor,
        target: torch.Tensor,
        target_mask: Optional[torch.Tensor],
        embeddings: nn.Embedding,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Performs single pass of decoder module.

        Args:
            source_encoded (torch.Tensor): encoded source sequence.
            source_mask (torch.Tensor): mask for source.
            target (torch.Tensor): current state of targets, which may be the
                full target or previous decoded, of shape
                B x seq_len x hidden_size.
            target_mask (torch.Tensor): mask for target.
            embeddings (nn.Embedding): embeddings.

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: decoder outputs and the
                embedded targets.
        """
        target_embedded = self.embed(target, embeddings)
        causal_mask = self._causal_mask(target_embedded.size(1))
        # -> B x seq_len x d_model.
        decoded = self.module(
            source_encoded,
            source_mask,
            target_embedded,
            target_mask,
            causal_mask,
        )
        return decoded, target_embedded

    def get_module(self) -> WrappedTransformerDecoder:
        decoder_layer = nn.TransformerDecoderLayer(
            d_model=self.decoder_input_size,
            dim_feedforward=self.hidden_size,
            nhead=self.attention_heads,
            dropout=self.dropout,
            activation="relu",
            norm_first=True,
            batch_first=True,
        )
        return WrappedTransformerDecoder(
            decoder_layer=decoder_layer,
            num_layers=self.layers,
            norm=nn.LayerNorm(self.embedding_size),
        )

    def _causal_mask(self, target_len: int) -> torch.Tensor:
        return nn.Transformer.generate_square_subsequent_mask(
            target_len,
            device=self.device,
            dtype=bool,
        )

    @property
    def name(self) -> str:
        return "transformer"

    @property
    def output_size(self) -> int:
        return self.embedding_size


class SeparateFeaturesTransformerDecoderLayer(nn.TransformerDecoderLayer):
    """Transformer decoder layer with separate features.

    Each decoding step gets a second multihead attention representation
    w.r.t. the encoded features. This and the original multihead attention
    representation w.r.t. the encoded symbols are then compressed in a
    linear layer and finally concatenated.

    The implementation is otherwise identical to nn.TransformerDecoderLayer.

    Args:
        *args: passed to superclass.
        *kwargs: passed to superclass.
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        factory_kwargs = {
            "device": kwargs.get("device"),
            "dtype": kwargs.get("dtype"),
        }
        d_model = kwargs["d_model"]
        self.features_multihead_attn = nn.MultiheadAttention(
            d_model,  # TODO: Separate feature embedding size?
            kwargs["nhead"],
            dropout=kwargs["dropout"],
            batch_first=kwargs["batch_first"],
            **factory_kwargs,
        )
        # If d_model is not even, an error will result. This is unlikely to
        # trigger since it must also be divisible by the number of attention
        # heads.
        if d_model % 2 != 0:
            raise Error(
                f"Feature-invariant transformer d_model ({d_model}) must be "
                "divisible by 2"
            )
        bias = kwargs.get("bias")
        self.source_linear = nn.Linear(
            d_model,
            d_model // 2,
            bias=bias,
            **factory_kwargs,
        )
        self.features_linear = nn.Linear(
            d_model,  # TODO: Separate feature embedding size?
            d_model // 2,
            bias=bias,
            **factory_kwargs,
        )

    def forward(
        self,
        source_encoded: torch.Tensor,
        source_mask: torch.Tensor,
        target: torch.Tensor,
        target_mask: Optional[torch.Tensor],
        features_encoded: torch.Tensor,
        features_mask: torch.Tensor,
        causal_mask: torch.Tensor,
    ) -> torch.Tensor:
        """Pass the inputs (and mask) through the decoder layer.

        This is based closely on the internals in torch.nn.transformer and
        follows the somewhat-inscrutable variable naming used there.

        Args:
            source_encoded (torch.Tensor): encoded source sequence.
            source_mask (torch.Tensor): mask for source.
            target (torch.Tensor): current embedded target, which
                may be the full target or previous decoded, of shape
                B x seq_len x hidden_size.
            target_mask (torch.Tensor): mask for target.
            features_encoded (torch.Tensor): encoded features.
            features_mask (torch.Tensor): mask for features.
            causal_mask (torch.Tensor).

        Returns:
            torch.Tensor.
        """
        # Self-attention (pre-norm).
        x = self.norm1(target)
        sa_output, _ = self.self_attn(
            x,
            x,
            x,
            attn_mask=causal_mask,
            key_padding_mask=target_mask,
        )
        output = target + self.dropout1(sa_output)
        # Cross-attention (pre-norm).
        x = self.norm2(output)
        # Cross-attends to source.
        source_attn_out, _ = self.multihead_attn(
            x, source_encoded, source_encoded, key_padding_mask=source_mask
        )
        source_attention = self.source_linear(self.dropout2(source_attn_out))
        # Cross-attends to features.
        feat_attn_out, _ = self.features_multihead_attn(
            x,
            features_encoded,
            features_encoded,
            key_padding_mask=features_mask,
        )
        features_attention = self.features_linear(self.dropout2(feat_attn_out))
        # Concatenates and adds residual connection.
        attn_output = torch.cat((source_attention, features_attention), dim=2)
        output = output + attn_output
        # Feed-forward.
        # Normalizes before the feed-forward network.
        output = output + self._ff_block(self.norm3(output))
        return output


class SeparateFeaturesTransformerDecoder(nn.TransformerDecoder):
    """Transformer decoder with separate features.

    Adding separate features into the transformer stack is implemented with
    SeparateFeaturesTransformerDecoderLayer.
    """

    def forward(
        self,
        source_encoded: torch.Tensor,
        source_mask: torch.Tensor,
        target: torch.Tensor,
        target_mask: Optional[torch.Tensor],
        features_encoded: torch.Tensor,
        features_mask: torch.Tensor,
        causal_mask: torch.Tensor,
    ) -> torch.Tensor:
        """Passes the inputs (and mask) through the decoder layer.

        Args:
            source_encoded (torch.Tensor): encoded source sequence.
            source_mask (torch.Tensor): mask for source.
            target (torch.Tensor): current embedded targets, which
                may be the full target or previous decoded, of shape
                B x seq_len x hidden_size.
            target_mask (torch.Tensor): causal mask for target.
            features_encoded (torch.Tensor): encoded features.
            features_mask (torch.Tensor): mask for source.
            causal_mask (torch.Tensor).

        Returns:
            torch.Tensor: Output tensor.
        """
        output = target
        for layer in self.layers:
            output = layer(
                source_encoded,
                source_mask,
                output,
                target_mask,
                features_encoded,
                features_mask,
                causal_mask,
            )
        if self.norm is not None:
            output = self.norm(output)
        return output


class PointerGeneratorTransformerDecoder(TransformerDecoder):
    """A transformer decoder which tracks the output of multihead attention.

    This is achieved with a hook into the forward pass.

    After:
        https://gist.github.com/airalcorn2/50ec06517ce96ecc143503e21fa6cb91

    Args:
        *args: passed to superclass.
        has_features_encoder (bool, optional).
        *kwargs: passed to superclass.
    """

    def __init__(
        self,
        *args,
        has_features_encoder: bool = False,
        **kwargs,
    ):
        self.has_features_encoder = has_features_encoder
        super().__init__(*args, **kwargs)
        # Stores the actual cross attentions.
        self.attention_output = AttentionOutput()
        # Refers to the attention from decoder to encoder.
        self._patch_attention(self.module.layers[-1].multihead_attn)
        self.hook_handle = self.module.layers[
            -1
        ].multihead_attn.register_forward_hook(self.attention_output)

    def forward(
        self,
        source_encoded: torch.Tensor,
        source_mask: torch.Tensor,
        target: torch.Tensor,
        target_mask: Optional[torch.Tensor],
        embeddings: nn.Embedding,
        *,
        features_encoded: Optional[torch.Tensor] = None,
        features_mask: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Performs single pass of decoder module.

        Args:
            source_encoded (torch.Tensor): encoded source sequence.
            source_mask (torch.Tensor): mask for source.
            target (torch.Tensor): current targets, which may be the full
                target or previous decoded, of shape B x seq_len x hidden_size.
            target_mask (torch.Tensor): mask for target.
            embeddings (nn.Embedding): embedding.
            features_encoded (Optional[torch.Tensor]): encoded features.
            features_mask (Optional[torch.Tensor]): mask for features.

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: decoder outputs and the
                embedded targets.
        """
        # FIXME: temporary.
        if target_mask is not None:
            assert target_mask.dim() == 2, target_mask.shape
        if features_mask is not None:
            assert features_mask.dim() == 2, features_mask.shape
        target_embedded = self.embed(target, embeddings)
        causal_mask = self._causal_mask(target_embedded.size(1))
        # -> B x seq_len x d_model.
        if self.has_features_encoder:
            decoded = self.module(
                source_encoded,
                source_mask,
                target_embedded,
                target_mask,
                features_encoded,
                features_mask,
                causal_mask,
            )
        else:
            decoded = self.module(
                source_encoded,
                source_mask,
                target_embedded,
                target_mask,
                causal_mask,
            )
        return decoded, target_embedded

    def get_module(self) -> nn.TransformerDecoder:
        if self.has_features_encoder:
            decoder_layer = SeparateFeaturesTransformerDecoderLayer(
                d_model=self.decoder_input_size,
                dim_feedforward=self.hidden_size,
                nhead=self.attention_heads,
                dropout=self.dropout,
                activation="relu",
                norm_first=True,
                batch_first=True,
            )
            return SeparateFeaturesTransformerDecoder(
                decoder_layer=decoder_layer,
                num_layers=self.layers,
                norm=nn.LayerNorm(self.embedding_size),
            )
        else:
            decoder_layer = nn.TransformerDecoderLayer(
                d_model=self.decoder_input_size,
                dim_feedforward=self.hidden_size,
                nhead=self.attention_heads,
                dropout=self.dropout,
                activation="relu",
                norm_first=True,
                batch_first=True,
            )
            return WrappedTransformerDecoder(
                decoder_layer=decoder_layer,
                num_layers=self.layers,
                norm=nn.LayerNorm(self.embedding_size),
            )

    def _patch_attention(self, attention_module: torch.nn.Module) -> None:
        """Wraps a module's forward pass such that `need_weights` is True.

        Args:
            attention_module (torch.nn.Module): the module from which we want
                to track multiheaded attention weights.
        """
        forward_orig = attention_module.forward

        def wrap(*args, **kwargs):
            kwargs["need_weights"] = True
            return forward_orig(*args, **kwargs)

        attention_module.forward = wrap
