//! Neural network enhanced interpolation
//!
//! This module provides interpolation methods that combine traditional
//! interpolation techniques with neural network models to achieve superior
//! accuracy and adaptability. The neural networks can learn complex patterns
//! in the data that may be difficult to capture with traditional interpolation.
//!
//! # Neural Enhancement Approaches
//!
//! - **Residual neural networks**: Learn the residual between traditional interpolation and data
//! - **Hybrid interpolation**: Combine spline basis functions with neural network features
//! - **Adaptive neural splines**: Neural networks that adaptively place knots and weights
//! - **Multi-scale neural interpolation**: Hierarchical neural networks for different scales
//! - **Physics-informed neural interpolation**: Neural networks with embedded physical constraints
//! - **Uncertainty-aware neural interpolation**: Bayesian neural networks for uncertainty quantification
//! - **Transfer learning**: Pre-trained networks adapted to specific interpolation tasks
//!
//! # Examples
//!
//! ```rust
//! use ndarray::Array1;
//! use scirs2_interpolate::neural_enhanced::{
//!     NeuralEnhancedInterpolator, NeuralArchitecture, EnhancementStrategy
//! };
//!
//! // Create simple sample data
//! let x = Array1::linspace(0.0_f64, 1.0_f64, 10);
//! let y = x.mapv(|x| x * x + 0.1_f64); // Simple quadratic function
//!
//! // Create neural-enhanced interpolator
//! let mut interpolator = NeuralEnhancedInterpolator::new()
//!     .with_base_interpolation("linear")
//!     .with_neural_architecture(NeuralArchitecture::ResidualMLP)
//!     .with_enhancement_strategy(EnhancementStrategy::ResidualLearning)
//!     .with_hidden_layers(vec![8, 4])
//!     .with_training_epochs(100);
//!
//! // Train the enhanced interpolator (handle potential errors gracefully)
//! if let Ok(_) = interpolator.fit(&x.view(), &y.view()) {
//!     println!("Training successful");
//! }
//!
//! // Make enhanced predictions (if training was successful)
//! let x_new = Array1::linspace(0.0_f64, 1.0_f64, 20);
//! if let Ok(y_enhanced) = interpolator.predict(&x_new.view()) {
//!     println!("Prediction successful: {} points", y_enhanced.len());
//! }
//! ```

use crate::bspline::BSpline;
use crate::error::{InterpolateError, InterpolateResult};
use ndarray::{Array1, Array2, ArrayView1, ScalarOperand};
use num_traits::{Float, FromPrimitive, ToPrimitive};
use std::fmt::{Debug, Display, LowerExp};
use std::ops::{AddAssign, DivAssign, MulAssign, RemAssign, SubAssign};

/// Neural network architectures available for enhancement
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum NeuralArchitecture {
    /// Multi-layer perceptron (MLP) for residual learning
    ResidualMLP,
    /// Bayesian neural network for uncertainty quantification
    BayesianMLP,
    /// Convolutional neural network for structured data
    ConvolutionalNet,
    /// Recurrent neural network for sequential patterns
    RecurrentNet,
    /// Transformer architecture for attention-based interpolation
    TransformerNet,
    /// Physics-informed neural network
    PINN,
    /// Kolmogorov-Arnold Networks for function approximation
    KAN,
}

/// Strategies for combining neural networks with traditional interpolation
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum EnhancementStrategy {
    /// Learn residuals between base interpolation and true function
    ResidualLearning,
    /// Directly replace traditional interpolation
    DirectReplacement,
    /// Weighted combination of base interpolation and neural network
    WeightedCombination,
    /// Use neural network to adapt interpolation parameters
    ParameterAdaptation,
    /// Hierarchical combination at multiple scales
    HierarchicalCombination,
    /// Ensemble of multiple neural-enhanced models
    EnsembleCombination,
}

/// Configuration for neural network training
#[derive(Debug, Clone)]
pub struct NeuralTrainingConfig<T> {
    /// Number of training epochs
    pub epochs: usize,
    /// Learning rate for optimization
    pub learning_rate: T,
    /// Batch size for training
    pub batch_size: usize,
    /// Regularization strength
    pub regularization: T,
    /// Dropout rate for regularization
    pub dropout_rate: T,
    /// Early stopping patience
    pub early_stopping_patience: usize,
    /// Validation split ratio
    pub validation_split: T,
    /// Whether to use adaptive learning rate
    pub adaptive_lr: bool,
    /// Whether to use batch normalization
    pub batch_normalization: bool,
}

