/// Model and runtime configuration for SleepFM inference.
///
/// `ModelConfig` mirrors the Python SleepFM hyperparameters from config.json.

#[derive(Debug, Clone, serde::Deserialize)]
pub struct ModelConfig {
    /// Embedding dimension (default 128).
    #[serde(default = "default_embed_dim")]
    pub embed_dim: usize,

    /// Number of attention heads in the transformer encoder (default 8).
    #[serde(default = "default_num_heads")]
    pub num_heads: usize,

    /// Number of transformer encoder layers (default 6).
    #[serde(default = "default_num_layers")]
    pub num_layers: usize,

    /// Number of heads in spatial/temporal attention pooling (default 8).
    #[serde(default = "default_pooling_head")]
    pub pooling_head: usize,

    /// Dropout rate (default 0.3, but set to 0.0 for inference).
    #[serde(default)]
    pub dropout: f64,

    /// Patch size in samples (default 640 = 5s * 128 Hz).
    #[serde(default = "default_patch_size")]
    pub patch_size: usize,

    /// Max sequence length for positional encoding (default 128).
    #[serde(default = "default_max_seq_length")]
    pub max_seq_length: usize,

    /// Sampling frequency in Hz (default 128).
    #[serde(default = "default_sampling_freq")]
    pub sampling_freq: usize,

    /// Sampling duration in minutes (default 5).
    #[serde(default = "default_sampling_duration")]
    pub sampling_duration: usize,
}

fn default_embed_dim()         -> usize { 128 }
fn default_num_heads()         -> usize { 8 }
fn default_num_layers()        -> usize { 6 }
fn default_pooling_head()      -> usize { 8 }
fn default_patch_size()        -> usize { 640 }
fn default_max_seq_length()    -> usize { 128 }
fn default_sampling_freq()     -> usize { 128 }
fn default_sampling_duration() -> usize { 5 }

impl Default for ModelConfig {
    fn default() -> Self {
        Self {
            embed_dim:         default_embed_dim(),
            num_heads:         default_num_heads(),
            num_layers:        default_num_layers(),
            pooling_head:      default_pooling_head(),
            dropout:           0.0,
            patch_size:        default_patch_size(),
            max_seq_length:    default_max_seq_length(),
            sampling_freq:     default_sampling_freq(),
            sampling_duration: default_sampling_duration(),
        }
    }
}

impl ModelConfig {
    /// Samples per 5-minute chunk: sampling_duration * 60 * sampling_freq
    pub fn samples_per_chunk(&self) -> usize {
        self.sampling_duration * 60 * self.sampling_freq
    }

    /// Number of time patches per chunk: samples_per_chunk / patch_size
    pub fn num_patches(&self) -> usize {
        self.samples_per_chunk() / self.patch_size
    }

    /// FFN hidden dimension inside TransformerEncoderLayer (PyTorch default = 2048 for d_model=128..512).
    /// PyTorch nn.TransformerEncoderLayer default: dim_feedforward = 2048.
    pub fn ffn_dim(&self) -> usize {
        2048
    }
}
