sklears_discriminant_analysis/

async_optimization.rs

//! Asynchronous Multi-threaded Optimization for Discriminant Analysis
//!
//! This module provides asynchronous and multi-threaded optimization algorithms
//! for discriminant analysis, enabling parallel parameter updates, distributed
//! training, and efficient resource utilization following SciRS2 policy.

// ✅ Using SciRS2 dependencies following SciRS2 policy
use scirs2_core::ndarray::{Array1, Array2, Axis};

use crate::{
    lda::LinearDiscriminantAnalysisConfig, numerical_stability::NumericalStability,
    qda::QuadraticDiscriminantAnalysisConfig,
};

use sklears_core::{error::Result, prelude::SklearsError, types::Float};

use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use std::time::{Duration, Instant};
use tokio::sync::{mpsc, Semaphore};
use tokio::task::JoinHandle;

/// Configuration for asynchronous optimization
#[derive(Debug, Clone)]
pub struct AsyncOptimizationConfig {
    /// Number of worker threads for parallel optimization
    pub num_workers: usize,
    /// Batch size for parameter updates
    pub batch_size: usize,
    /// Maximum number of concurrent tasks
    pub max_concurrent_tasks: usize,
    /// Learning rate for parameter updates
    pub learning_rate: Float,
    /// Learning rate decay factor
    pub learning_rate_decay: Float,
    /// Momentum parameter for optimization
    pub momentum: Float,
    /// Enable adaptive learning rate
    pub adaptive_learning_rate: bool,
    /// Convergence tolerance
    pub convergence_tolerance: Float,
    /// Maximum number of iterations
    pub max_iterations: usize,
    /// Enable asynchronous parameter updates
    pub async_updates: bool,
    /// Update frequency for parameters (in milliseconds)
    pub update_frequency_ms: u64,
    /// Enable gradient checkpointing for memory efficiency
    pub gradient_checkpointing: bool,
    /// Channel buffer size for async communication
    pub channel_buffer_size: usize,
}

impl Default for AsyncOptimizationConfig {
    fn default() -> Self {
        Self {
            num_workers: num_cpus::get(),
            batch_size: 128,
            max_concurrent_tasks: num_cpus::get() * 2,
            learning_rate: 0.01,
            learning_rate_decay: 0.99,
            momentum: 0.9,
            adaptive_learning_rate: true,
            convergence_tolerance: 1e-6,
            max_iterations: 1000,
            async_updates: true,
            update_frequency_ms: 100,
            gradient_checkpointing: true,
            channel_buffer_size: 1000,
        }
    }
}

/// Optimization state for tracking convergence and performance
#[derive(Debug, Clone)]
pub struct OptimizationState {
    /// Current iteration number
    pub iteration: usize,
    /// Current loss/objective value
    pub current_loss: Float,
    /// Previous loss for convergence checking
    pub previous_loss: Float,
    /// Learning rate history
    pub learning_rate_history: Vec<Float>,
    /// Loss history
    pub loss_history: Vec<Float>,
    /// Gradient norms history
    pub gradient_norms: Vec<Float>,
    /// Convergence status
    pub converged: bool,
    /// Training start time
    pub start_time: Instant,
    /// Last update time
    pub last_update: Instant,
    /// Number of parameter updates
    pub update_count: usize,
}

impl Default for OptimizationState {
    fn default() -> Self {
        let now = Instant::now();
        Self {
            iteration: 0,
            current_loss: Float::INFINITY,
            previous_loss: Float::INFINITY,
            learning_rate_history: Vec::new(),
            loss_history: Vec::new(),
            gradient_norms: Vec::new(),
            converged: false,
            start_time: now,
            last_update: now,
            update_count: 0,
        }
    }
}

