rlevo_evolution/algorithms/

es_classical.rs

//! Classical Evolution Strategies.
//!
//! Four canonical variants parameterized on a single [`EsConfig`]:
//!
//! - `(1+1)` — a single parent, a single offspring, 1/5th success-rule
//!   σ adaptation.
//! - `(1+λ)` — a single parent, λ offspring per generation; the best
//!   offspring replaces the parent iff its fitness improves. The
//!   underlying mutation/selection loop is also reused by Cartesian GP.
//! - `(μ,λ)` — μ parents, λ offspring; parents are discarded each
//!   generation.
//! - `(μ+λ)` — μ parents, λ offspring; survivors are the μ best of the
//!   combined pool.
//!
//! σ adaptation is by log-normal self-adaptation in the multi-parent
//! variants; `(1+1)` uses Rechenberg's 1/5th success rule.
//!
//! # References
//!
//! - Beyer & Schwefel (2002), *Evolution strategies: A comprehensive
//!   introduction*.

use std::marker::PhantomData;

use burn::tensor::{Tensor, TensorData, backend::Backend};
use rand::Rng;
use rand::RngExt;

use rlevo_core::bounds::Bounds;
use rlevo_core::config::{self, ConfigError, ConstraintKind, Validate};

use crate::ops::mutation::gaussian_mutation_per_row;
use crate::ops::replacement::{mu_comma_lambda, mu_plus_lambda};
use crate::ops::selection::argmax_host;
use crate::rng::{SeedPurpose, seed_stream};
use crate::strategy::{Strategy, StrategyMetrics};

/// Which selection scheme the ES uses.
#[derive(Debug, Clone, Copy)]
pub enum EsKind {
    /// `(1+1)` with 1/5-rule σ adaptation.
    OnePlusOne,
    /// `(1+λ)` with shared σ across offspring.
    OnePlusLambda { lambda: usize },
    /// `(μ,λ)` with log-normal per-individual σ adaptation.
    MuCommaLambda { mu: usize, lambda: usize },
    /// `(μ+λ)` with log-normal per-individual σ adaptation.
    MuPlusLambda { mu: usize, lambda: usize },
}

impl EsKind {
    /// Returns the effective offspring-population size for this variant.
    #[must_use]
    pub fn population_size(&self) -> usize {
        match self {
            EsKind::OnePlusOne => 1,
            EsKind::OnePlusLambda { lambda }
            | EsKind::MuCommaLambda { lambda, .. }
            | EsKind::MuPlusLambda { lambda, .. } => *lambda,
        }
    }
}

/// Default σ floor for the self-adaptive step size (see
/// [`EsConfig::sigma_min`]).
const DEFAULT_SIGMA_MIN: f32 = 1e-8;
/// Default σ ceiling for the self-adaptive step size (see
/// [`EsConfig::sigma_max`]).
const DEFAULT_SIGMA_MAX: f32 = 1e6;

/// Static configuration for an [`EvolutionStrategy`] run.
#[derive(Debug, Clone)]
pub struct EsConfig {
    /// Variant to run.
    pub kind: EsKind,
    /// Genome dimensionality.
    pub genome_dim: usize,
    /// Search-space bounds; used for initialization and clamping.
    pub bounds: Bounds,
    /// Initial σ (log-normal self-adaptation modifies it in state).
    pub initial_sigma: f32,
    /// Lower clamp for the self-adaptive σ.
    ///
    /// Both the log-normal update `σ' = σ · exp(τ · N(0,1))` (multi-parent
    /// variants) and the Rechenberg 1/5-rule (`(1+1)`) are unbounded
    /// multiplicative processes; without a floor σ can underflow toward `0`,
    /// collapsing the mutation amplitude so the search freezes. Must be
    /// strictly positive and `< sigma_max`. Default `DEFAULT_SIGMA_MIN`.
    pub sigma_min: f32,
    /// Upper clamp for the self-adaptive σ.
    ///
    /// Without a ceiling σ can overflow toward `+∞` (genes then saturate to a
    /// bound with no error). Default `DEFAULT_SIGMA_MAX` — far outside any
    /// practical step scale on the `[-5.12, 5.12]` benchmark domain, so it
    /// never binds in normal operation and only catches a runaway process.
    pub sigma_max: f32,
    /// Learning-rate scale for log-normal σ update. Standard default is
    /// `1.0 / sqrt(2 * sqrt(D))`.
    pub tau: f32,
}

