gentle-ga 0.1.0

//! Gentle, ergonomic GA framework with strong typing and static dispatch.
//!
//! Design goals:
//! - Minimal trait surface for users: implement `Problem` for your genome.
//! - Sensible defaults: mutate/crossover/repair are optional.
//! - Configurable selection policy (default tournament).
//! - Clear separation of concerns: Problem = domain logic; Engine = orchestration.
//! - Deterministic runs via seed; no RNG state stored in `Problem`.

use rand::{rngs::StdRng, Rng, SeedableRng};
use std::cmp::Ordering;

/// Direction of optimization for fitness comparison.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum Optimize {
    Maximize,
    Minimize,
}

/// A genome with its evaluated fitness.
#[derive(Clone, Debug)]
pub struct Evaluated<G, F> {
    pub genome: G,
    pub fitness: F,
}

/// Core problem definition. Implement this to adapt the GA to your domain.
///
/// You must supply:
/// - `random_genome` to create initial individuals
/// - `fitness` to evaluate individuals
///
/// You may optionally override:
/// - `mutate` (default: no-op)
/// - `crossover` (default: clone `a`)
/// - `repair` (default: no-op)
///
/// Methods accept an `rng` so all randomness flows from the engine, making
/// runs deterministic with a given seed.
pub trait Problem {
    type Genome: Clone + Send + Sync + 'static;
    type Fitness: PartialOrd + Copy + Send + Sync + 'static;

    fn random_genome<R: Rng + ?Sized>(&mut self, rng: &mut R) -> Self::Genome;

    fn fitness(&self, genome: &Self::Genome) -> Self::Fitness;

    fn mutate<R: Rng + ?Sized>(&mut self, _genome: &mut Self::Genome, _rng: &mut R) {
        // default no-op
    }

    fn crossover<R: Rng + ?Sized>(
        &mut self,
        a: &Self::Genome,
        _b: &Self::Genome,
        _rng: &mut R,
    ) -> Self::Genome {
        // default: clone a
        a.clone()
    }

    fn repair(&mut self, _genome: &mut Self::Genome) {
        // default no-op
    }
}

/// Selection strategy for choosing parents from an evaluated population.
///
/// Implementors can be stateless or carry configuration (e.g., tournament size).
pub trait Selection<P: Problem> {
    fn select<'a, R: Rng + ?Sized>(
        &self,
        rng: &mut R,
        population: &'a [Evaluated<P::Genome, P::Fitness>],
        optimize: Optimize,
    ) -> (&'a P::Genome, &'a P::Genome);
}

/// Default selection: tournament with configurable size `k` (>= 2).
#[derive(Clone, Copy, Debug)]
pub struct TournamentSelection {
    pub k: usize,
}

impl Default for TournamentSelection {
    fn default() -> Self {
        Self { k: 3 }
    }
}

impl<P: Problem> Selection<P> for TournamentSelection {
    fn select<'a, R: Rng + ?Sized>(
        &self,
        rng: &mut R,
        population: &'a [Evaluated<P::Genome, P::Fitness>],
        optimize: Optimize,
    ) -> (&'a P::Genome, &'a P::Genome) {
        debug_assert!(population.len() >= 2, "population must have at least 2 individuals");
        let pick_one = |rng: &mut R| {
            // Sample `k` competitors uniformly with replacement, pick the best by fitness.
            // For small populations, k will clamp to population size.
            let k = self.k.max(1).min(population.len());
            let mut best_idx = rng.gen_range(0..population.len());
            for _ in 1..k {
                let idx = rng.gen_range(0..population.len());
                best_idx = better_index(population, best_idx, idx, optimize);
            }
            &population[best_idx].genome
        };
        (pick_one(rng), pick_one(rng))
    }
}

/// Overall configuration for a GA run.
#[derive(Clone, Debug)]
pub struct Config {
    pub pop_size: usize,
    pub optimize: Optimize,
    pub crossover_rate: f64, // probability to apply crossover (0..=1)
    pub mutation_rate: f64,  // probability to apply mutation (0..=1)
    pub elitism: usize,      // number of best individuals to carry over
    pub seed: Option<u64>,   // deterministic seed; None => from entropy
}