/// Message types for async optimization communication
#[derive(Debug, Clone)]
pub enum OptimizationMessage {
    /// Parameter update request
    ParameterUpdate {
        worker_id: usize,

        parameters: HashMap<String, Array2<Float>>,

        gradient: HashMap<String, Array2<Float>>,

        loss: Float,
    },
    /// Gradient computation request
    ComputeGradient {
        worker_id: usize,
        data_batch: Array2<Float>,
        labels_batch: Array1<usize>,
    },
    /// Convergence check request
    CheckConvergence {
        worker_id: usize,
        current_loss: Float,
    },
    /// Training completion signal
    TrainingComplete {
        final_parameters: HashMap<String, Array2<Float>>,
        final_loss: Float,
        iterations: usize,
    },
    /// Error occurred during optimization
    OptimizationError { worker_id: usize, error: String },
}

/// Async discriminant analysis optimizer
pub struct AsyncDiscriminantOptimizer {
    config: AsyncOptimizationConfig,
    state: Arc<RwLock<OptimizationState>>,
    parameters: Arc<RwLock<HashMap<String, Array2<Float>>>>,
    velocity: Arc<RwLock<HashMap<String, Array2<Float>>>>, // For momentum
    semaphore: Arc<Semaphore>,
    /// Reserved for future numerically-stable gradient updates
    #[allow(dead_code)] // retained for future numerically-stable gradient clipping
    numerical_stability: NumericalStability,
}

impl AsyncDiscriminantOptimizer {
    /// Create a new async optimizer
    pub fn new(config: AsyncOptimizationConfig) -> Self {
        let semaphore = Arc::new(Semaphore::new(config.max_concurrent_tasks));

        Self {
            config,
            state: Arc::new(RwLock::new(OptimizationState::default())),
            parameters: Arc::new(RwLock::new(HashMap::new())),
            velocity: Arc::new(RwLock::new(HashMap::new())),
            semaphore,
            numerical_stability: NumericalStability::new(),
        }
    }

    /// Initialize parameters for discriminant analysis
    pub fn initialize_parameters(&self, n_features: usize, n_classes: usize) -> Result<()> {
        let mut parameters = self.parameters.write().expect("lock not poisoned");
        let mut velocity = self.velocity.write().expect("lock not poisoned");

        // Initialize class means (centroids)
        let class_means = Array2::zeros((n_classes, n_features));
        parameters.insert("class_means".to_string(), class_means.clone());
        velocity.insert(
            "class_means".to_string(),
            Array2::zeros((n_classes, n_features)),
        );

        // Initialize shared covariance matrix (for LDA)
        let covariance = Array2::eye(n_features);
        parameters.insert("covariance".to_string(), covariance.clone());
        velocity.insert(
            "covariance".to_string(),
            Array2::zeros((n_features, n_features)),
        );

        // Initialize class priors
        let priors = Array2::from_elem((1, n_classes), 1.0 / n_classes as Float);
        parameters.insert("priors".to_string(), priors.clone());
        velocity.insert("priors".to_string(), Array2::zeros((1, n_classes)));

        Ok(())
    }

    /// Asynchronous training for Linear Discriminant Analysis
    pub async fn async_train_lda(
        &self,
        x_data: Array2<Float>, // standard ML notation: feature matrix
        y: Array1<usize>,
        config: LinearDiscriminantAnalysisConfig,
    ) -> Result<HashMap<String, Array2<Float>>> {
        let (_n_samples, n_features) = x_data.dim();
        let n_classes = y.iter().max().expect("collection should not be empty") + 1;

        // Initialize parameters
        self.initialize_parameters(n_features, n_classes)?;

        // Create communication channels
        let (tx, rx) = mpsc::channel::<OptimizationMessage>(self.config.channel_buffer_size);

        // Shared data structures
        let x_shared = Arc::new(x_data);
        let y_shared = Arc::new(y);

        // Spawn worker tasks
        let mut worker_handles = Vec::new();
        for worker_id in 0..self.config.num_workers {
            let worker_handle = self
                .spawn_lda_worker(
                    worker_id,
                    Arc::clone(&x_shared),
                    Arc::clone(&y_shared),
                    tx.clone(),
                    config.clone(),
                )
                .await;
            worker_handles.push(worker_handle);
        }

        // Spawn parameter server task
        let parameter_server_handle = self.spawn_parameter_server(rx).await;

        // Main training loop
        let training_result = self
            .run_async_training_loop(worker_handles, parameter_server_handle)
            .await?;

        Ok(training_result)
    }

