"""Main GA engine coordinating the evolutionary workflow.

The engine encapsulates the high-level control flow for one GA "agent".
It is the *only* place that defines the order of operations:

    selection → variation (mutation/crossover + immigrants) → simulation
    → evaluation (features/objectives) → survivor update → convergence check

Design notes:
  * SOLID-friendly: all stages depend on narrow interfaces (see .interfaces).
  * Latency hiding: physics is executed asynchronously via a thread pool.
  * Fault tolerance: validation and de-duplication filter problematic outputs
    without aborting the run; invalid structures are penalized and logged.

Threading model:
  - A single engine thread orchestrates stages and submits physics jobs.
  - Physical-model work is queued on a ThreadPoolExecutor (non-blocking).
  - `ctx.prev_future` carries the previous generation's physics Future.
"""
from __future__ import annotations

from typing import Optional
from concurrent.futures import ThreadPoolExecutor
import multiprocessing as mp
import atexit
import os
import numpy as np

from .interfaces import (
    IPopulation,
    IThermostat,
    IEvaluator,
    ISelector,
    IVariation,
    IGenerative,
    ISimulator,
    IConvergence,
    ILogger,
    IPlotter,
    ITransitionKernel,
)

from ..io.resume import resume_if_possible
from .context import Context
from .config import GAConfig
from ..io.snapshot_recorder import GenerationSnapshotWriter
from ..utils.memory import cleanup_memory

_TIMEOUT = 1440000

