# Copyright 2021 UMD (<NAME>)

# Transformer encoder (without the decoder part), a wrapper for the multi-headed attention layers
# Modified from ESPnet so that it can be used without espnet dependency

from typing import Optional
from typing import Tuple
from abc import ABC
import torch
from qanta.embedding import PositionalEncoding
from qanta.attention import LayerNorm, PositionwiseFeedForward, EncoderLayer, MultiHeadedAttention, make_pad_mask


class MultiSequential(torch.nn.Sequential):
    """Multi-input multi-output torch.nn.Sequential."""

    def forward(self, *args):
        """Repeat."""
        for m in self:
            args = m(*args)
        return args


def repeat(N, fn):
    """Repeat module N times.

    Args:
        N (int): Number of repeat time.
        fn (Callable): Function to generate module.

    Returns:
        MultiSequential: Repeated model instance.

    """
    return MultiSequential(*[fn(n) for n in range(N)])


class TransformerEncoder(torch.nn.Module):
    """Transformer encoder

    Args:
        input_size: input dim
        output_size: dimension of attention
        attention_heads: the number of heads of multi head attention
        linear_units: the number of units of position-wise feed forward
        num_blocks: the number of decoder blocks
        dropout_rate: dropout rate
        attention_dropout_rate: dropout rate in attention
        positional_dropout_rate: dropout rate after adding positional encoding
        pos_enc_class: PositionalEncoding
        normalize_before: whether to use layer_norm before the first block
        padding_idx: padding_idx for input_layer
    """

    def __init__(
        self,
        input_size: int,
        output_size: int = 256,
        attention_heads: int = 4,
        linear_units: int = 2048,
        num_blocks: int = 6,
        dropout_rate: float = 0.1,
        positional_dropout_rate: float = 0.1,
        attention_dropout_rate: float = 0.0,
        pos_enc_class: PositionalEncoding = PositionalEncoding,
        normalize_before: bool = True,
        padding_idx: int = -1,
    ):
        super().__init__()
        self._output_size = output_size

        self.embed = torch.nn.Sequential(
            torch.nn.Linear(input_size, output_size),
            torch.nn.LayerNorm(output_size),
            torch.nn.Dropout(dropout_rate),
            torch.nn.ReLU(),
            pos_enc_class(output_size, positional_dropout_rate),
        )

        self.normalize_before = normalize_before

        positionwise_layer = PositionwiseFeedForward
        positionwise_layer_args = (
            output_size,
            linear_units,
            dropout_rate,
        )

        self.encoders = repeat(
            num_blocks,
            lambda lnum: EncoderLayer(
                output_size,
                MultiHeadedAttention(
                    attention_heads, output_size, attention_dropout_rate
                ),
                positionwise_layer(*positionwise_layer_args),
                dropout_rate,
                normalize_before,
                False,
            ),
        )
        if self.normalize_before:
            self.after_norm = LayerNorm(output_size)

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

    def forward(
        self,
        xs_pad: torch.Tensor,
        ilens: torch.Tensor,
        prev_states: torch.Tensor = None,
    ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
        """Embed positions in tensor.

        Args:
            xs_pad: input tensor (B, L, D)
            ilens: input length (B)
            prev_states: Not to be used now.
        Returns:
            position embedded tensor and mask
        """
        masks = (~make_pad_mask(ilens)[:, None, :]).to(xs_pad.device)
        xs_pad = self.embed(xs_pad)
        xs_pad, masks = self.encoders(xs_pad, masks)
        if self.normalize_before:
            xs_pad = self.after_norm(xs_pad)
        olens = masks.squeeze(1).sum(1)
        return xs_pad, olens, None

# Copyright 2021 UMD (<NAME>)

# Transformer encoder (without the decoder part), a wrapper for the multi-headed attention layers
# Modified from ESPnet so that it can be used without espnet dependency

from typing import Optional
from typing import Tuple
from abc import ABC
import torch
from qanta.embedding import PositionalEncoding
from qanta.attention import LayerNorm, PositionwiseFeedForward, EncoderLayer, MultiHeadedAttention, make_pad_mask


class MultiSequential(torch.nn.Sequential):
    """Multi-input multi-output torch.nn.Sequential."""

    def forward(self, *args):
        """Repeat."""
        for m in self:
            args = m(*args)
        return args


def repeat(N, fn):
    """Repeat module N times.

    Args:
        N (int): Number of repeat time.
        fn (Callable): Function to generate module.

    Returns:
        MultiSequential: Repeated model instance.

    """
    return MultiSequential(*[fn(n) for n in range(N)])


class TransformerEncoder(torch.nn.Module):
    """Transformer encoder

    Args:
        input_size: input dim
        output_size: dimension of attention
        attention_heads: the number of heads of multi head attention
        linear_units: the number of units of position-wise feed forward
        num_blocks: the number of decoder blocks
        dropout_rate: dropout rate
        attention_dropout_rate: dropout rate in attention
        positional_dropout_rate: dropout rate after adding positional encoding
        pos_enc_class: PositionalEncoding
        normalize_before: whether to use layer_norm before the first block
        padding_idx: padding_idx for input_layer
    """

    def __init__(
        self,
        input_size: int,
        output_size: int = 256,
        attention_heads: int = 4,
        linear_units: int = 2048,
        num_blocks: int = 6,
        dropout_rate: float = 0.1,
        positional_dropout_rate: float = 0.1,
        attention_dropout_rate: float = 0.0,
        pos_enc_class: PositionalEncoding = PositionalEncoding,
        normalize_before: bool = True,
        padding_idx: int = -1,
    ):
        super().__init__()
        self._output_size = output_size

        self.embed = torch.nn.Sequential(
            torch.nn.Linear(input_size, output_size),
            torch.nn.LayerNorm(output_size),
            torch.nn.Dropout(dropout_rate),
            torch.nn.ReLU(),
            pos_enc_class(output_size, positional_dropout_rate),
        )

        self.normalize_before = normalize_before

        positionwise_layer = PositionwiseFeedForward
        positionwise_layer_args = (
            output_size,
            linear_units,
            dropout_rate,
        )

        self.encoders = repeat(
            num_blocks,
            lambda lnum: EncoderLayer(
                output_size,
                MultiHeadedAttention(
                    attention_heads, output_size, attention_dropout_rate
                ),
                positionwise_layer(*positionwise_layer_args),
                dropout_rate,
                normalize_before,
                False,
            ),
        )
        if self.normalize_before:
            self.after_norm = LayerNorm(output_size)

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

    def forward(
        self,
        xs_pad: torch.Tensor,
        ilens: torch.Tensor,
        prev_states: torch.Tensor = None,
    ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
        """Embed positions in tensor.

        Args:
            xs_pad: input tensor (B, L, D)
            ilens: input length (B)
            prev_states: Not to be used now.
        Returns:
            position embedded tensor and mask
        """
        masks = (~make_pad_mask(ilens)[:, None, :]).to(xs_pad.device)
        xs_pad = self.embed(xs_pad)
        xs_pad, masks = self.encoders(xs_pad, masks)
        if self.normalize_before:
            xs_pad = self.after_norm(xs_pad)
        olens = masks.squeeze(1).sum(1)
        return xs_pad, olens, None

# Copyright 2021 UMD (<NAME>) # Transformer encoder (without the decoder part), a wrapper for the multi-headed attention layers # Modified from ESPnet so that it can be used without espnet dependency Multi-input multi-output torch.nn.Sequential. Repeat. Repeat module N times.      Args:         N (int): Number of repeat time.         fn (Callable): Function to generate module.      Returns:         MultiSequential: Repeated model instance. Transformer encoder      Args:         input_size: input dim         output_size: dimension of attention         attention_heads: the number of heads of multi head attention         linear_units: the number of units of position-wise feed forward         num_blocks: the number of decoder blocks         dropout_rate: dropout rate         attention_dropout_rate: dropout rate in attention         positional_dropout_rate: dropout rate after adding positional encoding         pos_enc_class: PositionalEncoding         normalize_before: whether to use layer_norm before the first block         padding_idx: padding_idx for input_layer Embed positions in tensor.          Args:             xs_pad: input tensor (B, L, D)             ilens: input length (B)             prev_states: Not to be used now.         Returns:             position embedded tensor and mask