# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/Split and shuffle quantization webdataset.ipynb.

# %% auto 0
__all__ = []

# %% ../nbs/Split and shuffle quantization webdataset.ipynb 1
import webdataset as wds
import pylab as plt
from pathlib import Path
import shutil
from fastprogress import progress_bar
from fastcore.script import call_parse
import numpy as np
import random
from collections import Counter
from whisperspeech import vad, vq_stoks, wh_transcribe

# %% ../nbs/Split and shuffle quantization webdataset.ipynb 2
@call_parse
def split_dataset(
    input:str,
    preproc_dir:str,
    dest_folder:Path,
    txt_tag="txt-medium",
    val_size=512,
    train_shards=400,
    drop_first_and_last=False,
    fix_dots_in_names=False,
):
    dest = Path(dest_folder)
    if dest.exists(): shutil.rmtree(dest)
    dest.mkdir(parents=True)

    if isinstance(input, (Path, str)):
        path = Path(input)
        if path.is_dir():
            glob = '*.tar.gz'
        else:
            glob = path.name
            path = path.parent
        input = Path(path).glob(glob)
    elif isinstance(input, list):
        pass
    else:
        raise ArgumentError("input should be either a list or a path with an optional glob specifier")
    shards = sorted([str(x) for x in input])

    def load(shards):
        ds = wds.WebDataset(shards, shardshuffle=True, rename_files=vad.fix_dots_in_names if fix_dots_in_names else None).compose(
            wds.decode(wds.torch_audio),
            wds.select(lambda x: 'wav' in x or 'flac' in x or 'mp3' in x or 'ogg' in x), # skip samples without audio
            wds.rename(audio="flac;mp3;wav;ogg"),
            vq_stoks.merge_in(vq_stoks.derived_dataset(preproc_dir, 'vad')),
            wds.map_dict(**{"vad.npy":wh_transcribe.chunk_merger}),
            wh_transcribe.split_to_chunks,
            vq_stoks.merge_in(vq_stoks.derived_dataset(preproc_dir, txt_tag)),
        )
        if drop_first_and_last:
            ds = ds.compose(
                # drop the first and last segment because they tend to be inaccurate
                # (the transcriptions don't have the "LibriVox" headers and "end of chapter" suffixes)
                wds.select(lambda x: x['i'] != 0 and x['i'] != x['imax']),
            )
        return ds
    
    ds = load([shards[0]])
    N = len([True for _ in ds]) * len(shards)
    
    ds = load(shards)
    dl = wds.WebLoader(ds, num_workers=32, batch_size=None).shuffle(200000)

    print(f"Writing the vq_stoks datasets ({N} samples):")
    with wds.TarWriter(str(Path(dest_folder)/"val.tar.gz")) as val_sink:
        with wds.ShardWriter(str(Path(dest_folder)/"train-%06d.tar.gz"), maxcount=N//(train_shards+1)) as train_sink:
            for i, s in enumerate(progress_bar(dl, total='noinfer')):
                s['samples'] = s['samples'].numpy()
                s = {k:s[k.split('.')[0]] for k in ['__key__', 'tend.pyd', 'tstart.pyd', 'samples.npy', 'txt']}
                if i < val_size:
                    val_sink.write(s)
                else:
                    train_sink.write(s)