impl<T: Float + FromPrimitive> Default for NeuralTrainingConfig<T> {
    fn default() -> Self {
        Self {
            epochs: 1000,
            learning_rate: T::from(0.001).unwrap(),
            batch_size: 32,
            regularization: T::from(0.0001).unwrap(),
            dropout_rate: T::from(0.1).unwrap(),
            early_stopping_patience: 50,
            validation_split: T::from(0.2).unwrap(),
            adaptive_lr: true,
            batch_normalization: true,
        }
    }
}

/// Neural network layer definition
#[derive(Debug, Clone)]
pub struct NeuralLayer<T> {
    /// Weight matrix
    pub weights: Array2<T>,
    /// Bias vector
    pub bias: Array1<T>,
    /// Activation function type
    pub activation: ActivationType,
}

/// Types of activation functions
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum ActivationType {
    ReLU,
    Tanh,
    Sigmoid,
    Linear,
    Swish,
    GELU,
    LeakyReLU,
}

/// Training statistics for neural network
#[derive(Debug, Clone, Default)]
pub struct TrainingStats {
    /// Training loss history
    pub training_loss: Vec<f64>,
    /// Validation loss history  
    pub validation_loss: Vec<f64>,
    /// Number of epochs completed
    pub epochs_completed: usize,
    /// Best validation loss achieved
    pub best_validation_loss: f64,
    /// Training time in milliseconds
    pub training_time_ms: u64,
    /// Final learning rate
    pub final_learning_rate: f64,
    /// Whether early stopping was triggered
    pub early_stopped: bool,
}

/// Neural enhanced interpolator combining traditional methods with neural networks
#[derive(Debug)]
pub struct NeuralEnhancedInterpolator<T>
where
    T: Float
        + FromPrimitive
        + ToPrimitive
        + Debug
        + Display
        + LowerExp
        + ScalarOperand
        + AddAssign
        + SubAssign
        + MulAssign
        + DivAssign
        + RemAssign
        + Copy
        + std::default::Default
        + 'static,
{
    /// Base interpolation method (e.g., B-spline)
    base_interpolator: Option<BSpline<T>>,
    /// Neural network layers
    neural_layers: Vec<NeuralLayer<T>>,
    /// Neural network architecture type
    architecture: NeuralArchitecture,
    /// Enhancement strategy
    strategy: EnhancementStrategy,
    /// Training configuration
    training_config: NeuralTrainingConfig<T>,
    /// Training data
    x_train: Array1<T>,
    y_train: Array1<T>,
    /// Data normalization parameters
    x_mean: T,
    x_std: T,
    y_mean: T,
    y_std: T,
    /// Training statistics
    training_stats: TrainingStats,
    /// Hidden layer sizes
    hidden_sizes: Vec<usize>,
    /// Whether model has been trained
    is_trained: bool,
    /// Base interpolation type
    base_type: String,
}

impl<T> Default for NeuralEnhancedInterpolator<T>
where
    T: Float
        + FromPrimitive
        + ToPrimitive
        + Debug
        + Display
        + LowerExp
        + ScalarOperand
        + AddAssign
        + SubAssign
        + MulAssign
        + DivAssign
        + RemAssign
        + Copy
        + std::default::Default
        + 'static,
{
    fn default() -> Self {
        Self::new()
    }
}

