# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/3B. Semantic token extraction.ipynb.

# %% auto 0
__all__ = []

# %% ../nbs/3B. Semantic token extraction.ipynb 2
import sys
import os
from os.path import expanduser
import itertools
from pathlib import Path

import numpy as np
import torch
import torchaudio
import torch.nn.functional as F
from torch.profiler import profile, record_function, ProfilerActivity

from fastprogress import progress_bar
from fastcore.script import *

from speechbrain.pretrained import EncoderClassifier
from . import vq_stoks, utils, vad_merge
import webdataset as wds

from .inference import get_compute_device

# %% ../nbs/3B. Semantic token extraction.ipynb 7
@call_parse
def prepare_stoks(
    input:str,  # audio file webdataset file path
    output:str, # output shard path
    vq_model:str="collabora/spear-tts-pytorch:whisper-vq-stoks-v2.model", # the model path (use repo_id:filename to download it from hugginface)
    n_samples:int=None, # process a limited amount of samples
    batch_size:int=64, # process several segments at once
    kind:str="max", # could be eqvad to get more uniform chunk lengths
    
):
    device = get_compute_device()
    vq_model = vq_stoks.RQBottleneckTransformer.load_model(vq_model).to(device)
    vq_model.ensure_whisper()
    
    spk_classifier = EncoderClassifier.from_hparams("speechbrain/spkrec-ecapa-voxceleb",
                                                    savedir=expanduser("~/.cache/speechbrain/"),
                                                    run_opts = {"device": device})
    
    total = n_samples//batch_size if n_samples else 'noinfer'

    if total == 'noinfer':
        import math, time
        start = time.time()
        ds = wds.WebDataset([utils.derived_name(input, 'mvad')]).decode()
        total = math.ceil(sum([len(x[f'{kind}.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], kind).compose(
        utils.resampler(16000, 'samples_16k'),
        wds.to_tuple('__key__', 'rpad_s', 'samples_16k'),
        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, samples16k in progress_bar(dl, total=total):
            with torch.no_grad():
                samples16k = samples16k.to(device).to(torch.float16)
                stoks = vq_model.encode_audio(samples16k).cpu().numpy().astype(np.int16)
                spk_embs = spk_classifier.encode_batch(
                   samples16k, wav_lens=torch.tensor(30 - rpad_ss, dtype=torch.float)/30)[:,0,:].cpu().numpy()
            for key, rpad_s, _stoks, spk_emb in zip(keys, rpad_ss, stoks, spk_embs):
                _stoks = _stoks[:int((30-rpad_s) * 25 + .5)]
                s = {
                    "__key__": key,
                    "stoks.npy": _stoks,
                }
                if spk_emb is not None: s["spk_emb.npy"] = spk_emb
                sink.write(s)
        sys.stdout.write("\n")
