optirs-core 0.3.1

//! Parallel optimizer operations using scirs2_core
//!
//! This module provides parallel processing capabilities for optimizers,
//! enabling efficient multi-core utilization for large-scale optimization.
//!
//! # Features
//!
//! - Parallel parameter group processing
//! - Parallel batch updates
//! - Automatic work distribution across CPU cores
//! - Zero-copy parameter handling
//!
//! # Performance
//!
//! Expected speedup: 4-8x on multi-core systems for multiple parameter groups

use scirs2_core::ndarray::{Array, Array1, Dimension, ScalarOperand};
use scirs2_core::numeric::Float;
use scirs2_core::parallel_ops::*;
use std::fmt::Debug;

use crate::error::Result;
use crate::optimizers::Optimizer;

/// Parallel optimizer wrapper for processing multiple parameter groups
///
/// This wrapper enables parallel processing of multiple parameter groups,
/// providing significant speedup on multi-core systems.
///
/// # Examples
///
/// ```
/// use scirs2_core::ndarray::Array1;
/// use optirs_core::optimizers::{SGD, Optimizer};
/// use optirs_core::parallel_optimizer::ParallelOptimizer;
///
/// // Create base optimizer
/// let optimizer = SGD::new(0.01);
///
/// // Wrap in parallel optimizer
/// let mut parallel_opt = ParallelOptimizer::new(optimizer);
///
/// // Process multiple parameter groups in parallel
/// let params_list = vec![
///     Array1::zeros(1000),
///     Array1::zeros(2000),
///     Array1::zeros(1500),
/// ];
/// let grads_list = vec![
///     Array1::from_elem(1000, 0.1),
///     Array1::from_elem(2000, 0.1),
///     Array1::from_elem(1500, 0.1),
/// ];
///
/// let updated = parallel_opt.step_parallel_groups(&params_list, &grads_list).expect("unwrap failed");
/// ```
#[derive(Debug)]
pub struct ParallelOptimizer<O, A, D>
where
    O: Optimizer<A, D> + Clone + Send + Sync,
    A: Float + ScalarOperand + Debug + Send + Sync,
    D: Dimension,
{
    base_optimizer: O,
    _phantom_a: std::marker::PhantomData<A>,
    _phantom_d: std::marker::PhantomData<D>,
}

impl<O, A, D> ParallelOptimizer<O, A, D>
where
    O: Optimizer<A, D> + Clone + Send + Sync,
    A: Float + ScalarOperand + Debug + Send + Sync,
    D: Dimension,
{
    /// Creates a new parallel optimizer wrapper
    ///
    /// # Arguments
    ///
    /// * `base_optimizer` - The base optimizer to parallelize
    pub fn new(base_optimizer: O) -> Self {
        Self {
            base_optimizer,
            _phantom_a: std::marker::PhantomData,
            _phantom_d: std::marker::PhantomData,
        }
    }

    /// Process multiple parameter groups in parallel
    ///
    /// This method distributes parameter groups across available CPU cores
    /// for parallel processing.
    ///
    /// # Arguments
    ///
    /// * `params_list` - List of parameter arrays
    /// * `grads_list` - List of gradient arrays
    ///
    /// # Returns
    ///
    /// Updated parameter arrays processed in parallel
    pub fn step_parallel_groups(
        &mut self,
        params_list: &[Array<A, D>],
        grads_list: &[Array<A, D>],
    ) -> Result<Vec<Array<A, D>>>
    where
        Array<A, D>: Clone + Send + Sync,
    {
        if params_list.len() != grads_list.len() {
            return Err(crate::error::OptimError::InvalidConfig(format!(
                "Parameter groups ({}) and gradient groups ({}) must have same length",
                params_list.len(),
                grads_list.len()
            )));
        }

        // Use parallel iterator from scirs2_core
        let results: Vec<Result<Array<A, D>>> = params_list
            .par_iter()
            .zip(grads_list.par_iter())
            .map(|(params, grads)| {
                let mut opt_clone = self.base_optimizer.clone();
                opt_clone.step(params, grads)
            })
            .collect();

        // Collect results and handle errors
        let mut updated_params = Vec::with_capacity(results.len());
        for result in results {
            updated_params.push(result?);
        }

        Ok(updated_params)
    }

    /// Get the underlying optimizer
    pub fn inner(&self) -> &O {
        &self.base_optimizer
    }

    /// Get mutable reference to underlying optimizer
    pub fn inner_mut(&mut self) -> &mut O {
        &mut self.base_optimizer
    }

    /// Get the current learning rate from the base optimizer
    pub fn get_learning_rate(&self) -> A {
        self.base_optimizer.get_learning_rate()
    }

    /// Set the learning rate on the base optimizer
    pub fn set_learning_rate(&mut self, learning_rate: A) {
        self.base_optimizer.set_learning_rate(learning_rate);
    }
}

