//! Phase 8.4 spike — download Moonshine-tiny from HuggingFace and dump
//! the safetensors layout. Used to design the candle module shape.
//!
//! Architecture (from HF config.json):
//!   encoder-decoder transformer, hidden_size=288, intermediate=1152,
//!   6 enc + 6 dec layers, 8 heads (head_dim=36 + pad to 40),
//!   partial_rotary_factor=0.9, vocab=32768, max_pos=194, audio input
//!   is raw waveform (no mel-spec preprocessing).
//!
//! Usage:
//! ```bash
//! cargo run -p rtx-csm --release --example moonshine_inspect
//! ```

use anyhow::{Context, Result};
use hf_hub::api::sync::Api;

fn main() -> Result<()> {
    let api = Api::new().context("hf_hub init")?;
    let repo = api.model("UsefulSensors/moonshine-tiny".to_string());
    let weights_path = repo
        .get("model.safetensors")
        .context("download model.safetensors")?;
    println!("downloaded: {}", weights_path.display());
    let bytes = std::fs::read(&weights_path)?;
    println!(
        "file size: {} bytes ({:.1} MB)",
        bytes.len(),
        bytes.len() as f64 / 1_000_000.0
    );

    // Parse the safetensors header to enumerate tensor names + shapes.
    let st = safetensors::SafeTensors::deserialize(&bytes)
        .context("parse safetensors header")?;
    let names = st.names();
    println!("\n=== {} tensors ===", names.len());

    // Group by prefix so the layered structure is visible.
    let mut grouped: std::collections::BTreeMap<String, Vec<(String, Vec<usize>, &str)>> =
        Default::default();
    for name in names {
        let tensor = st.tensor(name)?;
        let prefix = name
            .split('.')
            .take(2)
            .collect::<Vec<_>>()
            .join(".");
        let dtype = match tensor.dtype() {
            safetensors::Dtype::F32 => "f32",
            safetensors::Dtype::F16 => "f16",
            safetensors::Dtype::BF16 => "bf16",
            other => match other {
                _ => "?",
            },
        };
        grouped
            .entry(prefix)
            .or_default()
            .push((name.to_string(), tensor.shape().to_vec(), dtype));
    }

    for (prefix, items) in grouped.iter() {
        println!("\n[{prefix}]  ({} tensors)", items.len());
        for (name, shape, dtype) in items.iter().take(8) {
            println!("  {name:<60} {dtype} {shape:?}");
        }
        if items.len() > 8 {
            println!("  ... {} more", items.len() - 8);
        }

        // Special: dump everything matching `conv` since we need
        // exact strides for the audio stem.
        for (name, shape, dtype) in items.iter().filter(|(n, _, _)| n.contains("conv")) {
            if !items.iter().take(8).any(|(n2, _, _)| n2 == name) {
                println!("  {name:<60} {dtype} {shape:?}  (conv detail)");
            }
        }
    }

    let total_params: usize = st
        .names()
        .iter()
        .map(|n| {
            st.tensor(n)
                .map(|t| t.shape().iter().product::<usize>())
                .unwrap_or(0)
        })
        .sum();
    println!(
        "\ntotal parameters: {} ({:.1} M)",
        total_params,
        total_params as f64 / 1_000_000.0
    );

    Ok(())
}