from typing import Tuple
import torch
from torch import nn


class Wav2Vec2Processor(nn.Module):
    def __init__(self):
        """
        Convert tuple of waveforms whose length is different to a batch.

        Args:
            waveforms (Tuple[torch.Tensor]): The waveforms. Shape: (batch_size, wave_length).

        Returns:
            waveforms (torch.Tensor): The batched waveforms. Shape: (batch_size, max_wave_length).
            wave_lengths (torch.Tensor): The wave length of each waveform. Shape: (batch_size,).
        """
        super().__init__()

    def forward(self, waveforms: Tuple[torch.Tensor, ...]):
        device = waveforms[0].device
        wave_lengths = torch.tensor(
            tuple(waveform.size(0) for waveform in waveforms), device=device
        )

        max_length = wave_lengths.max().item()

        padded = []

        for waveform in waveforms:
            padded.append(
                nn.functional.pad(
                    waveform,
                    (0, max_length - waveform.size(0)),
                    mode="constant",
                    value=0.0,
                )
            )

        batched_waveforms = torch.stack(padded, dim=0)

        return batched_waveforms, wave_lengths