# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/1B. Voice activity detection.ipynb.

# %% auto 0
__all__ = []

# %% ../nbs/1B. Voice activity detection.ipynb 3
import os
import random
import torch
import torchaudio

from pathlib import Path
from fastprogress import progress_bar
from fastcore.script import call_parse

import numpy as np
import webdataset as wds
from huggingface_hub import hf_hub_download

import whisperx
from speechbrain.pretrained import EncoderClassifier
from clara import PLCLARA
import clara

# %% ../nbs/1B. Voice activity detection.ipynb 5
# some of the original file names have a dot in their name
# webdataset does not like it so let's patch it
def fix_dots_in_names(name):
    name, ext = name.rsplit('.', 1)
    return ".".join((name.replace('.', '_'), ext))

def load_dataset(url, decode=True, rename_files=None):
    ds = wds.WebDataset(url, rename_files=rename_files)
    if not decode: return ds
    return ds.decode(wds.torch_audio).compose(
        wds.rename(audio="flac;mp3;wav;ogg")
    )

# %% ../nbs/1B. Voice activity detection.ipynb 7
def extract_segments(vad_result, max_duration):
    binarize = whisperx.vad.Binarize(max_duration=max_duration)
    segments = binarize(vad_result)
    return [(x.start, x.end) for x in segments.get_timeline()]

def segment_audio(vad_model, audio, sr=16000):
    vad_result = vad_model({"waveform": audio, "sample_rate": sr})
    return extract_segments(vad_result, 30)

# %% ../nbs/1B. Voice activity detection.ipynb 25
def split_to_chunks(stream, pad_to_seconds=None):
    for s in stream:
        audio, sr = s['audio']
        imax = len(s['vad.npy']) - 1
        for i,(ts,te) in enumerate(s['vad.npy']):
            samples = audio[0,int(ts*sr):int(te*sr)]
            if pad_to_seconds is not None:
                padding = pad_to_seconds*sr-samples.shape[-1]
                lpad = random.randint(0, padding) if random_shift else 0
                samples = F.pad(samples, (lpad, padding-lpad))
            yield {"__key__": s['__key__'] + f"_{i:03d}",
                   "__url__": s['__url__'],
                   "i": i, "imax": imax,
                   "tstart": ts, "tend": te, "total_seconds": audio.shape[-1]/sr,
                   "lpad": lpad, "rpad": padding-lpad,
                   "lpad_s": lpad/sr, "rpad_s": (padding-lpad)/sr,
                   "samples": samples, "sample_rate": sr}

# %% ../nbs/1B. Voice activity detection.ipynb 26
def flac_to_vad_name(input, key='raw'):
    if '-flac-' in input:
        return input.rsplit("/", 1)[1].replace('flac', 'vad') + ".gz"
    else:
        return input.rsplit("/", 1)[1].replace(key, 'vad') + ".gz"

@call_parse
def process_shard(
    input:str,           # input shard URL/path
    output:str=None,     # output shard URL/path
    key:str='raw',       # string to replace with 'vad' in the shard name
):
    if output is None: output = flac_to_vad_name(input, key=key)
    
    rename_files=fix_dots_in_names if 'librilight' in input else None
    ds = load_dataset(input, rename_files=rename_files)
    dl = torch.utils.data.DataLoader(ds, num_workers=1, batch_size=None)
    vad_model = whisperx.vad.load_vad_model('cuda')
    
    tmp = output+".tmp"
    last_key = None
    with wds.TarWriter(tmp) as sink:
        for s in progress_bar(dl, total='noinfer'):
            audio, sr = s.get('audio', (None, None))
            if audio is None:
                print(f"warning: '{s['__key__']}' does not contain an audio file")
                continue
            sink.write({
                "__key__": s['__key__'],
                "vad.npy": np.array(segment_audio(vad_model, audio, sr=sr), dtype=np.float32),
            })
            last_key = s['__key__']
    os.rename(tmp, output)
