primordial_opt/

runner.rs

//! GA runner for executing genetic algorithm loops.
//!
//! Provides a high-level `run_ga` function that encapsulates the standard
//! genetic algorithm loop: evaluate → select → reproduce.

use rand::Rng;

use crate::{crossover, mutation, selection, Chromosome};

/// Configuration for the genetic algorithm.
#[derive(Debug, Clone)]
pub struct GaConfig {
    /// Number of chromosomes in the population.
    pub population_size: usize,
    /// Number of generations to run.
    pub generations: usize,
    /// Per-bit mutation probability (0.0 to 1.0).
    pub mutation_rate: f64,
    /// Number of elite chromosomes to preserve (ignored for multi-objective).
    pub elite_count: usize,
    /// Fraction of population used as parent pool (0.0 to 1.0).
    pub parent_fraction: f64,
}

impl Default for GaConfig {
    fn default() -> Self {
        Self {
            population_size: 100,
            generations: 1000,
            mutation_rate: 0.01,
            elite_count: 2,
            parent_fraction: 0.5,
        }
    }
}

impl GaConfig {
    /// Sets the population size.
    #[must_use]
    pub fn population_size(mut self, n: usize) -> Self {
        self.population_size = n;
        self
    }

    /// Sets the number of generations.
    #[must_use]
    pub fn generations(mut self, n: usize) -> Self {
        self.generations = n;
        self
    }

    /// Sets the mutation rate.
    #[must_use]
    pub fn mutation_rate(mut self, r: f64) -> Self {
        self.mutation_rate = r;
        self
    }

    /// Sets the elite count (single-objective only).
    #[must_use]
    pub fn elite_count(mut self, n: usize) -> Self {
        self.elite_count = n;
        self
    }

    /// Sets the parent pool fraction.
    #[must_use]
    pub fn parent_fraction(mut self, f: f64) -> Self {
        self.parent_fraction = f;
        self
    }
}

/// Progress snapshot passed to the callback each generation.
#[derive(Debug)]
pub struct Progress<'a> {
    /// Current generation number (0-indexed).
    pub generation: usize,
    /// Population sorted by selection criteria.
    pub population: &'a [Chromosome],
}

/// Result of running the genetic algorithm.
#[derive(Debug)]
pub struct GaResult {
    /// Final population, sorted by selection criteria.
    pub population: Vec<Chromosome>,
    /// Number of generations completed.
    pub generations_run: usize,
}

/// Runs a genetic algorithm loop.
///
/// # Arguments
///
/// * `config` - GA configuration parameters
/// * `chromosome_size` - Size in bytes for each chromosome
/// * `evaluate` - Fitness function: takes a chromosome, returns fitness values
/// * `on_progress` - Callback invoked each generation with current progress
/// * `rng` - Random number generator
///
/// # Example
///
/// ```ignore
/// let result = run_ga(
///     &GaConfig::default().population_size(100).generations(500),
///     16,
///     |c| vec![fitness(c)],
///     |p| println!("Gen {}: best = {:.4}", p.generation, p.population[0].fitness().unwrap()[0]),
///     &mut rng,
/// );
/// ```
pub fn run_ga<F, P, R>(
    config: &GaConfig,
    chromosome_size: usize,
    mut evaluate: F,
    mut on_progress: P,
    rng: &mut R,
) -> GaResult
where
    F: FnMut(&Chromosome) -> Vec<f64>,
    P: FnMut(&Progress),
    R: Rng,
{
    assert!(config.population_size > 0, "population_size must be > 0");
    assert!(config.generations > 0, "generations must be > 0");

    // Initialize population
    let mut population: Vec<Chromosome> = (0..config.population_size)
        .map(|_| Chromosome::new(chromosome_size, rng))
        .collect();

    for generation in 0..config.generations {
        // Evaluate fitness (skip if cached)
        for c in &mut population {
            if c.fitness().is_none() {
                let fitness = evaluate(c);
                c.set_fitness(fitness);
            }
        }

        // Detect single vs multi-objective
        let is_multi = population
            .first()
            .and_then(|c| c.fitness())
            .is_some_and(|f| f.len() > 1);

        // Select (sorts population by fitness)
        let sorted: Vec<Chromosome> = if is_multi {
            selection::nsga2(&population).collect()
        } else {
            selection::rank(&population).collect()
        };

        // Progress callback
        on_progress(&Progress {
            generation,
            population: &sorted,
        });

        // Last generation - return sorted population
        if generation == config.generations - 1 {
            return GaResult {
                population: sorted,
                generations_run: generation + 1,
            };
        }

        // Reproduce
        let elite_count = if is_multi {
            0
        } else {
            config.elite_count.min(config.population_size)
        };
        let pool_size =
            ((config.population_size as f64 * config.parent_fraction) as usize).max(2);

        let mut next_gen: Vec<Chromosome> = sorted[..elite_count].to_vec();

        while next_gen.len() < config.population_size {
            let p1 = &sorted[rng.random_range(0..pool_size)];
            let p2 = &sorted[rng.random_range(0..pool_size)];

            let (mut c1, mut c2) = crossover::single_point(p1, p2, rng);
            mutation::bit_flip(&mut c1, config.mutation_rate, rng);
            mutation::bit_flip(&mut c2, config.mutation_rate, rng);

            next_gen.push(c1);
            if next_gen.len() < config.population_size {
                next_gen.push(c2);
            }
        }

        population = next_gen;
    }

    GaResult {
        population,
        generations_run: config.generations,
    }
}

