//! Augmented Neural ODEs implementation
//!
//! Augmented Neural ODEs increase the expressivity of Neural ODEs by adding
//! extra dimensions to the state space. This helps overcome limitations of
//! standard Neural ODEs in modeling complex dynamics.
//!
//! Based on "Augmented Neural ODEs" (Dupont et al., 2019)

use crate::prelude::*;
use super::{ODEFunc, NeuralODE, NeuralODEConfig, Result, NeuralODEError};

/// Configuration for augmented Neural ODEs
#[derive(Debug, Clone)]
pub struct AugmentationConfig {
    /// Number of original dimensions
    pub original_dim: usize,
    /// Number of augmented dimensions to add
    pub augmented_dim: usize,
    /// Initialization strategy for augmented dimensions
    pub init_strategy: AugmentationInit,
    /// Whether to include augmented dimensions in output
    pub return_augmented: bool,
}

/// Initialization strategies for augmented dimensions
#[derive(Debug, Clone)]
pub enum AugmentationInit {
    /// Initialize to zeros
    Zeros,
    /// Initialize to small random values
    SmallRandom { scale: f32 },
    /// Initialize using a linear transformation of the original state
    Linear { matrix: Tensor },
    /// Initialize to constant values
    Constant { value: f32 },
}

impl Default for AugmentationConfig {
    fn default() -> Self {
        Self {
            original_dim: 0, // Must be set explicitly
            augmented_dim: 1,
            init_strategy: AugmentationInit::SmallRandom { scale: 0.01 },
            return_augmented: false,
        }
    }
}

/// Augmented Neural ODE that adds extra dimensions to increase expressivity
pub struct AugmentedNeuralODE {
    /// Underlying Neural ODE operating on augmented state
    neural_ode: NeuralODE,
    /// Augmentation configuration
    config: AugmentationConfig,
    /// Device for computations
    device: Device,
}

impl AugmentedNeuralODE {
    /// Create a new Augmented Neural ODE
    ///
    /// # Arguments
    /// - `dynamics`: ODE dynamics function for the augmented system
    /// - `ode_config`: Configuration for the underlying Neural ODE
    /// - `original_dim`: Number of original state dimensions
    /// - `augmented_dim`: Number of extra dimensions to add
    /// - `device`: Device for tensor operations
    pub fn new(
        dynamics: Box<dyn ODEFunc>,
        ode_config: NeuralODEConfig,
        original_dim: usize,
        augmented_dim: usize,
        device: &Device,
    ) -> Result<Self> {
        let augmentation_config = AugmentationConfig {
            original_dim,
            augmented_dim,
            init_strategy: AugmentationInit::SmallRandom { scale: 0.01 },
            return_augmented: false,
        };

        let neural_ode = NeuralODE::new(dynamics, ode_config, device)?;

        Ok(Self {
            neural_ode,
            config: augmentation_config,
            device: device.clone(),
        })
    }

    /// Create with custom augmentation configuration
    pub fn with_config(
        dynamics: Box<dyn ODEFunc>,
        ode_config: NeuralODEConfig,
        augmentation_config: AugmentationConfig,
        device: &Device,
    ) -> Result<Self> {
        let neural_ode = NeuralODE::new(dynamics, ode_config, device)?;

        Ok(Self {
            neural_ode,
            config: augmentation_config,
            device: device.clone(),
        })
    }

