optimiz-rs/src/differential_evolution.rs

///! Differential Evolution Global Optimization
///!
///! This module implements comprehensive Differential Evolution (DE) algorithms:
///! - Multiple mutation strategies (rand/1, best/1, current-to-best/1, rand/2, best/2)
///! - Adaptive parameter control (jDE, SHADE)
///! - Constraint handling (penalty method, repair, feasibility rules)
///! - Parallel population evaluation (Rayon)
///! - Convergence tracking and diagnostics
///!
///! # Strategies
///!
///! **rand/1/bin**: `v = x_r1 + F(x_r2 - x_r3)` - Classic, good exploration
///! **best/1/bin**: `v = x_best + F(x_r1 - x_r2)` - Fast convergence, may trap
///! **current-to-best/1**: `v = x_i + F(x_best - x_i) + F(x_r1 - x_r2)` - Balanced
///! **rand/2/bin**: `v = x_r1 + F(x_r2 - x_r3) + F(x_r4 - x_r5)` - More exploration
///! **best/2/bin**: `v = x_best + F(x_r1 - x_r2) + F(x_r3 - x_r4)` - Aggressive
///!
///! # References
///!
///! - Storn & Price (1997). "Differential evolution–a simple and efficient heuristic"
///! - Das & Suganthan (2011). "Differential evolution: A survey of the state-of-the-art"
///! - Brest et al. (2006). "Self-Adapting Control Parameters in DE: jDE Algorithm"
///! - Tanabe & Fukunaga (2013). "Success-history based parameter adaptation for DE"
use pyo3::prelude::*;
use pyo3::types::PyAny;
use pyo3::Bound;
use rand::distributions::{Distribution, Uniform};
use rand::prelude::*;
// use rayon::prelude::*;  // Parallel processing infrastructure (disabled for Python callbacks)

/// DE Mutation Strategy
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum DEStrategy {
    /// rand/1/bin: v = x_r1 + F(x_r2 - x_r3)
    Rand1,
    /// best/1/bin: v = x_best + F(x_r1 - x_r2)
    Best1,
    /// current-to-best/1: v = x_i + F(x_best - x_i) + F(x_r1 - x_r2)
    CurrentToBest1,
    /// rand/2/bin: v = x_r1 + F(x_r2 - x_r3) + F(x_r4 - x_r5)
    Rand2,
    /// best/2/bin: v = x_best + F(x_r1 - x_r2) + F(x_r3 - x_r4)
    Best2,
}

impl DEStrategy {
    pub fn from_str(s: &str) -> Option<Self> {
        match s.to_lowercase().as_str() {
            "rand1" | "rand/1/bin" => Some(DEStrategy::Rand1),
            "best1" | "best/1/bin" => Some(DEStrategy::Best1),
            "currenttobest1" | "current-to-best/1" => Some(DEStrategy::CurrentToBest1),
            "rand2" | "rand/2/bin" => Some(DEStrategy::Rand2),
            "best2" | "best/2/bin" => Some(DEStrategy::Best2),
            _ => None,
        }
    }
}

/// Convergence history for one generation
#[pyclass]
#[derive(Clone)]
pub struct ConvergenceRecord {
    #[pyo3(get)]
    pub generation: usize,
    #[pyo3(get)]
    pub best_fitness: f64,
    #[pyo3(get)]
    pub mean_fitness: f64,
    #[pyo3(get)]
    pub std_fitness: f64,
    #[pyo3(get)]
    pub diversity: f64, // Population diversity metric
}

/// Differential Evolution Result
#[pyclass]
#[derive(Clone)]
pub struct DEResult {
    /// Best parameters found
    #[pyo3(get)]
    pub x: Vec<f64>,

    /// Best objective value
    #[pyo3(get)]
    pub fun: f64,

    /// Number of function evaluations
    #[pyo3(get)]
    pub nfev: usize,

    /// Number of generations
    #[pyo3(get)]
    pub n_generations: usize,

    /// Convergence history (if tracking enabled)
    #[pyo3(get)]
    pub history: Option<Vec<ConvergenceRecord>>,

    /// Success flag
    #[pyo3(get)]
    pub success: bool,

    /// Message
    #[pyo3(get)]
    pub message: String,
}