/// Parallel batch processor for large parameter arrays
///
/// This processor splits large parameter arrays into chunks and processes
/// them in parallel, providing speedup for very large models.
pub struct ParallelBatchProcessor {
    /// Minimum chunk size for parallel processing
    min_chunk_size: usize,
    /// Number of threads to use (None = automatic)
    num_threads: Option<usize>,
}

impl ParallelBatchProcessor {
    /// Creates a new parallel batch processor
    ///
    /// # Arguments
    ///
    /// * `min_chunk_size` - Minimum size of each chunk (default: 1024)
    pub fn new(min_chunk_size: usize) -> Self {
        Self {
            min_chunk_size,
            num_threads: None,
        }
    }

    /// Set the number of threads to use
    ///
    /// # Arguments
    ///
    /// * `num_threads` - Number of threads (None for automatic)
    pub fn with_threads(mut self, num_threads: Option<usize>) -> Self {
        self.num_threads = num_threads;
        self
    }

    /// Determine if parallel processing should be used
    ///
    /// # Arguments
    ///
    /// * `size` - Size of the parameter array
    ///
    /// # Returns
    ///
    /// True if parallel processing would be beneficial
    pub fn should_use_parallel(&self, size: usize) -> bool {
        let num_cores = num_cpus::get();
        size >= self.min_chunk_size * num_cores
    }

    /// Get optimal chunk size for parallel processing
    ///
    /// # Arguments
    ///
    /// * `total_size` - Total size of the array
    ///
    /// # Returns
    ///
    /// Optimal chunk size for parallel processing
    pub fn optimal_chunk_size(&self, total_size: usize) -> usize {
        let num_cores = self.num_threads.unwrap_or_else(num_cpus::get);
        let chunk_size = total_size / num_cores;
        chunk_size.max(self.min_chunk_size)
    }
}

impl Default for ParallelBatchProcessor {
    fn default() -> Self {
        Self::new(1024)
    }
}

/// Helper function to process parameter groups in parallel
///
/// This is a convenience function for one-off parallel processing without
/// creating a ParallelOptimizer instance.
///
/// # Arguments
///
/// * `optimizer` - The optimizer to use (will be cloned for each group)
/// * `params_list` - List of parameter arrays
/// * `grads_list` - List of gradient arrays
///
/// # Returns
///
/// Updated parameter arrays processed in parallel
pub fn parallel_step<O, A, D>(
    optimizer: &mut O,
    params_list: &[Array<A, D>],
    grads_list: &[Array<A, D>],
) -> Result<Vec<Array<A, D>>>
where
    O: Optimizer<A, D> + Clone + Send + Sync,
    A: Float + ScalarOperand + Debug + Send + Sync,
    D: Dimension,
    Array<A, D>: Clone + Send + Sync,
{
    if params_list.len() != grads_list.len() {
        return Err(crate::error::OptimError::InvalidConfig(format!(
            "Parameter groups ({}) and gradient groups ({}) must have same length",
            params_list.len(),
            grads_list.len()
        )));
    }

    let results: Vec<Result<Array<A, D>>> = params_list
        .par_iter()
        .zip(grads_list.par_iter())
        .map(|(params, grads)| {
            let mut opt_clone = optimizer.clone();
            opt_clone.step(params, grads)
        })
        .collect();

    let mut updated_params = Vec::with_capacity(results.len());
    for result in results {
        updated_params.push(result?);
    }

    Ok(updated_params)
}

