//! Model configuration for CSM-1B.
//!
//! Mirrors `candle_transformers::models::csm::Config`; we keep our own type so
//! we can extend it (e.g., for the Stage 2 fine-tuning loop) without depending
//! on candle's internal layout.

use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub enum BackboneFlavor {
    /// Llama-3.2 1B: 16 layers, 2048 dim, 32 heads / 8 KV heads, FFN 8192.
    Llama1B,
    /// Reserved for the 3B and 8B sizes if Sesame ever ships them.
    Llama3B,
    Llama8B,
}

#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub enum DecoderFlavor {
    /// Llama-3.2 100M: 4 layers, 1024 dim, 8 heads / 2 KV heads, FFN 8192.
    Llama100M,
    Llama250M,
    Llama300M,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelConfig {
    pub backbone: BackboneFlavor,
    pub decoder: DecoderFlavor,
    pub text_vocab_size: usize,
    pub audio_vocab_size: usize,
    pub audio_num_codebooks: usize,
    pub max_seq_len: usize,
    pub sample_rate: u32,
    pub frame_rate_hz: f32,
}

impl Default for ModelConfig {
    fn default() -> Self {
        Self::csm_1b()
    }
}

impl ModelConfig {
    /// CSM-1B: the only variant Sesame released publicly.
    pub fn csm_1b() -> Self {
        Self {
            backbone: BackboneFlavor::Llama1B,
            decoder: DecoderFlavor::Llama100M,
            text_vocab_size: 128_256,
            audio_vocab_size: 2051,
            audio_num_codebooks: 32,
            max_seq_len: 2048,
            sample_rate: 24_000,
            frame_rate_hz: 12.5,
        }
    }

    pub fn frame_duration_ms(&self) -> f32 {
        1000.0 / self.frame_rate_hz
    }
}