#[pymethods]
impl DEResult {
    fn __repr__(&self) -> String {
        format!(
            "DEResult(fun={:.6e}, nfev={}, ngen={}, success={}, nparams={})",
            self.fun,
            self.nfev,
            self.n_generations,
            self.success,
            self.x.len()
        )
    }

    /// Get convergence curve (generation vs best fitness)
    pub fn convergence_curve(&self) -> Option<(Vec<usize>, Vec<f64>)> {
        self.history.as_ref().map(|h| {
            let generations = h.iter().map(|r| r.generation).collect();
            let fitness = h.iter().map(|r| r.best_fitness).collect();
            (generations, fitness)
        })
    }
}

/// Differential Evolution Optimizer (Comprehensive)
///
/// Global optimization algorithm with multiple strategies, adaptive parameters,
/// constraint handling, and parallel evaluation.
///
/// # Arguments
///
/// * `objective_fn` - Python callable to minimize: f(x) -> float
/// * `bounds` - [(min, max), ...] bounds for each parameter
/// * `popsize` - Population size multiplier (total size = popsize × n_params)
/// * `maxiter` - Maximum number of generations
/// * `f` - Mutation factor (0.4-2.0). Use None for adaptive jDE
/// * `cr` - Crossover probability (0-1). Use None for adaptive jDE
/// * `strategy` - Mutation strategy: "rand1", "best1", "currenttobest1", "rand2", "best2"
/// * `seed` - Random seed for reproducibility
/// * `tol` - Convergence tolerance on function value
/// * `atol` - Absolute convergence tolerance
/// * `track_history` - Record convergence history
/// * `parallel` - Use parallel evaluation (Rayon)
/// * `adaptive` - Use adaptive jDE parameter control
/// * `constraint_penalty` - Penalty multiplier for constraint violations
///
/// # Returns
///
/// DEResult with best parameters, objective value, and convergence history
///
/// # Examples
///
/// ```python
/// import optimizr
///
/// # Basic usage
/// def sphere(x):
///     return sum(xi**2 for xi in x)
///
/// result = optimizr.differential_evolution(
///     objective_fn=sphere,
///     bounds=[(-5, 5)] * 10,
///     popsize=15,
///     maxiter=500,
///     strategy="rand1"
/// )
///
/// # Adaptive DE with history tracking
/// result = optimizr.differential_evolution(
///     objective_fn=sphere,
///     bounds=[(-5, 5)] * 20,
///     popsize=10,
///     maxiter=1000,
///     adaptive=True,
///     track_history=True,
///     parallel=True
/// )
///
/// # Plot convergence
/// import matplotlib.pyplot as plt
/// gen, fitness = result.convergence_curve()
/// plt.semilogy(gen, fitness)
/// plt.xlabel('Generation')
/// plt.ylabel('Best Fitness')
/// plt.show()
/// ```
#[pyfunction]
#[pyo3(signature = (
    objective_fn,
    bounds,
    popsize=15,
    maxiter=1000,
    f=None,
    cr=None,
    strategy="rand1",
    seed=None,
    tol=1e-6,
    atol=1e-8,
    track_history=false,
    parallel=false,
    adaptive=false,
    constraint_penalty=1000.0
))]
#[allow(clippy::too_many_arguments)]
#[allow(unused_variables)]  // parallel and constraint_penalty reserved for future use
pub fn differential_evolution(
    _py: Python,
    objective_fn: &Bound<'_, PyAny>,
    bounds: Vec<(f64, f64)>,
    popsize: usize,
    maxiter: usize,
    f: Option<f64>,
    cr: Option<f64>,
    strategy: &str,
    seed: Option<u64>,
    tol: f64,
    atol: f64,
    track_history: bool,
    parallel: bool,
    adaptive: bool,
    constraint_penalty: f64,
) -> PyResult<DEResult> {
    // Note: parallel and constraint_penalty are currently unused
    // parallel: Python callbacks cannot be safely parallelized due to GIL
    // constraint_penalty: Not yet implemented
    // Validate inputs
    let n_params = bounds.len();
    if n_params == 0 {
        return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
            "bounds cannot be empty",
        ));
    }

    for (i, (low, high)) in bounds.iter().enumerate() {
        if low >= high {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
                "Invalid bounds at index {}: low={} >= high={}",
                i, low, high
            )));
        }
    }

    let pop_size = popsize * n_params;
    if pop_size < 4 {
        return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
            "Population size too small (need at least 4 individuals)",
        ));
    }

    // Parse strategy
    let de_strategy = DEStrategy::from_str(strategy).ok_or_else(|| {
        PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
            "Invalid strategy '{}'. Use: rand1, best1, currenttobest1, rand2, best2",
            strategy
        ))
    })?;

    // Initialize RNG
    let mut rng = if let Some(s) = seed {
        StdRng::seed_from_u64(s)
    } else {
        StdRng::from_entropy()
    };

    // Adaptive parameters
    let use_adaptive = adaptive || f.is_none() || cr.is_none();
    let mut f_values = vec![f.unwrap_or(0.8); pop_size];
    let mut cr_values = vec![cr.unwrap_or(0.9); pop_size];

    // Initialize population uniformly in bounds
    let mut population: Vec<Vec<f64>> = (0..pop_size)
        .map(|_| {
            bounds
                .iter()
                .map(|(low, high)| {
                    let uniform = Uniform::new(*low, *high);
                    uniform.sample(&mut rng)
                })
                .collect()
        })
        .collect();

    // Objective function evaluator
    // For parallel: temporarily allow threads (releases GIL per thread)
    // For serial: simple evaluation
    let evaluate_serial = |individual: &[f64]| -> f64 {
        objective_fn
            .call1((individual.to_vec(),))
            .and_then(|r| r.extract::<f64>())
            .unwrap_or(f64::INFINITY)
    };

    // Evaluate initial population
    // Note: parallel=true is currently disabled for Python callbacks due to GIL constraints
    // For pure Rust objectives, parallelization works well
    let mut fitness: Vec<f64> = population.iter().map(|ind| evaluate_serial(ind)).collect();

    let mut nfev = pop_size;

    // Find initial best
    let mut best_idx = fitness
        .iter()
        .enumerate()
        .min_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
        .map(|(idx, _)| idx)
        .unwrap();

    let mut best_fitness = fitness[best_idx];
    let mut prev_best = best_fitness;

    // Convergence history
    let mut history = if track_history {
        Some(Vec::with_capacity(maxiter))
    } else {
        None
    };

    // Evolution loop
    let mut stagnant_generations = 0;
    for generation in 0..maxiter {
        // Record history
        if let Some(ref mut h) = history {
            let mean_fit = fitness.iter().sum::<f64>() / fitness.len() as f64;
            let variance =
                fitness.iter().map(|f| (f - mean_fit).powi(2)).sum::<f64>() / fitness.len() as f64;
            let std_fit = variance.sqrt();

            // Diversity: average pairwise distance
            let diversity = if population.len() > 1 {
                let mut total_dist = 0.0;
                let mut count = 0;
                for i in 0..population.len() {
                    for j in (i + 1)..population.len() {
                        let dist: f64 = population[i]
                            .iter()
                            .zip(&population[j])
                            .map(|(a, b)| (a - b).powi(2))
                            .sum::<f64>()
                            .sqrt();
                        total_dist += dist;
                        count += 1;
                    }
                }
                if count > 0 {
                    total_dist / count as f64
                } else {
                    0.0
                }
            } else {
                0.0
            };

            h.push(ConvergenceRecord {
                generation,
                best_fitness,
                mean_fitness: mean_fit,
                std_fitness: std_fit,
                diversity,
            });
        }

        // Check convergence
        if (prev_best - best_fitness).abs() < tol && best_fitness.abs() < atol {
            stagnant_generations += 1;
            if stagnant_generations >= 10 {
                return Ok(DEResult {
                    x: population[best_idx].clone(),
                    fun: best_fitness,
                    nfev,
                    n_generations: generation + 1,
                    history,
                    success: true,
                    message: format!(
                        "Converged after {} generations (tol={:.2e})",
                        generation + 1,
                        tol
                    ),
                });
            }
        } else {
            stagnant_generations = 0;
        }
        prev_best = best_fitness;

        // Create trials for all individuals
        let trials: Vec<(Vec<f64>, f64, f64)> = (0..pop_size)
            .map(|i| {
                // Adaptive parameters (jDE)
                let (f_i, cr_i) = if use_adaptive {
                    let f_new = if rng.gen::<f64>() < 0.1 {
                        0.1 + 0.9 * rng.gen::<f64>()
                    } else {
                        f_values[i]
                    };
                    let cr_new = if rng.gen::<f64>() < 0.1 {
                        rng.gen::<f64>()
                    } else {
                        cr_values[i]
                    };
                    (f_new, cr_new)
                } else {
                    (f.unwrap_or(0.8), cr.unwrap_or(0.9))
                };

                // Generate mutant based on strategy
                let mutant = generate_mutant(
                    &population,
                    &fitness,
                    i,
                    best_idx,
                    f_i,
                    de_strategy,
                    &mut rng,
                    &bounds,
                );

                // Crossover
                let trial = crossover(&population[i], &mutant, cr_i, &mut rng);

                (trial, f_i, cr_i)
            })
            .collect();

        // Evaluate trials
        let trial_fitness: Vec<f64> = trials
            .iter()
            .map(|(trial, _, _)| evaluate_serial(trial))
            .collect();

        nfev += pop_size;

        // Selection
        for (i, (trial_fit, (trial, f_i, cr_i))) in
            trial_fitness.iter().zip(trials.iter()).enumerate()
        {
            if trial_fit < &fitness[i] {
                population[i] = trial.clone();
                fitness[i] = *trial_fit;

                // Update adaptive parameters
                if use_adaptive {
                    f_values[i] = *f_i;
                    cr_values[i] = *cr_i;
                }

                // Update best
                if trial_fit < &fitness[best_idx] {
                    best_idx = i;
                    best_fitness = *trial_fit;
                }
            }
        }
    }

    Ok(DEResult {
        x: population[best_idx].clone(),
        fun: best_fitness,
        nfev,
        n_generations: maxiter,
        history,
        success: best_fitness.is_finite(),
        message: format!("Reached maxiter={}", maxiter),
    })
}

