"""
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
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 select_multiobjective_iterative, evaluate_objectives, evaluate_features
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

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 = {
            'period': 40,
            'decay_rate': 0.1,
            'temperature': 1.0,
            'min_mutation_rate': 1,
            'initial_mutation_rate': 100.0,
            'min_prob': 0.1,
            'max_prob': 0.9,
        }
        
        self.multiobjective_params = {
            'weights':None,
            'repulsion_weight': 1.0,
            'temperature': 1.0,
            'num_select': 20,
            'repulsion_mode':'avg',
            'metric':'euclidean',
            'random_seed': 100.0,
        }

        self._convergence_params = {
            'objective_threshold':  0.01,
            'feature_threshold': 0.01,
            'stall_threshold': 5,  
            }
        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,

            'supercell':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 = {'generation':[]}
        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'])

        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=1e-2)

        

    # ----------------------------------------------------------------
    # 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):
        """ """
        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)
        
        # 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.supercell_steps = validated_params['supercell_steps']

        # 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()
        self.partitions['dataset'].read_files(
            file_location=dataset_path,
            source='xyz',
            verbose=True
        )

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

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

        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 generate_new_compositions(self, 
        generation:int, 
        foreigners: int, 
        features:np.array, 
        objectives:np.array,  
        temperature:float=None
        ):
        """
        Generate new structures based on composition or "foreigners" approach.

        Parameters
        ----------
        foreigners : int
            Number of new foreigner structures to generate.
        features : Any
            Feature set used to guide composition design if needed.
        templates : object
            Template structure(s) used as a basis for new compositions.
        Returns
        -------
        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:
                # Attempt to generate design points using the DoE method
                design_points = self.DoE.generate()
            except Exception as e:
                # Log the error details (including traceback) for further debugging
                logging.error(f"Error generating DoE design points: {e}", exc_info=True)
                return []

            dp_name = 'DoE'

        else:
            self.BO.fit(X=features, Y=objectives) # bounds=bounds, n_candidates=3, strategy="ei"
            design_points = self.BO.recommend_candidates( n_candidates=foreigners, temperature=temperature )
            dp_name = 'BO'

        new_structures = self.mutation_crossover_handler.foreigners_generation(
            structure=self.partitions['template'].containers[0],
            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(self.partitions[dp_name].containers)}).")
        return new_structures

    def export_structures(self, partition, file_path: str):
        """
        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.
        """
        os.makedirs(file_path, exist_ok=True)
        partition.export_files(
            file_location=f"{file_path}",
            source='xyz',
            label='enumerate',
            verbose=True
        )
        return True

    def _save_generation_data(
        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 number.
        conv_results : dict
            Dictionary with convergence results info.
        features : np.ndarray
            Array with computed features.
        objectives : np.ndarray
            Array with computed objectives.
        selected_indices : np.ndarray
            Indices of selected structures from multiobjective selection.
        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,
            "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
        }

        # 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.setdefault('save_generation_data', []).append(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
        -------
        None
        """
        '''
        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):
            self.generation_start = time.time()
            self.logger.info(f"--- Generation {generation} ---")

            if self.thermostat:
                self.thermostat.actualizate_temperature(generation=generation)
                temperature = self.thermostat.get_temperature()

            else:
                temperature = self.temperature

            # 1) Evaluate objectives and features
            t0 = time.time()
            features = evaluate_features(
                self.partitions['dataset'].containers, 
                self.features_funcs, 
            ) 
            if self.partitions['dataset'].size > self.min_size_for_filter or self.foreigners > 0: 
                objectives = evaluate_objectives(
                    self.partitions['dataset'].containers, 
                    self.objective_funcs, 
                )
                self.mutation_crossover_handler._adjust_mutation_probabilities( 
                    structures = self.partitions['dataset'].containers,
                    objectives = objectives,
                    features = features,
                )
            else:
                objectives = evaluate_objectives(
                    self.partitions['dataset'].containers, 
                    self.objective_funcs, 
                )

            self.time_log.setdefault('objectives_features', []).append(time.time() - t0)

            # 2) Example: multiobjective selection or quick fallback
            t0 = time.time()
            if self.partitions['dataset'].size > self.min_size_for_filter:
                selected_indices = select_multiobjective_iterative(
                    objectives,
                    features,
                    weights=self.multiobjective_params['weights'],
                    repulsion_weight=self.multiobjective_params['repulsion_weight'],
                    temperature=temperature,
                    num_select=self.multiobjective_params['num_select'],
                    repulsion_mode=self.multiobjective_params['repulsion_mode'],
                    metric=self.multiobjective_params['metric'],
                    random_seed=self.multiobjective_params['random_seed'],
                )   
                top_structures = copy.deepcopy( [self.partitions['dataset'].containers[i] for i in selected_indices] )
            else:     
                selected_indices = np.array(range( self.partitions['dataset'].size ))
                top_structures = copy.deepcopy( self.partitions['dataset'].containers )

            #ranking = np.argsort(np.argsort(objectives.ravel()))
            #print(111, ranking[selected_indices])

            self.time_log.setdefault('multiobjective_selection', []).append(time.time() - t0)

            # 3) 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,
            )
            mutated_crossover_structures = mutated_structures + crossed_structures
            self.mutation_rate_array = mutation_rate_array
            self.time_log.setdefault('Mutation_Crossover', []).append(time.time() - t0)

            # 4) Foreigners (new compositions)
            if self.foreigners > 0 or generation == 1:
                # Composition generation using a template structure (first container from template partition)
                new_structures = self.generate_new_compositions(
                    generation=generation,
                    foreigners=self.foreigners, 
                    features=features,
                    objectives=objectives, 
                    temperature=temperature,
                )
            else:
                new_structures = []
            self.time_log.setdefault('Foreigners', []).append(time.time() - t0)

            # 5) self Collision filtering
            t0 = time.time()
            no_collision_structures = []
            for s in mutated_crossover_structures + new_structures:
                if not s.AtomPositionManager.self_collision(factor=self.collision_factor):
                    no_collision_structures.append(s)
            self.time_log.setdefault('Self Collision', []).append(time.time() - t0)

            # 6) physical_model - simulations on the filtered structures
            t0 = time.time()
            self.partitions['generation'] = physical_model( 
                structures=no_collision_structures, 
                physical_model_func=self.physical_model_func, 
                temperature=temperature,
                logger=self.logger, 
                debug=self.debug
            )
            self.time_log.setdefault('physical_model', []).append(time.time() - t0)

            # 7) Hash Collision filtering
            t0 = time.time()
            unique_structures = []
            for structure in self.partitions['generation'].containers:
                # and False if the structure is a duplicate.
                if self.hash_map.add_structure(container=structure):
                    unique_structures.append(structure)
                else:
                    self.mutation_crossover_handler.hash_colition_penalization(container=structure)

            # Update the generation containers with the unique structures only.
            self.partitions['generation'].containers = unique_structures
            self.time_log.setdefault('Hash Collision', []).append(time.time() - t0)

            # 8) Export structures for this generation
            t0 = time.time()
            self.export_structures(self.partitions['generation'], f'{self.output_path}/{partition_path}/generation/gen{generation}')
            self.time_log.setdefault('export_structures', []).append(time.time() - t0)

            # 9) Update the main dataset
            t0 = time.time()
            self.partitions['dataset'].containers += self.partitions['generation'].containers
            structures_id = np.append(structures_id, [c.AtomPositionManager.metadata['id'] for c in self.partitions['generation'].containers ])
            generations = np.append(generations, np.ones(self.partitions['generation'].size)*generation)
            self.time_log.setdefault('update_dataset', []).append(time.time() - t0)

            # 10) Check convergence
            t0 = time.time()
            if self.partitions['dataset'].size >= self._min_size_for_filter and not features is None:
                conv_results = self.convergence_checker.check_convergence(
                    generation=generation,
                    objectives=objectives,
                    features=features,
                    debug=self.debug,
                    generation_start=self.generation_start,
                    time_log=self.time_log
                )

            else:
                conv_results = {}
            self.time_log.setdefault('check_convergence', []).append(time.time() - t0)

            # 11) SAVE GENERATION DATA
            self._save_generation_data(
                generation=generation,
                conv_results=conv_results,
                features=features,
                objectives=objectives,
                selected_indices=selected_indices,
                temperature=temperature,
                partition_path=partition_path
            )

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

        overall_end = time.time()  # End overall workflow timer
        total_time = overall_end - overall_start
        self.logger.info(f"Workflow completed in {total_time:.2f} seconds.")

        t0 = time.time()
        self.evolution_plotter.generate_all_plots(
            logger_dir=f'{self.output_path}/{partition_path}/logger', 
            output_dir=f'{self.output_path}/{partition_path}/plot'
        )
        self.time_log.setdefault('plot_evolution_data', []).append(time.time() - t0)

        return self.partitions['dataset']



