//! Load the converted WavLM-SV safetensors and compute speaker embeddings
//! for a pair of input WAVs. Reports cosine similarity (the standard
//! speaker verification score: > ~0.5 = same speaker).
//!
//! Usage:
//! ```
//! cargo run -p rtx-csm --release --example wavlm_sv_demo -- \
//!     --weights /tmp/wavlm_sv.safetensors \
//!     --a /path/to/utt_a.wav \
//!     --b /path/to/utt_b.wav
//! ```

use anyhow::{Context, Result};
use candle_core::{Device, DType, Tensor};
use clap::Parser;
use rtx_csm::{audio_io, wavlm_sv::WavLmSv};
use std::path::PathBuf;

const WAVLM_RATE: u32 = 16_000;

#[derive(Debug, Parser)]
#[command(name = "wavlm_sv_demo")]
struct Cli {
    /// Converted safetensors (run `wavlm_sv_convert` first).
    #[arg(long)]
    weights: PathBuf,
    /// First utterance.
    #[arg(long)]
    a: PathBuf,
    /// Second utterance.
    #[arg(long)]
    b: PathBuf,
    /// Force CPU device.
    #[arg(long)]
    cpu: bool,
}

fn main() -> Result<()> {
    tracing_subscriber::fmt().init();
    let cli = Cli::parse();
    let device = if cli.cpu {
        Device::Cpu
    } else if candle_core::utils::metal_is_available() {
        Device::new_metal(0)?
    } else {
        Device::Cpu
    };

    let vb = unsafe {
        candle_nn::VarBuilder::from_mmaped_safetensors(&[&cli.weights], DType::F32, &device)
    }
    .context("opening WavLM-SV safetensors")?;
    let model = WavLmSv::new(vb).context("WavLmSv::new")?;
    println!("loaded WavLM-SV from {}", cli.weights.display());

    // Load utterances at 16 kHz mono.
    let a_samples = audio_io::load_mono_at_rate(&cli.a, WAVLM_RATE)?;
    let b_samples = audio_io::load_mono_at_rate(&cli.b, WAVLM_RATE)?;
    println!(
        "a: {} samples ({:.2}s), b: {} samples ({:.2}s)",
        a_samples.len(),
        a_samples.len() as f32 / WAVLM_RATE as f32,
        b_samples.len(),
        b_samples.len() as f32 / WAVLM_RATE as f32,
    );

    let a_norm = WavLmSv::normalize_waveform(&a_samples);
    let b_norm = WavLmSv::normalize_waveform(&b_samples);
    let xs_a = Tensor::from_slice(&a_norm, (1, 1, a_norm.len()), &device)?;
    let xs_b = Tensor::from_slice(&b_norm, (1, 1, b_norm.len()), &device)?;

    let emb_a = model.embed(&xs_a).context("embed a")?;
    let emb_b = model.embed(&xs_b).context("embed b")?;
    println!(
        "embeddings: a={:?}, b={:?}",
        emb_a.dims(),
        emb_b.dims()
    );

    let sim = WavLmSv::cosine_similarity(&emb_a, &emb_b)?;
    // cosine_similarity returns (B,) for (B, D) inputs; take the first element.
    let sim_vec: Vec<f32> = sim.flatten_all()?.to_dtype(DType::F32)?.to_vec1()?;
    let sim = sim_vec[0];
    println!("cosine similarity = {sim:.4}");
    if sim > 0.5 {
        println!("  → likely same speaker");
    } else {
        println!("  → likely different speakers");
    }

    Ok(())
}