# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/3D. Split out validation.ipynb.

# %% auto 0
__all__ = []

# %% ../nbs/3D. Split out validation.ipynb 1
import os
import webdataset as wds
from pathlib import Path
from fastprogress import progress_bar
from fastcore.script import call_parse
import numpy as np
import random
from collections import Counter
from whisperspeech import utils

# %% ../nbs/3D. Split out validation.ipynb 4
@call_parse
def split_dataset(
    shard_spec:str,
    splits:str,
):
    shards = utils.shard_glob(shard_spec)
    splits = splits.split()
    
    bufs = {k:[] for k in splits}
    outputs = {k:wds.TarWriter(str(Path(shard_spec).parent/(Path(k).name+".tar.gz"))) for k in splits}
    needles = {k:bufs[split] for split in splits for k in utils.readlines(split)}
    
#     with open(Path(shard_spec).parent/"validation-samples", "w") as f:
#         for k in needles.keys():
#             f.write(k+'\n')

    print(f"Generating splits: {' '.join(outputs.keys())}")
    print(f"Looking for {len(needles)} samples...")
    
    ds = wds.WebDataset(shards).compose(
        wds.select(lambda x: x['__key__'] in needles),
    )
        
    dl = wds.WebLoader(ds, num_workers=32, batch_size=None)

    for s in progress_bar(dl, total='noinfer'):
        needles[s['__key__']].append(s)
        del needles[s['__key__']]
        pass

    for split,buf in bufs.items():
        for s in sorted(buf, key=lambda x: x['__key__']):
            outputs[split].write(s)
    
    if len(needles) > 0:
        print(f"Missed {len(needles)} samples!")
        os.exit(1)
