//! Per-pair TTS quality eval: speaker similarity + WER + amplitude.
//!
//! Reads a JSONL where each row has at minimum `ref_wav` and `gen_wav`,
//! optionally `ref_text` (used as the reference transcript for WER). The
//! reference text comes from the manifest because the reference audio
//! might not have a clean public transcript (e.g. context-conditioning
//! reference clips).
//!
//! Designed as the metric foundation for Sprint 1 of the Phase 8 roadmap.
//! Originally TTSDS2 (arXiv 2506.19441) was on the menu — its install on
//! Python 3.12 + torchaudio 2.x is broken (`torchaudio.sox_effects` was
//! removed; `pyannote.audio` 3.1 calls `set_audio_backend`; openai-whisper
//! pinned dep needs `pkg_resources` which requires `setuptools<81`). For
//! now we score with what we already own end-to-end in candle:
//!
//!   - **speaker_cosine**: WavLM-SV (microsoft/wavlm-base-plus-sv)
//!   - **wer**: Moonshine v2 transcribe vs `ref_text`
//!   - **peak_db / rms_db**: full-band amplitude of `gen_wav`
//!
//! Output: one JSONL line per input, with the input's fields plus a
//! `metrics` object. Stable column order = appendable to a CSV/jq pipeline.
//!
//! ```bash
//! target/release/examples/quality_eval \
//!     --in /tmp/i2d/eval_pairs.jsonl \
//!     --wavlm-sv /tmp/wavlm_sv.safetensors \
//!     --out /tmp/i2d/scores.jsonl
//! ```

use anyhow::{Context, Result};
use candle_core::{Device, Tensor};
use clap::Parser;
use hf_hub::api::sync::Api;
use rtx_csm::speaker_sim::{SpeakerSimilarity, WavLmSimilarity};
use rtx_csm::{audio_io, moonshine};
use serde::{Deserialize, Serialize};
use std::io::{BufRead, BufWriter, Write};
use std::path::PathBuf;

#[derive(Debug, Parser)]
struct Cli {
    /// Input JSONL with rows: `{"ref_wav": "...", "gen_wav": "...", "ref_text": "..."}`.
    /// `ref_text` is optional; without it WER is `null`.
    #[arg(long = "in")]
    input: PathBuf,
    /// Output JSONL. One line per input row, with the input's fields plus
    /// a `metrics` object.
    #[arg(long)]
    out: PathBuf,
    /// Converted WavLM-SV safetensors. See `examples/wavlm_sv_convert`.
    #[arg(long)]
    wavlm_sv: PathBuf,
    /// Max decoded tokens for Moonshine (capped at 194 by config).
    #[arg(long, default_value_t = 120)]
    max_tokens: usize,
    /// Optional path to the emotion2vec_plus_base `.pt` checkpoint. When
    /// set, each row also gets emotion-classifier metrics. Pair with
    /// `--target-emotion` to score the probability of a specific class.
    ///
    /// **Caveat (2026-04-30):** the emotion2vec_plus_base checkpoint as
    /// loaded through our pickle path produces saturated outputs —
    /// every audio sample classifies as "Surprised" with probability
    /// ≥ 0.99, INCLUDING ground-truth RAVDESS clips with explicit
    /// neutral / happy / angry / etc. labels. The metric is wired
    /// correctly (verified against the trait `EmotionDetector::classify`
    /// path) but the underlying classifier collapses to one dominant
    /// class on most input. Useful for verifying the emotion2vec
    /// pipeline works at all; not useful for verifying generation
    /// quality. Treat the output as diagnostic, not as ground truth.
    #[arg(long)]
    emotion2vec: Option<PathBuf>,
    /// Target emotion to score the probability of. One of: angry,
    /// disgusted, fearful, happy, neutral, excited, sad, surprised.
    /// Ignored unless `--emotion2vec` is also set. See caveat above.
    #[arg(long)]
    target_emotion: Option<String>,
    /// Force CPU device.
    #[arg(long)]
    cpu: bool,
}