impl EsConfig {
    /// Default configuration for a given ES variant and dimensionality.
    ///
    /// Sets `bounds = (-5.12, 5.12)` (the standard Rastrigin/sphere domain),
    /// `initial_sigma = 1.0`, and τ via the standard formula
    /// `1 / sqrt(2 · sqrt(D))` (Beyer & Schwefel 2002, eq. 12).
    #[must_use]
    pub fn default_for(kind: EsKind, genome_dim: usize) -> Self {
        #[allow(clippy::cast_precision_loss)]
        let d = genome_dim as f32;
        let tau = 1.0 / (2.0 * d.sqrt()).sqrt();
        Self {
            kind,
            genome_dim,
            bounds: Bounds::new(-5.12, 5.12),
            initial_sigma: 1.0,
            sigma_min: DEFAULT_SIGMA_MIN,
            sigma_max: DEFAULT_SIGMA_MAX,
            tau,
        }
    }
}

impl Validate for EsConfig {
    fn validate(&self) -> Result<(), ConfigError> {
        const C: &str = "EsConfig";
        config::nonzero(C, "genome_dim", self.genome_dim)?;
        config::positive(C, "initial_sigma", f64::from(self.initial_sigma))?;
        config::positive(C, "sigma_min", f64::from(self.sigma_min))?;
        config::ordered(
            C,
            "sigma_max",
            f64::from(self.sigma_min),
            f64::from(self.sigma_max),
        )?;
        config::positive(C, "tau", f64::from(self.tau))?;
        match self.kind {
            EsKind::OnePlusOne => {}
            EsKind::OnePlusLambda { lambda } => {
                config::at_least(C, "lambda", lambda, 1)?;
            }
            EsKind::MuPlusLambda { mu, lambda } => {
                config::at_least(C, "mu", mu, 1)?;
                config::at_least(C, "lambda", lambda, 1)?;
            }
            EsKind::MuCommaLambda { mu, lambda } => {
                config::at_least(C, "mu", mu, 1)?;
                config::at_least(C, "lambda", lambda, 1)?;
                if lambda < mu {
                    return Err(ConfigError {
                        config: C,
                        field: "lambda",
                        kind: ConstraintKind::Custom("(mu, lambda) requires lambda >= mu"),
                    });
                }
            }
        }
        Ok(())
    }
}

/// Generation state for [`EvolutionStrategy`].
#[derive(Debug, Clone)]
pub struct EsState<B: Backend> {
    /// Parent population. `(μ, D)` for μ-parent variants; `(1, D)` for
    /// (1+1) and (1+λ).
    parents: Tensor<B, 2>,
    /// Per-parent σ values.
    ///
    /// Shape between generations is `(μ,)` for log-normal adaptation and
    /// `(1,)` for `(1+1)`/`(1+λ)` with shared σ. Between an `ask` and the
    /// matching `tell` the tensor is temporarily `(μ + λ,)`: parent σ
    /// followed by per-offspring σ. See `ask` for the rationale.
    sigmas: Tensor<B, 1>,
    /// Parent fitnesses.
    parent_fitness: Vec<f32>,
    /// Best-so-far genome, shape `(1, D)`.
    best_genome: Option<Tensor<B, 2>>,
    /// Best-so-far fitness.
    best_fitness: f32,
    /// Completed-generation counter.
    generation: usize,
    /// (1+1) only: running success-rate counter for the 1/5th rule.
    successes_in_window: u32,
    /// (1+1) only: window length observed so far.
    window_len: u32,
}

impl<B: Backend> EsState<B> {
    /// Assembles an ES state, checking the parent fitness cache matches the
    /// parent count.
    ///
    /// # Errors
    ///
    /// Returns a [`ConfigError`] if `parents` has zero rows or if
    /// `parent_fitness` is non-empty with a length other than the parent
    /// count `μ` (`parents.dims()[0]`). The bootstrap state (empty
    /// `parent_fitness`) is accepted.
    #[allow(clippy::too_many_arguments)]
    pub fn try_new(
        parents: Tensor<B, 2>,
        sigmas: Tensor<B, 1>,
        parent_fitness: Vec<f32>,
        best_genome: Option<Tensor<B, 2>>,
        best_fitness: f32,
        generation: usize,
        successes_in_window: u32,
        window_len: u32,
    ) -> Result<Self, ConfigError> {
        let mu = parents.dims()[0];
        config::nonzero("EsState", "parents", mu)?;
        if !parent_fitness.is_empty() && parent_fitness.len() != mu {
            return Err(ConfigError {
                config: "EsState",
                field: "parent_fitness",
                kind: ConstraintKind::Custom("length must equal the parent count μ"),
            });
        }
        Ok(Self {
            parents,
            sigmas,
            parent_fitness,
            best_genome,
            best_fitness,
            generation,
            successes_in_window,
            window_len,
        })
    }

    /// Parent population, shape `(μ, D)` (or `(1, D)` for `(1+1)`/`(1+λ)`).
    #[must_use]
    pub fn parents(&self) -> &Tensor<B, 2> {
        &self.parents
    }

