From 8eb3e8946375c681e41203cce372a67f09b772ec Mon Sep 17 00:00:00 2001
From: Rutherther <rutherther@ditigal.xyz>
Date: Sun, 2 Nov 2025 12:23:53 +0100
Subject: [PATCH] feat(lib): add evaluation count to EA stats

---
 codes/eoa_lib/src/evolution.rs   | 90 +++++++++++++++++++++++---------
 codes/eoa_lib/src/replacement.rs | 18 +++++++
 2 files changed, 83 insertions(+), 25 deletions(-)

diff --git a/codes/eoa_lib/src/evolution.rs b/codes/eoa_lib/src/evolution.rs
index 6b52d3278b26e7b6d47280a1c52be581105be5d4..aefd767d96425dc34c7e7836b213dd95f632faa7 100644
--- a/codes/eoa_lib/src/evolution.rs
+++ b/codes/eoa_lib/src/evolution.rs
@@ -6,6 +6,7 @@ use crate::{comparison::BetterThanOperator, crossover::Crossover, fitness::Fitne
 #[derive(Clone, Debug)]
 pub struct EvolutionCandidate<TInput, TResult> {
     pub evaluated_chromosome: EvaluatedChromosome<TInput, TResult>,
+    pub evaluation: usize,
     pub iteration: usize
 }
 
@@ -49,7 +50,7 @@ pub fn evolution_algorithm
     // TODO: termination condition
     iterations: usize,
     rng: &mut dyn RngCore,
-    mut evolution: impl FnMut(
+    mut evolutionary_strategy: impl FnMut(
         usize,
         &EvolutionStats<TChromosome, TResult>,
         &EvaluatedPopulation<TChromosome, TResult>,
@@ -61,32 +62,62 @@ pub fn evolution_algorithm
         &mut TReplacement
     )
 ) -> Result<EvolutionResult<TChromosome, TResult>, Box<dyn Error>> {
-    let mut current_population = initial_population.evaluate(fitness)?;
+    let mut current_evaluation = 0;
 
-    let mut last_best_candidate = EvolutionCandidate {
-        evaluated_chromosome: current_population.best_candidate(better_than).clone(),
-        iteration: 0
-    };
+    let mut last_best_candidate: Option<EvolutionCandidate<TChromosome, TResult>> = None;
     let mut stats: EvolutionStats<TChromosome, TResult> = EvolutionStats {
         best_candidates: vec![]
     };
 
-    for iteration in 0..iterations {
-        // Figure out best candidate and save it if better than last time
-        let best_candidate = current_population.best_candidate(better_than);
-
-        if better_than.better_than(
-            &best_candidate.evaluation,
-            &last_best_candidate.evaluated_chromosome.evaluation
-        ) {
-            stats.best_candidates.push(last_best_candidate);
-
-            last_best_candidate = EvolutionCandidate {
-                evaluated_chromosome: best_candidate.clone(),
-                iteration
-            }
-        }
+    fn get_fitness_fn<TChromosome: Clone, TResult: Clone>(
+        current_evaluation: &mut usize,
+        better_than: &impl BetterThanOperator<TResult>,
+        fitness: &impl FitnessFunction<In = TChromosome, Out = TResult>,
+        current_iteration: &usize,
+        stats: &mut EvolutionStats<TChromosome, TResult>,
+        last_best_candidate: &mut Option<EvolutionCandidate<TChromosome, TResult>>
+    ) -> impl FnMut(&TChromosome) -> TResult {
+      |chromosome| {
+          let evaluation = fitness.fit(chromosome).unwrap();
+
+          if last_best_candidate.is_none() ||
+              better_than.better_than(
+                  &evaluation,
+                  &last_best_candidate.as_ref().unwrap().evaluated_chromosome.evaluation
+              ) {
+                  let previous_best = std::mem::replace(
+                      last_best_candidate,
+                      Some(EvolutionCandidate {
+                      evaluated_chromosome: EvaluatedChromosome {
+                          chromosome: chromosome.clone(),
+                          evaluation: evaluation.clone(),
+                      },
+                      evaluation: *current_evaluation,
+                      iteration: *current_iteration
+                  }));
+
+                  if let Some(previous_best) = previous_best {
+                      stats.best_candidates.push(previous_best);
+                  }
+              }
+          *current_evaluation += 1;
+
+          evaluation
+      }
+    }
 
+    let mut current_population = initial_population.evaluate_mut(
+        &mut get_fitness_fn(
+            &mut current_evaluation,
+            better_than,
+            fitness,
+            &0,
+            &mut stats,
+            &mut last_best_candidate
+        )
+    );
+
+    for iteration in 1..=iterations {
         // Selection
         let parents = selection.select(parents_count, &current_population, better_than, rng).collect::<Vec<_>>();
         let parent_pairings = pairing.pair(&current_population, parents.into_iter());
@@ -99,12 +130,21 @@ pub fn evolution_algorithm
             perturbation.perturb(offspring, rng);
         }
 
-        let evaluated_offsprings = offsprings.evaluate(fitness)?;
+        let evaluated_offsprings = offsprings.evaluate_mut(
+            &mut get_fitness_fn(
+                &mut current_evaluation,
+                better_than,
+                fitness,
+                &iteration,
+                &mut stats,
+                &mut last_best_candidate
+            )
+        );
 
         // Replace
         current_population = replacement.replace(current_population, evaluated_offsprings, better_than, rng);
 
-        evolution(
+        evolutionary_strategy(
             iteration,
             &stats,
             &current_population,
@@ -116,8 +156,8 @@ pub fn evolution_algorithm
         );
     }
 
-    let best_candidate = last_best_candidate.evaluated_chromosome.clone();
-    stats.best_candidates.push(last_best_candidate);
+    let best_candidate = last_best_candidate.as_ref().unwrap().evaluated_chromosome.clone();
+    stats.best_candidates.push(last_best_candidate.unwrap());
 
     Ok(EvolutionResult {
         population: current_population,
diff --git a/codes/eoa_lib/src/replacement.rs b/codes/eoa_lib/src/replacement.rs
index 853baf86092f6140fcb9ec2c0022e80a705ee2a2..5de7fae4b3b39260fd18a36be39fc901f3d8a5e4 100644
--- a/codes/eoa_lib/src/replacement.rs
+++ b/codes/eoa_lib/src/replacement.rs
@@ -51,6 +51,13 @@ impl<TChromosome> Population<TChromosome> {
         )
     }
 
+    pub fn evaluate_mut<TResult>(self, func: &mut dyn FnMut(&TChromosome) -> TResult) -> EvaluatedPopulation<TChromosome, TResult> {
+        EvaluatedPopulation::evaluate_mut(
+            self.population,
+            func
+        )
+    }
+
     pub fn into_iter(self) -> impl Iterator<Item = TChromosome> {
         self.population.into_iter()
     }
@@ -88,6 +95,17 @@ impl<TChromosome, TResult> EvaluatedPopulation<TChromosome, TResult> {
                 .collect::<Result<_, _>>()?))
     }
 
+    pub fn evaluate_mut(chromosomes: Vec<TChromosome>, func: &mut dyn FnMut(&TChromosome) -> TResult) -> Self {
+        EvaluatedPopulation::from_vec(
+            chromosomes.into_iter()
+                .map(|chromosome|
+                     EvaluatedChromosome {
+                         evaluation: func(&chromosome),
+                         chromosome
+                     })
+                .collect::<Vec<_>>())
+    }
+
     pub fn from_vec(vec: Vec<EvaluatedChromosome<TChromosome, TResult>>) -> Self {
         Self {
             population: vec