fn emotion_index(name: &str) -> Option<usize> {
    // Maps a CLI emotion name to emotion2vec class index. Order from
    // src/emotion2vec.rs::Classifier::tag_for_class:
    match name.trim().to_lowercase().as_str() {
        "angry" => Some(0),
        "disgust" | "disgusted" => Some(1),
        "fearful" => Some(2),
        "happy" => Some(3),
        "neutral" => Some(4),
        "excited" | "other" => Some(5),
        "sad" => Some(6),
        "surprised" | "calm" => Some(7), // emotion2vec lumps calm in here
        _ => None,
    }
}

fn softmax(logits: &[f32]) -> Vec<f32> {
    let max = logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    let exp: Vec<f32> = logits.iter().map(|&v| (v - max).exp()).collect();
    let sum: f32 = exp.iter().sum();
    exp.into_iter().map(|v| v / sum.max(1e-12)).collect()
}

#[derive(Debug, Deserialize, Serialize)]
struct PairRow {
    ref_wav: PathBuf,
    gen_wav: PathBuf,
    #[serde(skip_serializing_if = "Option::is_none", default)]
    ref_text: Option<String>,
    /// Free-form per-row metadata (iteration index, model id, etc).
    /// Round-tripped unchanged into the output row.
    #[serde(flatten, default)]
    extra: serde_json::Map<String, serde_json::Value>,
}

#[derive(Debug, Serialize)]
struct ScoredRow {
    ref_wav: PathBuf,
    gen_wav: PathBuf,
    #[serde(skip_serializing_if = "Option::is_none")]
    ref_text: Option<String>,
    #[serde(flatten)]
    extra: serde_json::Map<String, serde_json::Value>,
    metrics: Metrics,
}

#[derive(Debug, Serialize)]
struct Metrics {
    /// Cosine in [-1, 1]. Higher = more similar speaker identity.
    speaker_cosine: f32,
    /// Word error rate in [0, 1+]. None if no `ref_text`.
    wer: Option<f32>,
    /// Moonshine's transcript of `gen_wav` (always emitted; useful for
    /// debugging).
    transcript: String,
    /// Generated audio characteristics.
    gen_peak_db: f32,
    gen_rms_db: f32,
    /// Probability of the target emotion (None unless --emotion2vec +
    /// --target-emotion are passed).
    #[serde(skip_serializing_if = "Option::is_none")]
    target_emotion_prob: Option<f32>,
    /// Argmax emotion class label per emotion2vec.
    #[serde(skip_serializing_if = "Option::is_none")]
    top_emotion: Option<String>,
    /// Probability of the argmax emotion class.
    #[serde(skip_serializing_if = "Option::is_none")]
    top_emotion_prob: Option<f32>,
}