    /// Per-parent σ values (see the field docs for the transient `(μ + λ,)`
    /// shape held between `ask` and `tell`).
    #[must_use]
    pub fn sigmas(&self) -> &Tensor<B, 1> {
        &self.sigmas
    }

    /// Parent fitnesses (empty at bootstrap, else `μ` long).
    #[must_use]
    pub fn parent_fitness(&self) -> &[f32] {
        &self.parent_fitness
    }

    /// Best-so-far genome (shape `(1, D)`), or `None` before the first `tell`.
    #[must_use]
    pub fn best_genome(&self) -> Option<&Tensor<B, 2>> {
        self.best_genome.as_ref()
    }

    /// Best-so-far (canonical, maximise) fitness.
    #[must_use]
    pub fn best_fitness(&self) -> f32 {
        self.best_fitness
    }

    /// Completed-generation counter.
    #[must_use]
    pub fn generation(&self) -> usize {
        self.generation
    }

    /// `(1+1)` only: running success count for the 1/5th rule.
    #[must_use]
    pub fn successes_in_window(&self) -> u32 {
        self.successes_in_window
    }

    /// `(1+1)` only: window length observed so far.
    #[must_use]
    pub fn window_len(&self) -> u32 {
        self.window_len
    }
}

/// Classical Evolution Strategy.
///
/// # Example
///
/// ```no_run
/// use burn::backend::Flex;
/// use rlevo_evolution::algorithms::es_classical::{EsConfig, EsKind, EvolutionStrategy};
///
/// let strategy = EvolutionStrategy::<Flex>::new();
/// let params = EsConfig::default_for(EsKind::MuPlusLambda { mu: 5, lambda: 20 }, 10);
/// let _ = (strategy, params);
/// ```
#[derive(Debug, Clone, Copy, Default)]
pub struct EvolutionStrategy<B: Backend> {
    _backend: PhantomData<fn() -> B>,
}

impl<B: Backend> EvolutionStrategy<B> {
    /// Builds a new (stateless) strategy object.
    #[must_use]
    pub fn new() -> Self {
        Self {
            _backend: PhantomData,
        }
    }

    fn mu(kind: EsKind) -> usize {
        match kind {
            EsKind::OnePlusOne | EsKind::OnePlusLambda { .. } => 1,
            EsKind::MuCommaLambda { mu, .. } | EsKind::MuPlusLambda { mu, .. } => mu,
        }
    }

    fn sample_initial_parents(
        params: &EsConfig,
        rng: &mut dyn Rng,
        device: &<B as burn::tensor::backend::BackendTypes>::Device,
    ) -> (Tensor<B, 2>, Tensor<B, 1>) {
        let mu = Self::mu(params.kind);
        let (lo, hi): (f32, f32) = params.bounds.into();
        // Host-sample the initial parents from a deterministic `seed_stream`
        // rather than the process-wide Flex RNG (`B::seed` + `Tensor::random`),
        // whose draws interleave with sibling tests under the parallel runner
        // and are not reproducible across thread schedules.
        let genome_dim = params.genome_dim;
        let mut stream = seed_stream(rng.next_u64(), 0, SeedPurpose::Init);
        let mut parent_rows = Vec::with_capacity(mu * genome_dim);
        for _ in 0..mu * genome_dim {
            parent_rows.push(lo + (hi - lo) * stream.random::<f32>());
        }
        let parents =
            Tensor::<B, 2>::from_data(TensorData::new(parent_rows, [mu, genome_dim]), device);
        let sigmas = Tensor::<B, 1>::from_data(
            TensorData::new(vec![params.initial_sigma; mu], [mu]),
            device,
        );
        (parents, sigmas)
    }
}

