//! Shared types for NeuralOp Studio - Neural Operator Workbench.
//!
//! This crate defines the IPC types for training neural operators on PDEs.

use serde::{Deserialize, Serialize};

// ============================================================================
// PDE Types
// ============================================================================

/// Type of PDE to solve.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
pub enum PDEType {
    /// Poisson equation: -∇²u = f
    Poisson,
    /// Heat equation: ∂u/∂t = α∇²u
    Heat,
    /// Wave equation: ∂²u/∂t² = c²∇²u
    Wave,
    /// Burgers' equation: ∂u/∂t + u∂u/∂x = ν∇²u
    Burgers,
    /// Navier-Stokes equations
    NavierStokes,
    /// Advection equation: ∂u/∂t + c·∇u = 0
    Advection,
    /// Diffusion-reaction: ∂u/∂t = D∇²u + R(u)
    DiffusionReaction,
    /// Custom PDE defined by user
    Custom,
}

/// PDE definition.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PDEDefinition {
    /// Type of PDE.
    pub pde_type: PDEType,
    /// Equation string (LaTeX format).
    pub equation: String,
    /// Domain description.
    pub domain: Domain,
    /// Boundary conditions.
    pub boundary_conditions: Vec<BoundaryCondition>,
    /// Initial condition (for time-dependent PDEs).
    pub initial_condition: Option<InitialCondition>,
    /// Physical parameters.
    pub parameters: PDEParameters,
}

/// Domain definition.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Domain {
    /// Spatial dimensions (1D, 2D, 3D).
    pub dimensions: usize,
    /// Domain bounds: [(x_min, x_max), (y_min, y_max), ...]
    pub bounds: Vec<(f64, f64)>,
    /// Grid resolution per dimension.
    pub resolution: Vec<usize>,
    /// Time bounds (for time-dependent PDEs).
    pub time_bounds: Option<(f64, f64)>,
    /// Time resolution (for time-dependent PDEs).
    pub time_resolution: Option<usize>,
}

/// Boundary condition type.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub enum BCType {
    /// Dirichlet: u = g on boundary
    Dirichlet,
    /// Neumann: ∂u/∂n = g on boundary
    Neumann,
    /// Robin: au + b∂u/∂n = g on boundary
    Robin,
    /// Periodic: u(x) = u(x + L)
    Periodic,
}

/// Boundary condition.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BoundaryCondition {
    /// Boundary identifier (e.g., "left", "right", "top").
    pub boundary: String,
    /// Type of boundary condition.
    pub bc_type: BCType,
    /// Value or function description.
    pub value: String,
    /// Coefficients for Robin BC: (a, b).
    pub robin_coefficients: Option<(f64, f64)>,
}

/// Initial condition.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InitialCondition {
    /// Function description.
    pub function: String,
    /// Initial velocity (for wave equation).
    pub velocity: Option<String>,
}

/// Physical parameters for PDEs.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PDEParameters {
    /// Diffusion coefficient (α, D, ν).
    pub diffusion: Option<f64>,
    /// Wave speed (c).
    pub wave_speed: Option<f64>,
    /// Reaction rate.
    pub reaction_rate: Option<f64>,
    /// Reynolds number (for Navier-Stokes).
    pub reynolds_number: Option<f64>,
    /// Source term.
    pub source_term: Option<String>,
    /// Custom parameters.
    pub custom: Vec<(String, f64)>,
}

impl Default for PDEParameters {
    fn default() -> Self {
        Self {
            diffusion: Some(1.0),
            wave_speed: Some(1.0),
            reaction_rate: None,
            reynolds_number: None,
            source_term: None,
            custom: vec![],
        }
    }
}

// ============================================================================
// Neural Operator Types
// ============================================================================

/// Neural operator architecture type.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
pub enum OperatorType {
    /// Fourier Neural Operator.
    FNO,
    /// DeepONet (Deep Operator Network).
    DeepONet,
    /// Physics-Informed Neural Operator (PINO).
    PINO,
    /// Spectral-Inspired Neural Operator (SINO).
    SINO,
    /// Galerkin Transformer.
    GalerkinTransformer,
    /// Message Passing Neural Operator.
    MPNO,
}

/// Neural operator configuration.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OperatorConfig {
    /// Operator type.
    pub operator_type: OperatorType,
    /// Hidden dimension.
    pub hidden_dim: usize,
    /// Number of layers.
    pub num_layers: usize,
    /// Number of Fourier modes (for FNO).
    pub fourier_modes: Option<usize>,
    /// Branch network width (for DeepONet).
    pub branch_width: Option<usize>,
    /// Trunk network width (for DeepONet).
    pub trunk_width: Option<usize>,
    /// Physics loss weight (for PINO).
    pub physics_weight: Option<f64>,
    /// Activation function.
    pub activation: ActivationType,
    /// Use residual connections.
    pub residual: bool,
    /// Dropout rate.
    pub dropout: f64,
}

