# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/D. Common dataset utilities.ipynb.

# %% auto 0
__all__ = ['shard_glob', 'join_datasets', 'resampler', 'derived_name', 'derived_dataset', 'merge_in', 'AtomicTarWriter',
           'readlines']

# %% ../nbs/D. Common dataset utilities.ipynb 1
import os
import torch
import torchaudio
import torch.nn.functional as F
from pathlib import Path
from contextlib import contextmanager
import webdataset as wds

# %% ../nbs/D. Common dataset utilities.ipynb 2
import huggingface_hub

# adds a environment variable that disables downloads thought the Huggingface Hub API:

def wrap_downloader(old):
    def new(*args, **kwargs):
        if os.environ.get('HUGGINGFACE_LOCAL_ONLY', False):
            print(f"Enforcing local_files_only for {old.__qualname__}")
            kwargs['local_files_only'] = True
        return old(*args, **kwargs)
    return new

huggingface_hub.snapshot_download = wrap_downloader(huggingface_hub.snapshot_download)
huggingface_hub.hf_hub_download = wrap_downloader(huggingface_hub.hf_hub_download)

# %% ../nbs/D. Common dataset utilities.ipynb 3
def shard_glob(input):
    if isinstance(input, Path):
        input = str(input)
    if '{' in input:
        return wds.shardlists.expand_urls(input)
    if str:
        path = Path(input)
        if path.is_dir():
            glob = '*.tar.gz'
        else:
            glob = path.name
            path = path.parent
        input = Path(path).glob(glob)
    else:
        raise ArgumentError("input should be either a list or a path with an optional glob specifier")
    return [str(x) for x in input]

# %% ../nbs/D. Common dataset utilities.ipynb 7
class join_datasets(torch.utils.data.IterableDataset):
    def __init__(self, datasets):
        self.datasets = datasets
        self.iters = [iter(ds) for ds in self.datasets]
        
    def __iter__(self):
        probs = torch.tensor([getattr(ds, 'weight', 1) for ds in self.datasets], dtype=torch.float)
        while True:
            try:
                yield next(self.iters[torch.multinomial(probs, 1)])
            except StopIteration:
                return
    
    def __len__(self):
        return sum([ds.total_samples for ds in self.datasets])

# %% ../nbs/D. Common dataset utilities.ipynb 10
def resampler(newsr = 24000, key = 'samples_24k'):
    _last_sr = None
    tform = None
    
    def _resample(samples):
        for s in samples:
            sr = s['sample_rate']
            if sr != newsr:
                if sr != _last_sr: tform = torchaudio.transforms.Resample(sr, newsr)
                s[key] = tform(s['samples'])
            else:
                s[key] = s['samples']
            yield s
    
    return _resample

# %% ../nbs/D. Common dataset utilities.ipynb 11
def derived_name(url, kind, suffix=None):
    if suffix is None: suffix = '' if url.endswith('.gz') else ".gz"
    url = Path(url)
    return str(url.parent.parent/kind/url.name) + suffix

# %% ../nbs/D. Common dataset utilities.ipynb 12
def derived_dataset(kind, suffix=None, decoders=[]):
    def deriver(url):
        return wds.WebDataset(
            wds.SimpleShardList([derived_name(url, kind, suffix)])
        ).decode(*decoders)
    return deriver

# %% ../nbs/D. Common dataset utilities.ipynb 13
def merge_in(dataset_fun):
    """Merge a dataset into the current one returning samples with the union of keys. Pass in a function
    that takes a URL of a sample and returns a dataset for it (called everytime the URL changes).
    
    It requires (and validates) that both datasets have the same ordering of keys so you have
    to use it before any sample shuffling. Shard shuffling is ok.
    """
    def merge_loop(main_samples):
        #print("new merge loop:", dataset_fun)
        merged_samples = None
        cur_url = None
        i = None
        for s in main_samples:
            url = s['__url__']
            if url != cur_url:
                # this will open a new file when we get the first sample with a new __url__
                merged_samples = iter(dataset_fun(url))
                cur_url = url
            news = {}
            news.update(s)
            if '__skip_merge__' not in s:
                try:
                    merge_s = next(merged_samples)
                except StopIteration:
                    # if the original shard got repeated we won't observe a __url__ change
                    # in this case restart the dataset from the beginning
                    merged_samples = iter(dataset_fun(url))
                    merge_s = next(merged_samples)
                assert merge_s['__key__'] == s['__key__'], f"sample keys don't match: {merge_s['__key__']}, {s['__key__']} in file {s['__url__']}"
                news.update(merge_s)
            yield news
    return merge_loop

