//! Configuration for xLSTM (Extended Long Short-Term Memory).

use serde::{Deserialize, Serialize};

/// Configuration for an xLSTM network.
///
/// The xLSTM architecture interleaves sLSTM blocks (exponential gating, scalar memory)
/// and mLSTM blocks (matrix memory with key-value association).
/// `slstm_ratio` controls the proportion of sLSTM vs mLSTM blocks.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct XLSTMConfig {
    /// Input feature dimension.
    pub input_dim: usize,
    /// Hidden state dimension. OOM note: mLSTM matrix memory is (`head_dim` × `head_dim`) per head.
    pub hidden_dim: usize,
    /// Number of xLSTM blocks.
    pub num_blocks: usize,
    /// Number of attention heads for mLSTM blocks.
    pub num_heads: usize,
    /// Ratio of sLSTM blocks (0.0 = all mLSTM, 1.0 = all sLSTM).
    pub slstm_ratio: f64,
    /// Output dimension (prediction size).
    pub output_dim: usize,
    /// Dropout rate (applied between blocks during training).
    pub dropout: f64,
    /// Learning rate for `AdamW` optimizer.
    pub learning_rate: f64,
    /// Weight decay for regularization.
    pub weight_decay: f64,
    /// Gradient clipping max norm.
    pub grad_clip: f64,
}

impl Default for XLSTMConfig {
    fn default() -> Self {
        Self {
            input_dim: 51,
            hidden_dim: 64,
            num_blocks: 4,
            num_heads: 4,
            slstm_ratio: 0.5,
            output_dim: 1,
            dropout: 0.1,
            learning_rate: 1e-3,
            weight_decay: 1e-4,
            grad_clip: 1.0,
        }
    }
}