impl Default for OperatorConfig {
    fn default() -> Self {
        Self {
            operator_type: OperatorType::FNO,
            hidden_dim: 64,
            num_layers: 4,
            fourier_modes: Some(12),
            branch_width: None,
            trunk_width: None,
            physics_weight: None,
            activation: ActivationType::GELU,
            residual: true,
            dropout: 0.0,
        }
    }
}

/// Activation function type.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub enum ActivationType {
    /// ReLU.
    ReLU,
    /// GELU.
    GELU,
    /// Tanh.
    Tanh,
    /// Swish / SiLU.
    Swish,
    /// LeakyReLU.
    LeakyReLU,
}

// ============================================================================
// Training Types
// ============================================================================

/// Training configuration.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainingConfig {
    /// Number of training epochs.
    pub epochs: usize,
    /// Batch size.
    pub batch_size: usize,
    /// Learning rate.
    pub learning_rate: f64,
    /// Learning rate scheduler.
    pub scheduler: SchedulerType,
    /// Optimizer type.
    pub optimizer: OptimizerType,
    /// Weight decay.
    pub weight_decay: f64,
    /// Number of training samples.
    pub num_train_samples: usize,
    /// Number of validation samples.
    pub num_val_samples: usize,
    /// Number of test samples.
    pub num_test_samples: usize,
    /// Random seed.
    pub seed: Option<u64>,
}

impl Default for TrainingConfig {
    fn default() -> Self {
        Self {
            epochs: 100,
            batch_size: 32,
            learning_rate: 1e-3,
            scheduler: SchedulerType::CosineAnnealing,
            optimizer: OptimizerType::AdamW,
            weight_decay: 1e-4,
            num_train_samples: 1000,
            num_val_samples: 100,
            num_test_samples: 100,
            seed: Some(42),
        }
    }
}

/// Learning rate scheduler type.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub enum SchedulerType {
    /// Constant learning rate.
    Constant,
    /// Step decay.
    StepDecay,
    /// Exponential decay.
    ExponentialDecay,
    /// Cosine annealing.
    CosineAnnealing,
    /// Warmup with linear decay.
    WarmupLinearDecay,
}

/// Optimizer type.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub enum OptimizerType {
    /// SGD.
    SGD,
    /// Adam.
    Adam,
    /// AdamW.
    AdamW,
    /// RMSprop.
    RMSprop,
}

/// Training progress update.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainingProgress {
    /// Current epoch.
    pub epoch: usize,
    /// Total epochs.
    pub total_epochs: usize,
    /// Current batch.
    pub batch: usize,
    /// Total batches.
    pub total_batches: usize,
    /// Training loss.
    pub train_loss: f64,
    /// Validation loss.
    pub val_loss: Option<f64>,
    /// Physics loss (for PINO).
    pub physics_loss: Option<f64>,
    /// Relative L2 error.
    pub relative_error: Option<f64>,
    /// Current learning rate.
    pub learning_rate: f64,
    /// Elapsed time (seconds).
    pub elapsed_seconds: f64,
}

// ============================================================================
// Prediction and Results Types
// ============================================================================

/// Prediction request.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PredictionRequest {
    /// Input function samples.
    pub input_function: Vec<f64>,
    /// Query points.
    pub query_points: Vec<Vec<f64>>,
    /// Time point (for time-dependent PDEs).
    pub time: Option<f64>,
}

/// Prediction result.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PredictionResult {
    /// Predicted solution values.
    pub solution: Vec<f64>,
    /// Gradient (if computed).
    pub gradient: Option<Vec<Vec<f64>>>,
    /// Uncertainty estimate (if available).
    pub uncertainty: Option<Vec<f64>>,
    /// Inference time (milliseconds).
    pub inference_time_ms: f64,
}

/// Model evaluation metrics.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EvaluationMetrics {
    /// Mean squared error.
    pub mse: f64,
    /// Relative L2 error.
    pub relative_l2: f64,
    /// Maximum absolute error.
    pub max_error: f64,
    /// Physics residual (for PINO).
    pub physics_residual: Option<f64>,
    /// Number of test samples.
    pub num_samples: usize,
    /// Average inference time (ms).
    pub avg_inference_time_ms: f64,
}

/// Training result.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainingResult {
    /// Final training loss.
    pub final_train_loss: f64,
    /// Final validation loss.
    pub final_val_loss: f64,
    /// Best epoch.
    pub best_epoch: usize,
    /// Training loss history.
    pub train_loss_history: Vec<f64>,
    /// Validation loss history.
    pub val_loss_history: Vec<f64>,
    /// Evaluation metrics on test set.
    pub test_metrics: EvaluationMetrics,
    /// Total training time (seconds).
    pub total_time_seconds: f64,
    /// Number of parameters.
    pub num_parameters: usize,
}

