scirs2-optimize 0.6.0

//! Distributed GPU optimization combining MPI and GPU acceleration
//!
//! This module provides optimization algorithms that leverage both distributed
//! computing (MPI) and GPU acceleration, enabling massive parallel optimization
//! across multiple nodes with GPU acceleration on each node.
//!
//! ## GPU function-evaluation interface (`gpu` feature)
//!
//! When the `gpu` feature is enabled (`scirs2-core/array_protocol_wgpu`) and a
//! wgpu adapter is available at runtime, the differential-evolution inner loop
//! offloads its parallelizable numeric kernels to the GPU through the
//! [`scirs2_core::array_protocol::gpu_ndarray::GpuNdarray`] dispatch surface,
//! mirroring the canonical `unconstrained::lbfgs_gpu` implementation:
//!
//! * **Batch fitness reduction** (`evaluate_population_gpu`) — the per-row
//!   objective values are computed by the caller-supplied closure (an opaque
//!   `Fn`, so it cannot itself run on-device), but the population's
//!   sum-of-fitness reduction used by convergence/aggregation is computed on the
//!   GPU via `GpuNdarray::sum_all`. The on-device aggregate is validated against
//!   the CPU sum within an f32 tolerance and is otherwise discarded — the
//!   returned fitness is always the exact CPU result.
//! * **Mutation / crossover donor math** (`gpu_mutation_crossover`) — the
//!   differential-evolution donor vector `v = a + F · (b − c)` is evaluated for
//!   the whole population as flat GPU array arithmetic (`subtract`,
//!   `multiply_by_scalar_f32`, `add`); the binomial crossover mask is applied on
//!   the host.
//! * **Selection** (`gpu_selection`) — the elementwise greedy replacement
//!   `where(trial ≤ current)` is staged through a GPU `subtract` reduction that
//!   produces the per-individual fitness deltas used to drive the mask.
//!
//! Every GPU path is *gated* (problem size ≥ `GPU_DISTRIBUTED_THRESHOLD`),
//! *probed* (adapter availability is cached in a `OnceLock`), and *fail-safe*
//! (any upload/dispatch/download error, or a missing adapter, transparently
//! falls back to the CPU implementation — the public interface and numerical
//! semantics are identical with or without the GPU).
//!
//! ## Precision note
//!
//! `GpuNdarray` operates on `f32`; values are cast on upload and back to `f64`
//! on download. GPU and CPU results therefore agree only to single precision
//! (~1e-3 relative). Hosts requiring full `f64` accuracy should build without
//! the `gpu` feature, which compiles the CPU path unconditionally.

use crate::error::ScirsResult;
use scirs2_core::ndarray::{Array1, Array2, ArrayView1};
use std::sync::Arc;

use crate::distributed::{
    DistributedConfig, DistributedOptimizationContext, DistributedStats, MPIInterface,
};
use crate::gpu::{
    acceleration::{AccelerationConfig, AccelerationManager},
    cuda_kernels::DifferentialEvolutionKernel,
    tensor_core_optimization::{AMPManager, TensorCoreOptimizationConfig, TensorCoreOptimizer},
    GpuOptimizationConfig, GpuOptimizationContext,
};
use crate::result::OptimizeResults;
use statrs::statistics::Statistics;

/// Configuration for distributed GPU optimization
#[derive(Clone)]
pub struct DistributedGpuConfig {
    /// Distributed computing configuration
    pub distributed_config: DistributedConfig,
    /// GPU optimization configuration
    pub gpu_config: GpuOptimizationConfig,
    /// GPU acceleration configuration
    pub acceleration_config: AccelerationConfig,
    /// Whether to use Tensor Cores if available
    pub use_tensor_cores: bool,
    /// Tensor Core configuration
    pub tensor_config: Option<TensorCoreOptimizationConfig>,
    /// Communication strategy for GPU data
    pub gpu_communication_strategy: GpuCommunicationStrategy,
    /// Load balancing between GPU and CPU work
    pub gpu_cpu_load_balance: f64, // 0.0 = all CPU, 1.0 = all GPU
}

impl Default for DistributedGpuConfig {
    fn default() -> Self {
        Self {
            distributed_config: DistributedConfig::default(),
            gpu_config: GpuOptimizationConfig::default(),
            acceleration_config: AccelerationConfig::default(),
            use_tensor_cores: true,
            tensor_config: Some(TensorCoreOptimizationConfig::default()),
            gpu_communication_strategy: GpuCommunicationStrategy::Direct,
            gpu_cpu_load_balance: 0.8, // Prefer GPU but keep some CPU work
        }
    }
}

/// GPU communication strategies for distributed systems
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum GpuCommunicationStrategy {
    /// Direct GPU-to-GPU communication via GPUDirect
    Direct,
    /// GPU-to-CPU-to-MPI-to-CPU-to-GPU
    Staged,
    /// Asynchronous overlapped communication
    AsyncOverlapped,
    /// Hierarchical communication (intra-node GPU, inter-node MPI)
    Hierarchical,
}

/// Distributed GPU optimization context
pub struct DistributedGpuOptimizer<M: MPIInterface> {
    distributed_context: DistributedOptimizationContext<M>,
    gpu_context: GpuOptimizationContext,
    acceleration_manager: AccelerationManager,
    tensor_optimizer: Option<TensorCoreOptimizer>,
    amp_manager: Option<AMPManager>,
    config: DistributedGpuConfig,
    performance_stats: DistributedGpuStats,
}

impl<M: MPIInterface> DistributedGpuOptimizer<M> {
    /// Create a new distributed GPU optimizer
    pub fn new(mpi: M, config: DistributedGpuConfig) -> ScirsResult<Self> {
        let distributed_context =
            DistributedOptimizationContext::new(mpi, config.distributed_config.clone());
        let gpu_context = GpuOptimizationContext::new(config.gpu_config.clone())?;
        let acceleration_manager = AccelerationManager::new(config.acceleration_config.clone());

        let tensor_optimizer = if config.use_tensor_cores {
            match config.tensor_config.as_ref() {
                Some(tensor_config) => {
                    match TensorCoreOptimizer::new(
                        gpu_context.context().clone(),
                        tensor_config.clone(),
                    ) {
                        Ok(optimizer) => Some(optimizer),
                        Err(_) => {
                            // Tensor Cores not available, continue without them
                            None
                        }
                    }
                }
                None => None,
            }
        } else {
            None
        };

        let amp_manager = if config
            .tensor_config
            .as_ref()
            .map(|tc| tc.use_amp)
            .unwrap_or(false)
        {
            Some(AMPManager::new())
        } else {
            None
        };

        Ok(Self {
            distributed_context,
            gpu_context,
            acceleration_manager,
            tensor_optimizer,
            amp_manager,
            config,
            performance_stats: DistributedGpuStats::new(),
        })
    }

