# 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 *

import webdataset as wds
from . import utils, vad_merge
from .inference import get_compute_device

# %% ../nbs/3C. S2A acoustic tokens preparation.ipynb 4
def load_model():
    "Load the pretrained EnCodec model"
    from encodec.model import EncodecModel
    model = EncodecModel.encodec_model_24khz()
    model.set_target_bandwidth(1.5)
    model.to(get_compute_device()).eval()
    return model

# %% ../nbs/3C. S2A acoustic tokens preparation.ipynb 5
@call_parse
def prepare_atoks(
    input:str,  # audio file webdataset file path
    output:str, # output shard path
    n_samples:int=None, # process a limited amount of samples
    batch_size:int=4, # process several segments at once
    bandwidth:float=3,
):
    device = get_compute_device()
    amodel = load_model().to(device)  # Move model to computed device
    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)")

    if total == 'noinfer':
        import math, time
        start = time.time()
        ds = wds.WebDataset([utils.derived_name(input, 'mvad')]).decode()
        total = math.ceil(sum([len(x['max.spk_emb.npy']) for x in ds])/batch_size)
        print(f"Counting {total} batches: {time.time()-start:.2f}")

    ds = vad_merge.chunked_audio_dataset([input], 'max').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(output, throwaway=n_samples is not None) as sink:
        for keys, rpad_ss, samples in progress_bar(dl, total=total):
            csamples = samples.to(device).unsqueeze(1)  # Move tensors to computed device
            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,
                })