// ============================================================================
// Visualization Types
// ============================================================================

/// Solution field for visualization.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SolutionField {
    /// Grid coordinates (flattened).
    pub coordinates: Vec<Vec<f64>>,
    /// Solution values.
    pub values: Vec<f64>,
    /// Field name.
    pub name: String,
    /// Time point (for time-dependent solutions).
    pub time: Option<f64>,
}

/// Comparison between ground truth and prediction.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SolutionComparison {
    /// Ground truth solution.
    pub ground_truth: SolutionField,
    /// Predicted solution.
    pub prediction: SolutionField,
    /// Point-wise error.
    pub error: SolutionField,
    /// Relative error statistics.
    pub error_stats: ErrorStatistics,
}

/// Error statistics.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ErrorStatistics {
    /// Mean error.
    pub mean: f64,
    /// Standard deviation.
    pub std: f64,
    /// Maximum error.
    pub max: f64,
    /// Minimum error.
    pub min: f64,
    /// 95th percentile.
    pub p95: f64,
}

// ============================================================================
// Sample Data Functions
// ============================================================================

/// Create a sample Poisson problem.
#[must_use]
pub fn sample_poisson_problem() -> PDEDefinition {
    PDEDefinition {
        pde_type: PDEType::Poisson,
        equation: r"-\nabla^2 u = f".to_string(),
        domain: Domain {
            dimensions: 2,
            bounds: vec![(0.0, 1.0), (0.0, 1.0)],
            resolution: vec![64, 64],
            time_bounds: None,
            time_resolution: None,
        },
        boundary_conditions: vec![BoundaryCondition {
            boundary: "all".to_string(),
            bc_type: BCType::Dirichlet,
            value: "0".to_string(),
            robin_coefficients: None,
        }],
        initial_condition: None,
        parameters: PDEParameters {
            source_term: Some("sin(pi*x)*sin(pi*y)".to_string()),
            ..Default::default()
        },
    }
}

/// Create a sample heat equation problem.
#[must_use]
pub fn sample_heat_problem() -> PDEDefinition {
    PDEDefinition {
        pde_type: PDEType::Heat,
        equation: r"\frac{\partial u}{\partial t} = \alpha \nabla^2 u".to_string(),
        domain: Domain {
            dimensions: 2,
            bounds: vec![(0.0, 1.0), (0.0, 1.0)],
            resolution: vec![64, 64],
            time_bounds: Some((0.0, 1.0)),
            time_resolution: Some(100),
        },
        boundary_conditions: vec![BoundaryCondition {
            boundary: "all".to_string(),
            bc_type: BCType::Dirichlet,
            value: "0".to_string(),
            robin_coefficients: None,
        }],
        initial_condition: Some(InitialCondition {
            function: "sin(pi*x)*sin(pi*y)".to_string(),
            velocity: None,
        }),
        parameters: PDEParameters {
            diffusion: Some(0.01),
            ..Default::default()
        },
    }
}

/// Create a sample Burgers' equation problem.
#[must_use]
pub fn sample_burgers_problem() -> PDEDefinition {
    PDEDefinition {
        pde_type: PDEType::Burgers,
        equation:
            r"\frac{\partial u}{\partial t} + u\frac{\partial u}{\partial x} = \nu \nabla^2 u"
                .to_string(),
        domain: Domain {
            dimensions: 1,
            bounds: vec![(0.0, 2.0 * std::f64::consts::PI)],
            resolution: vec![256],
            time_bounds: Some((0.0, 1.0)),
            time_resolution: Some(100),
        },
        boundary_conditions: vec![BoundaryCondition {
            boundary: "left".to_string(),
            bc_type: BCType::Periodic,
            value: "periodic".to_string(),
            robin_coefficients: None,
        }],
        initial_condition: Some(InitialCondition {
            function: "sin(x)".to_string(),
            velocity: None,
        }),
        parameters: PDEParameters {
            diffusion: Some(0.01),
            ..Default::default()
        },
    }
}

/// Create a sample Navier-Stokes problem.
#[must_use]
pub fn sample_navier_stokes_problem() -> PDEDefinition {
    PDEDefinition {
        pde_type: PDEType::NavierStokes,
        equation: r"\frac{\partial \mathbf{u}}{\partial t} + (\mathbf{u} \cdot \nabla)\mathbf{u} = -\nabla p + \nu \nabla^2 \mathbf{u}".to_string(),
        domain: Domain {
            dimensions: 2,
            bounds: vec![(0.0, 2.0 * std::f64::consts::PI), (0.0, 2.0 * std::f64::consts::PI)],
            resolution: vec![64, 64],
            time_bounds: Some((0.0, 10.0)),
            time_resolution: Some(100),
        },
        boundary_conditions: vec![
            BoundaryCondition {
                boundary: "all".to_string(),
                bc_type: BCType::Periodic,
                value: "periodic".to_string(),
                robin_coefficients: None,
            },
        ],
        initial_condition: Some(InitialCondition {
            function: "vortex_pair".to_string(),
            velocity: None,
        }),
        parameters: PDEParameters {
            reynolds_number: Some(1000.0),
            diffusion: Some(0.001),
            ..Default::default()
        },
    }
}

