//! Sprint 2 Phase B — derive ActAdd-style steering vectors from an
//! emotion-labeled manifest by running teacher-forced forward passes
//! through the backbone and computing per-layer difference-of-means
//! between a target-emotion pool and a neutral-baseline pool.
//!
//! Inspired by EmoSteer-TTS (arXiv 2508.03543); see `src/steering.rs`
//! for the architectural deviation note (CSM is AR-token, the paper is
//! flow-matching DiT).
//!
//! Manifest format (JSONL): rows with `wav`, `transcript`, `emotion_tag`
//! (e.g. `[whisper]`, `[neutral]`, `[excited]`). Brackets are optional;
//! they're stripped for matching. `audio_to_manifest --auto-emotion-tag
//! --use-emotion2vec` produces this layout.
//!
//! Output: a safetensors file with keys `layer_<i>_steering`, shape
//! `(embed_dim,)`. Loadable by `examples/generate --steering-vec ...`.
//!
//! Usage:
//! ```bash
//! target/release/examples/steering_extract \
//!     --in /tmp/voice_corpus/manifest_xxx/manifest.jsonl \
//!     --target whisper \
//!     --baseline neutral \
//!     --out /tmp/whisper_steering.safetensors
//! ```

use anyhow::{Context, Result};
use candle_core::{DType, Device, Tensor};
use clap::Parser;
use rtx_csm::Generator;
use rtx_csm::audio_io;
use rtx_csm::prompt::{Segment, build_prompt};
use serde::Deserialize;
use std::collections::HashMap;
use std::path::PathBuf;

#[derive(Debug, Parser)]
struct Cli {
    /// Manifest JSONL with `wav`, `transcript`, `emotion_tag` rows. The
    /// manifest's wav paths can be absolute or relative to the manifest dir.
    #[arg(long = "in")]
    input: PathBuf,
    /// Target emotion (e.g. "whisper", "excited", "sad"). Brackets are
    /// optional — `whisper` and `[whisper]` both match.
    #[arg(long)]
    target: String,
    /// Baseline emotion to subtract (default: "neutral").
    #[arg(long, default_value = "neutral")]
    baseline: String,
    /// Output safetensors path.
    #[arg(long)]
    out: PathBuf,
    /// Maximum rows per pool to encode. 0 = no cap. Lower for quick smoke.
    #[arg(long, default_value_t = 50)]
    max_per_pool: usize,
    /// Which module's activations to capture: `backbone` (default; semantic
    /// codebook 0) or `decoder` (acoustic codebooks 1..N-1; smaller embed
    /// dim, designed to bias prosody/timbre without disturbing word
    /// content). The decoder path teacher-forces the ground-truth c0 and
    /// captures one frame per clip.
    #[arg(long, default_value = "backbone")]
    target_module: String,
    /// Force CPU.
    #[arg(long)]
    cpu: bool,
}

#[derive(Debug, Deserialize)]
struct Row {
    wav: PathBuf,
    transcript: String,
    #[serde(default)]
    emotion_tag: Option<String>,
    #[serde(default)]
    speaker: Option<u32>,
}