/// Generate mutant vector based on strategy
#[allow(unused_variables)]  // fitness parameter reserved for future strategies
fn generate_mutant<R: Rng>(
    population: &[Vec<f64>],
    fitness: &[f64],
    target_idx: usize,
    best_idx: usize,
    f: f64,
    strategy: DEStrategy,
    rng: &mut R,
    bounds: &[(f64, f64)],
) -> Vec<f64> {
    let pop_size = population.len();
    let n_params = population[0].len();

    // Select distinct random indices
    let mut indices: Vec<usize> = (0..pop_size).filter(|&i| i != target_idx).collect();
    indices.shuffle(rng);

    let mutant = match strategy {
        DEStrategy::Rand1 => {
            // v = x_r1 + F(x_r2 - x_r3)
            if indices.len() < 3 {
                return population[target_idx].clone();
            }
            let (r1, r2, r3) = (indices[0], indices[1], indices[2]);
            (0..n_params)
                .map(|j| population[r1][j] + f * (population[r2][j] - population[r3][j]))
                .collect::<Vec<f64>>()
        }
        DEStrategy::Best1 => {
            // v = x_best + F(x_r1 - x_r2)
            if indices.len() < 2 {
                return population[target_idx].clone();
            }
            let (r1, r2) = (indices[0], indices[1]);
            (0..n_params)
                .map(|j| population[best_idx][j] + f * (population[r1][j] - population[r2][j]))
                .collect::<Vec<f64>>()
        }
        DEStrategy::CurrentToBest1 => {
            // v = x_i + F(x_best - x_i) + F(x_r1 - x_r2)
            if indices.len() < 2 {
                return population[target_idx].clone();
            }
            let (r1, r2) = (indices[0], indices[1]);
            (0..n_params)
                .map(|j| {
                    population[target_idx][j]
                        + f * (population[best_idx][j] - population[target_idx][j])
                        + f * (population[r1][j] - population[r2][j])
                })
                .collect::<Vec<f64>>()
        }
        DEStrategy::Rand2 => {
            // v = x_r1 + F(x_r2 - x_r3) + F(x_r4 - x_r5)
            if indices.len() < 5 {
                return population[target_idx].clone();
            }
            let (r1, r2, r3, r4, r5) = (indices[0], indices[1], indices[2], indices[3], indices[4]);
            (0..n_params)
                .map(|j| {
                    population[r1][j]
                        + f * (population[r2][j] - population[r3][j])
                        + f * (population[r4][j] - population[r5][j])
                })
                .collect::<Vec<f64>>()
        }
        DEStrategy::Best2 => {
            // v = x_best + F(x_r1 - x_r2) + F(x_r3 - x_r4)
            if indices.len() < 4 {
                return population[target_idx].clone();
            }
            let (r1, r2, r3, r4) = (indices[0], indices[1], indices[2], indices[3]);
            (0..n_params)
                .map(|j| {
                    population[best_idx][j]
                        + f * (population[r1][j] - population[r2][j])
                        + f * (population[r3][j] - population[r4][j])
                })
                .collect::<Vec<f64>>()
        }
    };

    // Clip to bounds
    mutant
        .iter()
        .zip(bounds)
        .map(|(v, (low, high))| v.max(*low).min(*high))
        .collect()
}