impl<B: Backend> Strategy<B> for EvolutionStrategy<B>
where
    B::Device: Clone,
{
    type Params = EsConfig;
    type State = EsState<B>;
    type Genome = Tensor<B, 2>;

    /// Samples the initial parent population uniformly from `params.bounds`
    /// via a deterministic `seed_stream` (host-RNG convention) and
    /// initializes all σ values to `params.initial_sigma`.
    fn init(
        &self,
        params: &EsConfig,
        rng: &mut dyn Rng,
        device: &<B as burn::tensor::backend::BackendTypes>::Device,
    ) -> EsState<B> {
        debug_assert!(
            params.validate().is_ok(),
            "invalid EsConfig reached init: {params:?}"
        );
        let (parents, sigmas) = Self::sample_initial_parents(params, rng, device);
        EsState {
            parents,
            sigmas,
            parent_fitness: Vec::new(),
            best_genome: None,
            best_fitness: f32::NEG_INFINITY,
            generation: 0,
            successes_in_window: 0,
            window_len: 0,
        }
    }

    /// Generates the offspring population for the current generation.
    ///
    /// On the very first call (before any `tell`), returns the initial parents
    /// unchanged so that they can be fitness-evaluated as the seed population.
    /// On subsequent calls, duplicates parents by uniform random selection,
    /// applies log-normal σ adaptation (multi-parent variants) or inherits the
    /// shared σ (`(1+1)` / `(1+λ)`), then mutates via per-individual Gaussian
    /// noise. All stochastic draws go through `seed_stream`
    /// (host-RNG convention); offspring σ values are appended to
    /// `state.sigmas` for consumption by `tell`.
    fn ask(
        &self,
        params: &EsConfig,
        state: &EsState<B>,
        rng: &mut dyn Rng,
        device: &<B as burn::tensor::backend::BackendTypes>::Device,
    ) -> (Tensor<B, 2>, EsState<B>) {
        // First call: evaluate the initial parents as the "offspring"
        // so fitness is populated in the subsequent `tell`.
        if state.parent_fitness.is_empty() {
            return (state.parents.clone(), state.clone());
        }

        let lambda = params.kind.population_size();
        let mu = Self::mu(params.kind);

        let mut mutation_rng = seed_stream(
            rng.next_u64(),
            state.generation as u64,
            SeedPurpose::Mutation,
        );
        let mut sigma_rng =
            seed_stream(rng.next_u64(), state.generation as u64, SeedPurpose::Other);

        // Build an offspring population of size λ by sampling a parent
        // index per offspring and mutating. Uniform random parent
        // selection — no fitness pressure applied at this stage in
        // classical ES; survivor selection provides the pressure.
        let mut parent_indices: Vec<i64> = Vec::with_capacity(lambda);
        for _ in 0..lambda {
            #[allow(clippy::cast_possible_wrap)]
            parent_indices.push(sigma_rng.random_range(0..mu) as i64);
        }
        let idx_tensor = Tensor::<B, 1, burn::tensor::Int>::from_data(
            TensorData::new(parent_indices.clone(), [lambda]),
            device,
        );
        let duplicated_parents = state.parents.clone().select(0, idx_tensor.clone());
        let duplicated_sigmas = state.sigmas.clone().select(0, idx_tensor);

        // Apply log-normal σ adaptation (multi-parent case) or keep σ
        // shared (1+1 / 1+λ). Log-normal: σ' = σ * exp(τ · N(0,1)).
        let is_one_plus = matches!(
            params.kind,
            EsKind::OnePlusOne | EsKind::OnePlusLambda { .. }
        );
        let offspring_sigmas = if is_one_plus {
            duplicated_sigmas
        } else {
            // Host-sample the N(0,1) noise from the deterministic `sigma_rng`
            // so the log-normal σ update is reproducible across schedules.
            let mut noise_rows = Vec::with_capacity(lambda);
            for _ in 0..lambda {
                noise_rows.push(crate::sampling::standard_normal(&mut sigma_rng));
            }
            let noise = Tensor::<B, 1>::from_data(TensorData::new(noise_rows, [lambda]), device);
            // Clamp the log-normal random walk to `[sigma_min, sigma_max]` so σ
            // can neither underflow to 0 (search freezes) nor overflow to +∞
            // (genes saturate). Both bounds are construction-validated.
            (duplicated_sigmas * noise.mul_scalar(params.tau).exp())
                .clamp(params.sigma_min, params.sigma_max)
        };

        // Mutate parents by the per-offspring σ, drawing from the host
        // `mutation_rng`.
        let mutated = gaussian_mutation_per_row(
            duplicated_parents,
            offspring_sigmas.clone(),
            &mut mutation_rng,
            device,
        );

        // Clamp to bounds.
        let (lo, hi): (f32, f32) = params.bounds.into();
        let mutated = mutated.clamp(lo, hi);

        let mut state = state.clone();
        // Carry offspring σ to `tell` by appending them to `state.sigmas`.
        // After this point sigmas has shape `(μ + λ,)`: the first μ entries
        // are the unchanged parent σ, the last λ are the per-offspring σ.
        // `tell` slices both halves to align survivor σ with survivor genomes
        // (`(μ+λ)` selection draws from the union, `(μ,λ)` only from the λ
        // offspring slice). Folding the offspring σ into the existing field
        // avoids adding a transient pending-σ field to `EsState`.
        let combined_sigmas = Tensor::cat(vec![state.sigmas.clone(), offspring_sigmas], 0);
        state.sigmas = combined_sigmas;
        (mutated, state)
    }

    /// Applies variant-specific selection and σ adaptation, then returns the
    /// updated state and a per-generation metrics snapshot.
    ///
    /// Variant behaviour:
    /// - `(1+1)`: greedy replacement; σ updated by Rechenberg's 1/5th
    ///   success rule every `10·D` steps.
    /// - `(1+λ)`: best offspring replaces the parent only if it strictly
    ///   improves fitness; σ is carried over unchanged.
    /// - `(μ,λ)`: selects the μ best offspring; parent pool discarded.
    ///   Survivor σ values are gathered by the same truncation indices.
    /// - `(μ+λ)`: selects the μ best of the combined parent + offspring
    ///   pool. Survivor σ values are drawn from the concatenated σ vector
    ///   by the same indices.
    ///
    /// The first `tell` after `init` bootstraps `parent_fitness` from the
    /// initial-population evaluation rather than running selection.
    #[allow(clippy::too_many_lines)]
    fn tell(
        &self,
        params: &EsConfig,
        offspring: Tensor<B, 2>,
        fitness: Tensor<B, 1>,
        mut state: EsState<B>,
        _rng: &mut dyn Rng,
    ) -> (EsState<B>, StrategyMetrics) {
        let fitness_host = fitness
            .into_data()
            .into_vec::<f32>()
            .expect("fitness tensor must be readable as f32");

        // First `tell` after `init`: offspring here is actually the
        // initial parent population evaluated.
        if state.parent_fitness.is_empty() {
            state.parent_fitness.clone_from(&fitness_host);
            state.generation += 1;
            update_best(&mut state, &offspring, &fitness_host);
            let m = StrategyMetrics::from_host_fitness(
                state.generation,
                &fitness_host,
                state.best_fitness,
            );
            state.best_fitness = m.best_fitness_ever();
            state.parents = offspring;
            // Restore parent-count σ vector.
            let mu = Self::mu(params.kind);
            let device = state.parents.device();
            state.sigmas = Tensor::<B, 1>::from_data(
                TensorData::new(vec![params.initial_sigma; mu], [mu]),
                &device,
            );
            return (state, m);
        }

        let device = offspring.device();
        let mu = Self::mu(params.kind);
        // state.sigmas currently holds parent σ concatenated with
        // offspring σ, per `ask`'s scratchpad trick.
        let lambda = params.kind.population_size();
        #[allow(clippy::single_range_in_vec_init)]
        let parent_sigmas = state.sigmas.clone().slice([0..mu]);
        #[allow(clippy::single_range_in_vec_init)]
        let offspring_sigmas = state.sigmas.clone().slice([mu..(mu + lambda)]);

        match params.kind {
            EsKind::OnePlusOne => {
                // One parent, one offspring. Fitness[0] is the offspring.
                let parent_fit = state.parent_fitness[0];
                let offspring_fit = fitness_host[0];
                let success = offspring_fit > parent_fit;
                state.window_len += 1;
                if success {
                    state.successes_in_window += 1;
                    state.parents.clone_from(&offspring);
                    state.parent_fitness = vec![offspring_fit];
                }
                // Rechenberg 1/5-rule every 10 · D generations.
                #[allow(clippy::cast_precision_loss, clippy::cast_possible_truncation)]
                let window = 10_u32.saturating_mul(params.genome_dim as u32).max(1);
                if state.window_len >= window {
                    #[allow(clippy::cast_precision_loss)]
                    let rate = state.successes_in_window as f32 / state.window_len as f32;
                    let current_sigma = state
                        .sigmas
                        .clone()
                        .into_data()
                        .into_vec::<f32>()
                        .expect("sigma tensor must be readable as f32")[0];
                    // The 1/5-rule is also an unbounded multiplicative process;
                    // clamp to the same construction-validated window so σ can
                    // neither underflow to 0 nor overflow to +∞ over a long run.
                    let new_sigma = if rate > 0.2 {
                        current_sigma * 1.22
                    } else if rate < 0.2 {
                        current_sigma / 1.22
                    } else {
                        current_sigma
                    }
                    .clamp(params.sigma_min, params.sigma_max);
                    state.sigmas =
                        Tensor::<B, 1>::from_data(TensorData::new(vec![new_sigma], [1]), &device);
                    state.successes_in_window = 0;
                    state.window_len = 0;
                } else {
                    state.sigmas = parent_sigmas;
                }
            }
            EsKind::OnePlusLambda { .. } => {
                // Best of (parent, offspring pool).
                let best_off_idx = argmax_host(&fitness_host);
                let best_off_fit = fitness_host[best_off_idx];
                if best_off_fit > state.parent_fitness[0] {
                    #[allow(clippy::single_range_in_vec_init)]
                    let best_row = offspring.clone().slice([best_off_idx..best_off_idx + 1]);
                    state.parents = best_row;
                    state.parent_fitness = vec![best_off_fit];
                }
                state.sigmas = parent_sigmas;
            }
            EsKind::MuCommaLambda { mu, .. } => {
                let (survivors, survivor_f) =
                    mu_comma_lambda::<B>(offspring.clone(), &fitness_host, mu, &device);
                // Gather survivor σs matching the same indices.
                let survivor_idx =
                    crate::ops::selection::truncation_indices_host(&fitness_host, mu);
                let survivor_sigmas = offspring_sigmas.select(
                    0,
                    Tensor::<B, 1, burn::tensor::Int>::from_data(
                        TensorData::new(survivor_idx, [mu]),
                        &device,
                    ),
                );
                state.parents = survivors;
                state.parent_fitness = survivor_f;
                state.sigmas = survivor_sigmas;
            }
            EsKind::MuPlusLambda { mu, .. } => {
                let (survivors, survivor_f) = mu_plus_lambda::<B>(
                    state.parents.clone(),
                    &state.parent_fitness,
                    offspring.clone(),
                    &fitness_host,
                    mu,
                    &device,
                );
                // Survivor σ via truncation_indices_host on the combined fitness.
                let combined_f: Vec<f32> = state
                    .parent_fitness
                    .iter()
                    .chain(fitness_host.iter())
                    .copied()
                    .collect();
                let survivor_idx = crate::ops::selection::truncation_indices_host(&combined_f, mu);
                let combined_sigmas = Tensor::cat(vec![parent_sigmas, offspring_sigmas], 0);
                let survivor_sigmas = combined_sigmas.select(
                    0,
                    Tensor::<B, 1, burn::tensor::Int>::from_data(
                        TensorData::new(survivor_idx, [mu]),
                        &device,
                    ),
                );
                state.parents = survivors;
                state.parent_fitness = survivor_f;
                state.sigmas = survivor_sigmas;
            }
        }

        state.generation += 1;
        update_best(&mut state, &offspring, &fitness_host);
        let m =
            StrategyMetrics::from_host_fitness(state.generation, &fitness_host, state.best_fitness);
        state.best_fitness = m.best_fitness_ever();
        (state, m)
    }

    /// Returns the best-so-far genome and its fitness, or `None` before the
    /// first `tell` call.
    fn best(&self, state: &EsState<B>) -> Option<(Tensor<B, 2>, f32)> {
        state
            .best_genome
            .as_ref()
            .map(|g| (g.clone(), state.best_fitness))
    }
}

