"""Compare two voices with TitaNet.mlpackage (macOS, Python 3.9+).

    pip install coremltools kaldi-native-fbank soundfile numpy
    python example.py a.wav b.wav

The WAV files must be mono 16 kHz. Prints the cosine similarity of the two
speaker embeddings.
"""

import sys

import coremltools as ct
import kaldi_native_fbank as knf
import numpy as np
import soundfile as sf


def features(samples: np.ndarray) -> np.ndarray:
    """The features sherpa-onnx feeds TitaNet-L, shaped (1, 80, T)."""
    o = knf.FbankOptions()
    o.frame_opts.samp_freq = 16000
    o.frame_opts.frame_length_ms = 25
    o.frame_opts.frame_shift_ms = 10
    o.frame_opts.snip_edges = True
    o.frame_opts.dither = 0
    o.frame_opts.remove_dc_offset = False
    o.frame_opts.preemph_coeff = 0.97
    o.frame_opts.window_type = "hann"
    o.mel_opts.num_bins = 80
    o.mel_opts.low_freq = 0
    o.mel_opts.high_freq = -400  # 7600 Hz
    o.mel_opts.is_librosa = True
    fbank = knf.OnlineFbank(o)
    fbank.accept_waveform(16000, samples.tolist())
    fbank.input_finished()
    f = np.array([fbank.get_frame(i) for i in range(fbank.num_frames_ready)], dtype=np.float32)
    f = (f - f.mean(0)) / (f.std(0) + 1e-5)  # per-bin, population std
    return f.T[None].copy()


def embed(model: ct.models.MLModel, path: str) -> np.ndarray:
    samples, rate = sf.read(path, dtype="float32")
    assert rate == 16000 and samples.ndim == 1, "expected mono 16 kHz audio"
    return model.predict({"features": features(samples)})["embedding"][0]


if __name__ == "__main__":
    model = ct.models.MLModel("TitaNet.mlpackage", compute_units=ct.ComputeUnit.CPU_AND_GPU)
    a, b = embed(model, sys.argv[1]), embed(model, sys.argv[2])
    print(f"cosine similarity: {a @ b / np.linalg.norm(a) / np.linalg.norm(b):.3f}")