//! Configuration for Diffusion model (DDPM/DDIM) for price path generation.

use serde::{Deserialize, Serialize};

/// Noise schedule type for the diffusion process.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum NoiseSchedule {
    /// Linear beta schedule from `beta_start` to `beta_end`.
    Linear,
    /// Cosine schedule (Nichol & Dhariwal) -- smoother noise progression.
    Cosine,
}

impl Default for NoiseSchedule {
    fn default() -> Self {
        Self::Cosine
    }
}

/// Configuration for a Diffusion model.
///
/// OOM-safe defaults: small channels, short sequences, few res blocks.
/// RTX 3050 Ti (4GB) safe at `batch_size` <= 32.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DiffusionConfig {
    /// Number of diffusion timesteps (training noise levels).
    pub num_timesteps: usize,
    /// Number of DDIM sampling steps (inference, << `num_timesteps`).
    pub sampling_steps: usize,
    /// Sequence length of generated price paths.
    pub seq_len: usize,
    /// Feature dimension per timestep (1 = univariate price).
    pub feature_dim: usize,
    /// Hidden dimension for the denoiser network.
    pub hidden_dim: usize,
    /// Number of denoiser layers.
    pub num_layers: usize,
    /// Time embedding dimension.
    pub time_embed_dim: usize,
    /// Noise schedule type.
    pub schedule: NoiseSchedule,
    /// Learning rate.
    pub learning_rate: f64,
    /// Weight decay for regularization.
    pub weight_decay: f64,
    /// Gradient clipping max norm.
    pub grad_clip: f64,
}

impl Default for DiffusionConfig {
    fn default() -> Self {
        Self {
            num_timesteps: 1000,
            sampling_steps: 10,
            seq_len: 64,
            feature_dim: 1,
            hidden_dim: 128,
            num_layers: 3,
            time_embed_dim: 32,
            schedule: NoiseSchedule::Cosine,
            learning_rate: 1e-4,
            weight_decay: 1e-4,
            grad_clip: 1.0,
        }
    }
}

impl DiffusionConfig {
    /// Total data dimension (`seq_len` * `feature_dim`).
    pub const fn data_dim(&self) -> usize {
        self.seq_len * self.feature_dim
    }
}

#[cfg(test)]
#[allow(clippy::assertions_on_result_states)]
mod tests {
    use super::*;

    #[test]
    fn test_default_config() {
        let config = DiffusionConfig::default();
        assert_eq!(config.num_timesteps, 1000);
        assert_eq!(config.sampling_steps, 10);
        assert_eq!(config.data_dim(), 64);
    }

    #[test]
    fn test_data_dim_multivariate() {
        let config = DiffusionConfig {
            seq_len: 32,
            feature_dim: 4,
            ..Default::default()
        };
        assert_eq!(config.data_dim(), 128);
    }

    #[test]
    fn test_serde_roundtrip() {
        let config = DiffusionConfig::default();
        let json = serde_json::to_string(&config).unwrap_or_default();
        let restored: Result<DiffusionConfig, _> = serde_json::from_str(&json);
        assert!(restored.is_ok());
    }
}