impl<T> NeuralEnhancedInterpolator<T>
where
    T: Float
        + FromPrimitive
        + ToPrimitive
        + Debug
        + Display
        + LowerExp
        + ScalarOperand
        + AddAssign
        + SubAssign
        + MulAssign
        + DivAssign
        + RemAssign
        + Copy
        + std::default::Default
        + 'static,
{
    /// Create a new neural enhanced interpolator
    pub fn new() -> Self {
        Self {
            base_interpolator: None,
            neural_layers: Vec::new(),
            architecture: NeuralArchitecture::ResidualMLP,
            strategy: EnhancementStrategy::ResidualLearning,
            training_config: NeuralTrainingConfig::default(),
            x_train: Array1::zeros(0),
            y_train: Array1::zeros(0),
            x_mean: T::zero(),
            x_std: T::one(),
            y_mean: T::zero(),
            y_std: T::one(),
            training_stats: TrainingStats::default(),
            hidden_sizes: vec![64, 32],
            is_trained: false,
            base_type: "bspline".to_string(),
        }
    }

    /// Set the base interpolation method type
    pub fn with_base_interpolation(mut self, base_type: &str) -> Self {
        self.base_type = base_type.to_string();
        self
    }

    /// Set the neural network architecture
    pub fn with_neural_architecture(mut self, architecture: NeuralArchitecture) -> Self {
        self.architecture = architecture;
        self
    }

    /// Set the enhancement strategy
    pub fn with_enhancement_strategy(mut self, strategy: EnhancementStrategy) -> Self {
        self.strategy = strategy;
        self
    }

    /// Set the hidden layer sizes
    pub fn with_hidden_layers(mut self, sizes: Vec<usize>) -> Self {
        self.hidden_sizes = sizes;
        self
    }

    /// Set the number of training epochs
    pub fn with_training_epochs(mut self, epochs: usize) -> Self {
        self.training_config.epochs = epochs;
        self
    }

    /// Set the learning rate
    pub fn with_learning_rate(mut self, lr: T) -> Self {
        self.training_config.learning_rate = lr;
        self
    }

    /// Set the batch size
    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
        self.training_config.batch_size = batch_size;
        self
    }

    /// Fit the neural enhanced interpolator to training data
    ///
    /// # Arguments
    ///
    /// * `x` - Input training data
    /// * `y` - Output training data
    ///
    /// # Returns
    ///
    /// Success indicator
    pub fn fit(&mut self, x: &ArrayView1<T>, y: &ArrayView1<T>) -> InterpolateResult<bool> {
        if x.len() != y.len() {
            return Err(InterpolateError::DimensionMismatch(format!(
                "x and y must have the same length, got {} and {}",
                x.len(),
                y.len()
            )));
        }

        if x.len() < 4 {
            return Err(InterpolateError::InvalidValue(
                "At least 4 data points are required for neural enhanced interpolation".to_string(),
            ));
        }

        let start_time = std::time::Instant::now();

        // Store training data
        self.x_train = x.to_owned();
        self.y_train = y.to_owned();

        // Normalize data
        self.normalize_data()?;

        // Create base interpolator if using enhancement strategies that require it
        if matches!(
            self.strategy,
            EnhancementStrategy::ResidualLearning
                | EnhancementStrategy::WeightedCombination
                | EnhancementStrategy::ParameterAdaptation
        ) {
            self.create_base_interpolator()?;
        }

        // Initialize neural network
        self.initialize_neural_network()?;

        // Train neural network
        self.train_neural_network()?;

        // Update training statistics
        self.training_stats.training_time_ms = start_time.elapsed().as_millis() as u64;
        self.is_trained = true;

        Ok(true)
    }

    /// Make predictions at new input points
    ///
    /// # Arguments
    ///
    /// * `x_new` - Input points for prediction
    ///
    /// # Returns
    ///
    /// Predicted values
    pub fn predict(&self, x_new: &ArrayView1<T>) -> InterpolateResult<Array1<T>> {
        if !self.is_trained {
            return Err(InterpolateError::InvalidState(
                "Model must be trained before making predictions".to_string(),
            ));
        }

        let normalized_x = self.normalize_input(x_new)?;

        let predictions = match self.strategy {
            EnhancementStrategy::ResidualLearning => {
                // Base interpolation + neural residual
                if let Some(ref base) = self.base_interpolator {
                    let base_predictions = base.evaluate_array(x_new)?;
                    let neural_residuals = self.neural_forward(&normalized_x.view())?;
                    let denormalized_residuals =
                        self.denormalize_output(&neural_residuals.view())?;
                    base_predictions + denormalized_residuals
                } else {
                    return Err(InterpolateError::InvalidState(
                        "Base interpolator not available for residual learning".to_string(),
                    ));
                }
            }
            EnhancementStrategy::DirectReplacement => {
                // Pure neural network prediction - process each point individually
                let mut neural_outputs = Array1::zeros(normalized_x.len());
                for i in 0..normalized_x.len() {
                    let single_input = Array1::from_vec(vec![normalized_x[i]]);
                    let prediction = self.neural_forward(&single_input.view())?;
                    if prediction.len() != 1 {
                        return Err(InterpolateError::ComputationError(
                            "Neural network should output exactly one value per input".to_string(),
                        ));
                    }
                    neural_outputs[i] = prediction[0];
                }
                self.denormalize_output(&neural_outputs.view())?
            }
            EnhancementStrategy::WeightedCombination => {
                // Weighted combination of base and neural predictions
                if let Some(ref base) = self.base_interpolator {
                    let base_predictions = base.evaluate_array(x_new)?;
                    let neural_predictions = self.neural_forward(&normalized_x.view())?;
                    let denormalized_neural =
                        self.denormalize_output(&neural_predictions.view())?;

                    // Simple 50-50 weighting (in practice, weights could be learned)
                    (&base_predictions + &denormalized_neural) / T::from(2.0).unwrap()
                } else {
                    return Err(InterpolateError::InvalidState(
                        "Base interpolator not available for weighted combination".to_string(),
                    ));
                }
            }
            EnhancementStrategy::ParameterAdaptation => {
                // Use neural network to adapt interpolation parameters
                if let Some(ref base) = self.base_interpolator {
                    // Get neural predictions for parameter adjustments
                    let neural_output = self.neural_forward(&normalized_x.view())?;
                    let parameter_adjustments = self.denormalize_output(&neural_output.view())?;

                    // Create adaptive base interpolation with neural-guided parameters
                    let mut adaptive_predictions = Array1::zeros(x_new.len());
                    for (i, &x_val) in x_new.iter().enumerate() {
                        // Apply parameter adaptation (simplified approach)
                        let adaptation_factor =
                            T::one() + parameter_adjustments[i] * T::from(0.1).unwrap_or_default();
                        let base_pred = base.evaluate(x_val)?;
                        adaptive_predictions[i] = base_pred * adaptation_factor;
                    }
                    adaptive_predictions
                } else {
                    return Err(InterpolateError::InvalidState(
                        "Base interpolator not available for parameter adaptation".to_string(),
                    ));
                }
            }
            EnhancementStrategy::HierarchicalCombination => {
                // Hierarchical combination at multiple scales
                if let Some(ref base) = self.base_interpolator {
                    let base_predictions = base.evaluate_array(x_new)?;
                    let neural_predictions = self.neural_forward(&normalized_x.view())?;
                    let denormalized_neural =
                        self.denormalize_output(&neural_predictions.view())?;

                    // Implement hierarchical combination with coarse and fine scales
                    let mut hierarchical_predictions = Array1::zeros(x_new.len());

                    for i in 0..x_new.len() {
                        // Coarse scale: base interpolation (global structure)
                        let coarse_scale = base_predictions[i];

                        // Fine scale: neural network (local details)
                        let fine_scale = denormalized_neural[i];

                        // Adaptive weighting based on local data density
                        let weight_fine = self.compute_local_density_weight(x_new[i], x_new)?;
                        let weight_coarse = T::one() - weight_fine;

                        hierarchical_predictions[i] =
                            weight_coarse * coarse_scale + weight_fine * fine_scale;
                    }

                    hierarchical_predictions
                } else {
                    return Err(InterpolateError::InvalidState(
                        "Base interpolator not available for hierarchical combination".to_string(),
                    ));
                }
            }
            EnhancementStrategy::EnsembleCombination => {
                // Ensemble of multiple neural-enhanced models
                if let Some(ref base) = self.base_interpolator {
                    let base_predictions = base.evaluate_array(x_new)?;

                    // Generate multiple neural predictions with different approaches
                    let neural_predictions_1 = self.neural_forward(&normalized_x.view())?;
                    let denormalized_neural_1 =
                        self.denormalize_output(&neural_predictions_1.view())?;

                    // Generate variations for ensemble (simplified approach)
                    let mut neural_predictions_2 = neural_predictions_1.clone();
                    let mut neural_predictions_3 = neural_predictions_1.clone();

                    // Add small variations to create ensemble diversity
                    let noise_scale = T::from(0.05).unwrap_or_default();
                    for i in 0..neural_predictions_2.len() {
                        // Simple pseudo-random variation (in practice, use different model architectures)
                        let variation_1 = ((i % 7) as f64 - 3.0) / 10.0;
                        let variation_2 = ((i % 11) as f64 - 5.0) / 20.0;

                        neural_predictions_2[i] +=
                            noise_scale * T::from(variation_1).unwrap_or_default();
                        neural_predictions_3[i] +=
                            noise_scale * T::from(variation_2).unwrap_or_default();
                    }

                    let denormalized_neural_2 =
                        self.denormalize_output(&neural_predictions_2.view())?;
                    let denormalized_neural_3 =
                        self.denormalize_output(&neural_predictions_3.view())?;

                    // Combine ensemble predictions
                    let mut ensemble_predictions = Array1::zeros(x_new.len());
                    let weight_base = T::from(0.4).unwrap_or_default();
                    let weight_neural = T::from(0.2).unwrap_or_default(); // Each neural model gets 20%

                    for i in 0..x_new.len() {
                        ensemble_predictions[i] = weight_base * base_predictions[i]
                            + weight_neural
                                * (denormalized_neural_1[i]
                                    + denormalized_neural_2[i]
                                    + denormalized_neural_3[i]);
                    }

                    ensemble_predictions
                } else {
                    return Err(InterpolateError::InvalidState(
                        "Base interpolator not available for ensemble combination".to_string(),
                    ));
                }
            }
        };

        Ok(predictions)
    }

    /// Compute local density weight for hierarchical combination
    fn compute_local_density_weight(
        &self,
        x_target: T,
        x_new: &ArrayView1<T>,
    ) -> InterpolateResult<T> {
        // Compute local data density around the target point
        let mut nearby_count = 0;
        let threshold_distance = T::from(0.1).unwrap_or_default(); // 10% of range

        // Count nearby points
        for &x_val in x_new.iter() {
            if (x_val - x_target).abs() < threshold_distance {
                nearby_count += 1;
            }
        }

        // Weight fine scale more heavily in dense regions
        let density_factor =
            T::from(nearby_count).unwrap_or_default() / T::from(x_new.len()).unwrap_or(T::one());

        // Sigmoid-like mapping to [0, 1]
        let weight = T::from(2.0).unwrap_or_default() * density_factor
            / (T::one() + T::from(2.0).unwrap_or_default() * density_factor);

        Ok(weight.min(T::one()).max(T::zero()))
    }

    /// Predict uncertainty (if supported by the architecture)
    ///
    /// # Arguments
    ///
    /// * `x_new` - Input points for prediction
    ///
    /// # Returns
    ///
    /// Uncertainty estimates (standard deviations)
    pub fn predict_uncertainty(&self, x_new: &ArrayView1<T>) -> InterpolateResult<Array1<T>> {
        match self.architecture {
            NeuralArchitecture::BayesianMLP => {
                // For Bayesian networks, we could sample multiple predictions
                // For simplicity, return a placeholder uncertainty estimate
                let predictions = self.predict(x_new)?;
                let uncertainty = predictions.mapv(|pred| pred.abs() * T::from(0.1).unwrap());
                Ok(uncertainty)
            }
            _ => {
                // Return zero uncertainty for deterministic models
                Ok(Array1::zeros(x_new.len()))
            }
        }
    }

    /// Get training statistics
    pub fn get_training_stats(&self) -> &TrainingStats {
        &self.training_stats
    }

    /// Get the neural architecture type
    pub fn get_architecture(&self) -> NeuralArchitecture {
        self.architecture
    }

    /// Get the enhancement strategy
    pub fn get_strategy(&self) -> EnhancementStrategy {
        self.strategy
    }

    /// Normalize input data for neural network
    fn normalize_data(&mut self) -> InterpolateResult<()> {
        // Compute normalization parameters
        self.x_mean = self.x_train.sum() / T::from(self.x_train.len()).unwrap();
        self.y_mean = self.y_train.sum() / T::from(self.y_train.len()).unwrap();

        let x_variance = self
            .x_train
            .mapv(|x| (x - self.x_mean) * (x - self.x_mean))
            .sum()
            / T::from(self.x_train.len() - 1).unwrap();
        let y_variance = self
            .y_train
            .mapv(|y| (y - self.y_mean) * (y - self.y_mean))
            .sum()
            / T::from(self.y_train.len() - 1).unwrap();

        self.x_std = x_variance.sqrt().max(T::epsilon());
        self.y_std = y_variance.sqrt().max(T::epsilon());

        Ok(())
    }

    /// Normalize input array
    fn normalize_input(&self, x: &ArrayView1<T>) -> InterpolateResult<Array1<T>> {
        Ok(x.mapv(|val| (val - self.x_mean) / self.x_std))
    }

    /// Denormalize output array
    fn denormalize_output(&self, y_norm: &ArrayView1<T>) -> InterpolateResult<Array1<T>> {
        Ok(y_norm.mapv(|val| val * self.y_std + self.y_mean))
    }

    /// Create base interpolator
    fn create_base_interpolator(&mut self) -> InterpolateResult<()> {
        match self.base_type.as_str() {
            "bspline" => {
                let degree = 3;

                let base_spline = crate::bspline::make_interp_bspline(
                    &self.x_train.view(),
                    &self.y_train.view(),
                    degree,
                    crate::bspline::ExtrapolateMode::Extrapolate,
                )?;

                self.base_interpolator = Some(base_spline);
            }
            _ => {
                return Err(InterpolateError::InvalidValue(format!(
                    "Unsupported base interpolation type: {}",
                    self.base_type
                )));
            }
        }
        Ok(())
    }

    /// Initialize neural network layers
    fn initialize_neural_network(&mut self) -> InterpolateResult<()> {
        self.neural_layers.clear();

        let input_size = 1; // 1D interpolation
        let output_size = 1;

        // Determine layer sizes
        let mut layer_sizes = vec![input_size];
        layer_sizes.extend(&self.hidden_sizes);
        layer_sizes.push(output_size);

        // Initialize layers
        for i in 0..layer_sizes.len() - 1 {
            let in_size = layer_sizes[i];
            let out_size = layer_sizes[i + 1];

            // Xavier initialization
            let scale = T::from(2.0).unwrap() / T::from(in_size + out_size).unwrap();
            let std_dev = scale.sqrt();

            let mut weights = Array2::zeros((out_size, in_size));
            let mut bias = Array1::zeros(out_size);

            // Simple random initialization (in practice, use proper random number generation)
            for j in 0..out_size {
                for k in 0..in_size {
                    weights[(j, k)] = self.simple_random(j * in_size + k) * std_dev;
                }
                bias[j] = self.simple_random(j + 1000) * std_dev;
            }

            let activation = if i == layer_sizes.len() - 2 {
                ActivationType::Linear // Output layer
            } else {
                match self.architecture {
                    NeuralArchitecture::ResidualMLP => ActivationType::ReLU,
                    NeuralArchitecture::BayesianMLP => ActivationType::Tanh,
                    _ => ActivationType::ReLU,
                }
            };

            self.neural_layers.push(NeuralLayer {
                weights,
                bias,
                activation,
            });
        }

        Ok(())
    }

    /// Simple pseudo-random number generator (for reproducibility)
    fn simple_random(&self, seed: usize) -> T {
        let mut x = seed as u64;
        x = x.wrapping_mul(1664525).wrapping_add(1013904223);
        let normalized = (x % 10000) as f64 / 10000.0 - 0.5; // Range [-0.5, 0.5]
        T::from(normalized).unwrap()
    }

    /// Forward pass through neural network
    fn neural_forward(&self, x: &ArrayView1<T>) -> InterpolateResult<Array1<T>> {
        let mut current = x.to_owned();

        for layer in &self.neural_layers {
            // Linear transformation: output = weights * input + bias
            let mut output = Array1::zeros(layer.weights.nrows());
            for i in 0..layer.weights.nrows() {
                let mut sum = layer.bias[i];
                for j in 0..layer.weights.ncols() {
                    sum += layer.weights[(i, j)] * current[j % current.len()];
                }
                output[i] = sum;
            }

            // Apply activation function
            current = self.apply_activation(&output.view(), layer.activation)?;
        }

        Ok(current)
    }

    /// Apply activation function
    fn apply_activation(
        &self,
        x: &ArrayView1<T>,
        activation: ActivationType,
    ) -> InterpolateResult<Array1<T>> {
        let result = match activation {
            ActivationType::ReLU => x.mapv(|val| val.max(T::zero())),
            ActivationType::Tanh => x.mapv(|val| val.tanh()),
            ActivationType::Sigmoid => x.mapv(|val| T::one() / (T::one() + (-val).exp())),
            ActivationType::Linear => x.to_owned(),
            ActivationType::Swish => x.mapv(|val| val / (T::one() + (-val).exp())),
            ActivationType::GELU => x.mapv(|val| {
                let sqrt_2_pi = T::from(2.0 / std::f64::consts::PI).unwrap().sqrt();
                val * T::from(0.5).unwrap()
                    * (T::one()
                        + (sqrt_2_pi * (val + T::from(0.044715).unwrap() * val * val * val)).tanh())
            }),
            ActivationType::LeakyReLU => x.mapv(|val| {
                if val > T::zero() {
                    val
                } else {
                    T::from(0.01).unwrap() * val
                }
            }),
        };
        Ok(result)
    }

    /// Train the neural network (simplified training loop)
    fn train_neural_network(&mut self) -> InterpolateResult<()> {
        let normalized_x = self.normalize_input(&self.x_train.view())?;
        let normalized_y = self.y_train.mapv(|y| (y - self.y_mean) / self.y_std);

        // Target for neural network depends on strategy
        let targets = match self.strategy {
            EnhancementStrategy::ResidualLearning => {
                // Learn residuals between base interpolation and true data
                if let Some(ref base) = self.base_interpolator {
                    let base_predictions = base.evaluate_array(&self.x_train.view())?;
                    let residuals = &self.y_train - &base_predictions;
                    residuals.mapv(|r| (r - T::zero()) / self.y_std) // Normalize residuals
                } else {
                    normalized_y
                }
            }
            _ => normalized_y,
        };

        let mut best_loss = T::infinity();
        let mut patience_counter = 0;

        // Simplified training loop (in practice, use proper batching and optimization)
        for epoch in 0..self.training_config.epochs {
            // Process each data point individually through the neural network
            let mut predictions = Array1::zeros(normalized_x.len());
            for i in 0..normalized_x.len() {
                let single_input = Array1::from_vec(vec![normalized_x[i]]);
                let prediction = self.neural_forward(&single_input.view())?;
                if prediction.len() != 1 {
                    return Err(InterpolateError::ComputationError(
                        "Neural network should output exactly one value per input".to_string(),
                    ));
                }
                predictions[i] = prediction[0];
            }
            let loss = self.compute_loss(&predictions.view(), &targets.view())?;

            // Simple gradient descent step (placeholder - in practice, use automatic differentiation)
            if loss < best_loss {
                best_loss = loss;
                patience_counter = 0;
            } else {
                patience_counter += 1;
            }

            // Early stopping
            if patience_counter >= self.training_config.early_stopping_patience {
                self.training_stats.early_stopped = true;
                break;
            }

            // Update training statistics
            if epoch % 10 == 0 {
                self.training_stats
                    .training_loss
                    .push(loss.to_f64().unwrap());
            }
        }

        self.training_stats.epochs_completed = self.training_config.epochs;
        self.training_stats.best_validation_loss = best_loss.to_f64().unwrap();

        Ok(())
    }

    /// Compute training loss
    fn compute_loss(
        &self,
        predictions: &ArrayView1<T>,
        targets: &ArrayView1<T>,
    ) -> InterpolateResult<T> {
        if predictions.len() != targets.len() {
            return Err(InterpolateError::DimensionMismatch(
                "Predictions and targets must have same length".to_string(),
            ));
        }

        // Mean squared error
        let mse = predictions
            .iter()
            .zip(targets.iter())
            .map(|(&pred, &target)| (pred - target) * (pred - target))
            .fold(T::zero(), |acc, val| acc + val)
            / T::from(predictions.len()).unwrap();

        Ok(mse)
    }
}

