//! Evolutionary strategy optimizer.
//!
//! Production-grade genetic algorithm engine inspired by Darwin SDK and the
//! ProFiT framework (Siper et al., "Program Search for Financial Trading").
//!
//! Features:
//! - 6 selection methods: Tournament, Uniform, Roulette Wheel, Rank, UCB1, Epsilon-Greedy
//! - NSGA-II multi-objective optimization with Pareto fronts and crowding distance
//! - Island model with Ring and FullyConnected migration topologies
//! - Convergence detection with patience and diversity collapse
//! - Adaptive mutation rates via Rechenberg's 1/5 success rule
//! - Per-generation statistics tracking
//! - Population diversity measurement
//!
//! All genetic operators are deterministic given a seed and run in Rust.
//! Evolution loop orchestration is in Python (fitness calls hz.backtest()).

use pyo3::prelude::*;

// ---------------------------------------------------------------------------
// PRNG (deterministic LCG)
// ---------------------------------------------------------------------------

struct Rng {
    state: u64,
}

impl Rng {
    fn new(seed: u64) -> Self {
        Self { state: seed.wrapping_add(1) }
    }

    fn next_u64(&mut self) -> u64 {
        self.state = self.state.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
        self.state
    }

    fn next_f64(&mut self) -> f64 {
        (self.next_u64() >> 11) as f64 / (1u64 << 53) as f64
    }

    fn next_gaussian(&mut self) -> f64 {
        let u1 = self.next_f64().max(1e-15);
        let u2 = self.next_f64();
        (-2.0 * u1.ln()).sqrt() * (2.0 * std::f64::consts::PI * u2).cos()
    }
}

fn finite_or_zero(v: f64) -> f64 {
    if v.is_finite() { v } else { 0.0 }
}

// ---------------------------------------------------------------------------
// Selection Method Enum
// ---------------------------------------------------------------------------

/// Selection method for choosing parents during evolution.
///
/// Six methods from uniform exploration to UCB1 exploitation, matching the
/// selection strategies described in the ProFiT framework (Siper et al.).
#[pyclass(eq, eq_int, frozen, hash)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum SelectionMethod {
    /// Tournament: pick `tournament_size` individuals, keep the best.
    Tournament,
    /// Uniform: random selection with equal probability.
    Uniform,
    /// Roulette Wheel: probability proportional to fitness (shifted to positive).
    RouletteWheel,
    /// Rank: probability proportional to rank position.
    Rank,
    /// UCB1: Upper Confidence Bound 1 — balances exploitation with exploration.
    /// From Auer et al., "Finite-time Analysis of the Multiarmed Bandit Problem".
    /// Used in ProFiT (Siper et al.) for parent selection.
    Ucb1,
    /// Epsilon-Greedy: with probability epsilon pick random, else pick best.
    EpsilonGreedy,
}

#[pymethods]
impl SelectionMethod {
    fn __repr__(&self) -> String {
        format!("SelectionMethod.{:?}", self)
    }

    fn __str__(&self) -> String {
        format!("{:?}", self)
    }
}

// ---------------------------------------------------------------------------
// Migration Topology Enum
// ---------------------------------------------------------------------------

/// Topology for inter-island migration in the island model.
#[pyclass(eq, eq_int, frozen, hash)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum MigrationTopology {
    /// Ring: each island sends migrants to the next island (i -> i+1).
    Ring,
    /// FullyConnected: each island sends migrants to all other islands.
    FullyConnected,
}

#[pymethods]
impl MigrationTopology {
    fn __repr__(&self) -> String {
        format!("MigrationTopology.{:?}", self)
    }

    fn __str__(&self) -> String {
        format!("{:?}", self)
    }
}

// ---------------------------------------------------------------------------
// Objective Direction
// ---------------------------------------------------------------------------

/// Whether to maximize or minimize an objective in NSGA-II.
#[pyclass(eq, eq_int, frozen, hash)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum ObjectiveDirection {
    Maximize,
    Minimize,
}

#[pymethods]
impl ObjectiveDirection {
    fn __repr__(&self) -> String {
        format!("ObjectiveDirection.{:?}", self)
    }

    fn __str__(&self) -> String {
        format!("{:?}", self)
    }
}

// ---------------------------------------------------------------------------
// Core Types
// ---------------------------------------------------------------------------

/// Configuration for evolutionary optimization.
#[pyclass(frozen, get_all)]
#[derive(Debug, Clone)]
pub struct EvolutionConfig {
    pub pop_size: u32,
    pub generations: u32,
    pub mutation_rate: f64,
    pub crossover_rate: f64,
    pub tournament_size: u32,
    pub elite_count: u32,
    pub sigma: f64,
}

#[pymethods]
impl EvolutionConfig {
    #[new]
    #[pyo3(signature = (pop_size=50, generations=100, mutation_rate=0.1, crossover_rate=0.7, tournament_size=3, elite_count=2, sigma=0.1))]
    fn new(
        pop_size: u32, generations: u32, mutation_rate: f64, crossover_rate: f64,
        tournament_size: u32, elite_count: u32, sigma: f64,
    ) -> Self {
        Self { pop_size, generations, mutation_rate, crossover_rate, tournament_size, elite_count, sigma }
    }

    fn __repr__(&self) -> String {
        format!("EvolutionConfig(pop={}, gen={}, mut_rate={:.2}, cx_rate={:.2})",
            self.pop_size, self.generations, self.mutation_rate, self.crossover_rate)
    }
}

/// A single individual (genome + fitness).
#[pyclass(frozen, get_all)]
#[derive(Debug, Clone)]
pub struct Individual {
    pub genome: Vec<f64>,
    pub fitness: f64,
}

#[pymethods]
impl Individual {
    #[new]
    #[pyo3(signature = (genome, fitness=f64::NEG_INFINITY))]
    fn new(genome: Vec<f64>, fitness: f64) -> Self {
        Self { genome, fitness }
    }

    fn __repr__(&self) -> String {
        format!("Individual(dim={}, fitness={:.6})", self.genome.len(), self.fitness)
    }
}

/// Bounds for a single parameter.
#[pyclass(frozen, get_all)]
#[derive(Debug, Clone)]
pub struct ParamBound {
    pub name: String,
    pub min_val: f64,
    pub max_val: f64,
    pub is_discrete: bool,
}

#[pymethods]
impl ParamBound {
    #[new]
    #[pyo3(signature = (name, min_val, max_val, is_discrete=false))]
    fn new(name: String, min_val: f64, max_val: f64, is_discrete: bool) -> Self {
        Self { name, min_val, max_val, is_discrete }
    }

    fn __repr__(&self) -> String {
        format!("ParamBound('{}', [{:.2}, {:.2}], discrete={})",
            self.name, self.min_val, self.max_val, self.is_discrete)
    }
}