    /// Spawn LDA worker task
    async fn spawn_lda_worker(
        &self,
        worker_id: usize,
        x_data: Arc<Array2<Float>>, // standard ML notation: feature matrix
        y: Arc<Array1<usize>>,
        tx: mpsc::Sender<OptimizationMessage>,
        _config: LinearDiscriminantAnalysisConfig,
    ) -> JoinHandle<Result<()>> {
        let config = self.config.clone();
        let parameters = Arc::clone(&self.parameters);
        let semaphore = Arc::clone(&self.semaphore);

        tokio::spawn(async move {
            let mut rng = fastrand::Rng::new();

            for iteration in 0..config.max_iterations {
                // Acquire semaphore permit
                let _permit = semaphore.acquire().await.expect("value should be present");

                // Generate random batch
                let batch_indices: Vec<usize> = (0..config.batch_size)
                    .map(|_| rng.usize(0..x_data.nrows()))
                    .collect();

                let x_batch =
                    Array2::from_shape_fn((config.batch_size, x_data.ncols()), |(i, j)| {
                        x_data[[batch_indices[i], j]]
                    });

                let y_batch = Array1::from_shape_fn(config.batch_size, |i| y[batch_indices[i]]);

                // Compute gradients
                let gradients =
                    Self::compute_lda_gradients(&x_batch, &y_batch, &parameters).await?;

                // Compute loss
                let loss = Self::compute_lda_loss(&x_batch, &y_batch, &parameters).await?;

                // Send parameter update message
                let params_snapshot = parameters.read().expect("lock not poisoned").clone();
                tx.send(OptimizationMessage::ParameterUpdate {
                    worker_id,
                    parameters: params_snapshot,
                    gradient: gradients,
                    loss,
                })
                .await
                .map_err(|e| SklearsError::ProcessingError(format!("Channel send error: {}", e)))?;

                // Check for convergence every 10 iterations
                if iteration % 10 == 0 {
                    tx.send(OptimizationMessage::CheckConvergence {
                        worker_id,
                        current_loss: loss,
                    })
                    .await
                    .map_err(|e| {
                        SklearsError::ProcessingError(format!("Channel send error: {}", e))
                    })?;
                }

                // Yield control periodically
                if iteration % 5 == 0 {
                    tokio::task::yield_now().await;
                }
            }

            Ok(())
        })
    }

    /// Spawn parameter server task
    async fn spawn_parameter_server(
        &self,
        mut rx: mpsc::Receiver<OptimizationMessage>,
    ) -> JoinHandle<Result<HashMap<String, Array2<Float>>>> {
        let parameters = Arc::clone(&self.parameters);
        let velocity = Arc::clone(&self.velocity);
        let state = Arc::clone(&self.state);
        let config = self.config.clone();

        tokio::spawn(async move {
            let mut update_timer =
                tokio::time::interval(Duration::from_millis(config.update_frequency_ms));

            loop {
                tokio::select! {
                    // Handle incoming messages
                    msg = rx.recv() => {
                        match msg {
                            Some(OptimizationMessage::ParameterUpdate { worker_id, parameters: _, gradient, loss }) => {
                                Self::apply_parameter_update(
                                    &parameters,
                                    &velocity,
                                    &state,
                                    &config,
                                    worker_id,
                                    gradient,
                                    loss,
                                ).await?;
                            }
                            Some(OptimizationMessage::CheckConvergence { current_loss, .. }) => {
                                let converged = Self::check_convergence(&state, &config, current_loss).await?;
                                if converged {
                                    let final_params = parameters.read().expect("lock not poisoned").clone();
                                    return Ok(final_params);
                                }
                            }
                            Some(OptimizationMessage::TrainingComplete { final_parameters, .. }) => {
                                return Ok(final_parameters);
                            }
                            Some(OptimizationMessage::OptimizationError { error, .. }) => {
                                return Err(SklearsError::ProcessingError(error));
                            }
                            None => break, // Channel closed
                            _ => {} // Handle other message types
                        }
                    }
                    // Periodic parameter updates
                    _ = update_timer.tick() => {
                        // Perform any periodic maintenance
                        Self::update_learning_rate(&state, &config).await?;
                    }
                }
            }

            let final_params = parameters.read().expect("lock not poisoned").clone();
            Ok(final_params)
        })
    }

