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

# %% ../nbs/3B. Semantic token extraction.ipynb 7
@call_parse
def prepare_stoks(
    input:str,  # FLAC webdataset file path (or - to read the names from stdin)
    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="maxvad", # could be eqvad to get more uniform chunk lengths
    
):
    vq_model = vq_stoks.RQBottleneckTransformer.load_model(vq_model).cuda()
    vq_model.ensure_whisper()
#     vq_model.encode_mel = torch.compile(vq_model.encode_mel, mode="reduce-overhead", fullgraph=True)
    
    spk_classifier = EncoderClassifier.from_hparams("speechbrain/spkrec-ecapa-voxceleb",
                                                    savedir=f"{os.environ['HOME']}/.cache/speechbrain/",
                                                    run_opts={"device": "cuda"})
    
    total = n_samples//batch_size if n_samples else 'noinfer'

    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(utils.derived_name(input, f'{kind}-stoks', dir="."), 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.cuda().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")