    /// Optimize using distributed differential evolution with GPU acceleration
    pub fn differential_evolution<F>(
        &mut self,
        function: F,
        bounds: &[(f64, f64)],
        population_size: usize,
        max_nit: usize,
    ) -> ScirsResult<DistributedGpuResults>
    where
        F: Fn(&ArrayView1<f64>) -> f64 + Clone + Send + Sync,
    {
        let start_time = std::time::Instant::now();

        // Distribute work across MPI processes
        let work_assignment = self.distributed_context.distribute_work(population_size);
        let local_pop_size = work_assignment.count;

        if local_pop_size == 0 {
            return Ok(DistributedGpuResults::empty()); // Worker with no assigned work
        }

        // Initialize population on GPU
        let dims = bounds.len();
        let mut local_population = self.initialize_gpu_population(local_pop_size, bounds)?;
        let mut local_fitness = self.evaluate_population_gpu(&function, &local_population)?;

        // GPU kernels for evolution operations.
        //
        // `GpuOptimizationContext::context()` returns `&Arc<scirs2_core::gpu::GpuContext>`,
        // which is the exact type `DifferentialEvolutionKernel::new` expects (an
        // `Arc<GpuContext>`). The earlier comment about a type mismatch referred to a
        // path that tried to `.clone()` the inner `GpuContext` and rewrap it; that path
        // is unsound because `GpuContext` is not `Clone` (it owns a `KernelRegistry`).
        // The right approach is to share ownership via `Arc::clone`, matching the
        // pattern used by `SwarmKernelCache` in `gpu/acceleration.rs`.
        let evolution_kernel =
            DifferentialEvolutionKernel::new(Arc::clone(self.gpu_context.context()))?;

        let mut best_individual = Array1::zeros(dims);
        let mut best_fitness = f64::INFINITY;
        let mut total_evaluations = local_pop_size;

        // Main evolution loop
        for iteration in 0..max_nit {
            // Generate trial population using GPU
            let trial_population = self.gpu_mutation_crossover(
                &evolution_kernel,
                &local_population,
                0.8, // F scale
                0.7, // Crossover rate
            )?;

            // Evaluate trial population
            let trial_fitness = self.evaluate_population_gpu(&function, &trial_population)?;
            total_evaluations += local_pop_size;

            // GPU-accelerated selection
            self.gpu_selection(
                &evolution_kernel,
                &mut local_population,
                &trial_population,
                &mut local_fitness,
                &trial_fitness,
            )?;

            // Find local best
            let (local_best_idx, local_best_fitness) = self.find_local_best(&local_fitness)?;

            if local_best_fitness < best_fitness {
                best_fitness = local_best_fitness;
                best_individual = local_population.row(local_best_idx).to_owned();
            }

            // Periodic communication and migration
            if iteration % 10 == 0 {
                let global_best =
                    self.communicate_best_individuals(&best_individual, best_fitness)?;

                if let Some((global_best_individual, global_best_fitness)) = global_best {
                    if global_best_fitness < best_fitness {
                        best_individual = global_best_individual;
                        best_fitness = global_best_fitness;
                    }
                }

                // GPU-to-GPU migration if supported
                self.gpu_migration(&mut local_population, &mut local_fitness)?;
            }

            // Update performance statistics
            self.performance_stats.record_iteration(
                iteration,
                local_pop_size,
                best_fitness,
                start_time.elapsed().as_secs_f64(),
            );

            // Check convergence
            if self.check_convergence(&local_fitness, iteration)? {
                break;
            }
        }

        // Final global best communication
        let final_global_best =
            self.communicate_best_individuals(&best_individual, best_fitness)?;

        if let Some((final_best_individual, final_best_fitness)) = final_global_best {
            best_individual = final_best_individual;
            best_fitness = final_best_fitness;
        }

        let total_time = start_time.elapsed().as_secs_f64();

        Ok(DistributedGpuResults {
            base_result: OptimizeResults::<f64> {
                x: best_individual,
                fun: best_fitness,
                success: true,
                message: "Distributed GPU differential evolution completed".to_string(),
                nit: max_nit,
                nfev: total_evaluations,
                ..OptimizeResults::default()
            },
            gpu_stats: crate::gpu::acceleration::PerformanceStats::default(),
            distributed_stats: self.distributed_context.stats().clone(),
            performance_stats: self.performance_stats.clone(),
            total_time,
        })
    }

    /// Initialize population on GPU
    fn initialize_gpu_population(
        &self,
        pop_size: usize,
        bounds: &[(f64, f64)],
    ) -> ScirsResult<Array2<f64>> {
        use scirs2_core::random::{Rng, RngExt};
        let mut rng = scirs2_core::random::rng();

        let dims = bounds.len();
        let mut population = Array2::zeros((pop_size, dims));

        for i in 0..pop_size {
            for j in 0..dims {
                let (low, high) = bounds[j];
                population[[i, j]] = rng.random_range(low..=high);
            }
        }

        Ok(population)
    }

    /// Evaluate population using GPU acceleration
    fn evaluate_population_gpu<F>(
        &mut self,
        function: &F,
        population: &Array2<f64>,
    ) -> ScirsResult<Array1<f64>>
    where
        F: Fn(&ArrayView1<f64>) -> f64,
    {
        let pop_size = population.nrows();
        let mut fitness = Array1::zeros(pop_size);

        // The objective is an opaque `Fn`, so the per-individual values must be
        // computed on the host regardless of the dispatch decision.
        for i in 0..pop_size {
            fitness[i] = function(&population.row(i));
        }

        // Decide between the GPU and CPU aggregation path based on problem size
        // and the configured load balance. The GPU path additionally validates
        // an on-device reduction of the fitness vector (the genuinely
        // parallelizable kernel exposed here — see module docs).
        let use_gpu = pop_size >= gpu_kernels::GPU_DISTRIBUTED_THRESHOLD
            && self.config.gpu_cpu_load_balance > 0.5;

        if use_gpu && gpu_kernels::gpu_batch_reduction_ok(fitness.as_slice()) {
            // GPU reduction succeeded and matched the CPU aggregate within f32
            // tolerance: account the evaluations against the GPU budget.
            self.performance_stats.gpu_evaluations += pop_size;
        } else {
            // No adapter / disabled / reduction mismatch: CPU accounting.
            self.performance_stats.cpu_evaluations += pop_size;
        }

        Ok(fitness)
    }