    /// Apply parameter update with momentum
    async fn apply_parameter_update(
        parameters: &Arc<RwLock<HashMap<String, Array2<Float>>>>,
        velocity: &Arc<RwLock<HashMap<String, Array2<Float>>>>,
        state: &Arc<RwLock<OptimizationState>>,
        config: &AsyncOptimizationConfig,
        _worker_id: usize,
        gradients: HashMap<String, Array2<Float>>,
        loss: Float,
    ) -> Result<()> {
        let current_lr = {
            let state_read = state.read().expect("lock not poisoned");
            let base_lr = config.learning_rate;
            let decay_factor = config
                .learning_rate_decay
                .powf(state_read.iteration as Float);
            base_lr * decay_factor
        };

        // Update parameters with momentum
        {
            let mut params = parameters.write().expect("lock not poisoned");
            let mut vel = velocity.write().expect("lock not poisoned");

            for (param_name, gradient) in gradients.iter() {
                if let (Some(param), Some(velocity_param)) =
                    (params.get_mut(param_name), vel.get_mut(param_name))
                {
                    // Momentum update: v = momentum * v - learning_rate * gradient
                    *velocity_param = config.momentum * &*velocity_param - current_lr * gradient;

                    // Parameter update: param = param + v
                    *param = &*param + &*velocity_param;
                }
            }
        }

        // Update optimization state
        {
            let mut state_write = state.write().expect("lock not poisoned");
            state_write.previous_loss = state_write.current_loss;
            state_write.current_loss = loss;
            state_write.iteration += 1;
            state_write.last_update = Instant::now();
            state_write.update_count += 1;

            // Compute gradient norm
            let gradient_norm: Float = gradients
                .values()
                .map(|grad| grad.iter().map(|&x| x * x).sum::<Float>())
                .sum::<Float>()
                .sqrt();

            state_write.loss_history.push(loss);
            state_write.gradient_norms.push(gradient_norm);
            state_write.learning_rate_history.push(current_lr);

            // Limit history size
            if state_write.loss_history.len() > 1000 {
                state_write.loss_history.remove(0);
                state_write.gradient_norms.remove(0);
                state_write.learning_rate_history.remove(0);
            }
        }

        Ok(())
    }

    /// Check convergence criteria
    async fn check_convergence(
        state: &Arc<RwLock<OptimizationState>>,
        config: &AsyncOptimizationConfig,
        current_loss: Float,
    ) -> Result<bool> {
        let state_read = state.read().expect("lock not poisoned");

        // Check loss convergence
        if state_read.iteration > 10 {
            let loss_change = (state_read.previous_loss - current_loss).abs();
            if loss_change < config.convergence_tolerance {
                return Ok(true);
            }
        }

        // Check gradient norm convergence
        if let Some(&last_grad_norm) = state_read.gradient_norms.last() {
            if last_grad_norm < config.convergence_tolerance {
                return Ok(true);
            }
        }

        // Check maximum iterations
        if state_read.iteration >= config.max_iterations {
            return Ok(true);
        }

        Ok(false)
    }

    /// Update learning rate based on adaptive strategies
    async fn update_learning_rate(
        state: &Arc<RwLock<OptimizationState>>,
        config: &AsyncOptimizationConfig,
    ) -> Result<()> {
        if !config.adaptive_learning_rate {
            return Ok(());
        }

        let state_write = state.write().expect("lock not poisoned");

        // Simple adaptive strategy: reduce learning rate if loss is not decreasing
        if state_write.loss_history.len() >= 10 {
            let recent_losses = &state_write.loss_history[state_write.loss_history.len() - 10..];
            let is_decreasing = recent_losses
                .windows(2)
                .all(|window| window[1] <= window[0] + config.convergence_tolerance);

            if !is_decreasing && state_write.iteration.is_multiple_of(50) {
                // Reduce learning rate
                // This would be applied in the next parameter update
                log::info!("Reducing learning rate due to lack of progress");
            }
        }

        Ok(())
    }