impl Default for Config {
    fn default() -> Self {
        Self {
            pop_size: 100,
            optimize: Optimize::Maximize,
            crossover_rate: 0.9,
            mutation_rate: 0.1,
            elitism: 1,
            seed: None,
        }
    }
}

/// Builder for GA with ergonomic configuration and sensible defaults.
pub struct Builder<P: Problem, S: Selection<P> = TournamentSelection> {
    problem: P,
    config: Config,
    selection: S,
}

impl<P: Problem> Builder<P, TournamentSelection> {
    pub fn new(problem: P) -> Self {
        Self {
            problem,
            config: Config::default(),
            selection: TournamentSelection::default(),
        }
    }
}

impl<P: Problem, S: Selection<P>> Builder<P, S> {
    pub fn pop_size(mut self, n: usize) -> Self {
        self.config.pop_size = n.max(2);
        self
    }
    pub fn optimize(mut self, opt: Optimize) -> Self {
        self.config.optimize = opt;
        self
    }
    pub fn mutation_rate(mut self, p: f64) -> Self {
        self.config.mutation_rate = clamp01(p);
        self
    }
    pub fn crossover_rate(mut self, p: f64) -> Self {
        self.config.crossover_rate = clamp01(p);
        self
    }
    pub fn elitism(mut self, k: usize) -> Self {
        self.config.elitism = k;
        self
    }
    pub fn seed(mut self, seed: u64) -> Self {
        self.config.seed = Some(seed);
        self
    }
    pub fn selection<S2: Selection<P>>(self, selection: S2) -> Builder<P, S2> {
        Builder {
            problem: self.problem,
            config: self.config,
            selection,
        }
    }
    pub fn build(mut self) -> GA<P, S> {
        let mut rng = match self.config.seed {
            Some(s) => StdRng::seed_from_u64(s),
            None => StdRng::from_entropy(),
        };

        // Initialize and evaluate population
        let mut pop = Vec::with_capacity(self.config.pop_size);
        for _ in 0..self.config.pop_size {
            let g = self.problem.random_genome(&mut rng);
            let f = self.problem.fitness(&g);
            pop.push(Evaluated { genome: g, fitness: f });
        }
        GA {
            problem: self.problem,
            config: self.config,
            selection: self.selection,
            rng,
            population: pop,
            generation: 0,
        }
    }
}

/// Genetic Algorithm engine.
pub struct GA<P: Problem, S: Selection<P> = TournamentSelection> {
    problem: P,
    config: Config,
    selection: S,
    rng: StdRng,
    population: Vec<Evaluated<P::Genome, P::Fitness>>,
    generation: usize,
}

/// Summary of a completed run.
#[derive(Clone, Debug)]
pub struct RunResult<G, F> {
    pub generations: usize,
    pub best: Evaluated<G, F>,
    pub average_fitness: f64,
}

impl<P: Problem, S: Selection<P>> GA<P, S> {
    pub fn builder(problem: P) -> Builder<P, TournamentSelection> {
        Builder::new(problem)
    }

    /// Access the current generation count.
    pub fn generation(&self) -> usize {
        self.generation
    }

    /// Returns a reference to the best individual in the current population.
    pub fn best(&self) -> &Evaluated<P::Genome, P::Fitness> {
        best_of(&self.population, self.config.optimize)
            .expect("population must be non-empty")
    }

    /// Returns the mean fitness of the current population as f64.
    pub fn average_fitness(&self) -> f64
    where
        P::Fitness: Into<f64>,
    {
        let sum: f64 = self.population.iter().map(|e| e.fitness.into()).sum();
        sum / self.population.len() as f64
    }