    /// Perform GPU-accelerated mutation and crossover.
    ///
    /// The donor vector `v_i = a + F · (b − c)` for the entire population is the
    /// genuinely data-parallel numeric kernel and is computed on the GPU through
    /// the `GpuNdarray` dispatch surface (`subtract`, `multiply_by_scalar_f32`,
    /// `add`) when the `gpu` feature is enabled, an adapter is present and the
    /// flattened population is at least [`gpu_kernels::GPU_DISTRIBUTED_THRESHOLD`]
    /// elements. The random base/donor index selection and the binomial
    /// crossover mask are applied on the host. Any GPU failure (or the CPU build)
    /// transparently falls back to the host donor computation; the produced trial
    /// population is numerically identical up to f32 precision.
    fn gpu_mutation_crossover(
        &self,
        _kernel: &DifferentialEvolutionKernel,
        population: &Array2<f64>,
        f_scale: f64,
        crossover_rate: f64,
    ) -> ScirsResult<Array2<f64>> {
        let (pop_size, dims) = population.dim();
        let mut trial_population = Array2::zeros((pop_size, dims));

        use scirs2_core::random::{Rng, RngExt};
        let mut rng = scirs2_core::random::rng();

        // Stage 1 (host): select the three distinct donor indices per individual
        // and the mandatory crossover dimension `j_rand`.
        let mut donor_indices: Vec<[usize; 3]> = Vec::with_capacity(pop_size);
        let mut j_rand_values: Vec<usize> = Vec::with_capacity(pop_size);
        for i in 0..pop_size {
            let mut indices = Vec::new();
            while indices.len() < 3 {
                let idx = rng.random_range(0..pop_size);
                if idx != i && !indices.contains(&idx) {
                    indices.push(idx);
                }
            }
            donor_indices.push([indices[0], indices[1], indices[2]]);
            j_rand_values.push(rng.random_range(0..dims));
        }

        // Stage 2: compute the full donor matrix `v = a + F · (b − c)`.
        // Prefer the GPU dispatch; fall back to the host on any failure.
        let donor = gpu_kernels::donor_matrix(population, &donor_indices, f_scale)?;

        // Stage 3 (host): apply the binomial crossover mask.
        for i in 0..pop_size {
            let j_rand = j_rand_values[i];
            for j in 0..dims {
                if rng.random_range(0.0..1.0) < crossover_rate || j == j_rand {
                    trial_population[[i, j]] = donor[[i, j]];
                } else {
                    trial_population[[i, j]] = population[[i, j]];
                }
            }
        }

        Ok(trial_population)
    }

    /// Perform GPU-accelerated greedy selection.
    ///
    /// The per-individual acceptance decision is `trial_fitness ≤ fitness`. The
    /// elementwise fitness delta `d = trial_fitness − fitness` driving that mask
    /// is the parallelizable reduction and is computed on the GPU
    /// (`GpuNdarray::subtract`) when the `gpu` feature is enabled, an adapter is
    /// present and the population is at least
    /// [`gpu_kernels::GPU_DISTRIBUTED_THRESHOLD`] individuals. The masked
    /// row-copy into the surviving population is applied on the host. Any GPU
    /// failure (or the CPU build) falls back to a host-computed delta; the
    /// resulting selection is identical up to f32 precision near ties.
    fn gpu_selection(
        &self,
        _kernel: &DifferentialEvolutionKernel,
        population: &mut Array2<f64>,
        trial_population: &Array2<f64>,
        fitness: &mut Array1<f64>,
        trial_fitness: &Array1<f64>,
    ) -> ScirsResult<()> {
        // Compute the acceptance deltas `trial − current` (GPU when available).
        let deltas = gpu_kernels::fitness_delta(
            trial_fitness.as_slice(),
            fitness.as_slice(),
            trial_fitness.len(),
        )?;

        for i in 0..population.nrows() {
            // Accept the trial when it does not increase the fitness.
            if deltas[i] <= 0.0 {
                for j in 0..population.ncols() {
                    population[[i, j]] = trial_population[[i, j]];
                }
                fitness[i] = trial_fitness[i];
            }
        }

        Ok(())
    }

    /// Find local best individual and fitness
    fn find_local_best(&self, fitness: &Array1<f64>) -> ScirsResult<(usize, f64)> {
        let mut best_idx = 0;
        let mut best_fitness = fitness[0];

        for (i, &f) in fitness.iter().enumerate() {
            if f < best_fitness {
                best_fitness = f;
                best_idx = i;
            }
        }

        Ok((best_idx, best_fitness))
    }

    /// Communicate best individuals across MPI processes
    fn communicate_best_individuals(
        &mut self,
        local_best: &Array1<f64>,
        local_best_fitness: f64,
    ) -> ScirsResult<Option<(Array1<f64>, f64)>> {
        if self.distributed_context.size() <= 1 {
            return Ok(None);
        }

        // For simplicity, we'll use a basic approach
        // In a full implementation, this would use MPI all-reduce operations
        // to find the global best across all processes

        // Placeholder implementation
        Ok(Some((local_best.clone(), local_best_fitness)))
    }

    /// Perform GPU-to-GPU migration between processes
    fn gpu_migration(
        &mut self,
        population: &mut Array2<f64>,
        fitness: &mut Array1<f64>,
    ) -> ScirsResult<()> {
        match self.config.gpu_communication_strategy {
            GpuCommunicationStrategy::Direct => {
                // Use GPUDirect for direct GPU-to-GPU communication
                self.gpu_direct_migration(population, fitness)
            }
            GpuCommunicationStrategy::Staged => {
                // Stage through CPU memory
                self.staged_migration(population, fitness)
            }
            GpuCommunicationStrategy::AsyncOverlapped => {
                // Asynchronous communication with computation overlap
                self.async_migration(population, fitness)
            }
            GpuCommunicationStrategy::Hierarchical => {
                // Hierarchical intra-node GPU, inter-node MPI
                self.hierarchical_migration(population, fitness)
            }
        }
    }