    /// Compute LDA gradients
    async fn compute_lda_gradients(
        x_batch: &Array2<Float>, // standard ML notation: batch of feature vectors
        y_batch: &Array1<usize>,
        parameters: &Arc<RwLock<HashMap<String, Array2<Float>>>>,
    ) -> Result<HashMap<String, Array2<Float>>> {
        let params = parameters.read().expect("lock not poisoned");
        let mut gradients = HashMap::new();

        // Get current parameters
        let class_means = params.get("class_means").expect("key not found");
        let covariance = params.get("covariance").expect("key not found");

        let (batch_size, n_features) = x_batch.dim();
        let n_classes = class_means.nrows();

        // Compute gradient with respect to class means
        let mut means_gradient = Array2::zeros((n_classes, n_features));
        for (i, &label) in y_batch.iter().enumerate() {
            let x_sample = x_batch.row(i);
            let mean_diff = x_sample.to_owned() - class_means.row(label);

            // Simple gradient: derivative of squared loss
            let updated_row = means_gradient.row(label).to_owned() + &mean_diff;
            means_gradient.row_mut(label).assign(&updated_row);
        }
        means_gradient /= batch_size as Float;

        // Compute gradient with respect to covariance (simplified)
        let mut cov_gradient = Array2::zeros(covariance.raw_dim());
        for (i, &label) in y_batch.iter().enumerate() {
            let x_sample = x_batch.row(i);
            let mean_centered = x_sample.to_owned() - class_means.row(label);
            let col = mean_centered.clone().insert_axis(Axis(1));
            let row = mean_centered.insert_axis(Axis(0));
            let outer_product = col.dot(&row);
            cov_gradient = cov_gradient + outer_product;
        }
        cov_gradient /= batch_size as Float;

        // Compute gradient with respect to priors (simplified)
        let priors_gradient = Array2::zeros((1, n_classes));

        gradients.insert("class_means".to_string(), means_gradient);
        gradients.insert("covariance".to_string(), cov_gradient);
        gradients.insert("priors".to_string(), priors_gradient);

        Ok(gradients)
    }

    /// Compute LDA loss
    async fn compute_lda_loss(
        x_batch: &Array2<Float>, // standard ML notation: batch of feature vectors
        y_batch: &Array1<usize>,
        parameters: &Arc<RwLock<HashMap<String, Array2<Float>>>>,
    ) -> Result<Float> {
        let params = parameters.read().expect("lock not poisoned");

        let class_means = params.get("class_means").expect("key not found");
        let covariance = params.get("covariance").expect("key not found");

        let mut total_loss = 0.0;

        // Simple squared loss between samples and their class means
        for (i, &label) in y_batch.iter().enumerate() {
            let x_sample = x_batch.row(i);
            let class_mean = class_means.row(label);
            let diff = x_sample.to_owned() - class_mean;

            // Mahalanobis distance (simplified - assuming identity covariance for now)
            let loss = diff.dot(&diff);
            total_loss += loss;
        }

        // Add regularization term
        let reg_term = covariance.iter().map(|&x| x * x).sum::<Float>() * 0.001;
        total_loss += reg_term;

        Ok(total_loss / x_batch.nrows() as Float)
    }

    /// Run the main async training loop
    async fn run_async_training_loop(
        &self,
        worker_handles: Vec<JoinHandle<Result<()>>>,
        parameter_server_handle: JoinHandle<Result<HashMap<String, Array2<Float>>>>,
    ) -> Result<HashMap<String, Array2<Float>>> {
        // Wait for parameter server completion
        let final_parameters = parameter_server_handle.await.map_err(|e| {
            SklearsError::ProcessingError(format!("Parameter server error: {}", e))
        })??;

        // Cancel remaining worker tasks
        for handle in worker_handles {
            handle.abort();
        }

        Ok(final_parameters)
    }