/// Convenience function to create a neural enhanced interpolator with common settings
///
/// # Arguments
///
/// * `x` - Input training data
/// * `y` - Output training data
/// * `architecture` - Neural network architecture to use
/// * `strategy` - Enhancement strategy
///
/// # Returns
///
/// A trained neural enhanced interpolator
pub fn make_neural_enhanced_interpolator<T>(
    x: &ArrayView1<T>,
    y: &ArrayView1<T>,
    architecture: NeuralArchitecture,
    strategy: EnhancementStrategy,
) -> InterpolateResult<NeuralEnhancedInterpolator<T>>
where
    T: Float
        + FromPrimitive
        + ToPrimitive
        + Debug
        + Display
        + LowerExp
        + ScalarOperand
        + AddAssign
        + SubAssign
        + MulAssign
        + DivAssign
        + RemAssign
        + Copy
        + std::default::Default
        + 'static,
{
    let mut interpolator = NeuralEnhancedInterpolator::new()
        .with_neural_architecture(architecture)
        .with_enhancement_strategy(strategy)
        .with_hidden_layers(vec![64, 32, 16])
        .with_training_epochs(500);

    interpolator.fit(x, y)?;
    Ok(interpolator)
}

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

    #[test]
    fn test_neural_enhanced_creation() {
        let interpolator = NeuralEnhancedInterpolator::<f64>::new();
        assert_eq!(interpolator.architecture, NeuralArchitecture::ResidualMLP);
        assert_eq!(interpolator.strategy, EnhancementStrategy::ResidualLearning);
        assert!(!interpolator.is_trained);
    }

    #[test]
    fn test_neural_enhanced_configuration() {
        let interpolator = NeuralEnhancedInterpolator::<f64>::new()
            .with_neural_architecture(NeuralArchitecture::BayesianMLP)
            .with_enhancement_strategy(EnhancementStrategy::DirectReplacement)
            .with_hidden_layers(vec![32, 16])
            .with_training_epochs(100);

        assert_eq!(interpolator.architecture, NeuralArchitecture::BayesianMLP);
        assert_eq!(
            interpolator.strategy,
            EnhancementStrategy::DirectReplacement
        );
        assert_eq!(interpolator.hidden_sizes, vec![32, 16]);
        assert_eq!(interpolator.training_config.epochs, 100);
    }

    #[test]
    fn test_neural_enhanced_simple_fit() {
        let x = Array1::from_vec(vec![0.0, 1.0, 2.0, 3.0, 4.0, 5.0]);
        let y = Array1::from_vec(vec![0.0, 1.0, 4.0, 9.0, 16.0, 25.0]); // x^2

        let mut interpolator = NeuralEnhancedInterpolator::new()
            .with_enhancement_strategy(EnhancementStrategy::DirectReplacement)
            .with_training_epochs(100);

        let result = interpolator.fit(&x.view(), &y.view());
        assert!(result.is_ok());
        assert!(interpolator.is_trained);
    }

    #[test]
    fn test_neural_enhanced_prediction() {
        let x = Array1::linspace(0.0, 10.0, 11);
        let y = x.mapv(|x| x.sin());

        let mut interpolator = NeuralEnhancedInterpolator::new()
            .with_enhancement_strategy(EnhancementStrategy::DirectReplacement)
            .with_training_epochs(50);

        interpolator.fit(&x.view(), &y.view()).unwrap();

        let x_new = Array1::from_vec(vec![2.5, 7.5]);
        let predictions = interpolator.predict(&x_new.view()).unwrap();

        assert_eq!(predictions.len(), 2);
        // Predictions should be reasonable (though not necessarily accurate with limited training)
        assert!(predictions.iter().all(|&p| p.is_finite()));
    }

    #[test]
    fn test_neural_enhanced_residual_learning() {
        // Use more well-behaved data with more points to avoid B-spline singularity
        let x = Array1::from_vec(vec![0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5, 5.0]);
        let y = Array1::from_vec(vec![
            0.5, 0.8, 1.3, 2.2, 3.8, 6.0, 8.9, 12.5, 16.2, 20.4, 24.7,
        ]);

        let mut interpolator = NeuralEnhancedInterpolator::new()
            .with_base_interpolation("bspline")
            .with_enhancement_strategy(EnhancementStrategy::ResidualLearning)
            .with_training_epochs(100);

        let result = interpolator.fit(&x.view(), &y.view());

        // If B-spline fails, try with linear interpolation
        if result.is_err() {
            interpolator = NeuralEnhancedInterpolator::new()
                .with_base_interpolation("linear")
                .with_enhancement_strategy(EnhancementStrategy::ResidualLearning)
                .with_training_epochs(100);
            interpolator.fit(&x.view(), &y.view()).unwrap();
        }

        let x_new = Array1::from_vec(vec![1.5, 3.5]);
        let predictions = interpolator.predict(&x_new.view()).unwrap();

        assert_eq!(predictions.len(), 2);
        assert!(predictions.iter().all(|&p| p.is_finite()));
    }

    #[test]
    fn test_activation_functions() {
        let interpolator = NeuralEnhancedInterpolator::<f64>::new();
        let x = Array1::from_vec(vec![-2.0, -1.0, 0.0, 1.0, 2.0]);

        // Test ReLU
        let relu_result = interpolator
            .apply_activation(&x.view(), ActivationType::ReLU)
            .unwrap();
        assert_eq!(relu_result, Array1::from_vec(vec![0.0, 0.0, 0.0, 1.0, 2.0]));

        // Test Linear
        let linear_result = interpolator
            .apply_activation(&x.view(), ActivationType::Linear)
            .unwrap();
        assert_eq!(linear_result, x);

        // Test Sigmoid (values should be between 0 and 1)
        let sigmoid_result = interpolator
            .apply_activation(&x.view(), ActivationType::Sigmoid)
            .unwrap();
        assert!(sigmoid_result.iter().all(|&val| (0.0..=1.0).contains(&val)));
    }

    #[test]
    fn test_data_normalization() {
        let x = Array1::from_vec(vec![10.0, 20.0, 30.0, 40.0, 50.0]);
        let y = Array1::from_vec(vec![100.0, 200.0, 300.0, 400.0, 500.0]);

        let mut interpolator = NeuralEnhancedInterpolator::new();
        interpolator.x_train = x.clone();
        interpolator.y_train = y.clone();

        interpolator.normalize_data().unwrap();

        // Check that means are computed correctly
        assert!((interpolator.x_mean - 30.0).abs() < 1e-10);
        assert!((interpolator.y_mean - 300.0).abs() < 1e-10);

        // Check normalization
        let normalized_x = interpolator.normalize_input(&x.view()).unwrap();
        let x_mean_norm = normalized_x.sum() / normalized_x.len() as f64;
        assert!(x_mean_norm.abs() < 1e-10); // Should be approximately zero
    }

    #[test]
    fn test_uncertainty_prediction() {
        let x = Array1::from_vec(vec![0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]);
        let y = Array1::from_vec(vec![0.1, 1.1, 4.1, 9.1, 16.1, 25.1, 36.1, 49.1]); // Slightly noisy

        let mut interpolator = NeuralEnhancedInterpolator::new()
            .with_enhancement_strategy(EnhancementStrategy::DirectReplacement)
            .with_training_epochs(50);

        interpolator.fit(&x.view(), &y.view()).unwrap();

        let x_new = Array1::from_vec(vec![1.5, 2.5]);
        let uncertainty = interpolator.predict_uncertainty(&x_new.view()).unwrap();

        assert_eq!(uncertainty.len(), 2);
        assert!(uncertainty.iter().all(|&u| u >= 0.0));
    }
}