    /// Augment the initial state with extra dimensions
    fn augment_state(&self, y0: &Tensor) -> Result<Tensor> {
        let batch_size = if y0.ndim() == 2 {
            Some(y0.shape()[0])
        } else {
            None
        };

        let original_shape = if let Some(bs) = batch_size {
            [bs, self.config.original_dim]
        } else {
            [self.config.original_dim]
        };

        // Verify original state has correct dimensions
        let expected_shape = if batch_size.is_some() { &original_shape[..] } else { &original_shape[1..] };
        if y0.shape() != expected_shape {
            return Err(NeuralODEError::InvalidInput(format!(
                "Expected original state shape {:?}, got {:?}",
                expected_shape, y0.shape()
            )));
        }

        // Create augmented dimensions
        let augmented_part = match &self.config.init_strategy {
            AugmentationInit::Zeros => {
                if let Some(bs) = batch_size {
                    Tensor::zeros([bs, self.config.augmented_dim], &self.device)?
                } else {
                    Tensor::zeros([self.config.augmented_dim], &self.device)?
                }
            }
            AugmentationInit::SmallRandom { scale } => {
                let shape = if let Some(bs) = batch_size {
                    [bs, self.config.augmented_dim]
                } else {
                    [self.config.augmented_dim]
                };
                Tensor::randn(shape, &self.device)?.mul_scalar(*scale)?
            }
            AugmentationInit::Linear { matrix } => {
                // Apply linear transformation to original state
                if let Some(_bs) = batch_size {
                    y0.matmul(&matrix.t()?)?
                } else {
                    matrix.matmul(y0)?
                }
            }
            AugmentationInit::Constant { value } => {
                let shape = if let Some(bs) = batch_size {
                    [bs, self.config.augmented_dim]
                } else {
                    [self.config.augmented_dim]
                };
                Tensor::full(shape, *value, &self.device)?
            }
        };

        // Concatenate original and augmented parts
        let augmented_state = if batch_size.is_some() {
            Tensor::cat(&[y0.clone(), augmented_part], 1)?
        } else {
            Tensor::cat(&[y0.clone(), augmented_part], 0)?
        };

        Ok(augmented_state)
    }

    /// Extract original dimensions from augmented state
    fn extract_original(&self, augmented_state: &Tensor) -> Result<Tensor> {
        if augmented_state.ndim() == 1 {
            // Single trajectory: [total_dim] -> [original_dim]
            augmented_state.slice(0, 0..self.config.original_dim)
        } else if augmented_state.ndim() == 2 {
            if augmented_state.shape()[0] == self.config.original_dim + self.config.augmented_dim {
                // Single time series: [time_steps, total_dim] -> [time_steps, original_dim]
                augmented_state.slice(1, 0..self.config.original_dim)
            } else {
                // Batch: [batch_size, total_dim] -> [batch_size, original_dim]
                augmented_state.slice(1, 0..self.config.original_dim)
            }
        } else if augmented_state.ndim() == 3 {
            // Batch time series: [batch_size, time_steps, total_dim] -> [batch_size, time_steps, original_dim]
            augmented_state.slice(2, 0..self.config.original_dim)
        } else {
            Err(NeuralODEError::InvalidInput(format!(
                "Unsupported tensor dimensionality: {}",
                augmented_state.ndim()
            )))
        }
    }

    /// Forward pass with augmented state
    pub fn forward(&self, y0: &Tensor, t_span: &[f32]) -> Result<Tensor> {
        // Augment initial state
        let augmented_y0 = self.augment_state(y0)?;

        // Solve ODE in augmented space
        let augmented_solution = self.neural_ode.forward(&augmented_y0, t_span)?;

        // Return either full augmented state or just original dimensions
        if self.config.return_augmented {
            Ok(augmented_solution)
        } else {
            self.extract_original(&augmented_solution)
        }
    }

    /// Forward pass for batch of initial conditions
    pub fn forward_batch(&self, y0_batch: &Tensor, t_span: &[f32]) -> Result<Tensor> {
        if y0_batch.ndim() != 2 {
            return Err(NeuralODEError::InvalidInput(
                "Batch input must be 2D [batch_size, original_dim]".to_string()
            ));
        }

        // Augment batch
        let augmented_y0_batch = self.augment_state(y0_batch)?;

        // Solve in augmented space
        let augmented_solution = self.neural_ode.forward_batch(&augmented_y0_batch, t_span)?;

        // Extract original dimensions
        if self.config.return_augmented {
            Ok(augmented_solution)
        } else {
            self.extract_original(&augmented_solution)
        }
    }

    /// Get all parameters from underlying Neural ODE
    pub fn parameters(&self) -> Vec<Tensor> {
        self.neural_ode.parameters()
    }

    /// Get configuration
    pub fn config(&self) -> &AugmentationConfig {
        &self.config
    }

    /// Set whether to return augmented dimensions in output
    pub fn set_return_augmented(&mut self, return_augmented: bool) {
        self.config.return_augmented = return_augmented;
    }

    /// Get the underlying Neural ODE
    pub fn inner(&self) -> &NeuralODE {
        &self.neural_ode
    }