fn update_best<B: Backend>(state: &mut EsState<B>, pop: &Tensor<B, 2>, fitness: &[f32]) {
    if fitness.is_empty() {
        return;
    }
    let best_idx = argmax_host(fitness);
    let best_f = fitness[best_idx];
    if best_f > state.best_fitness {
        let device = pop.device();
        #[allow(clippy::cast_possible_wrap)]
        let idx = Tensor::<B, 1, burn::tensor::Int>::from_data(
            TensorData::new(vec![best_idx as i64], [1]),
            &device,
        );
        state.best_genome = Some(pop.clone().select(0, idx));
        state.best_fitness = best_f;
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::fitness::FromFitnessEvaluable;
    use crate::strategy::EvolutionaryHarness;
    use burn::backend::Flex;
    use rlevo_core::fitness::FitnessEvaluable;
    type TestBackend = Flex;

    #[test]
    fn try_new_checks_parent_fitness_length() {
        let device = Default::default();
        let parents = Tensor::<TestBackend, 2>::zeros([4, 2], &device);
        let sigmas = Tensor::<TestBackend, 1>::ones([4], &device);
        // Bootstrap (empty) and fully-populated caches both accept.
        assert!(
            EsState::try_new(
                parents.clone(),
                sigmas.clone(),
                vec![],
                None,
                f32::MIN,
                0,
                0,
                0
            )
            .is_ok()
        );
        assert!(
            EsState::try_new(
                parents.clone(),
                sigmas.clone(),
                vec![1.0; 4],
                None,
                1.0,
                1,
                0,
                0,
            )
            .is_ok()
        );
        // parent_fitness length 3 ≠ μ = 4.
        assert!(EsState::try_new(parents, sigmas, vec![1.0; 3], None, 1.0, 1, 0, 0).is_err());
    }

    #[test]
    fn default_config_validates() {
        let cfg = EsConfig::default_for(EsKind::MuPlusLambda { mu: 5, lambda: 20 }, 10);
        assert!(cfg.validate().is_ok());
    }

    #[test]
    fn rejects_comma_lambda_below_mu() {
        let cfg = EsConfig::default_for(EsKind::MuCommaLambda { mu: 10, lambda: 5 }, 10);
        assert_eq!(cfg.validate().unwrap_err().field, "lambda");
    }

    /// `genome_dim == 0` makes `tau = 1/sqrt(2·sqrt(0)) = +∞`; the config guard
    /// must reject it at construction (ADR 0026) so the non-finite τ never
    /// reaches the first `ask` (issue #132, `es_classical` §1.1 / `ep` §1.2).
    #[test]
    fn rejects_zero_genome_dim() {
        let cfg = EsConfig::default_for(EsKind::MuPlusLambda { mu: 5, lambda: 20 }, 0);
        assert!(
            !cfg.tau.is_finite(),
            "precondition: derived tau is non-finite for genome_dim == 0, got {}",
            cfg.tau
        );
        assert_eq!(
            cfg.validate().unwrap_err().field,
            "genome_dim",
            "genome_dim == 0 must be rejected before the non-finite tau can be used"
        );
    }

    /// An inverted σ window (`sigma_min >= sigma_max`) is rejected so the clamp
    /// bounds are always a valid interval (`es_classical` §1.1).
    #[test]
    fn rejects_inverted_sigma_window() {
        let mut cfg = EsConfig::default_for(EsKind::MuPlusLambda { mu: 5, lambda: 20 }, 10);
        cfg.sigma_min = 10.0;
        cfg.sigma_max = 1.0;
        assert_eq!(
            cfg.validate().unwrap_err().field,
            "sigma_max",
            "sigma_min >= sigma_max must be rejected"
        );
    }

    /// The log-normal σ of a multi-parent variant stays inside
    /// `[sigma_min, sigma_max]` across many generations even under an aggressive
    /// `tau` that would otherwise drive the walk to `0` or `+∞`
    /// (`es_classical` §1.1). Drives the strategy directly so the transient
    /// `(μ + λ,)` σ vector produced by `ask` is inspected.
    #[test]
    fn sigma_stays_within_bounds_across_updates() {
        use rand::SeedableRng;
        use rand::rngs::StdRng;

        let device = Default::default();
        let strategy = EvolutionStrategy::<TestBackend>::new();
        let mut params = EsConfig::default_for(EsKind::MuPlusLambda { mu: 4, lambda: 12 }, 3);
        params.tau = 5.0;
        params.sigma_min = 1e-4;
        params.sigma_max = 10.0;
        assert!(params.validate().is_ok(), "test config must be valid");

        let mut rng = StdRng::seed_from_u64(9);
        let mut state = strategy.init(&params, &mut rng, &device);
        for generation in 0..60 {
            let (offspring, next) = strategy.ask(&params, &state, &mut rng, &device);
            let sigmas: Vec<f32> = next
                .sigmas()
                .clone()
                .into_data()
                .into_vec::<f32>()
                .expect("sigma host-read of a tensor this test just built");
            for &s in &sigmas {
                assert!(
                    s.is_finite() && s >= params.sigma_min && s <= params.sigma_max,
                    "σ left [{}, {}] at gen {generation}: {s}",
                    params.sigma_min,
                    params.sigma_max
                );
            }
            let n = offspring.dims()[0];
            let fitness = Tensor::<TestBackend, 1>::from_data(
                TensorData::new(vec![1.0_f32; n], [n]),
                &device,
            );
            let (advanced, _) = strategy.tell(&params, offspring, fitness, next, &mut rng);
            state = advanced;
        }
    }

    /// Both the plain argmax and the truncation selector used by ES survivor
    /// selection must return valid, in-range, pairwise-distinct indices even
    /// when the fitness slice carries `NaN` values (fitness-hygiene: a `NaN`
    /// ranks as worst and never as a survivor). Guards against an
    /// out-of-bounds gather in `tell` (`es_classical` §7.2, valid index).
    #[test]
    fn selection_returns_in_range_indices() {
        let fitness = [1.0_f32, f32::NAN, 5.0, 2.0, -3.0, f32::NAN];
        let n = fitness.len();

        let amax = argmax_host(&fitness);
        assert!(amax < n, "argmax index {amax} out of range for len {n}");

        for mu in 1..=4 {
            let idx = crate::ops::selection::truncation_indices_host(&fitness, mu);
            assert_eq!(idx.len(), mu, "truncation must return exactly mu indices");
            for (a, &x) in idx.iter().enumerate() {
                assert!(
                    usize::try_from(x).is_ok_and(|xi| xi < n),
                    "truncation index {x} out of range for len {n}"
                );
                for &y in &idx[a + 1..] {
                    assert_ne!(x, y, "truncation indices must be pairwise distinct");
                }
            }
        }
    }

    /// Canonical (maximise) fitness `−Σ xᵢ²` read straight off a genome tensor,
    /// so tests can drive a strategy directly without the harness.
    fn neg_sphere(pop: &Tensor<TestBackend, 2>) -> Tensor<TestBackend, 1> {
        let device = pop.device();
        let [n, d] = pop.dims();
        let rows: Vec<f32> = pop
            .clone()
            .into_data()
            .into_vec::<f32>()
            .expect("population host-read of a tensor this test just built");
        #[allow(clippy::needless_range_loop)]
        let fit: Vec<f32> = (0..n)
            .map(|i| -(0..d).map(|j| rows[i * d + j].powi(2)).sum::<f32>())
            .collect();
        Tensor::<TestBackend, 1>::from_data(TensorData::new(fit, [n]), &device)
    }

    /// Drives an ES variant directly for `gens` generations against the
    /// canonical maximise `−sphere`, returning the `best_fitness_ever`
    /// trajectory reported by each `tell`.
    fn run_es_best_ever(kind: EsKind, dim: usize, gens: usize, seed: u64) -> Vec<f32> {
        use rand::SeedableRng;
        use rand::rngs::StdRng;

        let device = Default::default();
        let strategy = EvolutionStrategy::<TestBackend>::new();
        let params = EsConfig::default_for(kind, dim);
        let mut rng = StdRng::seed_from_u64(seed);
        let mut state = strategy.init(&params, &mut rng, &device);
        let mut traj = Vec::with_capacity(gens);
        for _ in 0..gens {
            let (offspring, next) = strategy.ask(&params, &state, &mut rng, &device);
            let fitness = neg_sphere(&offspring);
            let (advanced, m) = strategy.tell(&params, offspring, fitness, next, &mut rng);
            traj.push(m.best_fitness_ever());
            state = advanced;
        }
        traj
    }

    /// `best_fitness_ever` is a rolling maximum in canonical space, so on a
    /// maximise problem it must never decrease from one generation to the next
    /// across every ES variant (`es_classical` §7.2, monotone best-ever). This
    /// pins the invariant that the algorithm threads the rolling best through
    /// `tell` and never resets it — including the `(μ,λ)` variant that discards
    /// its parent pool each generation.
    #[test]
    fn best_fitness_ever_is_monotonic_on_maximize() {
        for kind in [
            EsKind::OnePlusOne,
            EsKind::OnePlusLambda { lambda: 6 },
            EsKind::MuPlusLambda { mu: 3, lambda: 8 },
            EsKind::MuCommaLambda { mu: 3, lambda: 8 },
        ] {
            let traj = run_es_best_ever(kind, 3, 40, 17);
            for w in traj.windows(2) {
                assert!(
                    w[1] >= w[0],
                    "best_fitness_ever decreased for {kind:?}: {} -> {}",
                    w[0],
                    w[1]
                );
            }
            assert!(
                traj.last().copied().unwrap().is_finite(),
                "rolling best must stay finite for {kind:?}"
            );
        }
    }

    struct Sphere;
    struct SphereFit;
    impl FitnessEvaluable for SphereFit {
        type Individual = Vec<f64>;
        type Landscape = Sphere;
        fn evaluate(&self, x: &Self::Individual, _: &Self::Landscape) -> f64 {
            x.iter().map(|v| v * v).sum()
        }
    }

    fn run_es(kind: EsKind, dim: usize, generations: usize, seed: u64) -> f32 {
        let device = Default::default();
        let strategy = EvolutionStrategy::<TestBackend>::new();
        let params = EsConfig::default_for(kind, dim);
        let fitness_fn = FromFitnessEvaluable::new(SphereFit, Sphere);
        let mut harness = EvolutionaryHarness::<TestBackend, _, _>::new(
            strategy,
            params,
            fitness_fn,
            seed,
            device,
            generations,
        )
        .expect("valid params");
        harness.reset();
        loop {
            let step = harness.step(());
            if step.done {
                break;
            }
        }
        harness.latest_metrics().unwrap().best_fitness_ever()
    }

    #[test]
    fn one_plus_lambda_converges_on_sphere_d2() {
        let best = run_es(EsKind::OnePlusLambda { lambda: 8 }, 2, 200, 7);
        assert!(best < 1e-2, "OnePlusLambda best={best}");
    }

    #[test]
    fn one_plus_one_converges_on_sphere_d2() {
        let best = run_es(EsKind::OnePlusOne, 2, 500, 11);
        assert!(best < 1e-2, "OnePlusOne best={best}");
    }

    #[test]
    fn mu_plus_lambda_converges_on_sphere_d2() {
        let best = run_es(EsKind::MuPlusLambda { mu: 3, lambda: 8 }, 2, 200, 7);
        assert!(best < 1e-2, "MuPlusLambda best={best}");
    }

    #[test]
    fn mu_comma_lambda_converges_on_sphere_d2() {
        let best = run_es(EsKind::MuCommaLambda { mu: 3, lambda: 8 }, 2, 200, 7);
        assert!(best < 1e-1, "MuCommaLambda best={best}");
    }

    #[test]
    fn mu_plus_lambda_converges_on_sphere_d10() {
        // Convergence on Sphere (D=10) to best_fitness < 1e-6 within
        // budget on Flex. We allow a generous budget because the
        // classical ES is slower than CMA-ES; the goal is to verify
        // convergence direction, not to optimize hyperparameters.
        let best = run_es(EsKind::MuPlusLambda { mu: 5, lambda: 20 }, 10, 1500, 42);
        assert!(best < 1e-6, "MuPlusLambda D10 best={best}");
    }
}