/// Result of evolution optimization.
#[pyclass(frozen, get_all)]
#[derive(Debug, Clone)]
pub struct EvolutionResult {
    pub best_genome: Vec<f64>,
    pub best_fitness: f64,
    pub generations_run: u32,
    pub fitness_history: Vec<f64>,
}

#[pymethods]
impl EvolutionResult {
    fn __repr__(&self) -> String {
        format!("EvolutionResult(fitness={:.6}, gens={})", self.best_fitness, self.generations_run)
    }
}

/// Per-generation statistics.
#[pyclass(frozen, get_all)]
#[derive(Debug, Clone)]
pub struct GenerationStats {
    pub generation: u32,
    pub population_size: u32,
    pub best_fitness: f64,
    pub median_fitness: f64,
    pub mean_fitness: f64,
    pub worst_fitness: f64,
    pub std_fitness: f64,
    pub diversity: f64,
}

#[pymethods]
impl GenerationStats {
    fn __repr__(&self) -> String {
        format!(
            "GenerationStats(gen={}, best={:.4}, median={:.4}, diversity={:.4})",
            self.generation, self.best_fitness, self.median_fitness, self.diversity,
        )
    }
}

/// Convergence detection state.
#[pyclass(frozen, get_all)]
#[derive(Debug, Clone)]
pub struct ConvergenceState {
    pub converged: bool,
    pub reason: String,
    pub generations_without_improvement: u32,
    pub best_fitness: f64,
}

#[pymethods]
impl ConvergenceState {
    fn __repr__(&self) -> String {
        if self.converged {
            format!("ConvergenceState(converged=True, reason='{}')", self.reason)
        } else {
            format!("ConvergenceState(converged=False, stale_gens={})", self.generations_without_improvement)
        }
    }
}

/// Adaptive mutation state tracking.
#[pyclass(frozen, get_all)]
#[derive(Debug, Clone)]
pub struct AdaptiveMutationState {
    pub sigma: f64,
    pub success_count: u32,
    pub total_count: u32,
    pub success_rate: f64,
}

#[pymethods]
impl AdaptiveMutationState {
    fn __repr__(&self) -> String {
        format!("AdaptiveMutationState(sigma={:.4}, success_rate={:.2})", self.sigma, self.success_rate)
    }
}

/// Island model configuration.
#[pyclass(frozen, get_all)]
#[derive(Debug, Clone)]
pub struct IslandConfig {
    pub n_islands: u32,
    pub migration_interval: u32,
    pub migration_count: u32,
    pub topology: MigrationTopology,
}

#[pymethods]
impl IslandConfig {
    #[new]
    #[pyo3(signature = (n_islands=4, migration_interval=5, migration_count=2, topology=MigrationTopology::Ring))]
    fn new(n_islands: u32, migration_interval: u32, migration_count: u32, topology: MigrationTopology) -> Self {
        Self { n_islands, migration_interval, migration_count, topology }
    }

    fn __repr__(&self) -> String {
        format!("IslandConfig(islands={}, interval={}, topology={:?})",
            self.n_islands, self.migration_interval, self.topology)
    }
}

/// NSGA-II individual with multi-objective fitness.
#[pyclass(frozen, get_all)]
#[derive(Debug, Clone)]
pub struct NsgaIndividual {
    pub index: u32,
    pub objectives: Vec<f64>,
    pub rank: u32,
    pub crowding_distance: f64,
}

#[pymethods]
impl NsgaIndividual {
    #[new]
    fn new(index: u32, objectives: Vec<f64>) -> Self {
        Self { index, objectives, rank: 0, crowding_distance: 0.0 }
    }

    fn __repr__(&self) -> String {
        format!("NsgaIndividual(idx={}, rank={}, crowd={:.4})", self.index, self.rank, self.crowding_distance)
    }
}

/// Result of NSGA-II selection.
#[pyclass(frozen, get_all)]
#[derive(Debug, Clone)]
pub struct NsgaResult {
    pub selected_indices: Vec<u32>,
    pub pareto_front_indices: Vec<u32>,
    pub fronts: Vec<Vec<u32>>,
}

#[pymethods]
impl NsgaResult {
    fn __repr__(&self) -> String {
        format!("NsgaResult(selected={}, pareto_size={})",
            self.selected_indices.len(), self.pareto_front_indices.len())
    }
}

// ---------------------------------------------------------------------------
// Selection Functions
// ---------------------------------------------------------------------------

/// Tournament selection: select `count` individuals from population.
#[pyfunction]
#[pyo3(signature = (fitnesses, tournament_size=3, count=1, seed=42))]
pub fn tournament_select(
    fitnesses: Vec<f64>,
    tournament_size: u32,
    count: u32,
    seed: u64,
) -> Vec<u32> {
    if fitnesses.is_empty() {
        return Vec::new();
    }

    let n = fitnesses.len();
    let ts = (tournament_size as usize).min(n);
    let mut rng = Rng::new(seed);
    let mut selected = Vec::with_capacity(count as usize);

    for _ in 0..count {
        let mut best_idx = (rng.next_u64() % n as u64) as usize;
        let mut best_fit = fitnesses[best_idx];

        for _ in 1..ts {
            let idx = (rng.next_u64() % n as u64) as usize;
            if fitnesses[idx] > best_fit {
                best_idx = idx;
                best_fit = fitnesses[idx];
            }
        }

        selected.push(best_idx as u32);
    }

    selected
}

/// Uniform selection: random with equal probability.
#[pyfunction]
#[pyo3(signature = (fitnesses, count=1, seed=42))]
pub fn uniform_select(
    fitnesses: Vec<f64>,
    count: u32,
    seed: u64,
) -> Vec<u32> {
    if fitnesses.is_empty() {
        return Vec::new();
    }

    let n = fitnesses.len();
    let mut rng = Rng::new(seed);
    let mut selected = Vec::with_capacity(count as usize);

    for _ in 0..count {
        selected.push((rng.next_u64() % n as u64) as u32);
    }

    selected
}

/// Roulette wheel selection: probability proportional to fitness.
///
/// Fitnesses are shifted so the minimum becomes a small positive value,
/// ensuring valid probabilities even with negative fitnesses.
#[pyfunction]
#[pyo3(signature = (fitnesses, count=1, seed=42))]
pub fn roulette_select(
    fitnesses: Vec<f64>,
    count: u32,
    seed: u64,
) -> Vec<u32> {
    if fitnesses.is_empty() {
        return Vec::new();
    }

    let n = fitnesses.len();
    let min_score = fitnesses.iter().cloned().fold(f64::INFINITY, f64::min);
    let shifted: Vec<f64> = fitnesses.iter().map(|s| s - min_score + 1e-6).collect();
    let total: f64 = shifted.iter().sum();

    let mut rng = Rng::new(seed);
    let mut selected = Vec::with_capacity(count as usize);

    for _ in 0..count {
        let mut pick = rng.next_f64() * total;
        let mut chosen = n - 1;
        for (i, &s) in shifted.iter().enumerate() {
            pick -= s;
            if pick <= 0.0 {
                chosen = i;
                break;
            }
        }
        selected.push(chosen as u32);
    }

    selected
}