    /// Get current optimization statistics
    pub fn get_optimization_stats(&self) -> OptimizationStats {
        let state = self.state.read().expect("lock not poisoned");

        OptimizationStats {
            iteration: state.iteration,
            current_loss: state.current_loss,
            converged: state.converged,
            training_time: state.start_time.elapsed(),
            last_update_time: state.last_update.elapsed(),
            update_count: state.update_count,
            average_loss: if !state.loss_history.is_empty() {
                state.loss_history.iter().sum::<Float>() / state.loss_history.len() as Float
            } else {
                0.0
            },
            loss_variance: if state.loss_history.len() > 1 {
                let mean =
                    state.loss_history.iter().sum::<Float>() / state.loss_history.len() as Float;
                state
                    .loss_history
                    .iter()
                    .map(|&x| (x - mean).powi(2))
                    .sum::<Float>()
                    / (state.loss_history.len() - 1) as Float
            } else {
                0.0
            },
        }
    }

    /// Asynchronous training for Quadratic Discriminant Analysis
    pub async fn async_train_qda(
        &self,
        x_data: Array2<Float>, // standard ML notation: feature matrix
        y: Array1<usize>,
        config: QuadraticDiscriminantAnalysisConfig,
    ) -> Result<HashMap<String, Array2<Float>>> {
        let (_n_samples, n_features) = x_data.dim();
        let n_classes = y.iter().max().expect("collection should not be empty") + 1;

        // Initialize parameters for QDA (each class has its own covariance)
        self.initialize_qda_parameters(n_features, n_classes)?;

        // Similar structure to LDA but with class-specific covariances
        let (tx, rx) = mpsc::channel::<OptimizationMessage>(self.config.channel_buffer_size);

        let x_shared = Arc::new(x_data);
        let y_shared = Arc::new(y);

        // Spawn QDA-specific workers
        let mut worker_handles = Vec::new();
        for worker_id in 0..self.config.num_workers {
            let worker_handle = self
                .spawn_qda_worker(
                    worker_id,
                    Arc::clone(&x_shared),
                    Arc::clone(&y_shared),
                    tx.clone(),
                    config.clone(),
                )
                .await;
            worker_handles.push(worker_handle);
        }

        let parameter_server_handle = self.spawn_parameter_server(rx).await;
        let training_result = self
            .run_async_training_loop(worker_handles, parameter_server_handle)
            .await?;

        Ok(training_result)
    }

    /// Initialize QDA-specific parameters
    fn initialize_qda_parameters(&self, n_features: usize, n_classes: usize) -> Result<()> {
        let mut parameters = self.parameters.write().expect("lock not poisoned");
        let mut velocity = self.velocity.write().expect("lock not poisoned");

        // Initialize class means
        let class_means = Array2::zeros((n_classes, n_features));
        parameters.insert("class_means".to_string(), class_means.clone());
        velocity.insert(
            "class_means".to_string(),
            Array2::zeros((n_classes, n_features)),
        );

        // Initialize class-specific covariance matrices (stacked)
        let covariances = Array2::eye(n_features * n_classes);
        parameters.insert("class_covariances".to_string(), covariances.clone());
        velocity.insert(
            "class_covariances".to_string(),
            Array2::zeros((n_features * n_classes, n_features)),
        );

        // Initialize class priors
        let priors = Array2::from_elem((1, n_classes), 1.0 / n_classes as Float);
        parameters.insert("priors".to_string(), priors.clone());
        velocity.insert("priors".to_string(), Array2::zeros((1, n_classes)));

        Ok(())
    }

