"""
evolutionary_structure_explorer.py
----------------------------------

This module provides the EvolutionaryStructureExplorer class, which integrates
the workflow for classification, selection, mutation, crossover, composition
generation, molecular dynamics (MD) simulation, and convergence checking.

The class implements:
    - Encapsulation of all workflow parameters with appropriate getters and setters.
    - Detailed inline documentation using docstrings.
    - Logging of workflow events and key steps to a file.

CHANGES (Lineage Tracking):
---------------------------
1) A new integer counter (self._structure_id_counter) is added in the constructor.
2) A private method (_assign_lineage_info) is introduced to:
   - Increment the ID counter and assign a unique ID to each new offspring.
   - Store metadata about the parents' IDs, generation index, and genetic operation type.
3) The apply_mutation_and_crossover method now calls _assign_lineage_info(...) immediately
   after each structure is generated by mutation or crossover, allowing you to keep track
   of any relationship or ancestry in the structures' metadata.

# multi agent coordinator
# Swarm
"""
import os
import time
import copy
import logging
import warnings
import numpy as np
import json  # Added for parameter logging

from sklearn.decomposition import PCA
from sklearn.metrics import pairwise_distances
from sklearn.linear_model import Ridge

from sage_lib.partition.Partition import Partition
from sage_lib.miscellaneous.data_mining import *

from ..config.params import validate_evolution_parameters  
from ..convergence.convergence_checker import ConvergenceChecker
from ..utils.logger import WorkflowLogger
from ..utils.lineage import LineageTracker
from ..utils.canonical_hash import CanonicalHashMap

from ..selection.multiobjective_selector import ProbabilisticParetoSampler 
from bansga.thermostat.thermostat import Thermostat
from ..selection.features import evaluate_features
from ..selection.objective import evaluate_objectives
from ..io.io_manager import save_generation_data  

from ..physical_model.physical_model import physical_model
from ..visualization.plot_evolution import EvolutionPlotter  
from ..mutation_crossover.mutation_crossover_handler import MutationCrossoverHandler
from ..classification.clustering import SOAPClusterAnalyzer
from ..metric.information_ensemble_metrics import InformationEnsambleMetric

DEBUG = False
verbose = False