    /// Advances the population by one generation.
    pub fn step(&mut self) {
        let n = self.population.len();
        let elitism = self.config.elitism.min(n);

        // 1) Elites
        let mut next: Vec<Evaluated<P::Genome, P::Fitness>> = if elitism > 0 {
            let mut sorted = self.population.clone();
            sorted.sort_by(|a, b| cmp_fit(self.config.optimize, &a.fitness, &b.fitness));
            sorted.into_iter().rev().take(elitism).collect()
        } else {
            Vec::with_capacity(n)
        };

        // 2) Offspring until next population full
        while next.len() < n {
            let (pa, pb) = self.selection.select(&mut self.rng, &self.population, self.config.optimize);

            // crossover
            let mut child = if chance(&mut self.rng, self.config.crossover_rate) {
                self.problem.crossover(pa, pb, &mut self.rng)
            } else {
                pa.clone()
            };

            // mutate
            if chance(&mut self.rng, self.config.mutation_rate) {
                self.problem.mutate(&mut child, &mut self.rng);
            }

            // repair and evaluate
            self.problem.repair(&mut child);
            let fit = self.problem.fitness(&child);
            next.push(Evaluated { genome: child, fitness: fit });
        }

        self.population = next;
        self.generation += 1;
    }

    /// Runs for a fixed number of generations and returns a summary.
    pub fn run_for(&mut self, generations: usize) -> RunResult<P::Genome, P::Fitness>
    where
        P::Fitness: Into<f64>,
    {
        for _ in 0..generations {
            self.step();
        }
        RunResult {
            generations: self.generation,
            best: self.best().clone(),
            average_fitness: self.average_fitness(),
        }
    }

    /// Runs until the predicate returns true; evaluated at the start of each generation.
    /// Returns when the predicate returns true or `max_generations` is reached.
    pub fn run_until<F>(
        &mut self,
        max_generations: usize,
        mut stop_when: F,
    ) -> RunResult<P::Genome, P::Fitness>
    where
        F: FnMut(usize, &Evaluated<P::Genome, P::Fitness>) -> bool,
        P::Fitness: Into<f64>,
    {
        for _ in 0..max_generations {
            if stop_when(self.generation, self.best()) {
                break;
            }
            self.step();
        }
        RunResult {
            generations: self.generation,
            best: self.best().clone(),
            average_fitness: self.average_fitness(),
        }
    }

    /// Immutable view of the population.
    pub fn population(&self) -> &[Evaluated<P::Genome, P::Fitness>] {
        &self.population
    }

    /// Mutable access to the underlying problem (e.g., to tweak parameters mid-run).
    pub fn problem_mut(&mut self) -> &mut P {
        &mut self.problem
    }
}

// ---------- Internals / helpers ----------

fn clamp01(p: f64) -> f64 {
    if p.is_nan() {
        0.0
    } else if p < 0.0 {
        0.0
    } else if p > 1.0 {
        1.0
    } else {
        p
    }
}

fn chance<R: Rng + ?Sized>(rng: &mut R, p: f64) -> bool {
    // Use Bernoulli sampling; also clamp to be robust against accidental out-of-range inputs.
    rng.gen_bool(clamp01(p))
}

fn cmp_fit<F: PartialOrd>(optimize: Optimize, a: &F, b: &F) -> Ordering {
    // Handle NaN as Equal to avoid panics; users should ensure Fitness is well-defined.
    let ord = a.partial_cmp(b).unwrap_or(Ordering::Equal);
    match optimize {
        Optimize::Maximize => ord,
        Optimize::Minimize => ord.reverse(),
    }
}

fn best_of<G: Clone, F: PartialOrd>(
    pop: &[Evaluated<G, F>],
    optimize: Optimize,
) -> Option<&Evaluated<G, F>> {
    pop.iter()
        .max_by(|a, b| cmp_fit(optimize, &a.fitness, &b.fitness))
}

fn better_index<G, F: PartialOrd>(
    pop: &[Evaluated<G, F>],
    i: usize,
    j: usize,
    optimize: Optimize,
) -> usize {
    match cmp_fit(optimize, &pop[i].fitness, &pop[j].fitness) {
        Ordering::Less => j,
        _ => i,
    }
}

/// Prelude for convenient imports.
pub mod prelude {
    pub use super::{
        Builder, Config, Evaluated, GA, Optimize, Problem, Selection, TournamentSelection,
    };
}