# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/4A. S2A dataset preparation-Copy1.ipynb.

# %% auto 0
__all__ = ['derived_name']

# %% ../nbs/4A. S2A dataset preparation-Copy1.ipynb 2
import sys
import os
import itertools
from pathlib import Path

import numpy as np
import torch
import torchaudio
import torch.nn.functional as F
from torch.profiler import profile, record_function, ProfilerActivity

from fastprogress import progress_bar
from fastcore.script import *

import whisper
from . import vad, wh_transcribe, vq_stoks
import webdataset as wds

# %% ../nbs/4A. S2A dataset preparation-Copy1.ipynb 5
def derived_name(input, kind, base="audio"):
    return input.replace(base, kind) + ".gz"

# %% ../nbs/4A. S2A dataset preparation-Copy1.ipynb 6
@call_parse
def prepare_mvad(
    input:str,  # FLAC webdataset file path (or - to read the names from stdin)
    output:str=None, # output file name
):    
    if output is None: output = derived_name(input, 'mvad')
            
    ds = wds.WebDataset([derived_name(input, 'vad')]).decode()
    dl = wds.WebLoader(ds, num_workers=min(4, len(input)), batch_size=None)

    tmp = output+".tmp"
    with wds.TarWriter(tmp) as sink:
        for s in progress_bar(dl, total='noinfer'):
            s["maxvad.npy"] = np.array(wh_transcribe.chunk_merger(s['vad.npy']))
            s["eqvad.npy"] = np.array(wh_transcribe.chunk_merger(s['vad.npy'], wh_transcribe.random_cutter))
            del s['vad.npy']
            sink.write(s)
    os.rename(tmp, output)