#[cfg(test)]
mod tests {
    use rand::SeedableRng;

    use super::*;

    fn rng() -> rand::rngs::StdRng {
        rand::rngs::StdRng::seed_from_u64(42)
    }

    #[test]
    fn run_ga_single_objective() {
        let config = GaConfig::default()
            .population_size(20)
            .generations(10)
            .mutation_rate(0.01);

        let mut last_gen = 0;
        let result = run_ga(
            &config,
            4,
            |c| {
                // Simple fitness: sum of bytes (minimize)
                let sum: u32 = c.data().iter().map(|&b| b as u32).sum();
                vec![sum as f64]
            },
            |p| {
                last_gen = p.generation;
            },
            &mut rng(),
        );

        assert_eq!(result.generations_run, 10);
        assert_eq!(result.population.len(), 20);
        assert_eq!(last_gen, 9); // 0-indexed
        assert!(result.population[0].fitness().is_some());
    }

    #[test]
    fn run_ga_multi_objective() {
        let config = GaConfig::default()
            .population_size(20)
            .generations(5)
            .mutation_rate(0.01);

        let result = run_ga(
            &config,
            4,
            |c| {
                // Two objectives
                let sum: u32 = c.data().iter().map(|&b| b as u32).sum();
                let max = c.data().iter().copied().max().unwrap_or(0) as f64;
                vec![sum as f64, max]
            },
            |_| {},
            &mut rng(),
        );

        assert_eq!(result.generations_run, 5);
        // Multi-objective should have 2 fitness values
        assert_eq!(result.population[0].fitness().map(|f| f.len()), Some(2));
    }

    #[test]
    fn config_builder() {
        let config = GaConfig::default()
            .population_size(50)
            .generations(100)
            .mutation_rate(0.05)
            .elite_count(5)
            .parent_fraction(0.3);

        assert_eq!(config.population_size, 50);
        assert_eq!(config.generations, 100);
        assert!((config.mutation_rate - 0.05).abs() < f64::EPSILON);
        assert_eq!(config.elite_count, 5);
        assert!((config.parent_fraction - 0.3).abs() < f64::EPSILON);
    }

    #[test]
    #[should_panic(expected = "population_size must be > 0")]
    fn run_ga_panics_on_zero_population() {
        let config = GaConfig::default().population_size(0);
        run_ga(&config, 4, |_| vec![0.0], |_| {}, &mut rng());
    }

    #[test]
    #[should_panic(expected = "generations must be > 0")]
    fn run_ga_panics_on_zero_generations() {
        let config = GaConfig::default().generations(0);
        run_ga(&config, 4, |_| vec![0.0], |_| {}, &mut rng());
    }
}