//! Diagnostic — feed REAL 16 kHz audio files into emotion2vec_plus_base
//! and dump per-class logits. Used to investigate why every clip in
//! audio_to_manifest --auto-emotion-tag --use-emotion2vec was getting
//! tagged the same emotion regardless of input.

use anyhow::{Context, Result};
use clap::Parser;
use hf_hub::api::sync::Api;
use rtx_csm::audio_io;
use rtx_csm::emotion2vec::{Classifier, Emotion2Vec};
use std::path::PathBuf;

#[derive(Debug, Parser)]
struct Cli {
    /// One or more audio files to classify.
    #[arg(long = "in")]
    inputs: Vec<PathBuf>,
}

fn main() -> Result<()> {
    tracing_subscriber::fmt().init();
    let cli = Cli::parse();
    if cli.inputs.is_empty() {
        anyhow::bail!("--in <path> required (one or more)");
    }
    let device = if candle_core::utils::metal_is_available() {
        candle_core::Device::new_metal(0)?
    } else {
        candle_core::Device::Cpu
    };

    let api = Api::new()?;
    let path = api
        .model("emotion2vec/emotion2vec_plus_base".into())
        .get("model.pt")
        .context("download emotion2vec_plus_base/model.pt")?;
    let model = Emotion2Vec::load_from_pickle(&path, &device)?;
    eprintln!("model loaded\n");

    let labels = [
        "angry",
        "disgusted",
        "fearful",
        "happy",
        "neutral",
        "other",
        "sad",
        "surprised",
        "<unk>",
    ];

    for input in cli.inputs.iter() {
        let pcm = audio_io::load_mono_at_rate(input, 16_000)?;
        let secs = pcm.len() as f32 / 16_000.0;
        // Per-utterance zero-mean unit-variance normalization
        // (config.yaml: `normalize: true`). data2vec2 / emotion2vec
        // expects this; the upstream FunASR pipeline applies it before
        // the local_encoder. Without it the model produces ~constant
        // output regardless of input.
        let mean = pcm.iter().sum::<f32>() / pcm.len().max(1) as f32;
        let var = pcm.iter().map(|x| (x - mean).powi(2)).sum::<f32>() / pcm.len().max(1) as f32;
        let std = var.sqrt().max(1e-7);
        let normed: Vec<f32> = pcm.iter().map(|x| (x - mean) / std).collect();
        let audio = candle_core::Tensor::from_vec(normed, (1, 1, pcm.len()), &device)?;
        let logits = model.forward(&audio)?;
        let v = logits.flatten_all()?.to_vec1::<f32>()?;
        let argmax = v
            .iter()
            .enumerate()
            .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
            .map(|(i, _)| i)
            .unwrap_or(0);
        println!(
            "=== {} ({:.2}s) ===",
            input.file_name().and_then(|s| s.to_str()).unwrap_or("?"),
            secs,
        );
        for (i, x) in v.iter().enumerate() {
            let marker = if i == argmax { "  ←" } else { "" };
            println!("  {} {:>10}  {:>9.4}{}", i, labels[i], x, marker);
        }
        println!(
            "  argmax: {} ({}) → 5-bucket {}\n",
            argmax,
            labels[argmax],
            Classifier::tag_for_class(argmax as u32).as_tag()
        );
    }
    Ok(())
}