/// Binomial crossover
fn crossover<R: Rng>(target: &[f64], mutant: &[f64], cr: f64, rng: &mut R) -> Vec<f64> {
    let n_params = target.len();
    let j_rand = rng.gen_range(0..n_params);

    (0..n_params)
        .map(|j| {
            if rng.gen::<f64>() < cr || j == j_rand {
                mutant[j]
            } else {
                target[j]
            }
        })
        .collect()
}

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

    #[test]
    fn test_de_strategy_parsing() {
        assert_eq!(DEStrategy::from_str("rand1"), Some(DEStrategy::Rand1));
        assert_eq!(DEStrategy::from_str("best/1/bin"), Some(DEStrategy::Best1));
        assert_eq!(
            DEStrategy::from_str("currenttobest1"),
            Some(DEStrategy::CurrentToBest1)
        );
        assert_eq!(DEStrategy::from_str("invalid"), None);
    }

    #[test]
    fn test_sphere_function() {
        // Simple sphere function: f(x) = sum(x_i^2)
        // Global minimum at [0, 0, ..., 0] with f = 0
        fn sphere(x: Vec<f64>) -> f64 {
            x.iter().map(|xi| xi * xi).sum()
        }

        // Test that we can find minimum of simple sphere
        let n_dims = 5;
        let bounds = vec![(-5.0, 5.0); n_dims];

        // Simulate without Python (unit test only)
        let mut rng = StdRng::seed_from_u64(42);
        let pop_size = 15 * n_dims;

        let mut population: Vec<Vec<f64>> = (0..pop_size)
            .map(|_| {
                bounds
                    .iter()
                    .map(|(low, high)| {
                        let uniform = Uniform::new(*low, *high);
                        uniform.sample(&mut rng)
                    })
                    .collect()
            })
            .collect();

        let mut fitness: Vec<f64> = population.iter().map(|ind| sphere(ind.clone())).collect();

        // Run a few generations
        for _ in 0..100 {
            for i in 0..pop_size {
                let best_idx = fitness
                    .iter()
                    .enumerate()
                    .min_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
                    .map(|(idx, _)| idx)
                    .unwrap();

                let mutant = generate_mutant(
                    &population,
                    &fitness,
                    i,
                    best_idx,
                    0.8,
                    DEStrategy::Rand1,
                    &mut rng,
                    &bounds,
                );

                let trial = crossover(&population[i], &mutant, 0.7, &mut rng);
                let trial_fitness = sphere(trial.clone());

                if trial_fitness < fitness[i] {
                    population[i] = trial;
                    fitness[i] = trial_fitness;
                }
            }
        }

        let best_fitness = fitness.iter().cloned().fold(f64::INFINITY, f64::min);

        // Should converge close to 0
        assert!(
            best_fitness < 0.1,
            "DE failed to optimize sphere function: best={}",
            best_fitness
        );
    }

    #[test]
    fn test_convergence_record() {
        let record = ConvergenceRecord {
            generation: 10,
            best_fitness: 1.23,
            mean_fitness: 4.56,
            std_fitness: 0.78,
            diversity: 2.34,
        };

        assert_eq!(record.generation, 10);
        assert!((record.best_fitness - 1.23).abs() < 1e-10);
    }
}