//! Demo: compute teacher-forced training loss for one frame on a real
//! (text, audio) pair using the loaded CSM-1B model.
//!
//! Pipeline:
//!   1. Load CSM-1B (FP path — required for backward through trainable params)
//!   2. Mimi-encode a reference WAV to get target codebook tokens
//!   3. Tokenize the matching transcript with Llama BPE
//!   4. Build the (1, S, 33) input tokens + mask
//!   5. Call `Model::forward_loss(tokens, mask, pos, target_codes)`
//!   6. Print the scalar cross-entropy
//!
//! This is the foundation primitive for LoRA fine-tuning. The next step is to
//! wrap the backbone q_proj/v_proj projections with `LoraLinear` (introducing
//! trainable params) and call `loss.backward()` followed by an AdamW step.
//!
//! Usage:
//!   cargo run -p rtx-csm --release --example forward_loss_demo -- \
//!       --wav /tmp/csm/hello.wav --text "Hello from Rust."

use anyhow::Result;
use clap::Parser;
use rtx_csm::{audio_io, Generator, Segment};
use std::path::PathBuf;

#[derive(Debug, Parser)]
struct Cli {
    #[arg(long)]
    wav: PathBuf,
    #[arg(long)]
    text: String,
    #[arg(long, default_value_t = 0)]
    speaker: u32,
    #[arg(long)]
    cpu: bool,
}

fn main() -> Result<()> {
    tracing_subscriber::fmt().init();
    let cli = Cli::parse();

    let device = if cli.cpu {
        candle_core::Device::Cpu
    } else {
        Generator::default_device()?
    };
    println!("device: {device:?}");

    // FP backbone — required because forward_loss isn't implemented on the
    // quantized backend (training-through-Q8 isn't a supported workflow).
    let mut generator = Generator::load_csm_1b(&device)?;

    // Audio side: load WAV, ensure 24 kHz mono, encode through Mimi.
    let audio = audio_io::load_mono_24k(&cli.wav)?;
    println!(
        "loaded {} samples (~{:.2}s)",
        audio.len(),
        audio.len() as f32 / generator.config.sample_rate as f32
    );
    let codes = generator.mimi.encode(&audio)?;
    let (b, cb, t_frames) = codes.dims3()?;
    println!("Mimi codes shape: ({b}, {cb}, {t_frames})");
    if t_frames < 1 {
        anyhow::bail!("audio too short; need at least one Mimi frame");
    }

    // Take frame 0 as the target. Shape after narrow + flatten: (cb,) u32.
    let target_codes = codes
        .narrow(2, 0, 1)?
        .squeeze(2)?
        .flatten_all()?
        .to_vec1::<u32>()?;
    println!("target_codes[0..8] = {:?}", &target_codes[..8.min(target_codes.len())]);

    // Build a "prompt + query" sequence, just like inference. The prompt is
    // the transcript; we ask the model to predict the first audio frame given
    // the text. Use Segment to leverage build_prompt.
    let current = Segment::new_text(cli.speaker, &cli.text);
    let prompt = rtx_csm::prompt::build_prompt(
        &[],
        &current,
        &generator.model,
        &mut generator.mimi,
        &generator.tokenizer,
    )?;
    println!("prompt tokens shape: {:?}", prompt.tokens.shape());

    // Forward-loss for the first frame.
    generator.model.clear_kv_cache();
    let loss = generator
        .model
        .inner
        .forward_loss(&prompt.tokens, &prompt.mask, 0, &target_codes)?;
    let loss_val = loss.to_scalar::<f32>()?;
    println!("\nteacher-forced cross-entropy loss for frame 0: {loss_val:.4}");

    // For context: a randomly initialized model would have CE ≈ ln(2051) ≈ 7.63.
    // A perfectly-predicting model would have CE ≈ 0. Pretrained CSM should be
    // somewhere in between for held-out audio.
    let random_baseline = (generator.config.audio_vocab_size as f32).ln();
    println!("random-baseline CE: {random_baseline:.4}");
    println!(
        "loss / random_baseline: {:.3} (lower = model is more confident)",
        loss_val / random_baseline
    );
    Ok(())
}