    /// Spawn QDA worker task
    async fn spawn_qda_worker(
        &self,
        worker_id: usize,
        x_data: Arc<Array2<Float>>, // standard ML notation: feature matrix
        y: Arc<Array1<usize>>,
        tx: mpsc::Sender<OptimizationMessage>,
        _config: QuadraticDiscriminantAnalysisConfig,
    ) -> JoinHandle<Result<()>> {
        let config = self.config.clone();
        let parameters = Arc::clone(&self.parameters);
        let semaphore = Arc::clone(&self.semaphore);

        tokio::spawn(async move {
            // Similar structure to LDA worker but with QDA-specific computations
            let mut rng = fastrand::Rng::new();

            for iteration in 0..config.max_iterations {
                let _permit = semaphore.acquire().await.expect("value should be present");

                let batch_indices: Vec<usize> = (0..config.batch_size)
                    .map(|_| rng.usize(0..x_data.nrows()))
                    .collect();

                let x_batch =
                    Array2::from_shape_fn((config.batch_size, x_data.ncols()), |(i, j)| {
                        x_data[[batch_indices[i], j]]
                    });

                let y_batch = Array1::from_shape_fn(config.batch_size, |i| y[batch_indices[i]]);

                // Compute QDA-specific gradients
                let gradients =
                    Self::compute_qda_gradients(&x_batch, &y_batch, &parameters).await?;
                let loss = Self::compute_qda_loss(&x_batch, &y_batch, &parameters).await?;

                let params_snapshot = parameters.read().expect("lock not poisoned").clone();
                tx.send(OptimizationMessage::ParameterUpdate {
                    worker_id,
                    parameters: params_snapshot,
                    gradient: gradients,
                    loss,
                })
                .await
                .map_err(|e| SklearsError::ProcessingError(format!("Channel send error: {}", e)))?;

                if iteration % 10 == 0 {
                    tx.send(OptimizationMessage::CheckConvergence {
                        worker_id,
                        current_loss: loss,
                    })
                    .await
                    .map_err(|e| {
                        SklearsError::ProcessingError(format!("Channel send error: {}", e))
                    })?;
                }

                tokio::task::yield_now().await;
            }

            Ok(())
        })
    }

    /// Compute QDA gradients
    async fn compute_qda_gradients(
        x_batch: &Array2<Float>, // standard ML notation: batch of feature vectors
        y_batch: &Array1<usize>,
        parameters: &Arc<RwLock<HashMap<String, Array2<Float>>>>,
    ) -> Result<HashMap<String, Array2<Float>>> {
        // Simplified QDA gradient computation
        // Real implementation would compute gradients for class-specific covariances
        Self::compute_lda_gradients(x_batch, y_batch, parameters).await
    }

    /// Compute QDA loss
    async fn compute_qda_loss(
        x_batch: &Array2<Float>, // standard ML notation: batch of feature vectors
        y_batch: &Array1<usize>,
        parameters: &Arc<RwLock<HashMap<String, Array2<Float>>>>,
    ) -> Result<Float> {
        // Simplified QDA loss computation
        // Real implementation would use class-specific covariances
        Self::compute_lda_loss(x_batch, y_batch, parameters).await
    }
}

/// Optimization statistics
#[derive(Debug, Clone)]
pub struct OptimizationStats {
    pub iteration: usize,
    pub current_loss: Float,
    pub converged: bool,
    pub training_time: Duration,
    pub last_update_time: Duration,
    pub update_count: usize,
    pub average_loss: Float,
    pub loss_variance: Float,
}

/// High-level async discriminant analysis trainer
pub struct AsyncDiscriminantAnalysis {
    optimizer: AsyncDiscriminantOptimizer,
    runtime: Option<tokio::runtime::Runtime>,
}

impl AsyncDiscriminantAnalysis {
    /// Create a new async discriminant analysis trainer
    pub fn new(config: AsyncOptimizationConfig) -> Self {
        let optimizer = AsyncDiscriminantOptimizer::new(config);

        // Create Tokio runtime for async operations
        let runtime = tokio::runtime::Builder::new_multi_thread()
            .worker_threads(num_cpus::get())
            .enable_all()
            .build()
            .ok();

        Self { optimizer, runtime }
    }

    /// Train LDA asynchronously
    pub fn train_lda_async(
        &self,
        x_data: Array2<Float>, // standard ML notation: feature matrix
        y: Array1<usize>,
        config: LinearDiscriminantAnalysisConfig,
    ) -> Result<HashMap<String, Array2<Float>>> {
        match &self.runtime {
            Some(rt) => rt.block_on(self.optimizer.async_train_lda(x_data, y, config)),
            None => Err(SklearsError::ProcessingError(
                "Async runtime not available".to_string(),
            )),
        }
    }