fn normalize_tag(s: &str) -> String {
    s.trim()
        .trim_start_matches('[')
        .trim_end_matches(']')
        .to_lowercase()
}

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

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

    let mut generator = Generator::load_csm_1b(&device)?;
    let module = cli.target_module.to_lowercase();
    if !["backbone", "decoder"].contains(&module.as_str()) {
        anyhow::bail!("--target-module must be 'backbone' or 'decoder' (got '{module}')");
    }
    let (num_layers, embed_dim): (usize, usize) = if module == "decoder" {
        let n = generator
            .decoder_num_layers()
            .ok_or_else(|| anyhow::anyhow!("decoder steering needs the FP backbone"))?;
        (n, 1024) // CSM-1B decoder = Llama100M
    } else {
        let n = generator
            .backbone_num_layers()
            .ok_or_else(|| anyhow::anyhow!("steering extraction requires the FP backbone"))?;
        (n, 2048) // CSM-1B backbone = Llama1B
    };
    tracing::info!("loaded CSM-1B; module={module} → {num_layers} layers × {embed_dim} embed");

    // Resolve manifest paths and normalize tags.
    let manifest_dir = cli
        .input
        .parent()
        .map(|p| p.to_path_buf())
        .unwrap_or_else(|| PathBuf::from("."));
    let target_key = normalize_tag(&cli.target);
    let baseline_key = normalize_tag(&cli.baseline);
    tracing::info!("target='{target_key}'  baseline='{baseline_key}'");

    let body = std::fs::read_to_string(&cli.input)
        .with_context(|| format!("read manifest {}", cli.input.display()))?;

    let mut pools: HashMap<String, (Vec<Vec<f32>>, usize)> = HashMap::new();
    pools.insert(
        target_key.clone(),
        (vec![vec![0.0f32; embed_dim]; num_layers], 0),
    );
    pools.insert(
        baseline_key.clone(),
        (vec![vec![0.0f32; embed_dim]; num_layers], 0),
    );

    let mut total_seen = 0usize;
    let mut total_used = 0usize;

    for line in body.lines() {
        if line.trim().is_empty() {
            continue;
        }
        let row: Row = match serde_json::from_str(line) {
            Ok(r) => r,
            Err(e) => {
                tracing::warn!("skipping unparseable row: {e}");
                continue;
            }
        };
        total_seen += 1;
        let tag = match &row.emotion_tag {
            Some(t) => normalize_tag(t),
            None => continue,
        };
        let pool = match pools.get_mut(&tag) {
            Some(p) => p,
            None => continue,
        };
        if cli.max_per_pool > 0 && pool.1 >= cli.max_per_pool {
            continue;
        }

        // Resolve WAV path.
        let wav = if row.wav.is_absolute() {
            row.wav.clone()
        } else {
            manifest_dir.join(&row.wav)
        };
        if !wav.exists() {
            tracing::warn!("missing wav, skipping: {}", wav.display());
            continue;
        }

        // Encode audio + build prompt. Mimi accumulates streaming state
        // across encode() calls and the transformer's internal frame
        // counter overflows its 8192-frame buffer after ~80 clips even
        // with reset_state(). Pattern used elsewhere (commit 0568dd3):
        // reset_state() per row, full reload() every 10.
        let audio = audio_io::load_mono_24k(&wav)?;
        let speaker = row.speaker.unwrap_or(0);

        if total_used > 0 && total_used % 10 == 0 {
            generator.mimi.reload()?;
        } else {
            generator.mimi.reset_state();
        }

        let captures = if module == "decoder" {
            // Decoder mode: text-only segment for the prompt; Mimi-encode
            // the audio separately to grab the middle frame's c0 as the
            // teacher-forced target. This isolates the decoder's response
            // to "given the text context and one audio token, what should
            // codebooks 1..N-1 be?" — pure acoustic, no semantic prefix.
            let segment = Segment::new_text(speaker, row.transcript.clone());
            let prompt = build_prompt(
                &[],
                &segment,
                &generator.model,
                &mut generator.mimi,
                &generator.tokenizer,
            )?;
            let codes = generator.mimi.encode(&audio)?;
            let (_b, _cb, num_frames) = codes.dims3()?;
            if num_frames == 0 {
                tracing::warn!("zero-frame encode, skipping: {}", wav.display());
                continue;
            }
            // Middle frame's c0 — past the audio onset, into steady state.
            let mid = num_frames / 2;
            let c0 = codes.narrow(2, mid, 1)?.narrow(1, 0, 1)?;
            let target_c0: u32 = c0.to_dtype(DType::U32)?.flatten_all()?.to_vec1::<u32>()?[0];
            generator.model.inner.capture_decoder_activations(
                &prompt.tokens,
                &prompt.mask,
                target_c0,
            )?
        } else {
            let segment = Segment::new(speaker, row.transcript.clone(), audio);
            let prompt = build_prompt(
                &[],
                &segment,
                &generator.model,
                &mut generator.mimi,
                &generator.tokenizer,
            )?;
            generator
                .model
                .inner
                .capture_backbone_activations(&prompt.tokens, &prompt.mask)?
        };
        if captures.len() != num_layers {
            anyhow::bail!(
                "captured {} layers, expected {} — model layout mismatch?",
                captures.len(),
                num_layers
            );
        }

        // Accumulate (move to CPU f32 for stable double-precision sum).
        for (l, t) in captures.iter().enumerate() {
            let v: Vec<f32> = t.to_dtype(DType::F32)?.to_vec1::<f32>()?;
            if v.len() != embed_dim {
                anyhow::bail!(
                    "layer {l} captured embed_dim {}, expected {embed_dim}",
                    v.len()
                );
            }
            for (acc, vi) in pool.0[l].iter_mut().zip(v.iter()) {
                *acc += *vi;
            }
        }
        pool.1 += 1;
        total_used += 1;

        if total_used % 5 == 0 {
            tracing::info!(
                "progress: seen={total_seen} used={total_used} target={} baseline={}",
                pools.get(&target_key).map(|p| p.1).unwrap_or(0),
                pools.get(&baseline_key).map(|p| p.1).unwrap_or(0),
            );
        }
    }

    let n_target = pools.get(&target_key).map(|p| p.1).unwrap_or(0);
    let n_base = pools.get(&baseline_key).map(|p| p.1).unwrap_or(0);
    if n_target == 0 || n_base == 0 {
        anyhow::bail!(
            "need ≥1 row per pool: target='{}'={}  baseline='{}'={}",
            target_key,
            n_target,
            baseline_key,
            n_base,
        );
    }

    // Compute per-layer (mean(target) - mean(baseline)).
    let mut diff_tensors: HashMap<String, Tensor> = HashMap::new();
    for layer_idx in 0..num_layers {
        let target_mean: Vec<f32> = pools[&target_key].0[layer_idx]
            .iter()
            .map(|v| v / n_target as f32)
            .collect();
        let base_mean: Vec<f32> = pools[&baseline_key].0[layer_idx]
            .iter()
            .map(|v| v / n_base as f32)
            .collect();
        let diff: Vec<f32> = target_mean
            .iter()
            .zip(base_mean.iter())
            .map(|(t, b)| t - b)
            .collect();
        let t = Tensor::from_vec(diff, (embed_dim,), &Device::Cpu)?;
        diff_tensors.insert(format!("layer_{layer_idx}_steering"), t);
    }

    if let Some(parent) = cli.out.parent() {
        std::fs::create_dir_all(parent).ok();
    }
    candle_core::safetensors::save(&diff_tensors, &cli.out)?;
    eprintln!("--- summary ---");
    eprintln!("target  '{target_key}': {n_target} rows");
    eprintln!("baseline '{baseline_key}': {n_base} rows");
    eprintln!("rows seen / used: {total_seen} / {total_used}");
    eprintln!("wrote {} layers to {}", num_layers, cli.out.display());
    Ok(())
}