//! 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::{DType, Device, 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,
    /// Optional path to dump a JSON fingerprint of the embeddings (for
    /// `scripts/wavlm_sv_parity.py` numerical comparison).
    #[arg(long)]
    parity_json: Option<PathBuf>,
}

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");
    }

    if let Some(parity_path) = cli.parity_json.as_ref() {
        let emb_a_vec: Vec<f32> = emb_a.flatten_all()?.to_dtype(DType::F32)?.to_vec1()?;
        let emb_b_vec: Vec<f32> = emb_b.flatten_all()?.to_dtype(DType::F32)?.to_vec1()?;
        let norm = |v: &[f32]| (v.iter().map(|x| x * x).sum::<f32>()).sqrt();
        let head = |v: &[f32], n: usize| v.iter().take(n).copied().collect::<Vec<_>>();
        let tail = |v: &[f32], n: usize| v.iter().rev().take(n).copied().collect::<Vec<_>>();
        let report = serde_json::json!({
            "model": "rtx-csm port of microsoft/wavlm-base-plus-sv",
            "wav_a": cli.a.display().to_string(),
            "wav_b": cli.b.display().to_string(),
            "embedding_dim": emb_a_vec.len(),
            "cosine_similarity_rust": sim,
            "embedding_a_norm": norm(&emb_a_vec),
            "embedding_b_norm": norm(&emb_b_vec),
            "embedding_a_head": head(&emb_a_vec, 8),
            "embedding_a_tail": tail(&emb_a_vec, 8).into_iter().rev().collect::<Vec<_>>(),
            "embedding_b_head": head(&emb_b_vec, 8),
            "embedding_b_tail": tail(&emb_b_vec, 8).into_iter().rev().collect::<Vec<_>>(),
        });
        std::fs::write(parity_path, serde_json::to_string_pretty(&report)?)?;
        println!("wrote parity fingerprint to {}", parity_path.display());
    }

    Ok(())
}