fn main() -> Result<()> {
    tracing_subscriber::fmt().with_target(false).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
    };
    tracing::info!("device: {device:?}");

    // Speaker similarity model.
    let sv = WavLmSimilarity::load(&cli.wavlm_sv, &device)?;
    tracing::info!("loaded WavLM-SV from {}", cli.wavlm_sv.display());

    // Moonshine encoder + decoder + tokenizer (all from HF cache).
    let api = Api::new()?;
    let repo = api.model("UsefulSensors/moonshine-tiny".to_string());
    let m_weights = repo.get("model.safetensors")?;
    let m_tok = repo.get("tokenizer.json")?;
    let m_cfg = moonshine::MoonshineConfig::tiny();
    let (m_enc, m_dec) = moonshine::load_full(&m_weights, &device, &m_cfg)?;
    let m_tokenizer =
        moonshine::load_tokenizer(&m_tok).map_err(|e| anyhow::anyhow!("tokenizer: {e}"))?;
    tracing::info!("loaded Moonshine-tiny");

    // Optional emotion2vec for emotion-classification metrics.
    let target_idx = cli.target_emotion.as_deref().and_then(emotion_index);
    if cli.emotion2vec.is_some() && cli.target_emotion.is_some() && target_idx.is_none() {
        anyhow::bail!(
            "--target-emotion '{}' is unrecognized — must be angry, disgusted, fearful, happy, neutral, excited, sad, or surprised",
            cli.target_emotion.as_deref().unwrap_or("")
        );
    }
    let emo_model = if let Some(path) = cli.emotion2vec.as_ref() {
        let m = rtx_csm::emotion2vec::Emotion2Vec::load_from_pickle(path, &device)?;
        tracing::info!(
            "loaded emotion2vec from {} (target='{}', class_idx={:?})",
            path.display(),
            cli.target_emotion.as_deref().unwrap_or("(none)"),
            target_idx,
        );
        Some(m)
    } else {
        None
    };

    // Stream the input JSONL line by line so big batches don't blow memory.
    let f = std::fs::File::open(&cli.input)
        .with_context(|| format!("open input {}", cli.input.display()))?;
    let reader = std::io::BufReader::new(f);

    if let Some(parent) = cli.out.parent() {
        std::fs::create_dir_all(parent).ok();
    }
    let out_file = std::fs::File::create(&cli.out)
        .with_context(|| format!("create out {}", cli.out.display()))?;
    let mut out_w = BufWriter::new(out_file);

    let mut n_rows = 0usize;
    let mut sum_cos = 0.0f64;
    let mut sum_wer = 0.0f64;
    let mut n_wer = 0usize;
    for (idx, line) in reader.lines().enumerate() {
        let line = line?;
        if line.trim().is_empty() {
            continue;
        }
        let row: PairRow =
            serde_json::from_str(&line).with_context(|| format!("parse row {idx}: {line}"))?;

        // Speaker similarity at 16 kHz. WavLM-SV's TDNN front-end has
        // kernel size 5 with stride 2 chains, so it requires at least a
        // few hundred samples of audio. Below ~0.25 s the model errors
        // out — score those as -1.0 (sentinel) so the eval still emits
        // a row instead of bailing on the whole batch.
        let ref_16k = audio_io::load_mono_at_rate(&row.ref_wav, 16_000)?;
        let gen_16k = audio_io::load_mono_at_rate(&row.gen_wav, 16_000)?;
        const MIN_SV_SAMPLES: usize = 4000; // 0.25 s at 16 kHz
        let speaker_cosine = if gen_16k.len() < MIN_SV_SAMPLES || ref_16k.len() < MIN_SV_SAMPLES {
            tracing::warn!(
                "row {idx}: gen={}s ref={}s — too short for WavLM-SV; cosine=-1",
                gen_16k.len() as f32 / 16_000.0,
                ref_16k.len() as f32 / 16_000.0,
            );
            -1.0f32
        } else {
            sv.score(&ref_16k, &gen_16k)? as f32
        };

        // Transcribe gen_wav via Moonshine (16 kHz).
        let pcm_t = Tensor::from_vec(gen_16k.clone(), (1, 1, gen_16k.len()), &device)?;
        let enc_out = m_enc.forward(&pcm_t)?;
        let token_ids = m_dec.generate_cached(&enc_out, &m_cfg, cli.max_tokens)?;
        let transcript = m_tokenizer
            .decode(&token_ids, /* skip_special */ true)
            .map_err(|e| anyhow::anyhow!("detok: {e}"))?
            .trim()
            .to_string();

        let wer = row
            .ref_text
            .as_ref()
            .map(|t| word_error_rate(t, &transcript));

        // Amplitude of gen_wav at native rate.
        let gen_native = audio_io::load_mono_24k(&row.gen_wav)?;
        let (peak_db, rms_db) = peak_rms_db(&gen_native);

        // Optional emotion classification on gen_wav at 16 kHz.
        let (target_emotion_prob, top_emotion, top_emotion_prob) = if let Some(emo) =
            emo_model.as_ref()
        {
            if gen_16k.len() < MIN_SV_SAMPLES {
                (None, None, None)
            } else {
                use rtx_csm::ser::EmotionDetector;
                // Run forward to get logits (the EmotionDetector trait
                // gives us only argmax via classify; we need the full
                // probability vector).
                let n = gen_16k.len();
                let mean = gen_16k.iter().sum::<f32>() / n as f32;
                let var = gen_16k.iter().map(|x| (x - mean).powi(2)).sum::<f32>() / n as f32;
                let std = var.sqrt().max(1e-7);
                let normed: Vec<f32> = gen_16k.iter().map(|x| (x - mean) / std).collect();
                let audio = Tensor::from_vec(normed, (1, 1, n), &device)?;
                let logits_t = emo.forward(&audio)?;
                let logits = logits_t.flatten_all()?.to_vec1::<f32>()?;
                let probs = softmax(&logits);
                let target_p = target_idx.and_then(|i| probs.get(i).copied());
                let (top_i, top_p) = probs
                    .iter()
                    .enumerate()
                    .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
                    .map(|(i, &p)| (i, p))
                    .unwrap_or((0, 0.0));
                // Use the trait's tag for human-readable label.
                let top_label = emo
                    .classify(&gen_16k)
                    .map(|l| format!("{l:?}"))
                    .unwrap_or_else(|_| format!("class_{top_i}"));
                let _ = top_i;
                (target_p, Some(top_label), Some(top_p))
            }
        } else {
            (None, None, None)
        };

        let scored = ScoredRow {
            ref_wav: row.ref_wav,
            gen_wav: row.gen_wav,
            ref_text: row.ref_text,
            extra: row.extra,
            metrics: Metrics {
                speaker_cosine,
                wer,
                transcript: transcript.clone(),
                gen_peak_db: peak_db,
                gen_rms_db: rms_db,
                target_emotion_prob,
                top_emotion,
                top_emotion_prob,
            },
        };
        writeln!(out_w, "{}", serde_json::to_string(&scored)?)?;
        out_w.flush()?;

        n_rows += 1;
        sum_cos += speaker_cosine as f64;
        if let Some(w) = wer {
            sum_wer += w as f64;
            n_wer += 1;
        }
        tracing::info!(
            "row {idx}: cos={speaker_cosine:.3} wer={} peak={peak_db:.2} rms={rms_db:.2}",
            wer.map(|w| format!("{w:.3}")).unwrap_or_else(|| "-".into())
        );
    }

    if n_rows == 0 {
        anyhow::bail!("no rows processed (empty input?)");
    }
    eprintln!("--- summary ---");
    eprintln!("rows           {n_rows}");
    eprintln!("speaker_cosine {:.3} (mean)", sum_cos / n_rows as f64);
    if n_wer > 0 {
        eprintln!(
            "wer            {:.3} (mean over {n_wer})",
            sum_wer / n_wer as f64
        );
    } else {
        eprintln!("wer            (no ref_text in any row)");
    }
    eprintln!("wrote          {}", cli.out.display());
    Ok(())
}