    /// Evaluate dynamics in augmented space
    pub fn evaluate_augmented_dynamics(&self, t: f32, augmented_y: &Tensor) -> Result<Tensor> {
        self.neural_ode.evaluate_dynamics(t, augmented_y)
    }

    /// Compare expressivity with standard Neural ODE
    ///
    /// This runs both augmented and standard versions and returns metrics
    pub fn compare_with_standard(
        &self,
        standard_dynamics: Box<dyn ODEFunc>,
        y0: &Tensor,
        t_span: &[f32],
    ) -> Result<ExpressionComparison> {
        // Run augmented version
        let augmented_result = self.forward(y0, t_span)?;

        // Run standard version
        let standard_config = NeuralODEConfig::default(); // Use same config as augmented
        let standard_ode = NeuralODE::new(standard_dynamics, standard_config, &self.device)?;
        let standard_result = standard_ode.forward(y0, t_span)?;

        // Compute comparison metrics
        let mse = augmented_result.sub(&standard_result)?.pow_scalar(2.0)?.mean()?;
        let max_diff = augmented_result.sub(&standard_result)?.abs()?.max()?;
        
        let mse_val = mse.to_scalar::<f32>()?;
        let max_diff_val = max_diff.to_scalar::<f32>()?;

        Ok(ExpressionComparison {
            mse: mse_val,
            max_absolute_diff: max_diff_val,
            augmented_trajectory: augmented_result,
            standard_trajectory: standard_result,
        })
    }
}

/// Results of comparing augmented vs standard Neural ODE
#[derive(Debug)]
pub struct ExpressionComparison {
    /// Mean squared error between trajectories
    pub mse: f32,
    /// Maximum absolute difference
    pub max_absolute_diff: f32,
    /// Trajectory from augmented Neural ODE
    pub augmented_trajectory: Tensor,
    /// Trajectory from standard Neural ODE
    pub standard_trajectory: Tensor,
}

/// Wrapper that automatically augments any ODEFunc
pub struct AugmentedODEFunc<F: ODEFunc> {
    inner: F,
    original_dim: usize,
    augmented_dim: usize,
}

impl<F: ODEFunc> AugmentedODEFunc<F> {
    pub fn new(inner: F, original_dim: usize, augmented_dim: usize) -> Self {
        Self {
            inner,
            original_dim,
            augmented_dim,
        }
    }
}

impl<F: ODEFunc> ODEFunc for AugmentedODEFunc<F> {
    fn forward(&self, t: f32, y: &Tensor) -> Result<Tensor> {
        // Extract original dimensions
        let y_original = if y.ndim() == 1 {
            y.slice(0, 0..self.original_dim)?
        } else {
            y.slice(1, 0..self.original_dim)?
        };

        // Apply original dynamics to original dimensions
        let dy_original = self.inner.forward(t, &y_original)?;

        // For augmented dimensions, we can use various strategies
        // Here we use a simple approach: augmented dimensions decay towards zero
        let y_augmented = if y.ndim() == 1 {
            y.slice(0, self.original_dim..)?
        } else {
            y.slice(1, self.original_dim..)?
        };

        let dy_augmented = y_augmented.mul_scalar(-0.1)?; // Slow decay

        // Combine derivatives
        let dy_combined = if y.ndim() == 1 {
            Tensor::cat(&[dy_original, dy_augmented], 0)?
        } else {
            Tensor::cat(&[dy_original, dy_augmented], 1)?
        };

        Ok(dy_combined)
    }

    fn parameters(&self) -> Vec<Tensor> {
        self.inner.parameters()
    }

    fn is_autonomous(&self) -> bool {
        self.inner.is_autonomous()
    }

    fn name(&self) -> String {
        format!("Augmented({})", self.inner.name())
    }
}

#[cfg(all(test, feature = "disabled_tests"))]
mod tests {
    use super::*;
    use crate::neural_ode::ode_func::LinearDynamics;

    #[test]
    fn test_augmentation_config() {
        let config = AugmentationConfig::default();
        assert_eq!(config.original_dim, 0);
        assert_eq!(config.augmented_dim, 1);
        assert!(!config.return_augmented);
        
        match config.init_strategy {
            AugmentationInit::SmallRandom { scale } => assert_eq!(scale, 0.01),
            _ => panic!("Expected SmallRandom initialization"),
        }
    }

