"""
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 ..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.
    md_params : dict
        Parameters for MD simulation or other property evaluations.
    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._convergence_params = {}
        self.information_metric = {}

        self._objectives_for_features_history = {}

        self._collision_factor = .4

        self.foreigners = 0

        # MD simulation parameters (if any additional parameters are needed, add here)
        self._md_params = {}

        # 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.thermostat = None

        self.physical_model_func = None

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

        self.setup_evolution_parameters(params)

        self.detailed_record = False

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

        self.convergence_checker = ConvergenceChecker(
            logger=self.logger, 
            detailed_record=self.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=1.0)

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

            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'],
        ) 

        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 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', 'convergence_params', 'collision_factor',
            'md_params', '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._md_params = validated_params['md_params']
        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_func = validated_params['physical_model_func']
        self.multiobjective_params = validated_params['multiobjective_params']
        self.information_metric = validated_params['information_metric']

        # 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:
            #MultiObjectiveSelector.temperature = temperature
            selected_indices = self.multiobjective_selector.select(
                objectives=objectives,
                features=features,
            )
            #print( objectives )
            #print( objectives[selected_indices] )
            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
        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.
        """
        # Evaluate feature descriptors for each candidate structure.
        features = evaluate_features(structures, self.features_funcs)
        
        # Filter candidates: Only retain those whose feature set is validated by self.BO.validate.
        valid_candidates = [
            structure
            for structure, feature in zip(structures, features)
            if self.BO.validate(features=feature)
        ]

        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
        -------
        no_collision_structures : list
            Filtered list of structures without self-collisions.
        """
        t0 = time.time()
        no_collision_structures = [
            s for s in structures
            if not s.AtomPositionManager.self_collision(factor=self._collision_factor)
        ]
        self.time_log['Self Collision'] = time.time() - t0
        return no_collision_structures

    def _execute_physical_model(self, no_collision_structures, temperature):
        """
        Executes the physical model (MD simulation or other) on filtered structures.

        Parameters
        ----------
        no_collision_structures : list
            List of structures that passed the self-collision check.
        temperature : float
            Current temperature, if needed by the physical model.

        Returns
        -------
        generation_partition : Partition
            Partition object containing the structures after the physical model was applied.
        """
        t0 = time.time()
        generation_partition = physical_model(
            structures=no_collision_structures, 
            physical_model_func=self.physical_model_func, 
            temperature=temperature,
            logger=self.logger, 
            debug=self.debug
        )
        self.time_log['physical_model'] = time.time() - t0
        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()
        unique_structures = []
        for structure in structures:
            if self.hash_map.add_structure(structure):
                unique_structures.append(structure)
            else:
                self.mutation_crossover_handler.hash_colition_penalization(container=structure)

        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()
        if self.partitions['dataset'].size >= self._min_size_for_filter and features is not None:

            if self.convergence_checker._information_driven: 
                cluster_array = self.information_ensamble_analyzer.get_cluster_counts(structures)
                information_novelty_has_improved = self.information_ensamble_metric.has_improved(cluster_array)
            else:
                information_novelty_has_improved = False

            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
            )
            
        else:
            conv_results = {}
        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,

            "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,

            "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):
        """
        Executes the entire evolutionary structure exploration workflow.

        This method loops over generations, performing classification,
        selection, mutation, crossover, composition generation, MD simulation,
        and convergence checking until either convergence is reached or the maximum
        number of generations is exceeded.

        Parameters
        ----------
        debug : bool, optional
            If True, enables debug mode for additional logging.

        Returns
        -------
        Partition
            The final dataset partition after running all generations.
        """

        '''
        generations_array, temperature_array = self.thermostat.evaluate_temperature_over_generations(max_generations=40)
        self.mutation_crossover_handler.evaluate_mutation_rate_over_generations( temperature=temperature_array )
        self.BO.evaluate_BO_over_generations( temperature=temperature_array )
        '''
        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, top_structures, objectives, features, temperature
            )

            # 6) Self-collision filtering
            no_collision_structures = self._filter_self_collisions(new_structures)

            # 7) Physical model execution
            self.partitions['generation'] = self._execute_physical_model(no_collision_structures, temperature)

            # 8) Hash collision filtering
            unique_structures = self._filter_hash_collisions(self.partitions['generation'].containers)
            unique_structures = self._filter_self_collisions(unique_structures)
            self.partitions['generation'].containers = unique_structures

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

            # 10) Update main dataset
            self._update_main_dataset(unique_structures, generation)

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

            # 12) Save generation data
            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()
        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:
            pass

        # === 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']


