"""LinearInput block for Transducer encoder."""

from typing import Optional, Tuple, Union

import torch

class LinearInput(torch.nn.Module):
    """ConvInput module definition.

    Args:
        input_size: Input size.
        conv_size: Convolution size.
        subsampling_factor: Subsampling factor.
        vgg_like: Whether to use a VGG-like network.
        output_size: Block output dimension.

    """

    def __init__(
        self,
        input_size: int,
        output_size: Optional[int] = None,
        subsampling_factor: int = 1,
    ) -> None:
        """Construct a ConvInput object."""
        super().__init__()
        self.embed = torch.nn.Sequential(
            torch.nn.Linear(input_size, output_size),
            torch.nn.LayerNorm(output_size),
            torch.nn.Dropout(0.1),
        )
        self.subsampling_factor = subsampling_factor
        self.min_frame_length = 1

    def forward(
        self, x: torch.Tensor, mask: Optional[torch.Tensor]
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        
        x = self.embed(x)
        return x, mask

    def get_size_before_subsampling(self, size: int) -> int:
        """Return the original size before subsampling for a given size.

        Args:
            size: Number of frames after subsampling.

        Returns:
            : Number of frames before subsampling.

        """
        return size