/// Rank-based selection: probability proportional to rank position.
///
/// Individuals are sorted by fitness. Rank 1 (worst) gets weight 1,
/// rank N (best) gets weight N. Higher-ranked individuals are more
/// likely to be selected.
#[pyfunction]
#[pyo3(signature = (fitnesses, count=1, seed=42))]
pub fn rank_select(
    fitnesses: Vec<f64>,
    count: u32,
    seed: u64,
) -> Vec<u32> {
    if fitnesses.is_empty() {
        return Vec::new();
    }

    let n = fitnesses.len();
    let mut indices: Vec<usize> = (0..n).collect();
    indices.sort_by(|&a, &b| {
        fitnesses[a].partial_cmp(&fitnesses[b]).unwrap_or(std::cmp::Ordering::Equal)
    });

    let total = (n * (n + 1) / 2) as f64;
    let mut rng = Rng::new(seed);
    let mut selected = Vec::with_capacity(count as usize);

    for _ in 0..count {
        let mut pick = rng.next_f64() * total;
        let mut chosen = indices[n - 1];
        for (rank, &idx) in indices.iter().enumerate() {
            pick -= (rank + 1) as f64;
            if pick <= 0.0 {
                chosen = idx;
                break;
            }
        }
        selected.push(chosen as u32);
    }

    selected
}

/// UCB1 selection: Upper Confidence Bound 1.
///
/// Balances exploitation (high fitness) with exploration (rarely selected).
/// Formula: score_i + sqrt(2 * ln(T) / n_i)
///
/// From Auer et al., "Finite-time Analysis of the Multiarmed Bandit Problem" (2002).
/// Used in the ProFiT evolutionary framework (Siper et al.) for parent selection
/// to balance exploring undersampled strategies with exploiting known good ones.
#[pyfunction]
#[pyo3(signature = (fitnesses, selection_counts, total_selections, count=1))]
pub fn ucb1_select(
    fitnesses: Vec<f64>,
    selection_counts: Vec<u64>,
    total_selections: u64,
    count: u32,
) -> Vec<u32> {
    if fitnesses.is_empty() {
        return Vec::new();
    }

    let n = fitnesses.len();
    let t = (total_selections + 1) as f64;
    let ln_t = t.ln();

    // Normalize fitnesses to [0, 1] for UCB1 to work correctly
    let min_f = fitnesses.iter().cloned().fold(f64::INFINITY, f64::min);
    let max_f = fitnesses.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
    let range = (max_f - min_f).max(1e-12);

    let mut ucb_scores: Vec<f64> = (0..n).map(|i| {
        let norm_score = (fitnesses[i] - min_f) / range;
        let n_i = if i < selection_counts.len() {
            (selection_counts[i] as f64).max(1.0)
        } else {
            1.0
        };
        let exploration = (2.0 * ln_t / n_i).sqrt();
        finite_or_zero(norm_score + exploration)
    }).collect();

    let mut selected = Vec::with_capacity(count as usize);
    for _ in 0..count {
        let best_idx = ucb_scores
            .iter()
            .enumerate()
            .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
            .map(|(i, _)| i)
            .unwrap_or(0);
        selected.push(best_idx as u32);
        // Mark selected individual so next iteration picks a different one
        ucb_scores[best_idx] = f64::NEG_INFINITY;
    }

    selected
}

/// Epsilon-greedy selection: with probability epsilon pick random, else pick best.
#[pyfunction]
#[pyo3(signature = (fitnesses, epsilon=0.1, count=1, seed=42))]
pub fn epsilon_greedy_select(
    fitnesses: Vec<f64>,
    epsilon: f64,
    count: u32,
    seed: u64,
) -> Vec<u32> {
    if fitnesses.is_empty() {
        return Vec::new();
    }

    let n = fitnesses.len();
    let best_idx = fitnesses
        .iter()
        .enumerate()
        .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
        .map(|(i, _)| i)
        .unwrap_or(0);

    let mut rng = Rng::new(seed);
    let mut selected = Vec::with_capacity(count as usize);

    for _ in 0..count {
        if rng.next_f64() < epsilon {
            selected.push((rng.next_u64() % n as u64) as u32);
        } else {
            selected.push(best_idx as u32);
        }
    }

    selected
}

// ---------------------------------------------------------------------------
// Genetic Operators
// ---------------------------------------------------------------------------

/// Gaussian mutation: mutate genome values with Gaussian noise.
#[pyfunction]
#[pyo3(signature = (genome, bounds_min, bounds_max, is_discrete, sigma=0.1, rate=0.1, seed=42))]
pub fn gaussian_mutate(
    genome: Vec<f64>,
    bounds_min: Vec<f64>,
    bounds_max: Vec<f64>,
    is_discrete: Vec<bool>,
    sigma: f64,
    rate: f64,
    seed: u64,
) -> Vec<f64> {
    let n = genome.len().min(bounds_min.len()).min(bounds_max.len());
    let mut rng = Rng::new(seed);
    let mut result = genome.clone();

    for i in 0..n {
        if rng.next_f64() < rate {
            let range = bounds_max[i] - bounds_min[i];
            let noise = rng.next_gaussian() * sigma * range;
            let mut val = result[i] + noise;
            val = val.clamp(bounds_min[i], bounds_max[i]);
            if i < is_discrete.len() && is_discrete[i] {
                val = val.round();
            }
            result[i] = val;
        }
    }

    result
}

/// Uniform crossover: combine two parent genomes.
#[pyfunction]
#[pyo3(signature = (parent_a, parent_b, rate=0.5, seed=42))]
pub fn uniform_crossover(
    parent_a: Vec<f64>,
    parent_b: Vec<f64>,
    rate: f64,
    seed: u64,
) -> Vec<f64> {
    let n = parent_a.len().min(parent_b.len());
    let mut rng = Rng::new(seed);
    let mut child = Vec::with_capacity(n);

    for i in 0..n {
        if rng.next_f64() < rate {
            child.push(parent_b[i]);
        } else {
            child.push(parent_a[i]);
        }
    }

    child
}