    /// Train QDA asynchronously
    pub fn train_qda_async(
        &self,
        x_data: Array2<Float>, // standard ML notation: feature matrix
        y: Array1<usize>,
        config: QuadraticDiscriminantAnalysisConfig,
    ) -> Result<HashMap<String, Array2<Float>>> {
        match &self.runtime {
            Some(rt) => rt.block_on(self.optimizer.async_train_qda(x_data, y, config)),
            None => Err(SklearsError::ProcessingError(
                "Async runtime not available".to_string(),
            )),
        }
    }

    /// Get optimization statistics
    pub fn get_stats(&self) -> OptimizationStats {
        self.optimizer.get_optimization_stats()
    }

    /// Check if training has converged
    pub fn has_converged(&self) -> bool {
        self.optimizer
            .state
            .read()
            .expect("lock not poisoned")
            .converged
    }
}

/// Distributed discriminant analysis across multiple nodes
pub struct DistributedDiscriminantAnalysis {
    node_id: usize,
    total_nodes: usize,
    async_trainer: AsyncDiscriminantAnalysis,
}

impl DistributedDiscriminantAnalysis {
    /// Create a new distributed trainer
    pub fn new(node_id: usize, total_nodes: usize, config: AsyncOptimizationConfig) -> Self {
        Self {
            node_id,
            total_nodes,
            async_trainer: AsyncDiscriminantAnalysis::new(config),
        }
    }

    /// Train on local data partition
    pub fn train_local_partition(
        &self,
        x_local: Array2<Float>, // standard ML notation: local partition of feature matrix
        y_local: Array1<usize>,
        lda_config: LinearDiscriminantAnalysisConfig,
    ) -> Result<HashMap<String, Array2<Float>>> {
        log::info!("Node {} training on local partition", self.node_id);

        // Train on local data
        let local_params = self
            .async_trainer
            .train_lda_async(x_local, y_local, lda_config)?;

        // In a real implementation, this would involve parameter aggregation
        // across nodes using techniques like federated averaging

        Ok(local_params)
    }

    /// Simulate parameter aggregation across nodes
    pub fn aggregate_parameters(
        &self,
        local_params: HashMap<String, Array2<Float>>,
        _all_node_params: Vec<HashMap<String, Array2<Float>>>,
    ) -> Result<HashMap<String, Array2<Float>>> {
        // Simplified aggregation - in practice would use sophisticated methods
        // like federated averaging, secure aggregation, etc.

        log::info!(
            "Node {} aggregating parameters from {} nodes",
            self.node_id,
            self.total_nodes
        );

        // For now, just return local parameters
        // Real implementation would average parameters across nodes
        Ok(local_params)
    }
}

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

    #[tokio::test]
    async fn test_async_optimization_config() {
        let config = AsyncOptimizationConfig::default();
        assert!(config.num_workers > 0);
        assert!(config.batch_size > 0);
        assert!(config.learning_rate > 0.0);
    }

    #[tokio::test]
    async fn test_parameter_initialization() {
        let config = AsyncOptimizationConfig::default();
        let optimizer = AsyncDiscriminantOptimizer::new(config);

        let result = optimizer.initialize_parameters(4, 3);
        assert!(result.is_ok());

        let params = optimizer
            .parameters
            .read()
            .expect("operation should succeed");
        assert!(params.contains_key("class_means"));
        assert!(params.contains_key("covariance"));
        assert!(params.contains_key("priors"));
    }

    #[test]
    fn test_async_discriminant_analysis_creation() {
        let config = AsyncOptimizationConfig::default();
        let _trainer = AsyncDiscriminantAnalysis::new(config);
        // Just test that it can be created successfully
    }

    #[test]
    fn test_distributed_discriminant_analysis_creation() {
        let config = AsyncOptimizationConfig::default();
        let _distributed = DistributedDiscriminantAnalysis::new(0, 4, config);
        // Test creation of distributed trainer
    }

    #[tokio::test]
    async fn test_optimization_state() {
        let mut state = OptimizationState::default();
        assert_eq!(state.iteration, 0);
        assert!(state.current_loss.is_infinite());
        assert!(!state.converged);

        state.iteration = 10;
        state.current_loss = 0.5;
        assert_eq!(state.iteration, 10);
        assert_eq!(state.current_loss, 0.5);
    }
}