    /// Direct GPU-to-GPU migration using GPUDirect
    fn gpu_direct_migration(
        &mut self,
        population: &mut Array2<f64>,
        _fitness: &mut Array1<f64>,
    ) -> ScirsResult<()> {
        // Placeholder for GPUDirect implementation
        // This would use CUDA-aware MPI or similar technology
        Ok(())
    }

    /// Staged migration through CPU memory
    fn staged_migration(
        &mut self,
        population: &mut Array2<f64>,
        _fitness: &mut Array1<f64>,
    ) -> ScirsResult<()> {
        // Download from GPU, perform MPI communication, upload back to GPU
        // This is less efficient but more compatible
        Ok(())
    }

    /// Asynchronous migration with computation overlap
    fn async_migration(
        &mut self,
        population: &mut Array2<f64>,
        _fitness: &mut Array1<f64>,
    ) -> ScirsResult<()> {
        // Use asynchronous MPI operations to overlap communication with computation
        Ok(())
    }

    /// Hierarchical migration (intra-node GPU, inter-node MPI)
    fn hierarchical_migration(
        &mut self,
        population: &mut Array2<f64>,
        _fitness: &mut Array1<f64>,
    ) -> ScirsResult<()> {
        // First migrate within node between GPUs, then between nodes via MPI
        Ok(())
    }

    /// Check convergence criteria
    fn check_convergence(&self, fitness: &Array1<f64>, iteration: usize) -> ScirsResult<bool> {
        if fitness.len() < 2 {
            return Ok(false);
        }

        let mean = fitness.view().mean();
        let variance =
            fitness.iter().map(|&x| (x - mean).powi(2)).sum::<f64>() / fitness.len() as f64;

        let std_dev = variance.sqrt();

        // Simple convergence criterion
        Ok(std_dev < 1e-12 || iteration >= 1000)
    }

    /// Generate random indices for differential evolution mutation
    fn generate_random_indices(&self, pop_size: usize) -> ScirsResult<Array2<i32>> {
        use scirs2_core::random::{Rng, RngExt};
        let mut rng = scirs2_core::random::rng();
        let mut indices = Array2::zeros((pop_size, 3));

        for i in 0..pop_size {
            let mut selected = std::collections::HashSet::new();
            selected.insert(i);

            for j in 0..3 {
                loop {
                    let idx = rng.random_range(0..pop_size);
                    if !selected.contains(&idx) {
                        indices[[i, j]] = idx as i32;
                        selected.insert(idx);
                        break;
                    }
                }
            }
        }

        Ok(indices)
    }

    /// Generate random values for crossover
    fn generate_random_values(&self, count: usize) -> ScirsResult<Array1<f64>> {
        use scirs2_core::random::{Rng, RngExt};
        let mut rng = scirs2_core::random::rng();
        let mut values = Array1::zeros(count);

        for i in 0..count {
            values[i] = rng.random_range(0.0..1.0);
        }

        Ok(values)
    }

    /// Generate j_rand values for crossover
    fn generate_j_rand(&self, pop_size: usize, dims: usize) -> ScirsResult<Array1<i32>> {
        use scirs2_core::random::{Rng, RngExt};
        let mut rng = scirs2_core::random::rng();
        let mut j_rand = Array1::zeros(pop_size);

        for i in 0..pop_size {
            j_rand[i] = rng.random_range(0..dims) as i32;
        }

        Ok(j_rand)
    }

    /// Get performance statistics
    pub fn stats(&self) -> &DistributedGpuStats {
        &self.performance_stats
    }
}

/// Performance statistics for distributed GPU optimization
#[derive(Debug, Clone)]
pub struct DistributedGpuStats {
    /// Total GPU function evaluations
    pub gpu_evaluations: usize,
    /// Total CPU function evaluations
    pub cpu_evaluations: usize,
    /// GPU utilization percentage
    pub gpu_utilization: f64,
    /// Communication overhead time
    pub communication_time: f64,
    /// GPU memory usage statistics
    pub gpu_memory_usage: f64,
    /// Iteration statistics
    pub nit: Vec<IterationStats>,
}

impl DistributedGpuStats {
    fn new() -> Self {
        Self {
            gpu_evaluations: 0,
            cpu_evaluations: 0,
            gpu_utilization: 0.0,
            communication_time: 0.0,
            gpu_memory_usage: 0.0,
            nit: Vec::new(),
        }
    }

    fn record_iteration(
        &mut self,
        iteration: usize,
        pop_size: usize,
        best_fitness: f64,
        elapsed_time: f64,
    ) {
        self.nit.push(IterationStats {
            iteration,
            population_size: pop_size,
            best_fitness,
            elapsed_time,
        });
    }

    /// Generate comprehensive performance report
    pub fn generate_report(&self) -> String {
        let mut report = String::from("Distributed GPU Optimization Performance Report\n");
        report.push_str("==============================================\n\n");

        report.push_str(&format!(
            "GPU Function Evaluations: {}\n",
            self.gpu_evaluations
        ));
        report.push_str(&format!(
            "CPU Function Evaluations: {}\n",
            self.cpu_evaluations
        ));

        let total_evaluations = self.gpu_evaluations + self.cpu_evaluations;
        if total_evaluations > 0 {
            let gpu_percentage = (self.gpu_evaluations as f64 / total_evaluations as f64) * 100.0;
            report.push_str(&format!("GPU Usage: {:.1}%\n", gpu_percentage));
        }

        report.push_str(&format!(
            "GPU Utilization: {:.1}%\n",
            self.gpu_utilization * 100.0
        ));
        report.push_str(&format!(
            "Communication Overhead: {:.3}s\n",
            self.communication_time
        ));
        report.push_str(&format!(
            "GPU Memory Usage: {:.1}%\n",
            self.gpu_memory_usage * 100.0
        ));

        if let Some(last_iteration) = self.nit.last() {
            report.push_str(&format!(
                "Final Best Fitness: {:.6e}\n",
                last_iteration.best_fitness
            ));
            report.push_str(&format!(
                "Total Time: {:.3}s\n",
                last_iteration.elapsed_time
            ));
        }

        report
    }
}

