import torch
from speechbrain.inference.interfaces import Pretrained


class CustomEncoderBestRQ(Pretrained):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def encode_batch(self, wavs, wav_lens=None, normalize=False):
        # Manage single waveforms in input
        if len(wavs.shape) == 1:
            wavs = wavs.unsqueeze(0)

        # Assign full length if wav_lens is not assigned
        if wav_lens is None:
            wav_lens = torch.ones(wavs.shape[0], device=self.device)

        # Storing waveform in the specified device
        wavs, wav_lens = wavs.to(self.device), wav_lens.to(self.device)
        wavs = wavs.float()

        feats = self.hparams.compute_features(wavs)
        feats = self.mods.normalizer(feats, wav_lens)
        src = self.mods.extractor(feats)
        enc_out = self.mods.encoder(src, wav_lens)
        return enc_out

    def encode_file(self, path, normalize=False):
        waveform = self.load_audio(path)
        # Fake a batch:
        batch = waveform.unsqueeze(0)
        rel_length = torch.tensor([1.0])
        outputs = self.encode_batch(batch, rel_length)
        return outputs

    def forward(self, wavs, wav_lens=None, normalize=False):
        return self.encode_batch(wavs=wavs, wav_lens=wav_lens, normalize=normalize)