/// Word error rate: Levenshtein on whitespace-split, lowercased,
/// punctuation-stripped tokens. Returns errors / max(ref_len, 1).
fn word_error_rate(reference: &str, hypothesis: &str) -> f32 {
    let r = tokenize(reference);
    let h = tokenize(hypothesis);
    if r.is_empty() {
        return if h.is_empty() { 0.0 } else { 1.0 };
    }
    let n = r.len();
    let m = h.len();
    let mut prev: Vec<usize> = (0..=m).collect();
    let mut curr = vec![0usize; m + 1];
    for i in 1..=n {
        curr[0] = i;
        for j in 1..=m {
            let cost = if r[i - 1] == h[j - 1] { 0 } else { 1 };
            curr[j] = (prev[j] + 1).min(curr[j - 1] + 1).min(prev[j - 1] + cost);
        }
        std::mem::swap(&mut prev, &mut curr);
    }
    prev[m] as f32 / n as f32
}

fn tokenize(s: &str) -> Vec<String> {
    let mut buf = String::with_capacity(s.len());
    for c in s.chars() {
        if c.is_alphanumeric() || c.is_whitespace() {
            for lc in c.to_lowercase() {
                buf.push(lc);
            }
        } else {
            buf.push(' ');
        }
    }
    buf.split_whitespace().map(String::from).collect()
}

fn peak_rms_db(pcm: &[f32]) -> (f32, f32) {
    if pcm.is_empty() {
        return (-100.0, -100.0);
    }
    let mut peak = 0.0f32;
    let mut sum_sq = 0.0f64;
    for &s in pcm {
        let a = s.abs();
        if a > peak {
            peak = a;
        }
        sum_sq += (s as f64) * (s as f64);
    }
    let rms = (sum_sq / pcm.len() as f64).sqrt() as f32;
    let peak_db = if peak < 1e-6 {
        -100.0
    } else {
        20.0 * peak.log10()
    };
    let rms_db = if rms < 1e-6 {
        -100.0
    } else {
        20.0 * rms.log10()
    };
    (peak_db, rms_db)
}