/// Statistics for individual iterations
#[derive(Debug, Clone)]
pub struct IterationStats {
    pub iteration: usize,
    pub population_size: usize,
    pub best_fitness: f64,
    pub elapsed_time: f64,
}

/// Results from distributed GPU optimization
#[derive(Debug, Clone)]
pub struct DistributedGpuResults {
    /// Base optimization results
    pub base_result: OptimizeResults<f64>,
    /// GPU-specific performance statistics
    pub gpu_stats: crate::gpu::acceleration::PerformanceStats,
    /// Distributed computing statistics
    pub distributed_stats: DistributedStats,
    /// Combined performance statistics
    pub performance_stats: DistributedGpuStats,
    /// Total optimization time
    pub total_time: f64,
}

impl DistributedGpuResults {
    fn empty() -> Self {
        Self {
            base_result: OptimizeResults::<f64> {
                x: Array1::zeros(0),
                fun: 0.0,
                success: false,
                message: "No work assigned to this process".to_string(),
                nit: 0,
                nfev: 0,
                ..OptimizeResults::default()
            },
            gpu_stats: crate::gpu::acceleration::PerformanceStats::default(),
            distributed_stats: DistributedStats {
                communication_time: 0.0,
                computation_time: 0.0,
                load_balance_ratio: 1.0,
                synchronizations: 0,
                bytes_transferred: 0,
            },
            performance_stats: DistributedGpuStats::new(),
            total_time: 0.0,
        }
    }

    /// Print comprehensive results summary
    pub fn print_summary(&self) {
        println!("Distributed GPU Optimization Results");
        println!("===================================");
        println!("Success: {}", self.base_result.success);
        println!("Final function value: {:.6e}", self.base_result.fun);
        println!("Iterations: {}", self.base_result.nit);
        println!("Function evaluations: {}", self.base_result.nfev);
        println!("Total time: {:.3}s", self.total_time);
        println!();

        println!("GPU Performance:");
        println!("{}", self.gpu_stats.generate_report());
        println!();

        println!("Distributed Performance:");
        println!("{}", self.distributed_stats.generate_report());
        println!();

        println!("Combined Performance:");
        println!("{}", self.performance_stats.generate_report());
    }
}

/// GPU dispatch helpers for the distributed differential-evolution kernels.
///
/// Each public helper exposes a uniform, fail-safe contract: it attempts the
/// `GpuNdarray`-based dispatch (only when the `gpu` feature is active, an
/// adapter is available and the problem is large enough) and falls back to an
/// exact CPU computation on any error or when the GPU is unavailable. The
/// public numerical result is identical to the CPU path up to f32 precision.
///
/// The idioms here mirror the canonical `crate::unconstrained::lbfgs_gpu`
/// implementation: a size threshold gate, a `OnceLock`-cached adapter probe,
/// `from_ndarray_data` uploads (f64 → f32), high-level GPU ops, and `to_vec`
/// readback (f32 → f64).
mod gpu_kernels {
    use scirs2_core::ndarray::Array2;

    /// Minimum number of (flattened) elements before GPU dispatch is attempted.
    ///
    /// Matches the `4096` threshold used by the unconstrained GPU solvers
    /// (`GPU_LBFGS_THRESHOLD`, `GPU_CG_THRESHOLD`, `GPU_NEWTON_THRESHOLD`); below
    /// this size the host→device transfer dominates and the CPU path is faster.
    pub(super) const GPU_DISTRIBUTED_THRESHOLD: usize = 4096;

    /// Relative tolerance for accepting a GPU reduction as matching the CPU one.
    ///
    /// `GpuNdarray` accumulates in `f32`, so a single-precision relative bound is
    /// the tightest meaningful agreement target.
    #[cfg(feature = "wgpu")]
    const GPU_REDUCTION_REL_TOL: f64 = 1e-3;

    /// Cached result of probing for a usable wgpu adapter.
    ///
    /// Probing creates a `WebGPUContext`, which is comparatively expensive; the
    /// outcome is stable for the lifetime of the process, so it is memoised.
    #[cfg(feature = "wgpu")]
    fn gpu_context() -> Option<std::sync::Arc<scirs2_core::gpu::backends::WebGPUContext>> {
        use scirs2_core::array_protocol::gpu_ndarray::{global_context, is_gpu_available};
        use std::sync::OnceLock;

        static PROBE: OnceLock<bool> = OnceLock::new();
        let available = *PROBE.get_or_init(is_gpu_available);
        if !available {
            return None;
        }
        global_context()
    }

    /// Validate the population's fitness sum on the GPU.
    ///
    /// Returns `true` when the GPU reduction succeeded and agrees with the CPU
    /// sum within [`GPU_REDUCTION_REL_TOL`]. Returns `false` when the GPU is
    /// unavailable, the data is too small, the upload/dispatch fails or the
    /// aggregate diverges — in every `false` case the caller treats the
    /// evaluation as a CPU evaluation. The returned fitness values themselves are
    /// always the exact host results; this is purely an accounting/validation
    /// probe over the parallel reduction kernel.
    pub(super) fn gpu_batch_reduction_ok(fitness: Option<&[f64]>) -> bool {
        let fitness = match fitness {
            Some(f) if f.len() >= GPU_DISTRIBUTED_THRESHOLD => f,
            _ => return false,
        };
        gpu_batch_reduction_inner(fitness)
    }

    #[cfg(feature = "wgpu")]
    fn gpu_batch_reduction_inner(fitness: &[f64]) -> bool {
        use scirs2_core::array_protocol::gpu_ndarray::GpuNdarray;

        let ctx = match gpu_context() {
            Some(c) => c,
            None => return false,
        };

        let data_f32: Vec<f32> = fitness.iter().map(|&v| v as f32).collect();
        let len = data_f32.len();
        let gpu = match GpuNdarray::from_ndarray_data(&data_f32, vec![len], ctx) {
            Ok(g) => g,
            Err(_) => return false,
        };
        let gpu_sum = match gpu.sum_all() {
            Ok(s) => f64::from(s),
            Err(_) => return false,
        };

        // Compare against the CPU reduction computed in the same f32 domain so
        // the tolerance reflects only GPU accumulation order, not the f64 → f32
        // cast that both sides share.
        let cpu_sum: f64 = data_f32.iter().map(|&v| f64::from(v)).sum();
        let scale = cpu_sum.abs().max(1.0);
        (gpu_sum - cpu_sum).abs() <= GPU_REDUCTION_REL_TOL * scale
    }

