//! Tiny-transformer hyperparameters. Host-only (no GPU), always compiled.

/// Architecture config for [`crate::TinyTransformer`]. Keep it tiny — T5 is a
/// correctness bring-up, not a real training run.
#[derive(Debug, Clone, Copy)]
pub struct Config {
    /// Vocabulary size (char-level in the bring-up).
    pub vocab: usize,
    /// Model / residual width. Must equal `n_heads * head_dim`.
    pub dim: usize,
    /// Number of decoder blocks.
    pub n_layers: usize,
    /// Number of attention heads.
    pub n_heads: usize,
    /// Per-head dimension (`dim / n_heads`).
    pub head_dim: usize,
    /// SwiGLU hidden width (gate/up project to this, down projects back).
    pub ffn_hidden: usize,
    /// RMSNorm epsilon.
    pub eps: f32,
    /// RoPE base frequency (theta).
    pub rope_theta: f32,
}

impl Config {
    /// A minimal config used by the bring-up / overfit test.
    pub fn tiny() -> Self {
        let n_heads = 2;
        let head_dim = 16;
        Config {
            vocab: 0, // set by the caller from the char vocab
            dim: n_heads * head_dim,
            n_layers: 2,
            n_heads,
            head_dim,
            ffn_hidden: 64,
            eps: 1e-5,
            rope_theta: 10000.0,
        }
    }

    /// Total learnable parameter count (for logging / sanity).
    pub fn num_params(&self) -> usize {
        let per_layer = 2 * self.dim                       // 2 rmsnorm gammas
            + 2 * self.head_dim                             // q/k per-head norm gammas
            + 3 * self.dim * self.dim                       // q/k/v proj
            + self.dim * self.dim                           // out proj
            + 2 * self.dim * self.ffn_hidden                // gate/up proj
            + self.ffn_hidden * self.dim; // down proj
        self.vocab * self.dim          // embedding
            + self.n_layers * per_layer
            + self.dim                 // final norm
            + self.dim * self.vocab // lm head
    }
}
