from torch.nn import Conv2d, Module, ReLU, Linear, functional
from torch import Tensor
from torchaudio.models import Conformer


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

        self.conv1 = Conv2d(in_channels=1, out_channels=16, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
        self.conv2 = Conv2d(in_channels=16, out_channels=32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
        self.relu = ReLU()

        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, mel_specs: Tensor, mel_lengths: Tensor) -> tuple[Tensor, Tensor]:
        # mel_specs: [B, n_mels, time]
        x: Tensor = mel_specs.unsqueeze(1)  # [B, 1, n_mels, time]

        x = self.relu(self.conv1(x))  # [batch, 16, n_mels/2, time/2]
        x = self.relu(self.conv2(x))  # [batch, 32, n_mels/4, time/4]

        # [B, channels, freq, time] → [B, time, channels*freq]
        batch, channels, freq, time = x.shape
        x = x.permute(0, 3, 1, 2).reshape(batch, time, channels * freq)
        # lengths = torch.tensor([time] * batch, dtype=torch.long, device=x.device)
        lengths = ((mel_lengths + 1) // 2 + 1) // 2  # 两层 stride=2

        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())