//! 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,
        }
    }

    /// Build a config from the architecture knobs, deriving `dim = n_heads *
    /// head_dim`. The scaling-run entry (`bin/train`) passes these from CLI so the
    /// model size is a tunable ladder rung (v1 = dim256/8L, v2/v3 scale further),
    /// instead of a hardcoded tiny config. `eps`/`rope_theta` keep the engine
    /// defaults (also what the xserv export reconciles against).
    pub fn from_arch(
        vocab: usize,
        n_heads: usize,
        head_dim: usize,
        n_layers: usize,
        ffn_hidden: usize,
    ) -> Self {
        Config {
            vocab,
            dim: n_heads * head_dim,
            n_layers,
            n_heads,
            head_dim,
            ffn_hidden,
            eps: 1e-5,
            rope_theta: 10000.0,
        }
    }

    /// Transformer-core parameter count: everything except the token embedding and
    /// the LM head (the two `vocab × dim` tables). This is the figure the scaling
    /// ladder is sized against — the 50257-vocab embed+lm_head adds a fixed ~25M on
    /// top that does not reflect model capacity. `num_params() = core + 2·vocab·dim`.
    pub fn core_params(&self) -> usize {
        self.num_params() - 2 * self.vocab * self.dim
    }

    /// 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
    }
}