class EvolutionaryStructureExplorer:
    """
    Encapsulates the workflow for an evolutionary structure explorer. This includes:
      - Reading and maintaining a dataset of structures.
      - Performing multi-objective selection on the dataset.
      - Generating new structures through mutation, crossover, and optional "foreigners" or
        composition-based methods.
      - Executing physical model evaluations (simulations or other computations).
      - Checking and recording convergence metrics over generations.
      - Tracking lineage (which structures came from which parents) via unique IDs.
      - Exporting generated structures and logging relevant data.

    Parameters
    ----------
    params : dict, optional
        A dictionary of parameters to override default evolutionary settings.
    dataset_path : str, optional
        Path to the initial structures file. Defaults to '.' if not provided.
    template_path : str, optional
        Path to the template structure file. Defaults to '.' if not provided.
    output_path : str, optional
        Directory to save output files and logs. Defaults to '.' if not provided.

    Attributes
    ----------
    dataset_path : str
        File path for the initial structures.
    template_path : str
        Template file path.
    output_path : str
        Output directory path.
    max_generations : int
        Maximum number of generations in the evolutionary process.
    min_size_for_filter : int
        Minimum dataset size required before applying classification/filtering.
    convergence_params : dict
        Convergence thresholds and related parameters.
    partitions : dict
        Holds references to various `Partition` objects for the dataset, templates, etc.
    collision_factor : float
        Threshold for collision checking between atoms in structures.
    foreigners : int
        Number of "foreign" or composition-based structures to generate each generation.
    lineage_tracking_enabled : bool
        Flag indicating if lineage tracking is active.
    lineage_tracker : LineageTracker
        The object managing lineage relationships.
    """

    def __init__(self, params:dict=None, dataset_path:str=None, template_path:str=None, output_path:str=None, debug=False):
        """
        Initializes the EvolutionaryStructureExplorer with file paths and default parameters.

        Parameters
        ----------
        dataset_path : str
            Path to the initial structures file.
        template_path : str
            Path to the template structure file.
        output_path : str
            Directory to save output files and logs.
        """
        self._dataset_path = '.' if dataset_path is None else dataset_path
        self._template_path = '.' if template_path is None else template_path 
        self._output_path = '.' if output_path is None else output_path

        self.debug = debug

        # Workflow parameters
        self._max_generations = 100
        self._min_size_for_filter = 10

        self.mutation_rate_params = {}
        self.multiobjective_params = {}
        self.thermostat_params = {}
        self._convergence_params = {}
        self.information_metric = {}

        self._objectives_for_features_history = {}

        self._collision_factor = .4
        self._filter_duplicates = True
        self.save_logs = True

        self.foreigners = 0

        # Placeholder for Partition objects
        self._partitions = {
            'template':None,
            'dataset':None,
            'generation':None,
        }  # Should be set by load_partitions()

        # Lists to keep track of convergence metrics over generations
        self._energy_prev = None
        self._distances_prev = None

        self.objective_funcs = None
        self.features_funcs = None

        self.mutation_funcs = None
        self.crossover_funcs = None

        self.DoE = None
        self.BO = None

        self.physical_model_func = None
        self.physical_model_mode = None
                

        # Setup logger
        self.time_log = {}
        self.logger = WorkflowLogger.setup_logger('EvolutionaryStructureExplorer', self._output_path)

        self.setup_evolution_parameters(params)

        # Lineage Tracking 
        self.lineage_tracking_enabled = True
        self.lineage_tracker = LineageTracker()

        self.thermostat = Thermostat(
            initial_temperature=self.thermostat_params['initial_temperature'],
            period=self.thermostat_params['period'],
            decay_rate=self.thermostat_params['decay_rate'],
            stall_growth_rate=self.thermostat_params['stall_growth_rate'],
            temperature_bounds=self.thermostat_params['temperature_bounds'],
            max_stall_offset=self.thermostat_params['max_stall_offset'],
            constant_temperature=self.thermostat_params['constant_temperature'],

        )

        self.convergence_checker = ConvergenceChecker(
            logger=self.logger, 
            detailed_record=self._convergence_params['detailed_record'],
            stall_threshold=self._convergence_params['stall_threshold'],
            information_driven=self._convergence_params['information_driven']
        )

        self.mutation_crossover_handler = MutationCrossoverHandler(
            mutation_funcs=self.mutation_funcs,
            crossover_funcs=self.crossover_funcs,
            lineage_tracker=self.lineage_tracker, 
            logger=self.logger, 
            mutation_rate_params=self.mutation_rate_params, 
            debug=self.debug
        )

        self.evolution_plotter = EvolutionPlotter()
        self.hash_map = CanonicalHashMap(symprec=0.1)

        self.multiobjective_selector = ProbabilisticParetoSampler(
            weights=self.multiobjective_params['weights'],
            repulsion_weight=self.multiobjective_params['repulsion_weight'],
            repetition_penalty=self.multiobjective_params['repetition_penalty'],
            objective_temperature=self.multiobjective_params['objective_temperature'],

            size=self.multiobjective_params['size'],
            repulsion_mode=self.multiobjective_params['repulsion_mode'],
            metric=self.multiobjective_params['metric'],
            random_seed= self.multiobjective_params['random_seed'],

            steepness=self.multiobjective_params['steepness'],
            max_count=self.multiobjective_params['max_count'],
            cooling_rate=self.multiobjective_params['cooling_rate'],

            selection_method=self.multiobjective_params['selection_method'],
            divisions=self.multiobjective_params['divisions'],
        ) 

        self.information_ensamble_metric = InformationEnsambleMetric(
            auto_percentile = self.information_metric['auto_percentile'],
            metric = self.information_metric['metric'],
        )
        self.information_ensamble_analyzer = SOAPClusterAnalyzer(
            n_components = self.information_metric['components'],
            r_cut = self.information_metric['r_cut'],
            n_max = self.information_metric['n_max'],
            l_max = self.information_metric['l_max'],
            sigma = self.information_metric['sigma'],
            max_clusters = self.information_metric['max_clusters'],
            )
                
        self._containers_preload = {
            'template':None,
            'dataset':None,
            'generation':None,
        } 

    # ----------------------------------------------------------------
    # GETTERS/SETTERS
    # ----------------------------------------------------------------
    @property
    def dataset_path(self):
        """Gets the file path for the initial structures."""
        return self._dataset_path

    @dataset_path.setter
    def dataset_path(self, value):
        """Sets the file path for the initial structures."""
        self._dataset_path = value

    @property
    def template_path(self):
        """Gets the template file path."""
        return self._template_path

    @template_path.setter
    def template_path(self, value):
        """Sets the template file path."""
        self._template_path = value

    @property
    def output_path(self):
        """Gets the output directory path."""
        return self._output_path

    @output_path.setter
    def output_path(self, value):
        """Sets the output directory path."""
        self._output_path = value

    @property
    def max_generations(self):
        """Gets the maximum number of generations."""
        return self._max_generations

    @max_generations.setter
    def max_generations(self, value):
        """Sets the maximum number of generations."""
        self._max_generations = value

    @property
    def convergence_params(self):
        """Gets the convergence parameters dictionary."""
        return self._convergence_params

    @convergence_params.setter
    def convergence_params(self, params):
        """Sets the convergence parameters dictionary."""
        if isinstance(params, dict):
            self._convergence_params = params
        else:
            raise ValueError("convergence_params must be a dictionary.")

    @property
    def min_size_for_filter(self):
        """ """
        return self._min_size_for_filter

    @min_size_for_filter.setter
    def min_size_for_filter(self, params):
        """Sets the convergence parameters dictionary."""
        if isinstance(params, int):
            self._min_size_for_filter = params
        else:
            raise ValueError("min_size_for_filter must be a int.")

    @property
    def collision_factor(self):
        """ """
        return self._collision_factor

    @collision_factor.setter
    def collision_factor(self, params):
        """Sets the convergence parameters dictionary."""
        if isinstance(params, float):
            self._collision_factor = params
        else:
            raise ValueError("collision_factor must be a float.")

    @property
    def filter_duplicates(self):
        """ """
        return self._filter_duplicates

    @filter_duplicates.setter
    def filter_duplicates(self, params):
        """Sets the convergence parameters dictionary."""
        if isinstance(params, bool):
            self._filter_duplicates = params
        else:
            raise ValueError("filter_duplicates must be a bool.")

    @property
    def partitions(self):
        """Gets the dictionary of partition objects."""
        return self._partitions

    # ----------------------------------------------------------------
    # PARAMETER SETUP
    # ----------------------------------------------------------------
    def setup_evolution_parameters(self, params: dict = None):
        """
        Validates and assigns evolutionary parameters to the corresponding class attributes.

        This method leverages the external 'validate_evolution_parameters' function to 
        ensure that all parameters are correctly set, issuing warnings and performing assertions
        as needed. The validated parameters are then assigned to the respective attributes of 
        the instance.

        Parameters
        ----------
        params : dict, optional
            Dictionary containing parameter overrides.
        """
        # Validate and correct parameters using the external function
        validated_params = validate_evolution_parameters(params)

        # Filter to JSON-serializable parameters
        serializable_keys = [
            'max_generations', 'min_size_for_filter', 'mutation_rate_params',
            'multiobjective_params', 'thermostat_params', 'convergence_params', 
            'collision_factor', 'filter_duplicates',
            'foreigners', 'dataset_path', 'template_path',
            'output_path', 'information_metric', 
        ]
        log_params = {key: validated_params[key] for key in serializable_keys if key in validated_params}
        params_str = json.dumps(log_params, indent=4)
                
        # Log parameters clearly and in order
        self.logger.info("Evolution parameters for this run:\n%s", params_str)
        
        # Assign validated parameters to class attributes
        self._max_generations = validated_params['max_generations']
        self._min_size_for_filter = validated_params['min_size_for_filter']
        self.mutation_rate_params = validated_params['mutation_rate_params']
        self._convergence_params = validated_params['convergence_params']
        self._collision_factor = validated_params['collision_factor']
        self._filter_duplicates = validated_params['filter_duplicates']
        self.foreigners = validated_params['foreigners']
        self._dataset_path = validated_params['dataset_path']
        self._template_path = validated_params['template_path']
        self._output_path = validated_params['output_path']
        self.objective_funcs = validated_params['objective_funcs']
        self.features_funcs = validated_params['features_funcs']
        self.mutation_funcs = validated_params['mutation_funcs']
        self.crossover_funcs = validated_params['crossover_funcs']
        self.DoE = validated_params['DoE']
        self.BO = validated_params['BO']
        self.thermostat = validated_params['thermostat']

        self.physical_model_mode = validated_params['physical_model_func']['T_mode']
        self.physical_model_func = validated_params['physical_model_func']['calculator']

        self.multiobjective_params = validated_params['multiobjective_params']
        self.thermostat_params = validated_params['thermostat_params']

        self.information_metric = validated_params['information_metric']
        self.save_logs = validated_params['save_logs']

        # Log the successful assignment of parameters
        self.logger.info("Evolution parameters have been successfully assigned from validated parameters.")
    
    # --- End of getters and setters ---
    def load_partitions(self, dataset_path: str, template_path: str = None):
        """
        Load the initial and template partitions from given file paths.

        Parameters
        ----------
        dataset_path : str
            Path to the dataset of initial structures.
        template_path : str, optional
            Path to template structures for generating foreigners or other derived structures.
        """
        self.logger.info("Loading partitions from files.")
        start_time = time.time()

        self.partitions['dataset'] = Partition()
        if isinstance(dataset_path, str):
            self.partitions['dataset'].read_files(
                file_location=dataset_path,
                source='xyz',
                verbose=True
            )
        if isinstance(self._containers_preload['dataset'], list):
            self.partitions['dataset'].add_container( self._containers_preload['dataset'] )

        for container in self.partitions['dataset'].containers:
            self.hash_map.add_structure( container )

        self.partitions['template'] = Partition()
        if isinstance(template_path, str):
            self.partitions['template'].read_files(
                file_location=template_path,
                source='xyz',
                verbose=True
            )
        else:
            self.logger.info("No template file provided, skipping template partition.")
        if isinstance(self._containers_preload['template'], list):
            self.partitions['template'].add_container( self._containers_preload['template'] )

        self.lineage_tracker.assign_lineage_info_par_partition(self.partitions['dataset'], 0, [], "initialization")

        elapsed = time.time() - start_time
        self.logger.info(f"Partitions loaded successfully. (Elapsed: {elapsed:.2f}s)")

    def export_structures(self, partition, file_path: str, sort: bool=False, key:int=0):
        """
        Export structures in the given partition to the specified file path.

        Parameters
        ----------
        partition : Partition
            Partition containing the structures to export.
        file_path : str
            Output directory to store the exported structure files.
        sort : bool
        """
        os.makedirs(file_path, exist_ok=True)

        if sort:
            objectives = evaluate_objectives(partition.containers, self.objective_funcs)
            for container, objective in zip(partition.containers, objectives):
                apm = container.AtomPositionManager
                if not isinstance(getattr(apm, 'metadata', None), dict):
                    apm.metadata = {}
                apm.metadata['objectives'] = list(objective)

            numeric_objectives = np.array([float(obj[key]) for obj in objectives])
            sorted_indices = np.argsort(numeric_objectives, kind='mergesort')
            partition.containers = [partition.containers[i] for i in sorted_indices]

        partition.export_files(
            file_location=f"{file_path}",
            source='xyz',
            label='enumerate',
            verbose=True
        )
        return True

    # ----------------------------------------------------------------
    #  PRIVATE METHODS 
    # ----------------------------------------------------------------
    def _evaluate_features_objectives(self, structures):
        """
        Evaluates features and objectives for a list of structures.

        Parameters
        ----------
        structures : list
            List of structure containers.

        Returns
        -------
        features : np.ndarray
            Computed features for each structure.
        objectives : np.ndarray
            Computed objectives for each structure.
        """
        t0 = time.time()

        # Evaluate features
        features = evaluate_features(structures, self.features_funcs)

        # Evaluate objectives
        objectives = evaluate_objectives(structures, self.objective_funcs)

        self.time_log['objectives_features'] = time.time() - t0
        return features, objectives

    def _perform_selection(self, objectives, features, temperature):
        """
        Performs multi-objective selection to find top structures.

        Parameters
        ----------
        objectives : np.ndarray
            Objectives array for the current dataset.
        features : np.ndarray
            Features array for the current dataset.
        temperature : float
            Current temperature for selection, if applicable.

        Returns
        -------
        selected_indices : np.ndarray
            Indices of selected structures.
        top_structures : list
            Deepcopy of the selected structures from the dataset partition.
        """
        t0 = time.time()
        dataset_size = self.partitions['dataset'].size

        if dataset_size > self.min_size_for_filter:
            self.multiobjective_selector.sampling_temperature = temperature
            selected_indices = self.multiobjective_selector.select(
                objectives=objectives,
                features=features,
            )
            top_structures = copy.deepcopy([self.partitions['dataset'].containers[i] for i in selected_indices])
        else:
            selected_indices = np.array(range(dataset_size))
            top_structures = copy.deepcopy(self.partitions['dataset'].containers)

        self.time_log['multiobjective_selection'] = time.time() - t0
        return selected_indices, top_structures

    def _apply_mutation_crossover_foreigners(self, generation, top_structures, objectives, features, temperature):
        """
        Applies mutation, crossover, and foreigners generation in one step.

        Parameters
        ----------
        generation : int
            Current generation index.
        top_structures : list
            List of selected top structures.
        objectives : np.ndarray
            Current objectives array.
        features : np.ndarray
            Current features array.
        temperature : float
            Current temperature for mutation/crossover parameter adjustments.

        Returns
        -------
        new_structures : list
            List of new structures produced by mutation, crossover, and foreigners.
        """
        # 1) Apply mutation/crossover
        t0 = time.time()
        mutated_structures, crossed_structures, mutation_rate_array = (
            self.mutation_crossover_handler.apply_mutation_and_crossover(
                structures=top_structures, 
                generation=generation, 
                objectives=objectives,
                temperature=temperature,
            )
        )
        self.mutation_rate_array = mutation_rate_array
        self.time_log['Mutation_Crossover'] = time.time() - t0

        # 2) Foreigners
        t0 = time.time()
        new_foreigners = []
        if self.foreigners > 0 or generation == 1:
            new_foreigners = self._generate_new_compositions(generation, self.foreigners, features, objectives, temperature)

        new_candidates = mutated_structures + crossed_structures + new_foreigners

        new_candidates_validated = self._validate_candidates(new_candidates)

        self.time_log['Foreigners'] = time.time() - t0

        return new_candidates_validated

    def _generate_new_compositions(self, 
        generation:int, 
        foreigners: int, 
        features:np.array, 
        objectives:np.array,  
        temperature:float=None
        ):
        """
        Generates new structures based on composition ("foreigners" approach)
        or using a DoE/BO pipeline.

        Parameters
        ----------
        generation : int
            Current generation index.
        features : np.ndarray
            Current features array.
        objectives : np.ndarray
            Current objectives array.
        temperature : float
            Current temperature, if needed for BO or sampling.

        Returns
        -------
        new_structures : list
            List of newly generated "foreigner" structures.
        """
        self.logger.info("Generating new composition structures.")
        if generation == 1:
            if not getattr(self, 'DoE', None):
                logging.error("DoE instance is not available. Skipping design point generation.")
                return []

            try:
                design_points = self.DoE.generate()
            except Exception as e:
                logging.error(f"Error generating DoE design points: {e}", exc_info=True)
                return []
            dp_name = 'DoE'
        else:
            self.BO.fit(X=features, Y=objectives)
            design_points = self.BO.recommend_candidates(n_candidates=self.foreigners, T=temperature)
            dp_name = 'BO'

        new_structures = self.mutation_crossover_handler.foreigners_generation(
            structures=self.partitions['template'].containers,
            feature_func=self.features_funcs,
            design_points=design_points,
        )

        self.partitions[dp_name] = Partition()
        self.partitions[dp_name].containers = new_structures
        self.lineage_tracker.assign_lineage_info_par_partition(self.partitions[dp_name], 0, [], dp_name)

        self.logger.info(f"New composition structures generated: {len(new_structures)}.")
        return new_structures

    def _validate_candidates(self, structures: list) -> list:
        """
        Validate candidate structures based on their composition in feature space.

        This method computes feature descriptors for each candidate structure using the provided
        feature functions (self.features_funcs) and then validates each candidate by applying the
        design validator (self.BO.validate). The validator returns True if the candidate's features
        satisfy the design requirements and False otherwise.

        Parameters
        ----------
        structures : list
            A list of candidate structure objects to be validated.

        Returns
        -------
        list
            A list of candidate structures that meet the feature space requirements.
        """
        # 1) Compute all feature descriptors in one bulk call
        features = evaluate_features(structures, self.features_funcs)
        
        # 2) Local bindings for speed
        validate = self.BO.validate if self.BO else lambda *args, **kwargs: True
        penalize = self.mutation_crossover_handler.penalization
        
        # 3) Single pass: build valid list and penalise invalid
        valid_candidates = []
        physically_invalid_candidates = []
        out_of_doe_candidates = []

        for struct, feat in zip(structures, features):
            # 1) DOE‐validation check
            if not validate(features=feat):
                out_of_doe_candidates.append(struct)
                penalize(container=struct, process='out_of_doe')
                continue  # skip to next structure

            # 2) Physical‐sanity check
            apm = struct.AtomPositionManager
            if apm.E and apm.E / apm.atomCount <= -10:
                physically_invalid_candidates.append(struct)
            else:
                valid_candidates.append(struct)

        # 3) Export any physically invalid structures
        if physically_invalid_candidates:
            invalid_partition = Partition()
            invalid_partition.containers = physically_invalid_candidates
            self.export_structures(
                invalid_partition,
                f'{self.output_path}/nonvalid'
            )

        return valid_candidates

    def _filter_self_collisions(self, structures):
        """
        Filters out structures that have self-collision above the collision factor threshold.

        Parameters
        ----------
        structures : list
            List of candidate structures to filter.

        Returns
        -------
        non_colliding : list
            Filtered list of structures without self-collisions.
        """
        t0 = time.time()
        if self._collision_factor < 1e-6:
            return structures

        # 1) Compute collision mask in one pass (C-optimized list comp)
        mask = [
            s.AtomPositionManager.self_collision(factor=self._collision_factor)
            for s in structures
        ]

        # 2) Use mask to filter in C-level loops
        non_colliding = [s for s, m in zip(structures, mask) if not m]
        colliding = [s for s, m in zip(structures, mask) if m]

        # 3) Export invalid structures, if any
        if colliding:
            partition = Partition()
            partition.containers = colliding
            self.export_structures(
                partition,
                f"{self.output_path}/colliding"
            )

        self.time_log['Self Collision'] = time.time() - t0
        return non_colliding

    def _execute_physical_model(self, structures, temperature, generation:int=None):
        """
        Run the physical model on structures, timing the call and filtering
        out both self-collisions and hash-collisions.

        Parameters
        ----------
        structures : Sequence[Structure]
            Candidate structures to process.
        temperature : float
            Temperature passed to the physical model.
        generation : int, optional
            Generation index (unused).

        Returns
        -------
        Partition
            The resulting partition with collisions removed.
        """

        # 1) Self-collision filtering
        no_collision_structures = self._filter_self_collisions(structures)

        # 2) Physical model execution
        t0 = time.time()

        generation_partition = physical_model(
            structures=no_collision_structures, 
            physical_model_func=self.physical_model_func, 
            temperature=temperature,
            T_mode=self.physical_model_mode,
            generation=generation,
            output_path=f"{self.output_path}",
            logger=self.logger, 
            debug=self.debug
        )

        self.time_log['physical_model'] = time.time() - t0

        # 3) Hash collision filtering
        valid_structures = self._validate_candidates(generation_partition.containers)
        unique_valid_structures = self._filter_hash_collisions(valid_structures)
        no_collision_unique_valid_structures = self._filter_self_collisions(unique_valid_structures)

        generation_partition.containers = no_collision_unique_valid_structures

        return generation_partition

    def _filter_hash_collisions(self, structures):
        """
        Filters out structures that collide in the hash map (i.e., duplicates).

        Parameters
        ----------
        structures : list
            List of candidate structures after the physical model.

        Returns
        -------
        unique_structures : list
            Structures that are unique (not in the hash map).
        """
        t0 = time.time()
        if not self._filter_duplicates:
            return structures

        unique_structures = []
        for structure in structures:
            if self.hash_map.add_structure(structure):
                unique_structures.append(structure)
            else:
                self.mutation_crossover_handler.penalization(container=structure, process='hash_colition')

        self.logger.info(f"Filter_hash_collisions completed. {len(structures)-len(unique_structures) } Structures filtered.") 

        self.time_log['Hash Collision'] = time.time() - t0
        return unique_structures

    def _update_main_dataset(self, unique_structures, generation):
        """
        Updates the main dataset with new structures and records the structure IDs and generations.

        Parameters
        ----------
        unique_structures : list
            Newly generated unique structures.
        generation : int
            Current generation index.
        """
        t0 = time.time()
        self.partitions['dataset'].containers += unique_structures
        self.time_log['update_dataset'] = time.time() - t0

    def _check_convergence_wrapper(self, generation:np.array, objectives:np.array, features:np.array, structures:list=None):
        """
        Checks convergence using the ConvergenceChecker, if conditions are met.

        Parameters
        ----------
        generation : int
            Current generation index.
        objectives : np.ndarray
            Objectives array.
        features : np.ndarray
            Features array.

        Returns
        -------
        conv_results : dict
            Dictionary containing convergence check results.
        """
        t0 = time.time()

        # Local binding for speed and brevity
        ds = self.partitions['dataset']
        size = ds.size

        # Early exit if not enough data or no features
        if features is None or size < self._min_size_for_filter:
            self.time_log['check_convergence'] = time.time() - t0
            return {}

        # Determine whether to do information‐driven novelty check
        info_driven = (
            self.convergence_checker._information_driven
            and size > self.information_ensamble_analyzer.n_components
        )

        # Compute novelty improvement only if needed
        information_novelty_has_improved = False
        if info_driven:
            clusters = self.information_ensamble_analyzer.get_cluster_counts(structures)
            information_novelty_has_improved = self.information_ensamble_metric.has_improved(clusters)

        # Perform the actual convergence check
        conv_results = self.convergence_checker.check_convergence(
            generation=generation,
            objectives=objectives,
            features=features,
            debug=self.debug,
            information_novelty_has_improved=information_novelty_has_improved,
            generation_start=self.generation_start,
            time_log=self.time_log
        )

        self.time_log['check_convergence'] = time.time() - t0
        return conv_results

    def _save_generation_data_wrapper(
        self,
        generation: int,
        conv_results: dict,
        features: np.ndarray,
        objectives: np.ndarray,
        selected_indices: np.ndarray,
        temperature:float,
        partition_path: str
    ):
        """
        Saves the key results for a given generation to disk
        and logs the elapsed time for this step.

        Parameters
        ----------
        generation : int
            Current generation index.
        conv_results : dict
            Convergence results for this generation.
        features : np.ndarray
            Features array.
        objectives : np.ndarray
            Objectives array.
        selected_indices : np.ndarray
            Indices of selected structures.
        temperature : float
            Temperature used in this generation.
        partition_path : str
            Subdirectory path for logging (e.g., "my_partition").
        """
        t0 = time.time()

        # Build the 'data' dict
        data_dict = {
            "convergence_results": conv_results,
            "features": features.tolist() if features is not None else None,

            "generation": generation,
            "num_structures_in_dataset": int(self.partitions['dataset'].size),
            "num_newstructures": int(self.partitions['generation'].size),
            "objectives": objectives.tolist() if objectives is not None else None,

            "selected_indices": (
                selected_indices.tolist()
                if (selected_indices is not None)
                else []
            ),

            "stall_count": self.convergence_checker._no_improvement_count,
            "stall_count_objetive": self.convergence_checker._no_improvement_count_objectives,

            "objectives_for_features_history": self.convergence_checker._objectives_for_features_history,
            "mutation_rate_history": self.mutation_rate_array,
            "mutation_probabilities": self.mutation_crossover_handler._mutation_probabilities,

            "mutation_attempt_counts": self.mutation_crossover_handler._mutation_attempt_counts,
            "mutation_fails_counts": self.mutation_crossover_handler._mutation_fails_counts,
            "mutation_success_counts": self.mutation_crossover_handler._mutation_success_counts,
            "mutation_unsuccess_counts": self.mutation_crossover_handler._mutation_unsuccess_counts,
            "mutation_hashcolition_counts": self.mutation_crossover_handler._mutation_hashcolition_counts,
            "mutation_outofdoe_counts": self.mutation_crossover_handler._mutation_outofdoe_counts,

            "crosvover_attempt_counts": self.mutation_crossover_handler._crossover_attempt_counts,
            "crossover_fails_counts": self.mutation_crossover_handler._crossover_fails_counts,
            "crossover_success_counts": self.mutation_crossover_handler._crossover_success_counts,
            "crossover_unsuccess_counts": self.mutation_crossover_handler._crossover_unsuccess_counts,
            "crossover_hashcolition_counts": self.mutation_crossover_handler._crossover_hashcolition_counts,
            "crossover_outofdoe_counts": self.mutation_crossover_handler._crossover_outofdoe_counts,

            "time_log": self.time_log,
            "T": temperature
            #"model_evolution_info": self.model_evolution_info  # additional logging
        }

        if self.convergence_checker._information_driven:
             data_dict["novelty_history"] = self.information_ensamble_metric.get_latest_novelty()
             data_dict["novelty_thresh_history"] = self.information_ensamble_metric.get_latest_novelty_thresh()
             data_dict["stall_count_information"] = self.convergence_checker._no_improvement_count_information

        # Call your existing save_generation_data utility
        save_generation_data(
            generation=generation,
            data=data_dict,
            output_directory=f"{self._output_path}/{partition_path}/logger"
        )

        # Log elapsed time
        self.time_log['save_generation_data'] = time.time() - t0
        if self.debug:
            self.logger.info(f"[DEBUG] Generation data saved for Gen={generation}.")

    def run(self, partition_path:str='.', debug=False):
        r"""
        Execute the full evolutionary structure exploration workflow.

        Let \(\mathcal{D}_0\) be the initial dataset partition of size \(N_0\).  
        We iterate for generations \(g = 1, 2, \dots, G_{\max}\), maintaining partitions
        \(\mathcal{D}_{g-1}\) of size \(N_{g-1} = |\mathcal{D}_{g-1}|\).  At each generation \(g\):

        1. **Temperature Scheduling**  
           Update the search temperature via the thermostat:
           .. math::
              T_g \;=\; \mathrm{Thermostat.update}(g,\;\tau_{g-1})
           where \(\tau_{g-1}\) is the previous stall‐count.

        2. **Feature & Objective Evaluation**  
           Compute the feature matrix
           \(\Phi_g \in \mathbb{R}^{N_{g-1}\times d}\) and objective matrix
           \(O_g \in \mathbb{R}^{N_{g-1}\times k}\) by
           .. math::
              \Phi_g = \bigl[\phi(s)\bigr]_{s\in\mathcal{D}_{g-1}}, \quad
              O_g    = \bigl[f(s)\bigr]_{s\in\mathcal{D}_{g-1}}.
           
        3. **Multi‐Objective Selection**  
           Compute selection probabilities via Boltzmann‐Pareto sampling:
           .. math::
              p_i = \frac{\exp\!\bigl(-O_{g,i}/T_g\bigr)}
                         {\sum_{j=1}^{N_{g-1}}\exp\!\bigl(-O_{g,j}/T_g\bigr)}, 
              \quad
              S_g = \mathrm{sample}\bigl(\mathcal{D}_{g-1},\,p\bigr).
           
        4. **Variation Operators**  
           Generate offspring by:
           - **Mutation**:
             .. math::
                \mathcal{M}_g = \{\mathrm{mutate}(s)\mid s\in S_g\}.
           - **Crossover**:
             .. math::
                \mathcal{C}_g = \{\mathrm{crossover}(s_i,s_j)\mid s_i,s_j\in S_g\}.
           Set \(\mathcal{E}_g = \mathcal{M}_g \cup \mathcal{C}_g\).

        5. **Composition (“Foreigners”) Generation**  
           Produce additional candidates
           .. math::
              \mathcal{F}_g = 
              \begin{cases}
                \mathrm{DoE}\bigl(\Phi_g,\,O_g\bigr), & g = 1,\\
                \mathrm{BO}\bigl(\Phi_g,\,O_g,\,T_g\bigr), & g > 1.
              \end{cases}

        6. **Candidate Validation & Collision Filtering**  
           Form the raw candidate set
           \(\mathcal{U}_g = \mathcal{E}_g\cup\mathcal{F}_g\).  Then:
           - **Feasibility**: retain \(x\) if \(\mathrm{validate}(x)=\mathrm{True}\).  
           - **Self‐Collision**: remove \(x\) if \(\mathrm{collision}(x)>\sigma_c\).  
           - **Hash‐Collision**: keep only the first occurrence of each canonical hash.

        7. **Physical‐Model Execution**  
           Evaluate valid candidates under the physical model:
           .. math::
              \mathcal{P}_g = \mathrm{physical\_model}(\mathcal{U}'_g,\,T_g),
           yielding a new partition of size \(N'_g\).

        8. **Dataset Update**  
           Merge new structures into the dataset:
           .. math::
              \mathcal{D}_g = \mathcal{D}_{g-1} \;\cup\;\mathcal{P}_g,
              \quad N_g = |\mathcal{D}_g|.

        9. **Convergence Checking**  
           Compute convergence flag
           .. math::
              c_g = \mathrm{check\_convergence}(g,\;O_g,\;\Phi_g),
           and if \(c_g=\mathrm{True}\), terminate the loop early.

        10. **Logging & Persistence**  
            If `save_logs=True`, call
            :meth:`_save_generation_data_wrapper` to record all metrics,
            histories, and structures for generation \(g\).

        11. **Post‐Processing**  
            After termination (by convergence or \(g=G_{\max}\)):
            - Export final `dataset` and `template` partitions (sorted by primary objective).
            - Generate evolution plots via :class:`EvolutionPlotter`.

        Parameters
        ----------
        partition_path : str, optional
            Subdirectory under `output_path` for per‐generation outputs and logs.
            Defaults to `'.'`.
        debug : bool, optional
            If True, enables verbose debug logging. Defaults to False.

        Returns
        -------
        Partition
            The final dataset partition \(\mathcal{D}_{g^*}\) containing all retained structures.

        Raises
        ------
        RuntimeError
            If a critical step (e.g., DoE generation, physical model) fails irrecoverably.
        """
        self.logger.info("Workflow started.")
        overall_start = time.time()  # Start overall workflow timer

        self.load_partitions(self.dataset_path, self.template_path)
        objectives, features = None, None
        generations, structures_id = np.zeros(self.partitions['dataset'].size), np.arange(1, self.partitions['dataset'].size+1)

        # Main workflow loop
        for generation in range(1, self._max_generations + 1):
            self.generation_start = time.time()
            self.logger.info(f"--- Generation {generation}/{self._max_generations} (Pop { len(self.partitions['dataset'].containers) }) ---")
            # 1) Get temperature for this generation (if thermostat is active)
            self.thermostat.actualizate_temperature(generation=generation, stall=self.convergence_checker._no_improvement_count)
            temperature = self.thermostat.get_temperature()

            # 2) Evaluate features/objectives
            features, objectives = self._evaluate_features_objectives(self.partitions['dataset'].containers)

            # 3) Adjust mutation probabilities (optional)
            if self.partitions['dataset'].size > self.min_size_for_filter or self.foreigners > 0:
                self.mutation_crossover_handler._adjust_mutation_probabilities( 
                    structures=self.partitions['dataset'].containers, objectives=objectives, features=features, 
                )

            # 4) Selection
            selected_indices, top_structures = self._perform_selection(objectives, features, temperature)

            # 5) Mutation / Crossover + Foreigners
            new_structures = self._apply_mutation_crossover_foreigners(
                generation=generation, top_structures=top_structures, objectives=objectives, features=features, temperature=temperature
            )

            # 6) Physical model execution 
            self.partitions['generation'] = self._execute_physical_model(structures=new_structures, temperature=temperature, generation=generation)

            # 7) Export structures
            self.export_structures(
                self.partitions['generation'],
                f'{self.output_path}/{partition_path}/generation/gen{generation}'
            )

            # 8) Update main dataset
            self._update_main_dataset(self.partitions['generation'].containers, generation)

            # 9) Check convergence
            conv_results = self._check_convergence_wrapper(generation, objectives, features, structures=self.partitions['dataset'].containers)

            # 10) Save generation data
            if self.save_logs:
                self._save_generation_data_wrapper(
                    generation=generation,
                    conv_results=conv_results,
                    features=features,
                    objectives=objectives,
                    selected_indices=selected_indices,
                    temperature=temperature,
                    partition_path=partition_path
                )

            # Check if workflow is converged
            if conv_results.get('converge', False):
                self.logger.info("Convergence reached. Terminating workflow.")
                break

        # Wrap-up
        total_time = time.time() - overall_start
        self.logger.info(f"Workflow completed in {total_time:.2f} seconds.")

        # Optionally generate final plots
        t0 = time.time()
        if self.save_logs:
            try:
                self.evolution_plotter.generate_all_plots(
                    logger_dir=f'{self.output_path}/{partition_path}/logger',
                    output_dir=f'{self.output_path}/{partition_path}/plot'
                )
            except Exception as e:
                # Log the exception with traceback so you know exactly what went wrong
                self.logger.exception(
                    "Failed to generate evolution plots in '%s/%s': %s",
                    self.output_path, partition_path, e
                )

        # === Final Evaluation === # 
        self.export_structures(partition=self.partitions['dataset'], file_path=f'{self.output_path}/{partition_path}/', sort=True, )
        self.export_structures(partition=self.partitions['template'], file_path=f'{self.output_path}/{partition_path}/generation' )

        self.time_log['plot_evolution_data'] = time.time() - t0

        return self.partitions['dataset']


