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

# %% auto 0
__all__ = []

# %% ../nbs/Split and shuffle 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

# %% ../nbs/Split and shuffle webdataset.ipynb 9
@call_parse
def split_dataset(
    input:str,
    dest_folder:Path,
    val_size=512,
    train_shards=400
):
    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 = [str(x) for x in input]

    ds = wds.WebDataset([shards[0]])
    N = len([True for _ in ds]) * len(shards)
    
    ds = wds.WebDataset(shards, shardshuffle=True)
    dl = wds.WebLoader(ds, num_workers=32, batch_size=None).shuffle(200000)

    print(f"Writing the 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')):
                if i < val_size:
                    val_sink.write(s)
                else:
                    train_sink.write(s)
