"""Conv1d block for Transducer encoder."""

from typing import Optional, Tuple, Union

import torch


class Conv1d(torch.nn.Module):
    """Conv1d module definition.

    Args:
        input_size: Input dimension.
        output_size: Output dimension.
        kernel_size: Size of the convolving kernel.
        stride: Stride of the convolution.
        dilation: Spacing between the kernel points.
        groups: Number of blocked connections from input channels to output channels.
        bias: Whether to add a learnable bias to the output.
        batch_norm: Whether to use batch normalization after convolution.
        relu: Whether to use a ReLU activation after convolution.
        causal: Whether to use causal convolution (set to True if streaming).
        dropout_rate: Dropout rate.

    """

    def __init__(
        self,
        input_size: int,
        output_size: int,
        kernel_size: Union[int, Tuple],
        stride: Union[int, Tuple] = 1,
        dilation: Union[int, Tuple] = 1,
        groups: Union[int, Tuple] = 1,
        bias: bool = True,
        batch_norm: bool = False,
        relu: bool = True,
        causal: bool = False,
        dropout_rate: float = 0.0,
    ) -> None:
        """Construct a Conv1d object."""
        super().__init__()

        if causal:
            self.lorder = kernel_size - 1
            stride = 1
        else:
            self.lorder = 0
            stride = stride

        self.conv = torch.nn.Conv1d(
            input_size,
            output_size,
            kernel_size,
            stride=stride,
            dilation=dilation,
            groups=groups,
            bias=bias,
        )

        self.dropout = torch.nn.Dropout(p=dropout_rate)

        if relu:
            self.relu_func = torch.nn.ReLU()

        if batch_norm:
            self.bn = torch.nn.BatchNorm1d(output_size)

        self.out_pos = torch.nn.Linear(input_size, output_size)

        self.input_size = input_size
        self.output_size = output_size

        self.relu = relu
        self.batch_norm = batch_norm
        self.causal = causal

        self.kernel_size = kernel_size
        self.padding = dilation * (kernel_size - 1)
        self.stride = stride

        self.cache = None

    def reset_streaming_cache(self, left_context: int, device: torch.device) -> None:
        """Initialize/Reset Conv1d cache for streaming.

        Args:
            left_context: Number of left frames during chunk-by-chunk inference.
            device: Device to use for cache tensor.

        """
        self.cache = torch.zeros(
            (1, self.input_size, self.kernel_size - 1), device=device
        )

    def forward(
        self,
        x: torch.Tensor,
        pos_enc: torch.Tensor,
        mask: Optional[torch.Tensor] = None,
        chunk_mask: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Encode input sequences.

        Args:
            x: Conv1d input sequences. (B, T, D_in)
            pos_enc: Positional embedding sequences. (B, 2 * (T - 1), D_in)
            mask: Source mask. (B, T)
            chunk_mask: Chunk mask. (T_2, T_2)

        Returns:
            x: Conv1d output sequences. (B, sub(T), D_out)
            mask: Source mask. (B, T) or (B, sub(T))
            pos_enc: Positional embedding sequences.
                       (B, 2 * (T - 1), D_att) or (B, 2 * (sub(T) - 1), D_out)

        """
        x = x.transpose(1, 2)

        if self.lorder > 0:
            x = torch.nn.functional.pad(x, (self.lorder, 0), "constant", 0.0)
        else:
            mask = self.create_new_mask(mask)
            pos_enc = self.create_new_pos_enc(pos_enc)

        x = self.conv(x)

        if self.batch_norm:
            x = self.bn(x)

        x = self.dropout(x)

        if self.relu:
            x = self.relu_func(x)

        x = x.transpose(1, 2)

        return x, mask, self.out_pos(pos_enc)

    def chunk_forward(
        self,
        x: torch.Tensor,
        pos_enc: torch.Tensor,
        mask: torch.Tensor,
        left_context: int = 0,
        right_context: int = 0,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Encode chunk of input sequence.

        Args:
            x: Conv1d input sequences. (B, T, D_in)
            pos_enc: Positional embedding sequences. (B, 2 * (T - 1), D_in)
            mask: Source mask. (B, T)
            left_context: Number of frames in left context.
            right_context: Number of frames in right context.

        Returns:
            x: Conv1d output sequences. (B, T, D_out)
            pos_enc: Positional embedding sequences. (B, 2 * (T - 1), D_out)

        """
        x = torch.cat([self.cache, x.transpose(1, 2)], dim=2)

        if right_context > 0:
            self.cache = x[:, :, -(self.lorder + right_context) : -right_context]
        else:
            self.cache = x[:, :, -self.lorder :]

        x = self.conv(x)

        if self.batch_norm:
            x = self.bn(x)

        x = self.dropout(x)

        if self.relu:
            x = self.relu_func(x)

        x = x.transpose(1, 2)

        return x, self.out_pos(pos_enc)

    def create_new_mask(self, mask: torch.Tensor) -> torch.Tensor:
        """Create new mask for output sequences.

        Args:
            mask: Mask of input sequences. (B, T)

        Returns:
            mask: Mask of output sequences. (B, sub(T))

        """
        if self.padding != 0:
            mask = mask[:, : -self.padding]

        return mask[:, :: self.stride]

    def create_new_pos_enc(self, pos_enc: torch.Tensor) -> torch.Tensor:
        """Create new positional embedding vector.

        Args:
            pos_enc: Input sequences positional embedding.
                     (B, 2 * (T - 1), D_in)

        Returns:
            pos_enc: Output sequences positional embedding.
                     (B, 2 * (sub(T) - 1), D_in)

        """
        pos_enc_positive = pos_enc[:, : pos_enc.size(1) // 2 + 1, :]
        pos_enc_negative = pos_enc[:, pos_enc.size(1) // 2 :, :]

        if self.padding != 0:
            pos_enc_positive = pos_enc_positive[:, : -self.padding, :]
            pos_enc_negative = pos_enc_negative[:, : -self.padding, :]

        pos_enc_positive = pos_enc_positive[:, :: self.stride, :]
        pos_enc_negative = pos_enc_negative[:, :: self.stride, :]

        pos_enc = torch.cat([pos_enc_positive, pos_enc_negative[:, 1:, :]], dim=1)

        return pos_enc