class GeneticAlgorithm:
    """High-level GA coordinator (single agent).

    This class wires the modular stages and executes the per-generation loop.
    It assumes the heavy-cost physical model runs inside `ISimulator.run()`,
    which is overlapped with GA bookkeeping via a thread pool.

    Attributes:
      population: Repository of individuals plus partitioned views.
      thermostat: Temperature scheduler controlling exploration/exploitation.
      evaluator: Computes features and objectives for a set of individuals.
      selector: Parent-selection policy (e.g., Pareto, Boltzmann).
      variation: Mutation/crossover engine, including operator learning.
      generative: Optional immigrant source (DoE/BO/generative model).
      simulator: Physical-model driver (MD/relaxation/etc.).
      convergence: Convergence monitor with stall/novelty logic.
      logger: Structured logger that persists progress and diagnostics.
      plotter: Post-run plotting utility.
      ctx: Mutable run context (timers, generation state, caches).
      cfg: Static configuration (population sizes, limits, paths).
      _executor: Thread pool for asynchronous physics submission.
    """
    def __init__(
        self,
        population: IPopulation,
        thermostat: IThermostat,
        evaluator: IEvaluator,
        selector: ISelector,
        
        variation: IVariation,
        generative: IGenerative,
        simulator : ISimulator,
        convergence: IConvergence,

        logger: ILogger,
        plotter: IPlotter,
        transition: ITransitionKernel | None = None, # NEW
        cfg: GAConfig | None = None,
        ctx: Context | None = None,
    ) -> None:
        """Initializes the GA engine with modular components.

        Args:
          population: Population storage and datasets manager.
          thermostat: Temperature controller.
          evaluator: Feature/objective evaluator.
          selector: Parent selector.
          variation: Mutation/crossover operator suite.
          generative: Immigrant generator (can be None if unused).
          simulator: Physical-model runner (MD + relaxation).
          convergence: Convergence detector.
          logger: Logger for progress and diagnostics.
          plotter: Plotting backend for final figures.
          cfg: Optional configuration; default `GAConfig()`.
          ctx: Optional context; default `Context()`.

        Side Effects:
          Creates a `ThreadPoolExecutor` for overlapping physics.

        Raises:
          ValueError: If mandatory dependencies are missing (defensive guard).
        """
        self.population = population
        self.thermostat = thermostat
        self.evaluator = evaluator
        self.selector = selector
        self.simulator = simulator
        self.variation = variation
        self.generative = generative

        self.convergence = convergence

        self.logger = logger
        self.plotter = plotter
        self.transition = transition

        self.ctx = ctx or Context()
        self.cfg = cfg or GAConfig()

        # NOTE: Choose `max_workers` according to available devices.
        # The engine submits a *single* physics job per generation; if you
        # split physics internally (e.g., per-individual parallelism),
        # increase workers or delegate concurrency to `ISimulator`.
        self._executor = ThreadPoolExecutor(max_workers=8)

    # ------------------------------
    # Private Cleanup memory
    # ------------------------------
    def _restart_executor(self, max_workers: int = 8) -> None:
        """Rotates the thread pool executor to prevent memory bloat/hangs."""
        self.logger.info("Restarting ThreadPoolExecutor...")
        if self._executor:
            # We don't use cancel_futures=True here because we assume it's only called when idle
            self._executor.shutdown(wait=True)
        self._executor = ThreadPoolExecutor(max_workers=max_workers)

    # ------------------------------
    # Private helpers (stage blocks)
    # ------------------------------
    def load_population(self, ) -> None:
        """Loads or initializes the starting dataset from `IPopulation`.

        Expected to populate the dataset and internal indices.

        Side Effects:
          - I/O: reads persisted individuals if present.
          - Logging: emits timing and status messages.
        """
        with self.ctx.timer("load data"):
            self.population.load_population(logger=self.logger)
            
            # Resume hook (keeps GA class clean)
            resume_if_possible(
                self.population, self.ctx, self.logger,
                enabled=getattr(self.cfg, "resume", True),
                mode=getattr(self.cfg, "resume_mode", "folders_all"),
            )

    def _evaluate(self, ) -> None:
        """Computes features and objectives for the current dataset.

        Sources:
          Uses the dataset as the evaluation input.

        Stores:
          - `ctx.features` and `ctx.objectives` for subsequent stages.

        Performance:
          This call should be lightweight compared to physics; it often
          reuses results produced by `ISimulator.run()`.

        Raises:
          RuntimeError: if the evaluator fails unexpectedly.
        """
        with self.ctx.timer("Evaluation"):
            features, objectives = self.evaluator.evaluate_features_objectives( self.population.get_dataset('dataset') )
            self.ctx.set_objectives(objectives)
            self.ctx.set_features(features)

    def _select(self, ) -> None:
        """Selects parents given current features/objectives and temperature.

        Stores:
          - Selection indices in `ctx.selection`.
          - 'selected' dataset with chosen individuals.
        """
        with self.ctx.timer("selection"):
            self.selector.set_temperature( self.ctx.get_temperature() )
            objectives, features = self.ctx.get_objectives(), self.ctx.get_features()
            sel_idx = self.selector.select( objectives=objectives, features=features, )

            self.ctx.set_selection(sel_idx)
            population_selected = self.population.filter(name='dataset', idx=sel_idx)
            self.population.set_population(name='selected', population=population_selected)

            population_offspring = self.population.filter(name='dataset', idx=sel_idx)
            self.population.set_population(name='offspring', population=population_offspring)

    def _generate(self) -> None:
        """Produces candidate offspring and immigrants, then validates.

        Pipeline:
          1) Variation (mutation + crossover) on 'offspring'.
          2) Optional immigration (DoE/BO/generative).
          3) Validation (schema/DoE/constraints).
          4) Self-collision pruning.

        Stores:
          - Overwrites 'offspring' with validated current candidates.

        Notes:
          - Penalizes invalids via `variation.penalization`.
          - Logs filtered counts for traceability.
        """
        with self.ctx.timer("variation"):
            offspring = self.population.get_containers('offspring')
            if not offspring:
                self.logger.warning("[variation] 'offspring' partition is empty. Skipping mutation and crossover.")
                mutated_individuals, crossed_individuals, mutation_rate_array = [], [], np.array([], dtype=int)
            else:
                mutated_individuals, crossed_individuals, mutation_rate_array = (
                    self.variation.apply_mutation_and_crossover(
                        individuals=offspring, 
                        generation=self.ctx.generation, 
                        objectives=self.ctx.top_objectives,
                        temperature=self.ctx.temperature,
                    )
                )
            current = mutated_individuals + crossed_individuals
        self.logger.info("Mutation and crossover completed. ")

        # Immigration (a.k.a. foreigners).
        with self.ctx.timer("Foreigners"):
            if self.generative is not None and self.generative.is_active(self.ctx.generation):

                targets = self.generative.generate_targets(
                    generation=self.ctx.generation,
                    features=self.ctx.get_features(),
                    objectives=self.ctx.get_objectives(),
                    temperature=self.ctx.temperature,
                )

                if targets is not None and len(targets) > 0:
                    immigrants = self.variation.guided_variation(
                        parents=self.population.get_dataset('offspring'),
                        targets=targets,
                        temperature=self.ctx.temperature,
                        tolerance=self.cfg.generative.tolerance,
                        max_iterations=self.cfg.generative.max_variation_iterations,
                        max_mutations=self.cfg.generative.max_mutations,
                    )

                    current.extend(immigrants.containers)
                    self.logger.info(
                        "[generative] injection completed: %d new individuals",
                        len(immigrants),
                    )
                #self.generative.bo.plot_model()
        
        with self.ctx.timer("validate"):
            self.population.set_population(name='current', population=current)
            features = self.evaluator.get_features(dataset=self.population.get_dataset('current') )
            current_valid, current_invalid = self.population.validate_candidates(individuals=current, features=features, remove=True)
            self.variation.penalization(individuals=current_invalid, process='out_of_doe')
            self.logger.info("[filter] Validate candidates completed. %i/%i Structures filtered.  (remove=True)", len(current_invalid), len(current) ) 

        with self.ctx.timer("selfcollision"):
            current_total = len(current_valid)
            current_valid, current_invalid = self.population.filter_selfcollision(current_valid, remove=True)
            self.variation.penalization(individuals=current_invalid, process='selfcollision')
            self.logger.info("[filter] Selfcollision check completed. %i/%i Structures filtered.  (remove=True)", len(current_invalid), current_total )

        try:
            self.population.set_population(name='offspring', population=current_valid)
        except Exception as e:
            if hasattr(self, "logger"):
                self.logger.warning(f"Failed to set 'offspring' population: {e}", exc_info=True)

    def _simulate(self, individuals: list = None):
        """Runs the physical model (MD/relax) on the provided individuals and re-validates.

        Args:
            individuals: List of individuals to simulate. If None, uses 'offspring' partition.

        Returns:
          List-like collection of *valid* post-physics individuals.
        """
        if individuals is None:
            current = self.population.get_containers('offspring')
        else:
            current = individuals

        with self.ctx.timer(f"simulator{self.ctx.generation}"):
            current = self.simulator.run(
                individuals=current,
                temperature=self.ctx.temperature,
                generation=self.ctx.generation,
                output_path=self.population.output_path,
            )

        with self.ctx.timer("validate"):
            # Temporarily extract unserializable metadata to protect auto-JSON export hook
            db_cache = {}
            for c in current:
                apm = getattr(c, "AtomPositionManager", None)
                if apm is not None and "detailed_balance" in apm.metadata:
                    db_cache[id(c)] = apm.metadata.pop("detailed_balance")

            self.population.set_population(name='current', population=current)
            
            # Restore metadata for the survive phase
            for c in current:
                if id(c) in db_cache:
                    getattr(c, "AtomPositionManager").metadata["detailed_balance"] = db_cache[id(c)]

            features = self.evaluator.get_features(dataset=self.population.get_dataset('current'))
            current_valid, current_invalid = self.population.validate_candidates(individuals=current, features=features, remove=False)
            self.variation.penalization(individuals=current_invalid, process='out_of_doe')
            self.logger.info("[filter] Validate candidates completed. %i/%i Structures filtered.", len(current_invalid), len(current) ) 

        with self.ctx.timer("hash_collisions"):
            current_total = len(current_valid)
            current_valid, current_invalid = self.population.filter_hash_collisions(current_valid)
            self.variation.penalization(individuals=current_invalid, process='hash_collision')
            self.logger.info("[filter] Hash collisions check completed. %i/%i Structures filtered.", len(current_invalid), current_total ) 

        with self.ctx.timer("selfcollision"):
            current_total = len(current_valid)
            current_valid, current_invalid = self.population.filter_selfcollision(current_valid)
            self.variation.penalization(individuals=current_invalid, process='selfcollision')
            self.logger.info("[filter] Selfcollision check completed. %i/%i Structures filtered.", len(current_invalid), current_total ) 

        return current_valid

    def _transition(self, parents: Any, candidates: list) -> list:
        """
        Unifies MH survival + ensemble moves via the TransitionKernel.
        """
        if not candidates:
            return []
        
        with self.ctx.timer("transition"):
            return self.transition.step(
                parents=parents,
                candidates=candidates,
            )

    def _agent_sync(self) -> int:
        """Pulls migrants from the shared mailbox and merges unique ones.

        Logic:
          - Pull candidate migrants via `population.agent_sync()`.
          - Drop duplicates via hash collision check (no rehashing).
          - Append uniques to 'current_valid'.
        Notes:
          Content-addressed digests ensure idempotent ingestion.
        
        Returns:
            int: Number of imported individuals.
        """
        agent = getattr(self.population, "agent", None)
        if not agent:
            return 0
            
        try:
            if hasattr(agent, "active") and not agent.active():
                return 0
        except Exception:
            return 0

        with self.ctx.timer("agent_sync"):
            self.logger.info("[Agentic_Sync] Exported %d structures to common genetic pool via HashBatchSync.", self.population.get_dataset('current_valid').size )

            # Clean up non-serializable callables before sync
            self._strip_detailed_balance(list(self.population.get_dataset('current_valid').containers))
            
            individuals_candidates = self.population.agent_sync( 
                dataset=self.population.get_dataset('current_valid'),
                objectives=self.ctx.get_objectives(),
                features=self.ctx.get_features(),
            )

            if individuals_candidates:
                policy_msg = getattr(agent, "policy_status", "Active")
                self.logger.info("[Agentic_Sync] Imported %d new structures from other agents. (Policy: %s)", len(individuals_candidates), policy_msg)
                
                # Append uniques directly to the valid dataset for the next evaluation pipeline
                self.population.dataset_add('current_valid', individuals_candidates)
                return len(individuals_candidates)
            
            return 0

    def _update_dataset(self, ):
        """Moves validated current individuals into the persistent dataset.

        Source:
          - 'current_valid' dataset.

        Destination:
          - 'dataset' dataset (archive for selection/evaluation/persistence).
        """
        with self.ctx.timer("Update_dataset"):
            # Clean up non-serializable callables before database insertion
            self._strip_detailed_balance(list(self.population.get_dataset('current_valid').containers))
            
            pruned = self.population.update_main_dataset(
                dataset=self.population.get_dataset('current_valid'),
                generation=self.ctx.generation,
                features_prev=self.ctx.get_features(),
                objectives_prev=self.ctx.get_objectives(),
            )

            if pruned > 0:
                self.logger.info(
                    "[Dataset] Pruned %d structures due to population cap.",
                    pruned
                )

    def _check_convergence(self, generation: Optional[int] = None):
        """Updates convergence status using the configured monitor.

        Populates in `ctx`:
          - `is_converge`: bool
          - `stagnation`: int (stall counter)

        Logging:
          Emits a single-line status with improvement flag, stall counters,
          convergence boolean, and generation elapsed time.
        """
        if generation is None:
            generation = self.ctx.generation

        with self.ctx.timer("convergence"):
            # The number of structures originally proposed by variation (before physics/selection)
            proposed_structures = self.population.size('offspring')
            
            # The number of structures that successfully passed physics and selection (the new valid ones)
            accepted_structures = self.population.size('current_valid')
            
            self.convergence.check(
                generation=generation,
                objectives=self.ctx.get_objectives(),
                features=self.ctx.get_features(),
                accepted_structures=accepted_structures,
                proposed_structures=proposed_structures,
            )
            self.ctx.is_converge = self.convergence.is_converge()
            self.ctx.stagnation = self.convergence.get_stagnation()

        self.logger.info(
            "[Gen=%d] improvement=%s, stall_count=%d/%s (obj), converged=%s, time=%.2f s",
            generation,
            self.convergence.improvement_found(),
            self.ctx.stagnation,
            self.convergence._stall_threshold,
            self.ctx.is_converge,
            self.ctx.elapsed('generation')
        )

    # --------------------------------------------------
    # Pipeline Orchestration Helpers
    # --------------------------------------------------
    def _setup_environment(self) -> None:
        """Tame native thread pools to reduce teardown/interference issues."""
        os.environ.setdefault("OMP_NUM_THREADS", "1")
        os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")
        os.environ.setdefault("MKL_NUM_THREADS", "1")
        os.environ.setdefault("VECLIB_MAXIMUM_THREADS", "1")
        os.environ.setdefault("NUMEXPR_NUM_THREADS", "1")

    def _strip_detailed_balance(self, individuals: List[Any]):
        """Remove non-serializable detailed_balance metadata from structural containers."""
        for ind in individuals:
            apm = getattr(ind, "AtomPositionManager", None)
            if apm is not None:
                meta = getattr(apm, "metadata", None)
                if meta is not None and "detailed_balance" in meta:
                    del apm.metadata["detailed_balance"]

    @property
    def _is_ensemble(self) -> bool:
        """Check if the current configuration is in ensemble/trajectory mode."""
        return getattr(self.cfg, "ensemble", None) is not None and self.cfg.ensemble.kind != "none"

    def _get_start_generation(self) -> int:
        """Compute effective start generation and log if resuming."""
        start_gen = max(
            getattr(self.ctx, "generation", self.cfg.initial_generation),
            self.cfg.initial_generation,
        )
        if start_gen != self.cfg.initial_generation:
            self.logger.info("Resuming at gen %d (initial_gen was %d).", start_gen, self.cfg.initial_generation)
        return start_gen

    def _execute_ensemble_step(self) -> None:
        """Synchronous execution flow for ensemble modes (Gibbs, Replica Exchange)."""
        # --- SYNCHRONOUS MODE ---
        offspring = self.population.get_containers('offspring')
        self.ctx.prev_future = self._executor.submit(self._simulate, individuals=offspring)
        with self.ctx.timer("simulator"):
            try:
                prev_result = self.ctx.prev_future.result(timeout=_TIMEOUT)
            except Exception as exc:
                import traceback
                self.logger.exception("Physics job from Gen %d failed: %s\n%s",
                                        self.ctx.generation, exc, traceback.format_exc())
                prev_result = []

        parents = self.population.get_dataset('selected')
        if self.transition is not None:
            prev_result = self._transition(parents=parents, candidates=prev_result)

        self.population.set_population(name='current_valid', population=prev_result)
        
        if self.cfg.save_generations:
            with self.ctx.timer("export_structures"):
                self.population.export_structures( 
                    dataset=self.population.get_dataset('current_valid'), 
                    file_path=f"{self.population.output_path}/generation/gen{self.ctx.generation}/config_g{self.ctx.generation}.xyz" 
                )

        with self.ctx.timer("cleanup_memory"):
            cleanup_memory(
                gen=self.ctx.generation,
                tag=f"post_gen_{self.ctx.generation}",
                ctx=self.ctx,
                logger=self.logger,
                mem_hiwater_mb=self.cfg.mem_hiwater_mb,
                restart_executor_fn=lambda gen: self._restart_executor(max_workers=getattr(self.cfg, "executor_max_workers", 8))
            )
        
        # Agent sync logic - delegated to kernel
        if self.transition.allows_agent_sync():
            self._agent_sync()
            
        self._update_dataset()
        self._check_convergence()

        self.ctx.stop_timer('generation')
        self._save_stats()
        self.ctx.start_timer('generation')

        if self.ctx.is_converge:
            self.logger.info(f"Convergence reached at generation {self.ctx.generation}.")
            self.ctx.break_early = True
        
        # Clear future so the finalizer block outside the loop skips it
        self.ctx.prev_future = None

    def _execute_async_step(self) -> None:
        """Asynchronous execution flow for standard GA mode (pipelining)."""
        if self.ctx.prev_future:
            gen_finalized = self.ctx.generation - 1
            with self.ctx.timer("simulator_wait"):
                try:
                    prev_result = self.ctx.prev_future.result(timeout=_TIMEOUT)
                except Exception as exc:
                    import traceback
                    self.logger.exception("Physics job from Gen %d failed: %s\n%s",
                                            gen_finalized, exc, traceback.format_exc())
                    prev_result = []

            # Complete eradication of un-serializable metadata
            for c in prev_result:
                apm_c = getattr(c, "AtomPositionManager", None)
                if apm_c is not None and "detailed_balance" in getattr(apm_c, "metadata", {}):
                    del apm_c.metadata["detailed_balance"]

            parents = getattr(self.ctx, "prev_parents", self.population.get_dataset('selected'))
            if self.transition is not None:
                prev_result = self._transition(parents=parents, candidates=prev_result)

            self.population.set_population(name='current_valid', population=prev_result)
            
            if self.cfg.save_generations:
                with self.ctx.timer("export_structures"):
                    self.population.export_structures( 
                        dataset=self.population.get_dataset('current_valid'), 
                        file_path=f"{self.population.output_path}/generation/gen{gen_finalized}/config_g{gen_finalized}.xyz" 
                    )

            with self.ctx.timer("cleanup_memory"):
                cleanup_memory(
                    gen=gen_finalized,
                    tag=f"post_gen_{gen_finalized}",
                    ctx=self.ctx,
                    logger=self.logger,
                    mem_hiwater_mb=self.cfg.mem_hiwater_mb,
                    restart_executor_fn=lambda gen: self._restart_executor(max_workers=getattr(self.cfg, "executor_max_workers", 8))
                )

            # Queue physics for current generation ASAP to overlap with bookkeeping
            self.ctx.prev_parents = self.population.get_dataset('selected')
            self.ctx.prev_future = self._executor.submit(self._simulate)

            # Post-physics bookkeeping for the generation we just finalized
            self._agent_sync()
            self._update_dataset()
            self._check_convergence(generation=gen_finalized)

            # Logging
            self.ctx.stop_timer('generation')
            self._save_stats()
            self.ctx.start_timer('generation')
            if self.ctx.is_converge:
                self.logger.info(f"Convergence reached at generation {gen_finalized}.")
                self.ctx.break_early = True
        else:
            # 2) Submit New Gen (Generation G+1)
            # We MUST capture the current offspring list now to avoid race conditions
            # when the main thread overwrites the 'offspring' partition in the next generation.
            current_offspring = self.population.get_containers('offspring')
            self.ctx.prev_future = self._executor.submit(self._simulate, individuals=current_offspring)
            self.ctx.stop_timer('generation')

    def _finalize_pending_physics(self) -> None:
        """Finalize the last pending physics job if any."""
        if self.ctx.prev_future and not self.ctx.prev_future.cancelled():
            try:
                final_result = self.ctx.prev_future.result(timeout=_TIMEOUT)
            except Exception as exc:
                import traceback
                self.logger.exception("Final physics job failed: %s\n%s", exc, traceback.format_exc())
                final_result = []
            
            self.population.set_population(name='current', population=final_result)

            # Unified sync check - delegated to kernel
            if self.transition.allows_agent_sync():
                self._agent_sync()
                
            self._update_dataset()
            self._check_convergence()

    def _finalize_run(self) -> None:
        """Finalize timers and export results."""
        self.ctx.stop_timer('engine')
        self._save_stats()
        
        # Final exports
        if getattr(self.cfg, "export", False):
            self.population.export_structures(
                dataset=self.population.get_dataset('dataset'),
                file_path=f"{self.population.output_path}/config_all.xyz"
            )
        self.logger.info("Workflow completed in %.2f s.", self.ctx.elapsed('engine'))

    def _save_stats(self, ):
        """Flushes timers/metrics; hook to persist run-time stats.

        Notes:
          Extend this method to persist per-generation KPIs, operator usage,
          temperature schedule, and feature/obj snapshots for post-hoc analysis.
        """
        with self.ctx.timer("Save_stats"):
            GenerationSnapshotWriter.save_generation(
                population=self.population,
                convergence=self.convergence,
                variation=self.variation,
                ctx=self.ctx,
                output_dir=self.population.output_path,
                tag="full",  # or "post-sim", "selection", etc. if you want multiple phases
            )
            self.ctx.clear_timers(clean_all=True, keep='engine')

    # --------------------------------------------------
    # Public API
    # --------------------------------------------------
    def run(self, ) -> IPopulation:
        """Executes the full GA workflow for one agent.

        Control flow:
          1) Initialize environment and load population.
          2) For each generation:
             a) Update temperature.
             b) Pipeline stages (Evaluation, Selection, Variation, Generation).
             c) Execution Mode (Sync Ensemble vs Async Pipelined).
          3) Finalize run and export final artifacts.
        """
        self.logger.info("Workflow started.")
        self.ctx.start_timer('engine')
        self._setup_environment()

        self.load_population()
        start_gen = self._get_start_generation()

        self.ctx.start_timer('generation')
        self.ctx.prev_future = None
        self.ctx.break_early = False

        try:
            for generation in range(start_gen, self.cfg.max_generations + 1):
                # 1) Evolution Stages
                self.logger.info(f"--- Generation {generation}/{self.cfg.max_generations} (Pop { self.population.size('dataset') }) ---")
                self.ctx.set_generation(generation)
                self.ctx.temperature = self.thermostat.update(
                    generation=generation,
                    stall=self.convergence.get_stagnation(),
                    signals=self.convergence.signals,
                )
                self._evaluate()
                
                #with self.ctx.timer("adjust_mutation_probs"):
                # !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
                #    self.variation.adjust_mutation_probabilities(dataset=self.population.get_dataset('dataset'), ctx=self.ctx)
                self._select()
                self._generate()

                # 2) Execution Mode
                if self._is_ensemble:
                    self._execute_ensemble_step()
                else:
                    self._execute_async_step()

                if self.ctx.break_early:
                    break

            # 3) Cleanup
            self._finalize_pending_physics()
            self._finalize_run()
            
            return self.population

        finally:
            # Ensure the executor is always cleaned up
            try:
                self._executor.shutdown(wait=True, cancel_futures=True)
            except Exception:
                pass

            # --- cleanup stray multiprocessing children ---
            try:
                for child in mp.active_children():
                    self.logger.warning("Cleaning up leaked child process: %s", child)
                    child.terminate()
                    child.join(timeout=5)
            except Exception as e:
                self.logger.error("Error during multiprocessing cleanup: %s", e)

        return True

# Global safeguard: ensure cleanup also runs on interpreter exit
def _cleanup_mp():
    for child in mp.active_children():
        child.terminate()
        child.join(timeout=5)

atexit.register(_cleanup_mp)