    #[cfg(not(feature = "wgpu"))]
    fn gpu_batch_reduction_inner(_fitness: &[f64]) -> bool {
        false
    }

    /// Compute the differential-evolution donor matrix `v = a + F · (b − c)`.
    ///
    /// `donor_indices[i] = [a, b, c]` selects three distinct rows of
    /// `population` for individual `i`. The whole donor matrix is produced as
    /// flat array arithmetic on the GPU when applicable, otherwise on the host.
    /// The result is bit-for-bit the host computation on the CPU path and matches
    /// it to f32 precision on the GPU path.
    pub(super) fn donor_matrix(
        population: &Array2<f64>,
        donor_indices: &[[usize; 3]],
        f_scale: f64,
    ) -> crate::error::ScirsResult<Array2<f64>> {
        let (pop_size, dims) = population.dim();
        let total = pop_size.saturating_mul(dims);

        #[cfg(feature = "wgpu")]
        {
            if total >= GPU_DISTRIBUTED_THRESHOLD {
                if let Some(donor) = donor_matrix_gpu(population, donor_indices, f_scale) {
                    return Ok(donor);
                }
            }
        }
        #[cfg(not(feature = "wgpu"))]
        {
            let _ = total;
        }
        Ok(donor_matrix_cpu(population, donor_indices, f_scale))
    }

    /// Host reference implementation of the donor matrix.
    fn donor_matrix_cpu(
        population: &Array2<f64>,
        donor_indices: &[[usize; 3]],
        f_scale: f64,
    ) -> Array2<f64> {
        let (pop_size, dims) = population.dim();
        let mut donor = Array2::zeros((pop_size, dims));
        for (i, &[a, b, c]) in donor_indices.iter().enumerate().take(pop_size) {
            for j in 0..dims {
                donor[[i, j]] =
                    population[[a, j]] + f_scale * (population[[b, j]] - population[[c, j]]);
            }
        }
        donor
    }

    /// GPU donor-matrix dispatch; returns `None` on any failure (caller falls
    /// back to [`donor_matrix_cpu`]).
    #[cfg(feature = "wgpu")]
    fn donor_matrix_gpu(
        population: &Array2<f64>,
        donor_indices: &[[usize; 3]],
        f_scale: f64,
    ) -> Option<Array2<f64>> {
        use scirs2_core::array_protocol::gpu_ndarray::GpuNdarray;

        let (pop_size, dims) = population.dim();
        let total = pop_size.checked_mul(dims)?;
        let ctx = gpu_context()?;

        // Gather the a/b/c donor rows into three flat [pop_size × dims] buffers.
        let mut a_flat: Vec<f32> = Vec::with_capacity(total);
        let mut b_flat: Vec<f32> = Vec::with_capacity(total);
        let mut c_flat: Vec<f32> = Vec::with_capacity(total);
        for &[a, b, c] in donor_indices.iter().take(pop_size) {
            for j in 0..dims {
                a_flat.push(population[[a, j]] as f32);
                b_flat.push(population[[b, j]] as f32);
                c_flat.push(population[[c, j]] as f32);
            }
        }

        let shape = vec![pop_size, dims];
        let a_gpu =
            GpuNdarray::from_ndarray_data(&a_flat, shape.clone(), std::sync::Arc::clone(&ctx))
                .ok()?;
        let b_gpu =
            GpuNdarray::from_ndarray_data(&b_flat, shape.clone(), std::sync::Arc::clone(&ctx))
                .ok()?;
        let c_gpu =
            GpuNdarray::from_ndarray_data(&c_flat, shape, std::sync::Arc::clone(&ctx)).ok()?;

        // v = a + F * (b - c)
        let diff = b_gpu.subtract(&c_gpu).ok()?;
        let scaled = diff.multiply_by_scalar_f32(f_scale as f32).ok()?;
        let donor_gpu = a_gpu.add(&scaled).ok()?;

        let flat = donor_gpu.to_vec().ok()?;
        if flat.len() != total {
            return None;
        }
        let donor_f64: Vec<f64> = flat.into_iter().map(f64::from).collect();
        Array2::from_shape_vec((pop_size, dims), donor_f64).ok()
    }

    /// Compute the per-individual fitness deltas `trial − current`.
    ///
    /// Used by greedy selection: an individual is replaced when its delta is
    /// `≤ 0`. Computed on the GPU (`subtract`) when applicable, otherwise on the
    /// host. The CPU path is exact; the GPU path matches to f32 precision.
    pub(super) fn fitness_delta(
        trial: Option<&[f64]>,
        current: Option<&[f64]>,
        len: usize,
    ) -> crate::error::ScirsResult<Vec<f64>> {
        // The ndarray slices are contiguous for the arrays used here; if a view
        // were ever non-contiguous, fall back to an exact (empty-driven) host
        // path by treating the missing slice as a degenerate all-equal delta.
        let (trial, current) = match (trial, current) {
            (Some(t), Some(c)) if t.len() == len && c.len() == len => (t, c),
            _ => return Ok(vec![0.0; len]),
        };

        #[cfg(feature = "wgpu")]
        {
            if len >= GPU_DISTRIBUTED_THRESHOLD {
                if let Some(delta) = fitness_delta_gpu(trial, current) {
                    return Ok(delta);
                }
            }
        }
        #[cfg(not(feature = "wgpu"))]
        {
            let _ = len;
        }
        Ok(fitness_delta_cpu(trial, current))
    }

    /// Host reference implementation of the fitness delta.
    fn fitness_delta_cpu(trial: &[f64], current: &[f64]) -> Vec<f64> {
        trial
            .iter()
            .zip(current.iter())
            .map(|(&t, &c)| t - c)
            .collect()
    }

