from torch.nn import Module, Linear, functional, GELU, Conv1d, GroupNorm, ModuleList
from torch import Tensor
from torchaudio.models import Conformer

class WaveformFilter(Module):
    def __init__(self):
        super().__init__()
        
        self.filters = ModuleList([
            Conv1d(in_channels=1, out_channels=512, kernel_size=10, stride=5),
            GroupNorm(num_channels=512, num_groups=32),
            GELU(),
            Conv1d(in_channels=512, out_channels=512, kernel_size=3, stride=2),
            GroupNorm(num_channels=512, num_groups=32),
            GELU(),
            Conv1d(in_channels=512, out_channels=512, kernel_size=3, stride=2),
            GroupNorm(num_channels=512, num_groups=32),
            GELU(),
            Conv1d(in_channels=512, out_channels=512, kernel_size=3, stride=2),
            GroupNorm(num_channels=512, num_groups=32),
            GELU(),
            Conv1d(in_channels=512, out_channels=512, kernel_size=3, stride=2),
            GroupNorm(num_channels=512, num_groups=32),
            GELU(),
            Conv1d(in_channels=512, out_channels=512, kernel_size=2, stride=2),
            GroupNorm(num_channels=512, num_groups=32),
            GELU(),
            Conv1d(in_channels=512, out_channels=512, kernel_size=2, stride=2),
            GroupNorm(num_channels=512, num_groups=32),
            GELU(),
        ])
    
    def compute_lengths(self, waveform_lengths: Tensor) -> Tensor:
        """Accurately compute output lengths after all convolutions."""
        lengths = waveform_lengths
        for module in self.filters:
            if isinstance(module, Conv1d):
                # Conv1d with padding=0, dilation=1:
                # output = floor((input - kernel_size) / stride) + 1
                lengths = (lengths - module.kernel_size[0]) // module.stride[0] + 1
        return lengths
    
    def forward(self, x: Tensor) -> Tensor:
        for filter in self.filters:
            x = filter(x)
        
        return x


class ASRModel(Module):
    def __init__(self, vocab_size: int, input_dim: int = 256, num_heads: int = 8, ffn_dim: int = 2048, num_layers: int = 6, dropout: float = 0.1) -> None:
        super().__init__()

        self.feature_extractor = WaveformFilter()
        self.proj = Linear(in_features=512, out_features=input_dim, bias=False)

        self.encoder = Conformer(input_dim=input_dim, num_heads=num_heads, ffn_dim=ffn_dim, num_layers=num_layers, depthwise_conv_kernel_size=31, dropout=dropout)

        self.ctc_head = Linear(in_features=input_dim, out_features=vocab_size, bias=False)

    def forward(self, waveforms: Tensor, waveform_lengths: Tensor) -> tuple[Tensor, Tensor]:
        assert len(waveform_lengths.shape) == 1, "The waveform_lengths tensor must be shape of [B] tensor."
        assert len(waveforms.shape) == 2, "The waveform tensor must be [B, time] tensor"
        # waveforms: [B, time]
        x: Tensor = waveforms.unsqueeze(1)  # [B, 1, time]
        #x: [B, 1, time]
        x: Tensor = self.feature_extractor(x)
        #x: [B, 512, time]
        x = x.permute(0, 2, 1)
        #x: [B, time, 512]
        x = self.proj(x)
        #x: [B, time, 256]
        # lengths = torch.tensor([time] * batch, dtype=torch.long, device=x.device)
        lengths = self.feature_extractor.compute_lengths(waveform_lengths)

        x, lengths = self.encoder(x, lengths)

        logits = self.ctc_head(x) # [B, time, vocab_size]
        log_probs = functional.log_softmax(logits, dim=-1)

        return log_probs, lengths
    
    def get_num_params(self) -> int:
        return sum(p.numel() for p in self.parameters())