/// Create a sample FNO configuration.
#[must_use]
pub fn sample_fno_config() -> OperatorConfig {
    OperatorConfig {
        operator_type: OperatorType::FNO,
        hidden_dim: 64,
        num_layers: 4,
        fourier_modes: Some(12),
        activation: ActivationType::GELU,
        residual: true,
        dropout: 0.0,
        ..Default::default()
    }
}

/// Create a sample DeepONet configuration.
#[must_use]
pub fn sample_deeponet_config() -> OperatorConfig {
    OperatorConfig {
        operator_type: OperatorType::DeepONet,
        hidden_dim: 100,
        num_layers: 6,
        branch_width: Some(100),
        trunk_width: Some(100),
        activation: ActivationType::Tanh,
        residual: false,
        dropout: 0.0,
        ..Default::default()
    }
}

/// Create a sample PINO configuration.
#[must_use]
pub fn sample_pino_config() -> OperatorConfig {
    OperatorConfig {
        operator_type: OperatorType::PINO,
        hidden_dim: 64,
        num_layers: 4,
        fourier_modes: Some(12),
        physics_weight: Some(0.1),
        activation: ActivationType::GELU,
        residual: true,
        dropout: 0.0,
        ..Default::default()
    }
}

/// Create a sample training configuration.
#[must_use]
pub fn sample_training_config() -> TrainingConfig {
    TrainingConfig {
        epochs: 100,
        batch_size: 32,
        learning_rate: 1e-3,
        scheduler: SchedulerType::CosineAnnealing,
        optimizer: OptimizerType::AdamW,
        weight_decay: 1e-4,
        num_train_samples: 1000,
        num_val_samples: 100,
        num_test_samples: 100,
        seed: Some(42),
    }
}

// ============================================================================
// Tests
// ============================================================================

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

    #[test]
    fn test_pde_types() {
        assert_eq!(PDEType::Poisson, PDEType::Poisson);
        assert_ne!(PDEType::Heat, PDEType::Wave);
    }

    #[test]
    fn test_sample_poisson() {
        let problem = sample_poisson_problem();
        assert_eq!(problem.pde_type, PDEType::Poisson);
        assert_eq!(problem.domain.dimensions, 2);
        assert!(problem.initial_condition.is_none());
    }

    #[test]
    fn test_sample_heat() {
        let problem = sample_heat_problem();
        assert_eq!(problem.pde_type, PDEType::Heat);
        assert!(problem.initial_condition.is_some());
        assert!(problem.domain.time_bounds.is_some());
    }

    #[test]
    fn test_sample_burgers() {
        let problem = sample_burgers_problem();
        assert_eq!(problem.pde_type, PDEType::Burgers);
        assert_eq!(problem.domain.dimensions, 1);
    }

    #[test]
    fn test_sample_navier_stokes() {
        let problem = sample_navier_stokes_problem();
        assert_eq!(problem.pde_type, PDEType::NavierStokes);
        assert!(problem.parameters.reynolds_number.is_some());
    }

    #[test]
    fn test_operator_configs() {
        let fno = sample_fno_config();
        assert_eq!(fno.operator_type, OperatorType::FNO);
        assert!(fno.fourier_modes.is_some());

        let deeponet = sample_deeponet_config();
        assert_eq!(deeponet.operator_type, OperatorType::DeepONet);
        assert!(deeponet.branch_width.is_some());

        let pino = sample_pino_config();
        assert_eq!(pino.operator_type, OperatorType::PINO);
        assert!(pino.physics_weight.is_some());
    }

    #[test]
    fn test_training_config() {
        let config = sample_training_config();
        assert_eq!(config.epochs, 100);
        assert_eq!(config.optimizer, OptimizerType::AdamW);
    }

    #[test]
    fn test_serialization() {
        let problem = sample_poisson_problem();
        let json = serde_json::to_string(&problem).unwrap();
        assert!(json.contains("Poisson"));

        let parsed: PDEDefinition = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed.pde_type, problem.pde_type);
    }

    #[test]
    fn test_operator_config_serialization() {
        let config = sample_fno_config();
        let json = serde_json::to_string(&config).unwrap();
        assert!(json.contains("FNO"));

        let parsed: OperatorConfig = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed.operator_type, config.operator_type);
    }
}