/// Initialize a random population within bounds.
#[pyfunction]
#[pyo3(signature = (bounds_min, bounds_max, is_discrete, pop_size=50, seed=42))]
pub fn init_population(
    bounds_min: Vec<f64>,
    bounds_max: Vec<f64>,
    is_discrete: Vec<bool>,
    pop_size: u32,
    seed: u64,
) -> Vec<Vec<f64>> {
    let dim = bounds_min.len().min(bounds_max.len());
    let mut rng = Rng::new(seed);
    let mut population = Vec::with_capacity(pop_size as usize);

    for _ in 0..pop_size {
        let mut genome = Vec::with_capacity(dim);
        for j in 0..dim {
            let range = bounds_max[j] - bounds_min[j];
            let mut val = bounds_min[j] + rng.next_f64() * range;
            if j < is_discrete.len() && is_discrete[j] {
                val = val.round();
            }
            genome.push(val);
        }
        population.push(genome);
    }

    population
}

// ---------------------------------------------------------------------------
// NSGA-II Multi-Objective Optimization
// ---------------------------------------------------------------------------

/// Dominates check: a dominates b iff a >= b on all objectives and > on at least one.
fn dominates(a: &[f64], b: &[f64]) -> bool {
    let mut strictly_better = false;
    for (ai, bi) in a.iter().zip(b.iter()) {
        if ai < bi {
            return false;
        }
        if ai > bi {
            strictly_better = true;
        }
    }
    strictly_better
}

/// Fast non-dominated sorting (NSGA-II).
///
/// Returns fronts where fronts[0] is the Pareto front (non-dominated set),
/// fronts[1] is dominated only by front 0, etc.
///
/// Based on Deb et al., "A Fast and Elitist Multiobjective Genetic Algorithm: NSGA-II" (2002).
#[pyfunction]
#[pyo3(signature = (objectives))]
pub fn non_dominated_sort(objectives: Vec<Vec<f64>>) -> Vec<Vec<u32>> {
    let n = objectives.len();
    if n == 0 {
        return Vec::new();
    }

    let mut domination_count = vec![0usize; n];
    let mut dominated_set: Vec<Vec<usize>> = vec![Vec::new(); n];
    let mut front0 = Vec::new();

    for i in 0..n {
        for j in 0..n {
            if i == j {
                continue;
            }
            if dominates(&objectives[i], &objectives[j]) {
                dominated_set[i].push(j);
            } else if dominates(&objectives[j], &objectives[i]) {
                domination_count[i] += 1;
            }
        }
        if domination_count[i] == 0 {
            front0.push(i);
        }
    }

    let mut fronts: Vec<Vec<u32>> = vec![front0.iter().map(|&i| i as u32).collect()];

    loop {
        let prev_front = fronts.last().unwrap();
        let mut next_front = Vec::new();
        for &p in prev_front {
            for &q in &dominated_set[p as usize] {
                domination_count[q] -= 1;
                if domination_count[q] == 0 {
                    next_front.push(q as u32);
                }
            }
        }
        if next_front.is_empty() {
            break;
        }
        fronts.push(next_front);
    }

    fronts
}

/// Compute crowding distance for individuals in a front.
///
/// Returns crowding distances for each individual in the front.
/// Boundary individuals get infinite distance. Interior individuals get
/// distance proportional to the objective space around them.
#[pyfunction]
#[pyo3(signature = (objectives, front_indices))]
pub fn crowding_distance_assign(
    objectives: Vec<Vec<f64>>,
    front_indices: Vec<u32>,
) -> Vec<f64> {
    let n = front_indices.len();
    if n == 0 {
        return Vec::new();
    }
    if n <= 2 {
        return vec![f64::INFINITY; n];
    }

    let n_objectives = objectives.get(front_indices[0] as usize)
        .map(|o| o.len())
        .unwrap_or(0);

    let mut distances = vec![0.0f64; n];

    for m in 0..n_objectives {
        let mut sorted_pos: Vec<usize> = (0..n).collect();
        sorted_pos.sort_by(|&a, &b| {
            let obj_a = objectives.get(front_indices[a] as usize)
                .and_then(|o| o.get(m))
                .copied()
                .unwrap_or(0.0);
            let obj_b = objectives.get(front_indices[b] as usize)
                .and_then(|o| o.get(m))
                .copied()
                .unwrap_or(0.0);
            obj_a.partial_cmp(&obj_b).unwrap_or(std::cmp::Ordering::Equal)
        });

        distances[sorted_pos[0]] = f64::INFINITY;
        distances[sorted_pos[n - 1]] = f64::INFINITY;

        let first_obj = objectives.get(front_indices[sorted_pos[0]] as usize)
            .and_then(|o| o.get(m))
            .copied()
            .unwrap_or(0.0);
        let last_obj = objectives.get(front_indices[sorted_pos[n - 1]] as usize)
            .and_then(|o| o.get(m))
            .copied()
            .unwrap_or(0.0);
        let obj_range = last_obj - first_obj;

        if obj_range.abs() < 1e-12 {
            continue;
        }

        for i in 1..n - 1 {
            let next_obj = objectives.get(front_indices[sorted_pos[i + 1]] as usize)
                .and_then(|o| o.get(m))
                .copied()
                .unwrap_or(0.0);
            let prev_obj = objectives.get(front_indices[sorted_pos[i - 1]] as usize)
                .and_then(|o| o.get(m))
                .copied()
                .unwrap_or(0.0);
            distances[sorted_pos[i]] += (next_obj - prev_obj) / obj_range;
        }
    }

    distances
}

/// NSGA-II selection: select N individuals using front rank + crowding distance.
///
/// Multi-objective evolutionary optimization. Individuals are sorted by
/// Pareto dominance rank (lower is better). Within the same rank,
/// individuals with higher crowding distance are preferred to maintain diversity.
#[pyfunction]
#[pyo3(signature = (objectives, n_select))]
pub fn nsga2_select(objectives: Vec<Vec<f64>>, n_select: u32) -> NsgaResult {
    let n = objectives.len();
    if n == 0 {
        return NsgaResult {
            selected_indices: Vec::new(),
            pareto_front_indices: Vec::new(),
            fronts: Vec::new(),
        };
    }

    let fronts = non_dominated_sort(objectives.clone());
    let pareto_front = fronts.first().cloned().unwrap_or_default();

    let mut selected = Vec::with_capacity(n_select as usize);

    for front in &fronts {
        if selected.len() + front.len() <= n_select as usize {
            selected.extend_from_slice(front);
        } else {
            let distances = crowding_distance_assign(objectives.clone(), front.clone());
            let mut front_with_dist: Vec<(u32, f64)> = front.iter()
                .enumerate()
                .map(|(i, &idx)| (idx, distances.get(i).copied().unwrap_or(0.0)))
                .collect();
            front_with_dist.sort_by(|a, b|
                b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)
            );
            let remaining = n_select as usize - selected.len();
            for (idx, _) in front_with_dist.iter().take(remaining) {
                selected.push(*idx);
            }
            break;
        }
    }

    NsgaResult {
        selected_indices: selected,
        pareto_front_indices: pareto_front,
        fronts,
    }
}