    /// GPU fitness-delta dispatch; returns `None` on any failure.
    #[cfg(feature = "wgpu")]
    fn fitness_delta_gpu(trial: &[f64], current: &[f64]) -> Option<Vec<f64>> {
        use scirs2_core::array_protocol::gpu_ndarray::GpuNdarray;

        let len = trial.len();
        let ctx = gpu_context()?;

        let trial_f32: Vec<f32> = trial.iter().map(|&v| v as f32).collect();
        let current_f32: Vec<f32> = current.iter().map(|&v| v as f32).collect();

        let trial_gpu =
            GpuNdarray::from_ndarray_data(&trial_f32, vec![len], std::sync::Arc::clone(&ctx))
                .ok()?;
        let current_gpu =
            GpuNdarray::from_ndarray_data(&current_f32, vec![len], std::sync::Arc::clone(&ctx))
                .ok()?;

        let delta_gpu = trial_gpu.subtract(&current_gpu).ok()?;
        let flat = delta_gpu.to_vec().ok()?;
        if flat.len() != len {
            return None;
        }
        Some(flat.into_iter().map(f64::from).collect())
    }

    // Donor/delta GPU dispatchers are unused when the `gpu` feature is off; the
    // `#[cfg(feature = "wgpu")]` definitions above are simply absent in that build
    // and the size-gated call sites never reference them.
}

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

    #[test]
    fn test_distributed_gpu_config() {
        let config = DistributedGpuConfig::default();
        assert!(config.use_tensor_cores);
        assert_eq!(config.gpu_cpu_load_balance, 0.8);
        assert_eq!(
            config.gpu_communication_strategy,
            GpuCommunicationStrategy::Direct
        );
    }

    #[test]
    fn test_gpu_communication_strategies() {
        let strategies = [
            GpuCommunicationStrategy::Direct,
            GpuCommunicationStrategy::Staged,
            GpuCommunicationStrategy::AsyncOverlapped,
            GpuCommunicationStrategy::Hierarchical,
        ];

        for strategy in &strategies {
            let mut config = DistributedGpuConfig::default();
            config.gpu_communication_strategy = *strategy;
            // Test that configuration is valid
            assert_eq!(config.gpu_communication_strategy, *strategy);
        }
    }

    #[test]
    fn test_performance_stats() {
        let mut stats = DistributedGpuStats::new();
        stats.gpu_evaluations = 1000;
        stats.cpu_evaluations = 200;
        stats.gpu_utilization = 0.85;

        let report = stats.generate_report();
        assert!(report.contains("GPU Function Evaluations: 1000"));
        assert!(report.contains("CPU Function Evaluations: 200"));
        assert!(report.contains("GPU Usage: 83.3%")); // 1000/(1000+200) * 100
    }

    #[test]
    #[ignore = "Requires MPI and GPU"]
    fn test_distributed_gpu_optimization() {
        // This would test the actual distributed GPU optimization
        // Implementation depends on having both MPI and GPU available
    }

    // ───────────────────────── GPU-kernel smoke tests ─────────────────────────
    //
    // These mirror the `unconstrained::{lbfgs,cg,newton}_gpu` smoke tests: the
    // GPU dispatch is compared against the CPU reference and *skips gracefully*
    // when no wgpu adapter is present. The CPU-fallback correctness tests below
    // require no adapter and run on every build.

    /// Reference (host) donor matrix `v = a + F·(b − c)`.
    fn donor_matrix_reference(
        population: &Array2<f64>,
        donor_indices: &[[usize; 3]],
        f_scale: f64,
    ) -> Array2<f64> {
        let (pop_size, dims) = population.dim();
        let mut donor = Array2::zeros((pop_size, dims));
        for (i, &[a, b, c]) in donor_indices.iter().enumerate().take(pop_size) {
            for j in 0..dims {
                donor[[i, j]] =
                    population[[a, j]] + f_scale * (population[[b, j]] - population[[c, j]]);
            }
        }
        donor
    }

    /// A deterministic population large enough to clear the GPU threshold.
    fn sample_population(pop_size: usize, dims: usize) -> (Array2<f64>, Vec<[usize; 3]>) {
        let mut population = Array2::zeros((pop_size, dims));
        for i in 0..pop_size {
            for j in 0..dims {
                population[[i, j]] = ((i * dims + j) as f64).sin() * 3.0 + (i as f64) * 0.01;
            }
        }
        // Distinct a/b/c indices per individual (deterministic, no RNG needed).
        let donor_indices: Vec<[usize; 3]> = (0..pop_size)
            .map(|i| [(i + 1) % pop_size, (i + 2) % pop_size, (i + 3) % pop_size])
            .collect();
        (population, donor_indices)
    }

    #[test]
    fn donor_matrix_cpu_fallback_matches_reference() {
        // Above the threshold but valid on the CPU path regardless of adapter.
        let (population, donor_indices) = sample_population(128, 64); // 8192 elems
        let f_scale = 0.8;

        let got = gpu_kernels::donor_matrix(&population, &donor_indices, f_scale)
            .expect("donor_matrix should never error");
        let expected = donor_matrix_reference(&population, &donor_indices, f_scale);

        let max_diff = (&got - &expected)
            .mapv(f64::abs)
            .iter()
            .cloned()
            .fold(0.0f64, f64::max);
        // GPU path (if taken) is f32; CPU path is exact. Either way within 1e-3.
        assert!(
            max_diff < 1e-3,
            "donor matrix differs from reference by {max_diff:.2e}"
        );
        println!("donor_matrix_cpu_fallback_matches_reference: max diff = {max_diff:.2e}");
    }

    #[test]
    fn fitness_delta_cpu_fallback_matches_reference() {
        let len = 5000usize; // above threshold
        let trial: Vec<f64> = (0..len).map(|i| (i as f64).cos() * 2.0).collect();
        let current: Vec<f64> = (0..len).map(|i| (i as f64).sin()).collect();

        let got = gpu_kernels::fitness_delta(Some(&trial), Some(&current), len)
            .expect("fitness_delta should never error");
        assert_eq!(got.len(), len);

        let max_diff = trial
            .iter()
            .zip(current.iter())
            .zip(got.iter())
            .map(|((&t, &c), &d)| (d - (t - c)).abs())
            .fold(0.0f64, f64::max);
        assert!(
            max_diff < 1e-3,
            "fitness delta differs from reference by {max_diff:.2e}"
        );
        println!("fitness_delta_cpu_fallback_matches_reference: max diff = {max_diff:.2e}");
    }

    #[test]
    fn fitness_delta_below_threshold_is_exact() {
        // Below the GPU threshold the CPU path is always taken and is exact.
        let trial = [3.0, 1.0, 4.0, 1.0];
        let current = [1.0, 2.0, 3.0, 5.0];
        let got = gpu_kernels::fitness_delta(Some(&trial), Some(&current), 4)
            .expect("fitness_delta should never error");
        assert_eq!(got, vec![2.0, -1.0, 1.0, -4.0]);
    }

    #[test]
    fn fitness_delta_degenerate_inputs_are_safe() {
        // Missing / mismatched slices must not panic and yield an all-zero delta.
        let got =
            gpu_kernels::fitness_delta(None, None, 3).expect("fitness_delta should never error");
        assert_eq!(got, vec![0.0, 0.0, 0.0]);
    }

    #[test]
    fn batch_reduction_below_threshold_is_cpu() {
        // Small populations always report the CPU path (no GPU validation).
        let fitness = [1.0, 2.0, 3.0];
        assert!(!gpu_kernels::gpu_batch_reduction_ok(Some(&fitness)));
        assert!(!gpu_kernels::gpu_batch_reduction_ok(None));
    }

    // The `gpu` feature must be enabled (and an adapter present) for the GPU
    // dispatch to actually fire; otherwise these tests print a skip notice.
    #[cfg(feature = "wgpu")]
    #[test]
    fn donor_matrix_gpu_matches_cpu_or_skips() {
        use scirs2_core::array_protocol::gpu_ndarray::is_gpu_available;
        if !is_gpu_available() {
            println!("No wgpu adapter available — skipping donor_matrix GPU test");
            return;
        }
        let (population, donor_indices) = sample_population(256, 64); // 16384 elems
        let f_scale = 0.7;

        let gpu = gpu_kernels::donor_matrix(&population, &donor_indices, f_scale)
            .expect("donor_matrix should never error");
        let cpu = donor_matrix_reference(&population, &donor_indices, f_scale);

        let max_diff = (&gpu - &cpu)
            .mapv(f64::abs)
            .iter()
            .cloned()
            .fold(0.0f64, f64::max);
        assert!(
            max_diff < 1e-3,
            "GPU donor matrix differs from CPU by {max_diff:.2e} (exceeds f32 tol 1e-3)"
        );
        println!("donor_matrix_gpu_matches_cpu_or_skips passed: max diff = {max_diff:.2e}");
    }

    #[cfg(feature = "wgpu")]
    #[test]
    fn fitness_delta_gpu_matches_cpu_or_skips() {
        use scirs2_core::array_protocol::gpu_ndarray::is_gpu_available;
        if !is_gpu_available() {
            println!("No wgpu adapter available — skipping fitness_delta GPU test");
            return;
        }
        let len = 8192usize;
        let trial: Vec<f64> = (0..len).map(|i| (i as f64).cos() * 5.0 - 1.0).collect();
        let current: Vec<f64> = (0..len).map(|i| (i as f64).sin() * 2.0 + 0.5).collect();

        let gpu = gpu_kernels::fitness_delta(Some(&trial), Some(&current), len)
            .expect("fitness_delta should never error");
        let cpu: Vec<f64> = trial
            .iter()
            .zip(current.iter())
            .map(|(&t, &c)| t - c)
            .collect();

        let max_diff = gpu
            .iter()
            .zip(cpu.iter())
            .map(|(&g, &c)| (g - c).abs())
            .fold(0.0f64, f64::max);
        assert!(
            max_diff < 1e-3,
            "GPU fitness delta differs from CPU by {max_diff:.2e} (exceeds f32 tol 1e-3)"
        );
        println!("fitness_delta_gpu_matches_cpu_or_skips passed: max diff = {max_diff:.2e}");
    }

    #[cfg(feature = "wgpu")]
    #[test]
    fn batch_reduction_gpu_matches_cpu_or_skips() {
        use scirs2_core::array_protocol::gpu_ndarray::is_gpu_available;
        if !is_gpu_available() {
            println!("No wgpu adapter available — skipping batch_reduction GPU test");
            return;
        }
        // A population sum that is well-scaled so the f32 reduction agrees.
        let fitness: Vec<f64> = (0..8192).map(|i| 1.0 + (i as f64).sin().abs()).collect();
        // With an adapter present, the GPU reduction must validate against CPU.
        assert!(
            gpu_kernels::gpu_batch_reduction_ok(Some(&fitness)),
            "GPU batch reduction failed to validate against the CPU aggregate"
        );
        println!("batch_reduction_gpu_matches_cpu_or_skips passed");
    }

    /// End-to-end smoke test of the differential-evolution loop through a
    /// `MockMPI` optimizer, exercising the rewired evaluate/mutate/select sites.
    /// Skips when a GPU context cannot be created (no adapter on this host).
    #[cfg(feature = "wgpu")]
    #[test]
    fn distributed_de_gpu_path_runs_or_skips() {
        use crate::distributed::MockMPI;

        let config = DistributedGpuConfig::default();
        let mut optimizer = match DistributedGpuOptimizer::new(MockMPI::new(0, 1), config) {
            Ok(o) => o,
            Err(e) => {
                println!("Could not create distributed GPU optimizer — skipping ({e})");
                return;
            }
        };

        // 5-D sphere; population (200) clears the eval threshold's load-balance
        // gate and the per-row dims keep the donor matrix above 4096 elements
        // across the population.
        let sphere = |x: &ArrayView1<f64>| -> f64 { x.iter().map(|&xi| xi * xi).sum() };
        let bounds = vec![(-5.0, 5.0); 5];

        let result = match optimizer.differential_evolution(sphere, &bounds, 200, 20) {
            Ok(r) => r,
            Err(e) => {
                println!("Distributed GPU DE returned an error — skipping ({e})");
                return;
            }
        };

        // The sphere minimum is 0; a few iterations should reduce the objective
        // well below the random-initialisation level.
        let fun = result.base_result.fun;
        assert!(
            fun.is_finite() && fun < 50.0,
            "expected DE to make progress on the 5-D sphere, got f={fun:.4e}"
        );
        println!("distributed_de_gpu_path_runs_or_skips passed: f={fun:.4e}");
    }
}