//! Step A: Mimi codec round-trip parity check.
//!
//! Synthesizes a pseudo-speech signal at 24 kHz, encodes through Mimi to
//! discrete codes (32 codebooks), decodes back to waveform, asserts the RMS
//! reconstruction error is within bound.
//!
//! Marked `#[ignore]` because it requires ~300 MB of weights from
//! `kyutai/mimi`. Run manually:
//!
//! ```
//! HF_HUB_ENABLE_HF_TRANSFER=0 cargo test -p rtx-csm --test mimi_roundtrip --release -- --ignored --nocapture
//! ```

use rtx_csm::{
    audio_io::{TARGET_SAMPLE_RATE, write_wav_24k_mono},
    error::Result,
    hub,
    mimi::{Mimi, NUM_CODEBOOKS},
    util::{pick_device, rms},
};
use std::f32::consts::TAU;

/// Build a speech-shaped 24 kHz mono signal: voiced fundamental + first three
/// formants, amplitude-modulated to simulate syllable rate. Roughly 3 seconds.
fn synth_speech_like(seconds: f32) -> Vec<f32> {
    let n = (seconds * TARGET_SAMPLE_RATE as f32) as usize;
    let mut out = Vec::with_capacity(n);
    let f0 = 130.0_f32; // male-ish fundamental
    let formants = [700.0_f32, 1220.0, 2600.0];
    let formant_gains = [1.0_f32, 0.6, 0.3];
    let syllable_rate = 4.0_f32; // Hz
    for i in 0..n {
        let t = i as f32 / TARGET_SAMPLE_RATE as f32;
        // amplitude envelope (avoid 0 for stability)
        let env = 0.5 * (1.0 + (TAU * syllable_rate * t).sin()).max(0.05);
        let voiced = (TAU * f0 * t).sin();
        let mut formant_sum = 0.0;
        for (f, g) in formants.iter().zip(formant_gains.iter()) {
            formant_sum += g * (TAU * f * t).sin();
        }
        let s = 0.5 * env * (0.6 * voiced + 0.4 * formant_sum / formant_gains.iter().sum::<f32>());
        out.push(s);
    }
    out
}

#[test]
#[ignore]
fn mimi_roundtrip_rms_bound() -> Result<()> {
    let _ = tracing_subscriber::fmt::try_init();

    let device = pick_device()?;
    println!("device: {device:?}");

    let mimi_weights = hub::resolve_mimi()?;
    println!("mimi weights: {}", mimi_weights.display());

    let mut mimi = Mimi::load(&mimi_weights, &device)?;

    let original = synth_speech_like(3.0);
    println!(
        "input samples: {} (~{:.1}s)",
        original.len(),
        original.len() as f32 / TARGET_SAMPLE_RATE as f32
    );

    let codes = mimi.encode(&original)?;
    let dims = codes.dims();
    println!("codes shape: {dims:?}");
    assert_eq!(dims.len(), 3, "codes must be (B, C, T)");
    assert_eq!(dims[0], 1, "batch=1");
    assert_eq!(dims[1], NUM_CODEBOOKS, "codebooks=32");

    let decoded = mimi.decode(&codes)?;
    println!("decoded samples: {}", decoded.len());

    // Truncate to common length — Mimi may pad to a multiple of 1920 (24kHz/12.5Hz).
    let n = original.len().min(decoded.len());
    let err = rms(&original[..n], &decoded[..n]);
    println!("rms error: {err:.5}");

    // Save artifacts for ear-test debugging when -- --nocapture is used.
    if let Ok(dir) = std::env::var("CSM_TEST_OUT") {
        let _ = std::fs::create_dir_all(&dir);
        write_wav_24k_mono(format!("{dir}/mimi_in.wav"), &original)?;
        write_wav_24k_mono(format!("{dir}/mimi_out.wav"), &decoded[..n])?;
        println!("artifacts written to {dir}");
    }

    // Synthesized non-speech: looser bound than the 0.05 target for real speech.
    // Real-speech bound will be re-checked when a fixture WAV is added in Step B.
    assert!(
        err < 0.20,
        "Mimi round-trip RMS {err} exceeds 0.20 bound on synthesized signal"
    );
    Ok(())
}