/// Extract the Pareto front (non-dominated set) indices.
#[pyfunction]
#[pyo3(signature = (objectives))]
pub fn pareto_front(objectives: Vec<Vec<f64>>) -> Vec<u32> {
    let fronts = non_dominated_sort(objectives);
    fronts.into_iter().next().unwrap_or_default()
}

// ---------------------------------------------------------------------------
// Island Model
// ---------------------------------------------------------------------------

/// Perform island migration: move top genomes between islands.
///
/// Each island is a list of (genome, fitness) pairs. Migration copies the
/// top-K genomes from source islands to destination islands based on the topology.
#[pyfunction]
#[pyo3(signature = (islands, fitnesses, config))]
pub fn island_migrate(
    islands: Vec<Vec<Vec<f64>>>,
    fitnesses: Vec<Vec<f64>>,
    config: &IslandConfig,
) -> Vec<Vec<Vec<f64>>> {
    let n = islands.len();
    if n < 2 {
        return islands;
    }

    let mut result = islands.clone();
    let k = config.migration_count as usize;

    // Get top-K indices per island
    let mut migrants: Vec<Vec<Vec<f64>>> = Vec::with_capacity(n);
    for (island_idx, island_fit) in fitnesses.iter().enumerate() {
        let mut ranked: Vec<usize> = (0..island_fit.len()).collect();
        ranked.sort_by(|&a, &b| {
            island_fit[b].partial_cmp(&island_fit[a]).unwrap_or(std::cmp::Ordering::Equal)
        });
        let top_k: Vec<Vec<f64>> = ranked.iter()
            .take(k.min(islands[island_idx].len()))
            .map(|&i| islands[island_idx][i].clone())
            .collect();
        migrants.push(top_k);
    }

    match config.topology {
        MigrationTopology::Ring => {
            for (i, island_migrants) in migrants.iter().enumerate() {
                let target = (i + 1) % n;
                for migrant in island_migrants {
                    result[target].push(migrant.clone());
                }
            }
        }
        MigrationTopology::FullyConnected => {
            for (i, island_migrants) in migrants.iter().enumerate() {
                for j in 0..n {
                    if i == j {
                        continue;
                    }
                    for migrant in island_migrants {
                        result[j].push(migrant.clone());
                    }
                }
            }
        }
    }

    result
}

/// Get the default selection method for each island.
///
/// Assigns different selection methods to each island for diversity,
/// following the ProFiT framework's approach of per-island specialization.
/// Methods cycle through: RouletteWheel, Rank, UCB1, EpsilonGreedy.
#[pyfunction]
#[pyo3(signature = (n_islands))]
pub fn island_selection_methods(n_islands: u32) -> Vec<SelectionMethod> {
    let methods = [
        SelectionMethod::RouletteWheel,
        SelectionMethod::Rank,
        SelectionMethod::Ucb1,
        SelectionMethod::EpsilonGreedy,
    ];
    (0..n_islands)
        .map(|i| methods[i as usize % methods.len()])
        .collect()
}

// ---------------------------------------------------------------------------
// Convergence Detection
// ---------------------------------------------------------------------------

/// Check if evolution has converged.
///
/// Returns convergence state based on:
/// 1. Patience: no improvement for `patience` generations
/// 2. Diversity collapse: population diversity below threshold
#[pyfunction]
#[pyo3(signature = (best_fitness, diversity, prev_best_fitness, gens_without_improvement, patience=20, min_improvement=1e-4, diversity_threshold=0.01))]
pub fn check_convergence(
    best_fitness: f64,
    diversity: f64,
    prev_best_fitness: f64,
    gens_without_improvement: u32,
    patience: u32,
    min_improvement: f64,
    diversity_threshold: f64,
) -> ConvergenceState {
    let improvement = best_fitness - prev_best_fitness;
    let new_gens = if improvement > min_improvement {
        0
    } else {
        gens_without_improvement + 1
    };

    let new_best = if improvement > min_improvement {
        best_fitness
    } else {
        prev_best_fitness
    };

    // Check patience
    if new_gens >= patience {
        return ConvergenceState {
            converged: true,
            reason: format!("No improvement for {} generations (patience={})", new_gens, patience),
            generations_without_improvement: new_gens,
            best_fitness: new_best,
        };
    }

    // Check diversity collapse
    if diversity < diversity_threshold && diversity >= 0.0 {
        return ConvergenceState {
            converged: true,
            reason: format!("Diversity collapsed to {:.6} (threshold={:.6})", diversity, diversity_threshold),
            generations_without_improvement: new_gens,
            best_fitness: new_best,
        };
    }

    ConvergenceState {
        converged: false,
        reason: String::new(),
        generations_without_improvement: new_gens,
        best_fitness: new_best,
    }
}

// ---------------------------------------------------------------------------
// Adaptive Mutation (Rechenberg's 1/5 Rule)
// ---------------------------------------------------------------------------

/// Adapt mutation sigma using Rechenberg's 1/5 success rule.
///
/// If the success rate is above the target (default 0.2), increase sigma
/// to explore more broadly. If below, decrease sigma to fine-tune.
///
/// Reference: Rechenberg, I. (1973), "Evolutionsstrategie".
/// Also used in the Continuous Program Search framework (Siper et al.)
/// for adapting mutation strength based on acceptance rates.
#[pyfunction]
#[pyo3(signature = (current_sigma, success_count, total_count, target_success_rate=0.2, adaptation_factor=1.5, sigma_min=0.001, sigma_max=2.0))]
pub fn adapt_mutation_sigma(
    current_sigma: f64,
    success_count: u32,
    total_count: u32,
    target_success_rate: f64,
    adaptation_factor: f64,
    sigma_min: f64,
    sigma_max: f64,
) -> AdaptiveMutationState {
    if total_count == 0 {
        return AdaptiveMutationState {
            sigma: current_sigma,
            success_count,
            total_count,
            success_rate: 0.0,
        };
    }

    let success_rate = success_count as f64 / total_count as f64;
    let new_sigma = if success_rate > target_success_rate {
        (current_sigma * adaptation_factor).min(sigma_max)
    } else if success_rate < target_success_rate {
        (current_sigma / adaptation_factor).max(sigma_min)
    } else {
        current_sigma
    };

    AdaptiveMutationState {
        sigma: new_sigma,
        success_count,
        total_count,
        success_rate,
    }
}

