//! Smoke-test helper: write a random `LayerSteering` safetensors so the
//! apply hook in `LlamaModel` can be exercised end-to-end before the real
//! emotion-extraction binary lands (Sprint 2 Phase B).
//!
//! The vectors are random Gaussian noise — they do NOT encode any
//! meaningful direction. Output should sound DIFFERENT from the unsteered
//! generation but no quality claim. This binary exists purely so the
//! plumbing is proven before we invest in corpus extraction.
//!
//! Usage:
//! ```bash
//! target/release/examples/steering_random \
//!     --num-layers 16 --embed-dim 2048 \
//!     --out /tmp/steering_random.safetensors
//!
//! target/release/examples/generate \
//!     --text "Activation steering smoke test." \
//!     --speaker 0 --max-audio-ms 3000 --seed 42 \
//!     --steering-vec /tmp/steering_random.safetensors \
//!     --steering-scale 0.5 \
//!     --out /tmp/steered.wav
//! ```

use anyhow::Result;
use candle_core::{DType, Device, Tensor};
use clap::Parser;

#[derive(Debug, Parser)]
struct Cli {
    #[arg(long)]
    out: std::path::PathBuf,
    /// CSM-1B has 16 backbone layers (BackboneFlavor::Llama1B).
    #[arg(long, default_value_t = 16)]
    num_layers: usize,
    /// CSM-1B's Llama backbone embed_dim = 2048.
    #[arg(long, default_value_t = 2048)]
    embed_dim: usize,
    /// Stddev of the per-layer Gaussian. Smaller = subtler steering.
    #[arg(long, default_value_t = 0.05)]
    stddev: f32,
    #[arg(long, default_value_t = 42)]
    seed: u64,
}

fn main() -> Result<()> {
    let cli = Cli::parse();
    let dev = Device::Cpu;

    // Box-Muller from a seeded uniform RNG so re-runs are reproducible
    // without an extra `rand_distr` dep.
    use rand::{Rng, SeedableRng};
    let mut rng = rand::rngs::StdRng::seed_from_u64(cli.seed);
    let mut gauss = || -> f32 {
        let u1: f32 = rng.gen_range(1e-9..1.0);
        let u2: f32 = rng.r#gen();
        let r = (-2.0 * u1.ln()).sqrt();
        let theta = 2.0 * std::f32::consts::PI * u2;
        r * theta.cos() * cli.stddev
    };

    let mut tensors: Vec<(String, Tensor)> = Vec::with_capacity(cli.num_layers);
    for i in 0..cli.num_layers {
        let v: Vec<f32> = (0..cli.embed_dim).map(|_| gauss()).collect();
        let t = Tensor::from_vec(v, (cli.embed_dim,), &dev)?.to_dtype(DType::F32)?;
        tensors.push((format!("layer_{i}_steering"), t));
    }

    let map: std::collections::HashMap<String, Tensor> = tensors.into_iter().collect();
    candle_core::safetensors::save(&map, &cli.out)?;
    println!(
        "wrote {} layer vectors (embed_dim={}, stddev={}) to {}",
        cli.num_layers,
        cli.embed_dim,
        cli.stddev,
        cli.out.display(),
    );
    Ok(())
}