# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/3C. S2A acoustic tokens preparation.ipynb.

# %% auto 0
__all__ = []

# %% ../nbs/3C. S2A acoustic tokens preparation.ipynb 2
import sys
import os
import itertools
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F

from fastprogress import progress_bar
from fastcore.script import *

from . import utils, vad_merge, extract_acoustic
import webdataset as wds

# %% ../nbs/3C. S2A acoustic tokens preparation.ipynb 4
@call_parse
def prepare_atoks(
    input:str,  # FLAC webdataset file path (or - to read the names from stdin)
    n_samples:int=None, # process a limited amount of samples
    batch_size:int=1, # process several segments at once
    bandwidth:float=3,
):
    amodel = extract_acoustic.load_model()
    amodel.set_target_bandwidth(bandwidth)

    total = n_samples//batch_size if n_samples else 'noinfer'
    if n_samples: print(f"Benchmarking run of {n_samples} samples ({total} batches)")

    ds = vad_merge.chunked_audio_dataset([input], 'maxvad').compose(
        utils.resampler(24000, 'samples_24k'),
        wds.to_tuple('__key__', 'rpad_s', 'samples_24k'),
        wds.batched(64),
    )

    dl = wds.WebLoader(ds, num_workers=1, batch_size=None).unbatched().batched(batch_size)

    with utils.AtomicTarWriter(utils.derived_name(input, f'atoks-{bandwidth}kbps', dir="."), throwaway=n_samples is not None) as sink:
        for keys, rpad_ss, samples in progress_bar(dl, total=total):
            csamples = samples.cuda().unsqueeze(1)
            atokss = amodel.encode(csamples)[0][0]
            atokss = atokss.cpu().numpy().astype(np.int16)
            for key, rpad_s, atoks in zip(keys, rpad_ss, atokss):
                atoks = atoks[:,:int((30-rpad_s) * 75 + 0.5)]
                sink.write({
                    "__key__": key,
                    "atoks.npy": atoks,
                })