# %% ../nbs/D. Common dataset utilities.ipynb 14
def split_to_chunks(stream, ikey='vad.npy', copy_keys=[], split_keys=[], pad_to_seconds=30, random_shift=False):
    for s in stream:
        audio, sr = s['audio']
        chunks = s[ikey]
        imax = len(chunks) - 1
        for i,(ts,te) in enumerate(chunks):
            if 'mask.npy' in s and not s['mask.npy'][i]:
                # used for fishing out samples in validation sets, see also "3D. Split out validation"
                continue
            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))
            subs = {"__key__": s['__key__'] + f"_{i:03d}",
                    "src_key": s['__key__'],
                    "__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,
                    "src_sample": s}
            for k in copy_keys:
                subs[k] = s[k]
            for k in split_keys:
                subs[k] = s[k][i]
            yield subs

# %% ../nbs/D. Common dataset utilities.ipynb 15
import re
import tempfile

# %% ../nbs/D. Common dataset utilities.ipynb 16
# a patch to ignore invalid utf-8 metadata
import torio.io._streaming_media_decoder
def new_parse_si(i):
    media_type = i.media_type
    try:
        metadata = i.metadata
    except UnicodeDecodeError:
        metadata = {}
    if media_type == "audio":
        return torio.io._streaming_media_decoder.SourceAudioStream(
            media_type=i.media_type,
            codec=i.codec_name,
            codec_long_name=i.codec_long_name,
            format=i.format,
            bit_rate=i.bit_rate,
            num_frames=i.num_frames,
            bits_per_sample=i.bits_per_sample,
            metadata=metadata,
            sample_rate=i.sample_rate,
            num_channels=i.num_channels,
        )
    if media_type == "video":
        return torio.io._streaming_media_decoder.SourceVideoStream(
            media_type=i.media_type,
            codec=i.codec_name,
            codec_long_name=i.codec_long_name,
            format=i.format,
            bit_rate=i.bit_rate,
            num_frames=i.num_frames,
            bits_per_sample=i.bits_per_sample,
            metadata=metadata,
            width=i.width,
            height=i.height,
            frame_rate=i.frame_rate,
        )
    return torio.io._streaming_media_decoder.SourceStream(
        media_type=i.media_type,
        codec=i.codec_name,
        codec_long_name=i.codec_long_name,
        format=None,
        bit_rate=None,
        num_frames=None,
        bits_per_sample=None,
        metadata=metadata,
    )
torio.io._streaming_media_decoder._parse_si = new_parse_si

def torch_audio_opus(key, data):
    """Decode audio using the torchaudio library.

    :param key: file name extension
    :param data: data to be decoded
    """
    extension = re.sub(r".*[.]", "", key)
    if extension not in ["flac", "mp3", "sox", "wav", "m4a", "ogg", "wma", "opus"]:
        return None

    import torchaudio

    with tempfile.TemporaryDirectory() as dirname:
        fname = os.path.join(dirname, f"file.{extension}")
        with open(fname, "wb") as stream:
            stream.write(data)
        return torchaudio.load(fname, backend='soundfile' if extension == "mp3" else None)

# %% ../nbs/D. Common dataset utilities.ipynb 17
def find_audio(stream, okey='audio', ikeys='flac;mp3;sox;wav;m4a;ogg;wma;opus'):
    ikeys = ikeys.split(';')
    for s in stream:
        for ikey in ikeys:
            if ikey in s:
                s[okey] = s[ikey]
                yield s
                break
            # implicitly skips elements without any audio

# %% ../nbs/D. Common dataset utilities.ipynb 18
def vad_dataset(shards, ikey='vad.npy', kind='vad'):
    return wds.WebDataset(shards).compose(
        wds.decode(torch_audio_opus),
        find_audio,
        merge_in(derived_dataset(kind)),
        lambda x: split_to_chunks(x, ikey=ikey),
    )

# %% ../nbs/D. Common dataset utilities.ipynb 19
@contextmanager
def AtomicTarWriter(name, throwaway=False):
    Path(name).parent.mkdir(exist_ok=True, parents=True)
    tmp = name+".tmp"
    with wds.TarWriter(tmp, compress=name.endswith('gz')) as sink:
        yield sink
    if not throwaway:
        os.rename(tmp, name)

# %% ../nbs/D. Common dataset utilities.ipynb 20
def readlines(fname):
    with open(fname) as file:
        return [line.rstrip() for line in file]