    #[test]
    fn test_augmented_neural_ode_creation() {
        let device = Device::cpu();
        let dynamics = LinearDynamics::decay(1.0, 4, device.clone()).unwrap(); // 4D total (2 original + 2 augmented)
        let ode_config = NeuralODEConfig::default();
        
        let augmented_ode = AugmentedNeuralODE::new(
            Box::new(dynamics),
            ode_config,
            2, // original_dim
            2, // augmented_dim
            &device,
        ).unwrap();
        
        assert_eq!(augmented_ode.config.original_dim, 2);
        assert_eq!(augmented_ode.config.augmented_dim, 2);
    }

    #[test]
    fn test_state_augmentation() {
        let device = Device::cpu();
        let dynamics = LinearDynamics::decay(1.0, 3, device.clone()).unwrap(); // 3D total
        let ode_config = NeuralODEConfig::default();
        
        let augmented_ode = AugmentedNeuralODE::new(
            Box::new(dynamics),
            ode_config,
            2, // original_dim
            1, // augmented_dim
            &device,
        ).unwrap();
        
        let y0 = Tensor::ones([2], &device).unwrap();
        let augmented_y0 = augmented_ode.augment_state(&y0).unwrap();
        
        assert_eq!(augmented_y0.shape(), &[3]); // Should be [original_dim + augmented_dim]
        
        let aug_slice = augmented_y0.to_vec::<f32>().unwrap();
        // First two should be original state
        assert_eq!(aug_slice[0], 1.0);
        assert_eq!(aug_slice[1], 1.0);
        // Third should be small random value (approximately 0 for small scale)
        assert!(aug_slice[2].abs() < 0.1);
    }

    #[test]
    fn test_batch_augmentation() {
        let device = Device::cpu();
        let dynamics = LinearDynamics::decay(1.0, 4, device.clone()).unwrap(); // 4D total
        let ode_config = NeuralODEConfig::default();
        
        let augmented_ode = AugmentedNeuralODE::new(
            Box::new(dynamics),
            ode_config,
            2, // original_dim
            2, // augmented_dim
            &device,
        ).unwrap();
        
        let y0_batch = Tensor::ones([3, 2], &device).unwrap(); // Batch of 3
        let augmented_y0_batch = augmented_ode.augment_state(&y0_batch).unwrap();
        
        assert_eq!(augmented_y0_batch.shape(), &[3, 4]); // [batch_size, total_dim]
    }

    #[test]
    fn test_original_extraction() {
        let device = Device::cpu();
        let dynamics = LinearDynamics::decay(1.0, 3, device.clone()).unwrap();
        let ode_config = NeuralODEConfig::default();
        
        let augmented_ode = AugmentedNeuralODE::new(
            Box::new(dynamics),
            ode_config,
            2, // original_dim
            1, // augmented_dim
            &device,
        ).unwrap();
        
        // Test extraction from time series: [time_steps, total_dim]
        let augmented_solution = Tensor::randn(&[5, 3], &device).unwrap();
        let original_part = augmented_ode.extract_original(&augmented_solution).unwrap();
        
        assert_eq!(original_part.shape(), &[5, 2]); // [time_steps, original_dim]
    }

    #[test]
    fn test_augmented_forward() {
        let device = Device::cpu();
        let dynamics = LinearDynamics::decay(1.0, 3, device.clone()).unwrap();
        let ode_config = NeuralODEConfig::default();
        
        let augmented_ode = AugmentedNeuralODE::new(
            Box::new(dynamics),
            ode_config,
            2, // original_dim
            1, // augmented_dim
            &device,
        ).unwrap();
        
        let y0 = Tensor::ones([2], &device).unwrap();
        let t_span = vec![0.0, 0.5, 1.0];
        
        let result = augmented_ode.forward(&y0, &t_span).unwrap();
        
        // Should return only original dimensions
        assert_eq!(result.shape(), &[3, 2]); // [time_steps, original_dim]
        
        let result_slice = result.to_vec::<f32>().unwrap();
        // Should start with initial conditions
        assert!((result_slice[0] - 1.0).abs() < 1e-6);
        assert!((result_slice[1] - 1.0).abs() < 1e-6);
    }