/// Parallel processing for Array1 specifically (optimized path)
pub fn parallel_step_array1<O, A>(
    optimizer: &mut O,
    params_list: &[Array1<A>],
    grads_list: &[Array1<A>],
) -> Result<Vec<Array1<A>>>
where
    O: Optimizer<A, scirs2_core::ndarray::Ix1> + Clone + Send + Sync,
    A: Float + ScalarOperand + Debug + Send + Sync,
{
    if params_list.len() != grads_list.len() {
        return Err(crate::error::OptimError::InvalidConfig(format!(
            "Parameter groups ({}) and gradient groups ({}) must have same length",
            params_list.len(),
            grads_list.len()
        )));
    }

    let results: Vec<Result<Array1<A>>> = params_list
        .par_iter()
        .zip(grads_list.par_iter())
        .map(|(params, grads)| {
            let mut opt_clone = optimizer.clone();
            opt_clone.step(params, grads)
        })
        .collect();

    let mut updated_params = Vec::with_capacity(results.len());
    for result in results {
        updated_params.push(result?);
    }

    Ok(updated_params)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::optimizers::SGD;
    use approx::assert_relative_eq;

    #[test]
    fn test_parallel_optimizer_basic() {
        let optimizer = SGD::new(0.1);
        let mut parallel_opt = ParallelOptimizer::new(optimizer);

        let params_list = vec![
            Array1::from_vec(vec![1.0f32, 2.0, 3.0]),
            Array1::from_vec(vec![4.0, 5.0, 6.0]),
        ];
        let grads_list = vec![
            Array1::from_vec(vec![0.1, 0.2, 0.3]),
            Array1::from_vec(vec![0.1, 0.2, 0.3]),
        ];

        let results = parallel_opt
            .step_parallel_groups(&params_list, &grads_list)
            .expect("unwrap failed");

        assert_eq!(results.len(), 2);
        assert_relative_eq!(results[0][0], 0.99, epsilon = 1e-6);
        assert_relative_eq!(results[1][0], 3.99, epsilon = 1e-6);
    }

    #[test]
    fn test_parallel_optimizer_multiple_groups() {
        let optimizer = SGD::new(0.01);
        let mut parallel_opt = ParallelOptimizer::new(optimizer);

        // Create 10 parameter groups
        let params_list: Vec<Array1<f32>> =
            (0..10).map(|i| Array1::from_elem(100, i as f32)).collect();
        let grads_list: Vec<Array1<f32>> = (0..10).map(|_| Array1::from_elem(100, 0.1)).collect();

        let results = parallel_opt
            .step_parallel_groups(&params_list, &grads_list)
            .expect("unwrap failed");

        assert_eq!(results.len(), 10);
        // Verify first group was updated correctly
        assert_relative_eq!(results[0][0], 0.0 - 0.01 * 0.1, epsilon = 1e-6);
    }

    #[test]
    fn test_parallel_step_function() {
        let mut optimizer = SGD::new(0.1);

        let params_list = vec![
            Array1::from_vec(vec![1.0f32, 2.0]),
            Array1::from_vec(vec![3.0, 4.0]),
        ];
        let grads_list = vec![
            Array1::from_vec(vec![0.1, 0.2]),
            Array1::from_vec(vec![0.3, 0.4]),
        ];

        let results =
            parallel_step(&mut optimizer, &params_list, &grads_list).expect("unwrap failed");

        assert_eq!(results.len(), 2);
        assert_relative_eq!(results[0][0], 0.99, epsilon = 1e-6);
        assert_relative_eq!(results[1][0], 2.97, epsilon = 1e-6);
    }

    #[test]
    fn test_parallel_batch_processor() {
        let processor = ParallelBatchProcessor::new(1024);

        // Small array - should not use parallel
        assert!(!processor.should_use_parallel(100));

        // Large array - should use parallel
        let num_cores = num_cpus::get();
        assert!(processor.should_use_parallel(1024 * num_cores * 2));

        // Test optimal chunk size calculation
        let chunk_size = processor.optimal_chunk_size(10000);
        assert!(chunk_size >= 1024);
    }

    #[test]
    fn test_parallel_batch_processor_threads() {
        let processor = ParallelBatchProcessor::new(1024).with_threads(Some(4));

        let chunk_size = processor.optimal_chunk_size(10000);
        // With 4 threads, chunk size should be around 10000/4 = 2500
        assert!(chunk_size >= 1024);
        assert!(chunk_size <= 10000);
    }

    #[test]
    fn test_parallel_optimizer_learning_rate() {
        let optimizer = SGD::new(0.1);
        let mut parallel_opt: ParallelOptimizer<_, f64, scirs2_core::ndarray::Ix1> =
            ParallelOptimizer::new(optimizer);

        assert_relative_eq!(parallel_opt.get_learning_rate(), 0.1, epsilon = 1e-6);

        parallel_opt.set_learning_rate(0.2);
        assert_relative_eq!(parallel_opt.get_learning_rate(), 0.2, epsilon = 1e-6);
    }

    #[test]
    fn test_parallel_step_array1() {
        let mut optimizer = SGD::new(0.1);

        let params_list = vec![
            Array1::from_vec(vec![1.0f32, 2.0, 3.0]),
            Array1::from_vec(vec![4.0, 5.0, 6.0]),
        ];
        let grads_list = vec![
            Array1::from_vec(vec![0.1, 0.2, 0.3]),
            Array1::from_vec(vec![0.1, 0.2, 0.3]),
        ];

        let results =
            parallel_step_array1(&mut optimizer, &params_list, &grads_list).expect("unwrap failed");

        assert_eq!(results.len(), 2);
        assert_relative_eq!(results[0][0], 0.99, epsilon = 1e-6);
        assert_relative_eq!(results[1][0], 3.99, epsilon = 1e-6);
    }
}