import torch
import numpy as np
from .utils import record_user


class VoiceActivityDetection:
    def __init__(self, sampling_rate=16000):
        self.model, utils = torch.hub.load(repo_or_dir='snakers4/silero-vad',
                                           model='silero_vad',
                                           force_reload=False)

        (self.get_speech_timestamps,
         self.save_audio,
         self.read_audio,
         self.VADIterator,
         self.collect_chunks) = utils

        self.sampling_rate = sampling_rate

    def contains_speech(self, audio):
        frames = np.frombuffer(b''.join(audio), dtype=np.int16)

        # normalization see https://discuss.pytorch.org/t/torchaudio-load-normalization-question/71470
        frames = frames / (1 << 15)

        audio = torch.tensor(frames.astype(np.float32))
        speech_timestamps = self.get_speech_timestamps(audio, self.model,
                                                       sampling_rate=self.sampling_rate)
        return len(speech_timestamps) > 0


if __name__ == "__main__":
    vad = VoiceActivityDetection()
    audio = record_user(3, vad)
