//! Configuration types for PIDDM demo.

use serde::{Deserialize, Serialize};

/// PDE type for diffusion model training/inference.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
#[serde(rename_all = "snake_case")]
pub enum PdeType {
    /// Poisson equation: ∇²u = f
    #[default]
    Poisson,
    /// Heat equation: ∂u/∂t = α∇²u
    Heat,
    /// Darcy flow: -∇·(k∇p) = f
    Darcy,
    /// Burgers equation: ∂u/∂t + u·∇u = ν∇²u
    Burgers,
}

impl PdeType {
    /// Get display name for PDE type.
    #[must_use]
    pub fn display_name(&self) -> &'static str {
        match self {
            Self::Poisson => "Poisson Equation",
            Self::Heat => "Heat Equation",
            Self::Darcy => "Darcy Flow",
            Self::Burgers => "Burgers Equation",
        }
    }
}

/// Noise scheduler type.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
#[serde(rename_all = "snake_case")]
pub enum SchedulerType {
    /// Denoising Diffusion Probabilistic Models
    #[default]
    DDPM,
    /// Denoising Diffusion Implicit Models
    DDIM,
}

/// Configuration for PIDDM training.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PiddmTrainingConfig {
    /// PDE type to solve
    pub pde_type: PdeType,
    /// Grid resolution
    pub resolution: u32,
    /// Number of training epochs
    pub epochs: u32,
    /// Learning rate
    pub learning_rate: f64,
    /// Batch size
    pub batch_size: u32,
    /// Number of diffusion timesteps
    pub diffusion_steps: u32,
    /// Physics loss weight
    pub physics_weight: f32,
    /// Scheduler type
    pub scheduler: SchedulerType,
    /// `UNet` hidden dimension
    pub hidden_dim: usize,
    /// Number of `UNet` layers
    pub num_layers: usize,
}

impl Default for PiddmTrainingConfig {
    fn default() -> Self {
        Self {
            pde_type: PdeType::Poisson,
            resolution: 64,
            epochs: 100,
            learning_rate: 1e-4,
            batch_size: 16,
            diffusion_steps: 1000,
            physics_weight: 0.1,
            scheduler: SchedulerType::DDPM,
            hidden_dim: 64,
            num_layers: 4,
        }
    }
}

/// Configuration for PIDDM sampling/inference.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PiddmSamplingConfig {
    /// Number of samples to generate
    pub num_samples: u32,
    /// Grid resolution
    pub resolution: u32,
    /// Number of denoising steps (can be less than training steps for DDIM)
    pub sampling_steps: u32,
    /// Physics guidance strength (0 = no guidance, 1 = full guidance)
    pub guidance_strength: f32,
    /// Scheduler type
    pub scheduler: SchedulerType,
}

impl Default for PiddmSamplingConfig {
    fn default() -> Self {
        Self {
            num_samples: 4,
            resolution: 64,
            sampling_steps: 50,
            guidance_strength: 0.5,
            scheduler: SchedulerType::DDIM,
        }
    }
}

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

    #[test]
    fn test_pde_type_display() {
        assert_eq!(PdeType::Poisson.display_name(), "Poisson Equation");
        assert_eq!(PdeType::Heat.display_name(), "Heat Equation");
    }

    #[test]
    fn test_training_config_default() {
        let config = PiddmTrainingConfig::default();
        assert_eq!(config.resolution, 64);
        assert_eq!(config.diffusion_steps, 1000);
    }

    #[test]
    fn test_sampling_config_default() {
        let config = PiddmSamplingConfig::default();
        assert_eq!(config.sampling_steps, 50);
        assert_eq!(config.guidance_strength, 0.5);
    }
}