from .explorer import EvolutionaryStructureExplorer
import os
import copy
import pickle
from ..utils.logger import WorkflowLogger
from typing import Any, Dict, Tuple, List, Optional

# =============================================================================
# Higher-Level Manager to Run Multiple Supercell Instances
# =============================================================================
class SupercellEvolutionManager:
    """
    Manages multiple evolutionary structure exploration instances across
    a series of supercell dimensions. For each defined supercell dimension,
    an instance of EvolutionaryStructureExplorer is created, executed, and
    its results stored for further analysis.
    
    Additional functionality includes result retrieval, listing processed
    supercells, saving state to disk, and resetting the manager state.
    """
    def __init__(self,                  
                 params_base: dict, 
                 dataset_path: Optional[str] = None, 
                 template_path: Optional[str] = None,
                 output_path: str = '.', 
                 supercell_list: List[Tuple[int, int, int]] = [(1, 1, 1)], 
                 params_supercell: Optional[Dict[str, List[Any]]] = None,
                 debug: bool = False):
        """
        Initializes the SupercellEvolutionManager with the necessary parameters.
        
        Parameters
        ----------
        params_base : dict
            Base evolutionary parameters that will be shared and updated for each explorer instance.
        dataset_path : str, optional
            File path to the dataset of initial structures.
        template_path : str, optional
            File path to the template structures.
        output_path : str
            Directory where outputs and logs are stored.
        supercell_list : list of tuple(int, int, int)
            List of supercell dimensions (e.g., [(1,1,1), (2,1,1), (2,2,1), ...]).
        params_supercell : dict, optional
            Parameters specific to each supercell. For each key, a list corresponding to each supercell dimension is expected.
        debug : bool, optional
            If True, enables debug logging.
            
        Raises
        ------
        AssertionError
            If the provided parameters do not conform to the expected types or lengths.
        ValueError
            If any of the supercell dimensions is not a tuple of three integers.
        """
        # Basic type validations
        assert isinstance(params_base, dict), "params_base must be a dictionary."
        assert isinstance(supercell_list, list), "supercell_list must be a list of 3-tuples."
        for dim in supercell_list:
            assert isinstance(dim, tuple) and len(dim) == 3, f"Each supercell dimension must be a tuple of 3 integers. Found: {dim}"

        # If supercell-specific parameters are provided, ensure they match the number of supercells.
        if params_supercell is not None:
            assert isinstance(params_supercell, dict), "params_supercell must be a dictionary if provided."
            for key, value in params_supercell.items():
                if not isinstance(value, list):
                    raise ValueError(f"Value for key '{key}' in params_supercell must be a list.")
                if len(value) < len(supercell_list):
                    raise ValueError(
                        f"Invalid params_supercell for key '{key}': received list length {len(value)}; expected {len(supercell_list)}."
                    )

        self.params_base = params_base

        for params_key, params_item in params_supercell.items():
            if len(params_item) < len(supercell_list):
                raise ValueError(f"Invalid params_supercell for key {params_key}: key len {len(params_item)}. Expected {len(supercell_list)} (same dim as supercell_list).")

        self.params_supercell = params_supercell
        
        self.dataset_path = dataset_path
        self.template_path = template_path
        self.output_path = output_path
        self.supercell_list = supercell_list
        self.debug = debug
        self.logger = WorkflowLogger.setup_logger('SupercellEvolutionManager', self.output_path)

        # Dictionaries to store each explorer instance and its results.
        # Dictionaries to store each explorer instance and its results.
        self.explorer_instances: Dict[Tuple[int, int, int], EvolutionaryStructureExplorer] = {}
        self.results: Dict[Tuple[int, int, int], Any] = {}
        
    def run_all_supercells(self):
        """
        Executes the evolutionary exploration for each defined supercell dimension.
        
        For each supercell:
          1. Validates the dimension tuple.
          2. Searches within previous results to locate the best candidate whose
             dimensions are submultiples of the current dimensions (favoring larger volumes).
          3. Creates a unique output partition name and updates the base parameters.
          4. Instantiates an EvolutionaryStructureExplorer.
          5. Transfers candidate structure data if a compatible candidate exists.
          6. Runs the explorer instance, logging and storing outcomes.
          
        Returns
        -------
        Dict[Tuple[int, int, int], Any]
            A mapping from each supercell dimension to its exploration result partition.
        
        Raises
        ------
        ValueError
            If any supercell dimension is incorrectly specified.
        """
        total_supercells = len(self.supercell_list)
        self.logger.info(f"Starting evolutionary exploration for {total_supercells} supercell configuration(s).")
        
        for sc_dim_i, sc_dim in enumerate(self.supercell_list):
            self.logger.info(f"Processing supercell {sc_dim_i + 1}/{total_supercells}: Dimension {sc_dim}")

            # Validate that each supercell dimension is a 3-tuple.
            if not isinstance(sc_dim, tuple) or len(sc_dim) != 3:
                error_message = f"Invalid supercell dimensions: {sc_dim}. Expected a tuple of 3 integers."
                self.logger.error(error_message)
                raise ValueError(error_message)

            # Determine the best candidate from previously computed results.
            best_candidate = None
            best_candidate_size = 0
            self.logger.info("Searching for a compatible candidate among previous results.")
            # Iterate through available candidate dimensions in self.results to find a suitable candidate.
            for candidate_dim in self.results.keys():
                if (sc_dim[0] % candidate_dim[0] == 0 and
                    sc_dim[1] % candidate_dim[1] == 0 and
                    sc_dim[2] % candidate_dim[2] == 0):
                    candidate_size = candidate_dim[0] * candidate_dim[1] * candidate_dim[2]
                    if candidate_size > best_candidate_size:
                        best_candidate = candidate_dim
                        best_candidate_size = candidate_size

            if best_candidate is not None:
                self.logger.info(f"Candidate found: {best_candidate} as a submultiple of {sc_dim} with volume {best_candidate_size}.")
            else:
                self.logger.info("No compatible candidate found for preloading data.")

            # Construct a partition identifier incorporating the supercell dimensions.
            partition_path = f'supercell_{sc_dim[0]}{sc_dim[1]}{sc_dim[2]}'
            self.logger.info(f"Output partition set to: {partition_path}.")
            # Note: partition_path is defined here and can be used later for saving outputs.

            # Instantiate a new evolutionary explorer instance for the current supercell.
            params_updated = self.params_base.copy()
            if isinstance(self.dataset_path, list):
                params_updated['dataset_path'] = self.dataset_path[sc_dim_i]
                self.logger.info(f"Using dataset path from list: {params_updated['dataset_path']}.")
            elif self.dataset_path is not None:
                params_updated['dataset_path'] = self.dataset_path
                self.logger.info(f"Using provided dataset path: {self.dataset_path}.")

            if isinstance(self.template_path, list):
                params_updated['template_path'] = self.template_path[sc_dim_i]
                self.logger.info(f"Using template path from list: {params_updated['template_path']}.")
            elif self.template_path is not None:
                params_updated['template_path'] = self.template_path
                self.logger.info(f"Using provided template path: {self.template_path}.")
            
            # Set output path: ensure outputs are stored in a subdirectory specific to the supercell.
            if isinstance(self.output_path, str):
                params_updated['output_path'] = os.path.join(self.output_path, partition_path)
            else:
                params_updated['output_path'] = partition_path

            # Update any supercell-specific parameters.
            for key, value_list in self.params_supercell.items():
                params_updated[key] = value_list[sc_dim_i]
                self.logger.info(f"Parameter '{key}' set to {value_list[sc_dim_i]} for current supercell.")

            # Instantiate the evolutionary explorer.
            self.logger.info("Instantiating EvolutionaryStructureExplorer instance.")
            explorer = EvolutionaryStructureExplorer(params=params_updated, debug=self.debug)
            explorer.logger.info(f"Starting exploration for supercell dimension: {sc_dim}")
            
            # If a best candidate is found, attempt to transfer its data.
            if best_candidate is not None:
                explorer.logger.info(
                    f"Found previously explored supercell {best_candidate} as a submultiple of {sc_dim}. "
                    "Transferring candidate structure data."
                )
                try:
                    # Deep copy the candidate results to avoid unintended side effects.
                    best_candidate_results = copy.deepcopy(self.results[best_candidate])
                    # For every container in the candidate's results, update atom positions to reflect the new supercell.
                    for container in best_candidate_results.containers:
                        container.AtomPositionManager.generate_supercell( repeat= tuple( int(sc/bc+.1) for sc, bc in zip(sc_dim, best_candidate)) )
                    # Preload the dataset for the explorer with the updated candidate containers.
                    explorer._containers_preload['dataset'] = best_candidate_results.containers

                    # Deep copy the candidate results to avoid unintended side effects.
                    best_explorer_instances = copy.deepcopy(self.explorer_instances[best_candidate])
                    # For every container in the candidate's results, update atom positions to reflect the new supercell.
                    for container in best_explorer_instances._partitions['template'].containers:
                        container.AtomPositionManager.generate_supercell( repeat= tuple( int(sc/bc+.1) for sc, bc in zip(sc_dim, best_candidate)) )
                    # Preload the dataset for the explorer with the updated candidate containers.
                    explorer._containers_preload['template'] = best_explorer_instances._partitions['template'].containers

                except Exception as e:
                    explorer.logger.error(f"Error transferring candidate results for candidate {best_candidate}: {e}")
            else:
                explorer.logger.info("No compatible candidate found; starting exploration without preloaded data.")

            # Run the evolutionary exploration.
            try:
                self.logger.info(f"Running evolutionary exploration for supercell dimension {sc_dim}.")
                result_partition = explorer.run(debug=self.debug)
                self.logger.info(f"Exploration completed successfully for supercell {sc_dim}.")
            except Exception as e:
                explorer.logger.error(f"Exploration failed for supercell {sc_dim}: {e}", exc_info=True)
                continue

            # Save the current explorer instance and its result.
            self.explorer_instances[sc_dim] = explorer
            self.results[sc_dim] = result_partition
            self.logger.info(f"Results for supercell {sc_dim} stored successfully. ( {result_partition.N} new strucutres )")

        self.logger.info("Completed exploration for all supercells.")
        return self.results