    #[test]
    fn test_return_augmented_flag() {
        let device = Device::cpu();
        let dynamics = LinearDynamics::decay(1.0, 3, device.clone()).unwrap();
        let ode_config = NeuralODEConfig::default();
        
        let mut augmented_ode = AugmentedNeuralODE::new(
            Box::new(dynamics),
            ode_config,
            2, // original_dim
            1, // augmented_dim
            &device,
        ).unwrap();
        
        let y0 = Tensor::ones([2], &device).unwrap();
        let t_span = vec![0.0, 1.0];
        
        // Test with return_augmented = false (default)
        let result_original = augmented_ode.forward(&y0, &t_span).unwrap();
        assert_eq!(result_original.shape(), &[2, 2]); // Original dimensions only
        
        // Test with return_augmented = true
        augmented_ode.set_return_augmented(true);
        let result_augmented = augmented_ode.forward(&y0, &t_span).unwrap();
        assert_eq!(result_augmented.shape(), &[2, 3]); // All dimensions
    }

    #[test]
    fn test_different_initialization_strategies() {
        let device = Device::cpu();
        
        // Test zeros initialization
        let config_zeros = AugmentationConfig {
            original_dim: 2,
            augmented_dim: 2,
            init_strategy: AugmentationInit::Zeros,
            return_augmented: true,
        };
        
        let dynamics = LinearDynamics::decay(1.0, 4, device.clone()).unwrap();
        let ode_config = NeuralODEConfig::default();
        
        let augmented_ode = AugmentedNeuralODE::with_config(
            Box::new(dynamics),
            ode_config,
            config_zeros,
            &device,
        ).unwrap();
        
        let y0 = Tensor::ones([2], &device).unwrap();
        let augmented_y0 = augmented_ode.augment_state(&y0).unwrap();
        
        let aug_slice = augmented_y0.to_vec::<f32>().unwrap();
        assert_eq!(aug_slice[0], 1.0); // Original
        assert_eq!(aug_slice[1], 1.0); // Original
        assert_eq!(aug_slice[2], 0.0); // Augmented (zeros)
        assert_eq!(aug_slice[3], 0.0); // Augmented (zeros)
        
        // Test constant initialization
        let config_const = AugmentationConfig {
            original_dim: 2,
            augmented_dim: 1,
            init_strategy: AugmentationInit::Constant { value: 0.5 },
            return_augmented: true,
        };
        
        let dynamics2 = LinearDynamics::decay(1.0, 3, device.clone()).unwrap();
        let augmented_ode2 = AugmentedNeuralODE::with_config(
            Box::new(dynamics2),
            NeuralODEConfig::default(),
            config_const,
            &device,
        ).unwrap();
        
        let augmented_y0_2 = augmented_ode2.augment_state(&y0).unwrap();
        let aug_slice_2 = augmented_y0_2.to_vec::<f32>().unwrap();
        assert_eq!(aug_slice_2[2], 0.5); // Should be constant value
    }

    #[test]
    fn test_augmented_ode_func_wrapper() {
        let device = Device::cpu();
        let inner_dynamics = LinearDynamics::decay(1.0, 2, device.clone()).unwrap();
        let augmented_dynamics = AugmentedODEFunc::new(inner_dynamics, 2, 1);
        
        let y = Tensor::ones([3], &device).unwrap(); // [2 original + 1 augmented]
        let dy_dt = augmented_dynamics.forward(0.0, &y).unwrap();
        
        assert_eq!(dy_dt.shape(), &[3]);
        
        let dy_slice = dy_dt.to_vec::<f32>().unwrap();
        // Original dimensions should follow inner dynamics: -1 * y
        assert!((dy_slice[0] + 1.0).abs() < 1e-6);
        assert!((dy_slice[1] + 1.0).abs() < 1e-6);
        // Augmented dimension should decay slowly: -0.1 * y
        assert!((dy_slice[2] + 0.1).abs() < 1e-6);
    }

    #[test]
    fn test_parameters_passthrough() {
        let device = Device::cpu();
        let dynamics = LinearDynamics::decay(1.0, 3, device.clone()).unwrap();
        let ode_config = NeuralODEConfig::default();
        
        let augmented_ode = AugmentedNeuralODE::new(
            Box::new(dynamics),
            ode_config,
            2,
            1,
            &device,
        ).unwrap();
        
        let params = augmented_ode.parameters();
        assert!(!params.is_empty()); // Should have parameters from underlying dynamics
    }
}