// ---------------------------------------------------------------------------
// Population Diversity
// ---------------------------------------------------------------------------

/// Compute population diversity as the mean pairwise distance between genomes.
///
/// Uses normalized L2 distance. Returns 0.0 for populations with fewer than
/// 2 individuals. Higher values indicate more diverse populations.
#[pyfunction]
#[pyo3(signature = (population))]
pub fn population_diversity(population: Vec<Vec<f64>>) -> f64 {
    let n = population.len();
    if n < 2 {
        return 0.0;
    }

    let dim = population[0].len();
    if dim == 0 {
        return 0.0;
    }

    let mut total_dist = 0.0;
    let mut pair_count = 0u64;

    for i in 0..n {
        for j in (i + 1)..n {
            let dist: f64 = population[i].iter()
                .zip(population[j].iter())
                .map(|(a, b)| (a - b) * (a - b))
                .sum::<f64>()
                .sqrt()
                / dim as f64;
            total_dist += dist;
            pair_count += 1;
        }
    }

    if pair_count == 0 { 0.0 } else { finite_or_zero(total_dist / pair_count as f64) }
}

// ---------------------------------------------------------------------------
// Generation Statistics
// ---------------------------------------------------------------------------

/// Compute per-generation statistics from fitness values.
#[pyfunction]
#[pyo3(signature = (generation, fitnesses))]
pub fn compute_generation_stats(generation: u32, fitnesses: Vec<f64>) -> GenerationStats {
    if fitnesses.is_empty() {
        return GenerationStats {
            generation,
            population_size: 0,
            best_fitness: 0.0,
            median_fitness: 0.0,
            mean_fitness: 0.0,
            worst_fitness: 0.0,
            std_fitness: 0.0,
            diversity: 0.0,
        };
    }

    let n = fitnesses.len() as f64;
    let mean = fitnesses.iter().sum::<f64>() / n;
    let variance = fitnesses.iter().map(|f| (f - mean) * (f - mean)).sum::<f64>() / n;
    let std_dev = variance.sqrt();

    let mut sorted = fitnesses.clone();
    sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));

    let best = *sorted.last().unwrap();
    let worst = sorted[0];
    let median = sorted[sorted.len() / 2];

    // Diversity as coefficient of variation
    let diversity = if mean.abs() > 1e-12 {
        finite_or_zero(std_dev / mean.abs())
    } else {
        0.0
    };

    GenerationStats {
        generation,
        population_size: fitnesses.len() as u32,
        best_fitness: finite_or_zero(best),
        median_fitness: finite_or_zero(median),
        mean_fitness: finite_or_zero(mean),
        worst_fitness: finite_or_zero(worst),
        std_fitness: finite_or_zero(std_dev),
        diversity,
    }
}

// ---------------------------------------------------------------------------
// Module Registration
// ---------------------------------------------------------------------------

pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
    // Enums
    m.add_class::<SelectionMethod>()?;
    m.add_class::<MigrationTopology>()?;
    m.add_class::<ObjectiveDirection>()?;

    // Types
    m.add_class::<EvolutionConfig>()?;
    m.add_class::<Individual>()?;
    m.add_class::<ParamBound>()?;
    m.add_class::<EvolutionResult>()?;
    m.add_class::<GenerationStats>()?;
    m.add_class::<ConvergenceState>()?;
    m.add_class::<AdaptiveMutationState>()?;
    m.add_class::<IslandConfig>()?;
    m.add_class::<NsgaIndividual>()?;
    m.add_class::<NsgaResult>()?;

    // Selection
    m.add_function(wrap_pyfunction!(tournament_select, m)?)?;
    m.add_function(wrap_pyfunction!(uniform_select, m)?)?;
    m.add_function(wrap_pyfunction!(roulette_select, m)?)?;
    m.add_function(wrap_pyfunction!(rank_select, m)?)?;
    m.add_function(wrap_pyfunction!(ucb1_select, m)?)?;
    m.add_function(wrap_pyfunction!(epsilon_greedy_select, m)?)?;

    // Genetic operators
    m.add_function(wrap_pyfunction!(gaussian_mutate, m)?)?;
    m.add_function(wrap_pyfunction!(uniform_crossover, m)?)?;
    m.add_function(wrap_pyfunction!(init_population, m)?)?;

    // NSGA-II
    m.add_function(wrap_pyfunction!(non_dominated_sort, m)?)?;
    m.add_function(wrap_pyfunction!(crowding_distance_assign, m)?)?;
    m.add_function(wrap_pyfunction!(nsga2_select, m)?)?;
    m.add_function(wrap_pyfunction!(pareto_front, m)?)?;

    // Island model
    m.add_function(wrap_pyfunction!(island_migrate, m)?)?;
    m.add_function(wrap_pyfunction!(island_selection_methods, m)?)?;

    // Convergence
    m.add_function(wrap_pyfunction!(check_convergence, m)?)?;

    // Adaptive mutation
    m.add_function(wrap_pyfunction!(adapt_mutation_sigma, m)?)?;

    // Diversity & stats
    m.add_function(wrap_pyfunction!(population_diversity, m)?)?;
    m.add_function(wrap_pyfunction!(compute_generation_stats, m)?)?;

    Ok(())
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    // -- PRNG --

    #[test]
    fn test_rng_deterministic() {
        let mut rng1 = Rng::new(42);
        let mut rng2 = Rng::new(42);
        for _ in 0..100 {
            assert_eq!(rng1.next_u64(), rng2.next_u64());
        }
    }

    #[test]
    fn test_rng_range() {
        let mut rng = Rng::new(42);
        for _ in 0..1000 {
            let v = rng.next_f64();
            assert!(v >= 0.0 && v < 1.0);
        }
    }

    // -- Tournament Select --

    #[test]
    fn test_tournament_select() {
        let fitnesses = vec![0.1, 0.5, 0.3, 0.9, 0.2];
        let selected = tournament_select(fitnesses, 3, 5, 42);
        assert_eq!(selected.len(), 5);
        for &idx in &selected {
            assert!((idx as usize) < 5);
        }
    }

    #[test]
    fn test_tournament_select_empty() {
        assert!(tournament_select(vec![], 3, 5, 42).is_empty());
    }

    // -- Uniform Select --

    #[test]
    fn test_uniform_select_basic() {
        let fitnesses = vec![0.1, 0.5, 0.3];
        let selected = uniform_select(fitnesses, 100, 42);
        assert_eq!(selected.len(), 100);
        // Should have variety
        let unique: std::collections::HashSet<u32> = selected.iter().copied().collect();
        assert!(unique.len() > 1);
    }

    #[test]
    fn test_uniform_select_empty() {
        assert!(uniform_select(vec![], 5, 42).is_empty());
    }

    // -- Roulette Select --

    #[test]
    fn test_roulette_select_biased() {
        let fitnesses = vec![0.0, 0.0, 0.0, 100.0, 0.0];
        let selected = roulette_select(fitnesses, 100, 42);
        let count_3 = selected.iter().filter(|&&x| x == 3).count();
        assert!(count_3 > 50, "Roulette should favor high fitness: {}", count_3);
    }

    #[test]
    fn test_roulette_select_negative() {
        let fitnesses = vec![-10.0, -5.0, 0.0, 5.0];
        let selected = roulette_select(fitnesses, 100, 42);
        assert_eq!(selected.len(), 100);
        // Should still work with negative values
        let count_3 = selected.iter().filter(|&&x| x == 3).count();
        assert!(count_3 > 20);
    }

    // -- Rank Select --

    #[test]
    fn test_rank_select_biased() {
        let fitnesses = vec![0.1, 0.5, 0.3, 0.9, 0.2];
        let selected = rank_select(fitnesses, 1000, 42);
        let count_3 = selected.iter().filter(|&&x| x == 3).count();
        // Rank 5 (best) should be selected ~5/15 = 33% of the time
        assert!(count_3 > 200, "Rank selection should favor best: {}", count_3);
    }

    // -- UCB1 Select --

    #[test]
    fn test_ucb1_select_exploration() {
        let fitnesses = vec![1.0, 1.0, 1.0, 1.0, 1.0];
        let selection_counts = vec![100, 100, 100, 100, 1]; // index 4 barely explored
        let selected = ucb1_select(fitnesses, selection_counts, 500, 1);
        // Index 4 should be selected (high exploration bonus)
        assert_eq!(selected[0], 4);
    }

    #[test]
    fn test_ucb1_select_exploitation() {
        let fitnesses = vec![0.0, 0.0, 0.0, 10.0, 0.0];
        let selection_counts = vec![10, 10, 10, 10, 10]; // equal exploration
        let selected = ucb1_select(fitnesses, selection_counts, 50, 1);
        // Index 3 should be selected (highest fitness with equal exploration)
        assert_eq!(selected[0], 3);
    }

    // -- Epsilon Greedy --

    #[test]
    fn test_epsilon_greedy_exploitation() {
        let fitnesses = vec![0.1, 0.5, 0.3, 0.9, 0.2];
        let selected = epsilon_greedy_select(fitnesses, 0.0, 10, 42);
        // With epsilon=0.0, should always select best (index 3)
        for &idx in &selected {
            assert_eq!(idx, 3);
        }
    }

    #[test]
    fn test_epsilon_greedy_exploration() {
        let fitnesses = vec![0.1, 0.5, 0.3, 0.9, 0.2];
        let selected = epsilon_greedy_select(fitnesses, 1.0, 100, 42);
        // With epsilon=1.0, should be fully random
        let unique: std::collections::HashSet<u32> = selected.iter().copied().collect();
        assert!(unique.len() > 1);
    }

    // -- Gaussian Mutate --

    #[test]
    fn test_gaussian_mutate() {
        let genome = vec![0.5, 0.5, 0.5];
        let mutated = gaussian_mutate(
            genome.clone(),
            vec![0.0, 0.0, 0.0],
            vec![1.0, 1.0, 1.0],
            vec![false, false, false],
            0.1, 1.0, 42,
        );
        assert_eq!(mutated.len(), 3);
        for v in &mutated {
            assert!(*v >= 0.0 && *v <= 1.0);
        }
    }

    #[test]
    fn test_gaussian_mutate_discrete() {
        let genome = vec![5.0];
        let mutated = gaussian_mutate(
            genome,
            vec![0.0],
            vec![10.0],
            vec![true],
            0.5, 1.0, 42,
        );
        assert!((mutated[0] - mutated[0].round()).abs() < 1e-10);
    }

    // -- Crossover --

    #[test]
    fn test_uniform_crossover() {
        let a = vec![0.0, 0.0, 0.0, 0.0];
        let b = vec![1.0, 1.0, 1.0, 1.0];
        let child = uniform_crossover(a, b, 0.5, 42);
        assert_eq!(child.len(), 4);
    }

    // -- Init Population --

    #[test]
    fn test_init_population() {
        let pop = init_population(
            vec![0.0, 0.0],
            vec![1.0, 10.0],
            vec![false, true],
            20, 42,
        );
        assert_eq!(pop.len(), 20);
        for genome in &pop {
            assert_eq!(genome.len(), 2);
            assert!(genome[0] >= 0.0 && genome[0] <= 1.0);
            assert!(genome[1] >= 0.0 && genome[1] <= 10.0);
            assert!((genome[1] - genome[1].round()).abs() < 1e-10);
        }
    }

    #[test]
    fn test_init_population_deterministic() {
        let pop1 = init_population(vec![0.0], vec![1.0], vec![false], 10, 42);
        let pop2 = init_population(vec![0.0], vec![1.0], vec![false], 10, 42);
        assert_eq!(pop1, pop2);
    }

    // -- NSGA-II --

    #[test]
    fn test_dominates_fn() {
        assert!(dominates(&[2.0, 2.0], &[1.0, 1.0]));
        assert!(!dominates(&[2.0, 1.0], &[1.0, 2.0]));
        assert!(!dominates(&[1.0, 1.0], &[2.0, 2.0]));
        assert!(!dominates(&[1.0, 1.0], &[1.0, 1.0]));
    }

    #[test]
    fn test_non_dominated_sort_basic() {
        let objectives = vec![
            vec![3.0, 1.0],
            vec![1.0, 3.0],
            vec![2.0, 2.0],
            vec![1.0, 1.0],
        ];
        let fronts = non_dominated_sort(objectives);
        assert_eq!(fronts.len(), 2);
        assert_eq!(fronts[0].len(), 3); // indices 0, 1, 2 are non-dominated
        assert_eq!(fronts[1].len(), 1); // index 3 is dominated
        assert!(fronts[1].contains(&3));
    }

    #[test]
    fn test_non_dominated_sort_empty() {
        let fronts = non_dominated_sort(vec![]);
        assert!(fronts.is_empty());
    }

    #[test]
    fn test_crowding_distance() {
        let objectives = vec![
            vec![1.0, 3.0],
            vec![2.0, 2.0],
            vec![3.0, 1.0],
        ];
        let front = vec![0, 1, 2];
        let distances = crowding_distance_assign(objectives, front);
        assert_eq!(distances.len(), 3);
        // Boundary points should have infinite distance
        assert!(distances[0].is_infinite());
        assert!(distances[2].is_infinite());
        // Interior point should have finite distance
        assert!(distances[1].is_finite());
        assert!(distances[1] > 0.0);
    }

    #[test]
    fn test_nsga2_select_prefers_pareto() {
        let objectives = vec![
            vec![3.0, 1.0],
            vec![1.0, 3.0],
            vec![2.0, 2.0],
            vec![1.0, 1.0],
            vec![0.5, 0.5],
        ];
        let result = nsga2_select(objectives, 3);
        assert_eq!(result.selected_indices.len(), 3);
        // Dominated individuals (3, 4) should not be selected
        assert!(!result.selected_indices.contains(&4));
    }

    #[test]
    fn test_pareto_front_fn() {
        let objectives = vec![
            vec![3.0, 1.0],
            vec![1.0, 3.0],
            vec![2.0, 2.0],
            vec![1.0, 1.0],
        ];
        let front = pareto_front(objectives);
        assert_eq!(front.len(), 3);
        assert!(!front.contains(&3));
    }

    // -- Island Model --

    #[test]
    fn test_island_migrate_ring() {
        let islands = vec![
            vec![vec![1.0], vec![2.0]],
            vec![vec![3.0], vec![4.0]],
        ];
        let fitnesses = vec![
            vec![0.5, 1.0],
            vec![0.3, 0.8],
        ];
        let config = IslandConfig::new(2, 5, 1, MigrationTopology::Ring);
        let result = island_migrate(islands, fitnesses, &config);
        // Island 0's best (2.0) should migrate to island 1
        assert_eq!(result[1].len(), 3);
        // Island 1's best (4.0) should migrate to island 0
        assert_eq!(result[0].len(), 3);
    }

    #[test]
    fn test_island_migrate_fully_connected() {
        let islands = vec![
            vec![vec![1.0]],
            vec![vec![2.0]],
            vec![vec![3.0]],
        ];
        let fitnesses = vec![vec![1.0], vec![1.0], vec![1.0]];
        let config = IslandConfig::new(3, 5, 1, MigrationTopology::FullyConnected);
        let result = island_migrate(islands, fitnesses, &config);
        // Each island should receive migrants from the other 2
        for island in &result {
            assert_eq!(island.len(), 3);
        }
    }

    #[test]
    fn test_island_selection_methods() {
        let methods = island_selection_methods(6);
        assert_eq!(methods.len(), 6);
        assert_eq!(methods[0], SelectionMethod::RouletteWheel);
        assert_eq!(methods[1], SelectionMethod::Rank);
        assert_eq!(methods[2], SelectionMethod::Ucb1);
        assert_eq!(methods[3], SelectionMethod::EpsilonGreedy);
        assert_eq!(methods[4], SelectionMethod::RouletteWheel); // wraps around
    }

    // -- Convergence --

    #[test]
    fn test_convergence_patience() {
        let state = check_convergence(1.0, 1.0, 1.0, 19, 20, 1e-4, 0.01);
        assert!(state.converged);
        assert!(state.reason.contains("No improvement"));
    }

    #[test]
    fn test_convergence_improvement_resets() {
        let state = check_convergence(2.0, 1.0, 1.0, 19, 20, 1e-4, 0.01);
        assert!(!state.converged);
        assert_eq!(state.generations_without_improvement, 0);
    }

    #[test]
    fn test_convergence_diversity_collapse() {
        let state = check_convergence(1.0, 0.005, 1.0, 0, 20, 1e-4, 0.01);
        assert!(state.converged);
        assert!(state.reason.contains("Diversity"));
    }

    // -- Adaptive Mutation --

    #[test]
    fn test_adapt_sigma_increase() {
        let state = adapt_mutation_sigma(0.1, 30, 100, 0.2, 1.5, 0.001, 2.0);
        // 30% success > 20% target -> sigma should increase
        assert!(state.sigma > 0.1);
        assert!((state.sigma - 0.15).abs() < 1e-10);
    }

    #[test]
    fn test_adapt_sigma_decrease() {
        let state = adapt_mutation_sigma(0.1, 10, 100, 0.2, 1.5, 0.001, 2.0);
        // 10% success < 20% target -> sigma should decrease
        assert!(state.sigma < 0.1);
    }

    #[test]
    fn test_adapt_sigma_bounds() {
        let state = adapt_mutation_sigma(1.9, 100, 100, 0.2, 1.5, 0.001, 2.0);
        assert!(state.sigma <= 2.0);
        let state2 = adapt_mutation_sigma(0.002, 0, 100, 0.2, 1.5, 0.001, 2.0);
        assert!(state2.sigma >= 0.001);
    }

    #[test]
    fn test_adapt_sigma_zero_total() {
        let state = adapt_mutation_sigma(0.1, 0, 0, 0.2, 1.5, 0.001, 2.0);
        assert!((state.sigma - 0.1).abs() < 1e-10);
    }

    // -- Population Diversity --

    #[test]
    fn test_diversity_identical() {
        let pop = vec![vec![0.5, 0.5], vec![0.5, 0.5], vec![0.5, 0.5]];
        let d = population_diversity(pop);
        assert!((d - 0.0).abs() < 1e-10);
    }

    #[test]
    fn test_diversity_spread() {
        let pop = vec![vec![0.0, 0.0], vec![1.0, 1.0]];
        let d = population_diversity(pop);
        assert!(d > 0.0);
    }

    #[test]
    fn test_diversity_single() {
        let pop = vec![vec![0.5, 0.5]];
        let d = population_diversity(pop);
        assert!((d - 0.0).abs() < 1e-10);
    }

    #[test]
    fn test_diversity_empty() {
        let d = population_diversity(vec![]);
        assert!((d - 0.0).abs() < 1e-10);
    }

    // -- Generation Stats --

    #[test]
    fn test_generation_stats() {
        let stats = compute_generation_stats(5, vec![0.1, 0.5, 0.3, 0.8, 0.2]);
        assert_eq!(stats.generation, 5);
        assert_eq!(stats.population_size, 5);
        assert!((stats.best_fitness - 0.8).abs() < 1e-10);
        assert!((stats.worst_fitness - 0.1).abs() < 1e-10);
        assert!(stats.mean_fitness > 0.0);
        assert!(stats.std_fitness > 0.0);
    }

    #[test]
    fn test_generation_stats_empty() {
        let stats = compute_generation_stats(0, vec![]);
        assert_eq!(stats.population_size, 0);
        assert!((stats.best_fitness - 0.0).abs() < 1e-10);
    }

    #[test]
    fn test_generation_stats_identical() {
        let stats = compute_generation_stats(0, vec![1.0, 1.0, 1.0]);
        assert!((stats.std_fitness - 0.0).abs() < 1e-10);
        assert!((stats.diversity - 0.0).abs() < 1e-10);
    }
}
