from __future__ import annotations

import os
import time
import copy
import logging
import warnings
import numpy as np
from typing import List, Optional, Tuple, Dict, Any, Type, Union, Sequence, Iterator
from pathlib import Path
from tqdm import tqdm

from sage_lib.partition.Partition import Partition

from ezga.core.interfaces import (
    IPopulation, ILineage, IHash, IAgent, IIndividual
)

from ezga.utils.lineage import LineageTracker  # type: ignore
from ezga.utils.structure_hash_map import Structure_Hash_Map # type: ignore
from ezga.sync.agenticsync import Agentic_Sync
from ezga.utils.molecule_blacklist import BlacklistDetector

class Population(IPopulation):
    """
    Main repository for evolutionary structures and datasets.
    
    Handles loading, exporting, synchronizing, and filtering individuals
    using the 'sage_lib.partition' backend.
    """
    def __init__(
        self,
        *,
        dataset_path: Optional[Union[str, Path]] = None,
        template_path: Optional[Union[str, Path]] = None,
        output_path: Optional[Union[str, Path]] = None,
        collision_factor: float = 0.0,
        filter_duplicates: bool = True,
        export_hash_collisions: bool = True,
        size_limit: Optional[int] = None,
        fetch_limit: Optional[int] = None,

        db_path: Optional[Union[str, Path]] = None,
        db_ro_path: Optional[Union[str, Path]] = None,
        skip_hashing: bool = False,
        skip_lineage: bool = False,

        debug: bool = False,
        hash_map: Type[IHash] = Structure_Hash_Map,
        agent: Type[IAgent] = Agentic_Sync,
        lineage: Type[ILineage] = LineageTracker,

        constraints: Optional[List[Any]] = None,
        ef_bounds: Optional[Tuple[float, float]] = None,
        blacklist: Optional[List[str]] = None,
    ):
        """
        Initializes the Population.

        Args:
            dataset_path: Path to the initial structures file.
            template_path: Path to the template structure file.
            output_path: Directory to save output files.
            collision_factor: Factor for collision detection (van der Waals overlap).
            filter_duplicates: Whether to filter duplicate structures via hashing.
            size_limit: Maximum number of structures to keep in the main dataset.
            fetch_limit: Limit on the number of structures to fetch from sync agent.
            db_path: Path to local writable database (typically HDF5/ASE).
            db_ro_path: Path to read-only source database.
            skip_hashing: If True, skip registration of structures in the hash map.
            skip_lineage: If True, skip assigning lineage (parent/child) metadata.
            debug: Enable verbose debugging output.
            hash_map: Strategy for computing and storing structure hashes.
            agent: Multi-agent synchronization backend.
            lineage: Metadata manager for evolutionary genealogy.
            constraints: Optimization constraints for DoE.
            ef_bounds: Energy-per-atom bounds for physical sanity filtering.
            blacklist: List of SMARTS/patterns to reject.
        """
        self.logger = logging.getLogger(__name__)
        self._dataset_path = dataset_path
        self._template_path = template_path
        self._output_path = str(output_path) if output_path else '.'
        
        self.collision_factor = collision_factor
        self._filter_duplicates = filter_duplicates
        self.export_hash_collisions = export_hash_collisions
        self.size_limit = size_limit
        self.fetch_limit = fetch_limit
        
        self.skip_hashing = skip_hashing
        self.skip_lineage = skip_lineage

        self._db_path: Optional[str] = None
        self.db_path = db_path

        self._db_ro_path: Optional[str] = None
        self.db_ro_path = db_ro_path

        self.ef_bounds = ef_bounds
        self.blacklist_detector = BlacklistDetector(blacklist or [])
        self.debug = debug

        self._datasets: Dict[str, Optional[Partition]] = {
            'template': None,
            'dataset': None,
            'generation': None,
            'template_preload': None,
            'dataset_preload': None,
            'selected': None,
            'offspring': None,
        }

        self.hash_map = hash_map
        self.agent = agent
        self.lineage = lineage
        
        # We avoid direct import cycle by using a late import or checking type
        from ezga.DoE.DoE import DesignOfExperiments
        self.DoE = DesignOfExperiments(constraints=constraints)

    # ----------------------------------------------------------------
    # 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

    # --------------------------------------------------------------
    # Writable database path
    # --------------------------------------------------------------
    @property
    def db_path(self) -> Path:
        """Path to the writable database (as a Path object)."""
        return Path(self._db_path)

    @db_path.setter
    def db_path(self, value: str | Path | None):
        """Set writable database path; create directory if needed."""
        value = Path(value or ".")
        value.mkdir(parents=True, exist_ok=True)
        self._db_path = str(value)

    # --------------------------------------------------------------
    # Read-only database path
    # --------------------------------------------------------------
    @property
    def db_ro_path(self) -> Path:
        """Path to the read-only database (as a Path object)."""
        return Path(self._db_ro_path)

    @db_ro_path.setter
    def db_ro_path(self, value: str | Path | None):
        """Set read-only database path; create directory if needed."""
        value = Path(value or "db_ro")
        value.mkdir(parents=True, exist_ok=True)
        self._db_ro_path = str(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 datasets(self):
        """Gets the dictionary of partition objects."""
        return self._datasets

    def get_dataset(self, name: str) -> Partition: 
        """Gets a specific partition by name, creating it if it doesn't exist."""
        ds = self._datasets.get(name)
        if ds is None:
            return self.set_dataset(name)
        return ds

    def get_containers(self, name: str) -> List[IIndividual]:
        """Gets the list of structure containers from a partition."""
        ds = self.get_dataset(name)
        return list(ds.containers)

    def set_dataset(self, name: str, dataset: Optional[Union[Partition, List[IIndividual]]] = None) -> Partition:
        """
        Creates or overwrites a partition.
        
        Args:
            name: Label for the partition.
            dataset: Optional initial data (Partition or list of individuals).
            
        Returns:
            The created/updated Partition object.
        """
        if dataset is None:
            part = Partition()
        elif isinstance(dataset, list):
            part = Partition()
            part.add(dataset)
        else:
            part = dataset
            
        self._datasets[name] = part
        return part

    def set_population(self, name: str, population: List[IIndividual]) -> bool:
        """Alias for set_dataset to comply with IPopulation protocol."""
        self.set_dataset(name, dataset=population)
        return True

    def dataset_add(self, name: str, dataset: Union[Partition, List[IIndividual], IIndividual]) -> bool: 
        """Adds data to an existing partition."""
        self.get_dataset(name).add(dataset)
        return True

    def size(self, name: Optional[str] = None) -> int:
        """Returns the number of structures in a partition (default: 'dataset')."""
        target = name or 'dataset'
        ds = self._datasets.get(target)
        return ds.size if ds else 0
 
    def filter(self, name: str, idx: List[int]) -> Partition:
        """
        Filter and materialize a subset of individuals from a partition efficiently.
        
        This method is robust against missing HDF5 datasets (KeyError). If bulk 
        materialization fails, it attempts to recover as many valid structures as possible.

        Args:
            name: Partition name to filter (e.g., 'dataset', 'offspring').
            idx: List of structure indices or IDs to fetch.
            
        Returns:
            A new Partition containing the materialized structures.
        """
        dataset = self.get_dataset(name)
        
        # 1) Attempt multi-threaded bulk materialization (efficient)
        try:
            return dataset.materialize_by_ids(
                idx,
                dedup=True,
                sort_backend_reads=True,
                independent_duplicates=True
            )
        except (KeyError, Exception) as e:
            self.logger.warning(f"[Population] Bulk materialization for '{name}' failed: {e}. Attempting robust recovery.")

        # 2) Graceful fallback: Materialize one-by-one and skip broken IDs
        recovered_partition = Partition()
        failed_ids = []
        
        # We use tqdm to show progress if there are many IDs to recover
        for i in idx:
            try:
                # Load individual ID
                item = dataset.materialize_by_ids([i], dedup=True, independent_duplicates=True)
                recovered_partition.add(list(item.containers))
            except (KeyError, Exception):
                failed_ids.append(i)
        
        if failed_ids:
            self.logger.error(
                f"[Population] Failed to recover {len(failed_ids)} structures from '{name}'. "
                f"Missing IDs: {failed_ids[:10]}{'...' if len(failed_ids) > 10 else ''}"
            )

        return recovered_partition

    # ----------------------------------------------------------------
    # LOAD
    # ----------------------------------------------------------------
    def load_population(self, logger: Optional[Any] = None): 
        """
        Loads the initial population and template from disk or sync agent.
        """

        # ------------------------------------------------------------------
        # 0)  Book‑keeping: log start‑up and start a wall‑clock timer
        # ------------------------------------------------------------------
        if logger:
            self.logger = logger  # Use engine-provided logger if available
        self.logger.info("Loading dataset from files.")

        # ------------------------------------------------------------------
        # 1)  DATASET PARTITION ------------------------------------------------
        # ------------------------------------------------------------------
        # 1.1  Always create an *empty* Partition object first so that we have a
        #      valid container even if subsequent file I/O fails.  This avoids
        #      having to repeat *None* checks later on.
        if not isinstance(self.db_ro_path, (str, Path) ):
            self.db_ro_path = 'db_ro'
        if not isinstance(self.db_path, (str, Path)):
            self.db_path = '.'

        self.set_dataset( 
            'dataset', 
            Partition(
                storage='composite', # composite # hybrid
                base_root=self.db_ro_path,
                local_root=self.db_path,
            )
        )

        # 1.2  If the caller provided a valid *path* (and not, e.g., ``None`` or
        #      an already‑opened file object), read the structure files.
        if isinstance(self.dataset_path, (str, os.PathLike)):
            ds = self.get_dataset('dataset')
            if self.size_limit:
                ds.read_files(
                    file_location=str(self.dataset_path),
                    verbose=True,
                    sampling='random',
                    n_samples=self.size_limit
                )
            else:
                ds.read_files(
                    file_location=str(self.dataset_path),
                    verbose=True,
                )


                
        # 1.3  Merge any *pre‑loaded* containers (prepared by the caller before
        #      invoking this method) into the freshly created partition.
        if isinstance(self.get_dataset('dataset_preload'), list):
            self.dataset_add(name='dataset', dataset=self.get_dataset('dataset_preload'))

        # 1.4  Feed every single *Structure* object into the global hash map so
        #      that the rest of the code base can perform O(1) look‑ups for
        #      de‑duplication or collision tracking.
        if not self.skip_hashing and False:
            for container in tqdm(
                self.get_dataset('dataset').containers,
                desc="[load] Registering structures in hash map",
                unit="struct"
            ):
                self.hash_map.add_structure( container, force_rehash=False )

        # 1.5  If we are running in *synchronised* mode (i.e., multiple workers
        #      on different nodes), pull the current batch of structures from
        #      the synchronisation backend and add only the ones that do *not*
        #      produce a hash collision.
        agent_is_active = False
        if self.agent is not None and not isinstance(self.agent, type):
            try:
                agent_is_active = self.agent.active()
            except Exception:
                agent_is_active = False

        if agent_is_active:
            try: 
                sync_new = list(self.agent.get_batch(prune=False, fetch_limit=self.fetch_limit))
                containers_list = [];  [containers_list.extend(sub.containers) for _, sub in sync_new]
                unique_structures, not_unique_structures = self.filter_hash_collisions(individuals=containers_list, force_rehash=False)
                self.dataset_add(name='dataset', dataset=unique_structures )

                if logger:
                    logger.info(
                        "Genetic injection: %d structures added to dataset (Gen=%d).",
                        len(unique_structures),         
                        0,    
                    )
            except Exception as exc:
                if hasattr(self, "logger"):
                    self.logger.warning(
                        "[Agentic_Sync] Failed during load_population step: %s", exc, exc_info=True
                    )

        # ------------------------------------------------------------------
        # 2)  TEMPLATE PARTITION ---------------------------------------------
        # ------------------------------------------------------------------
        # 2.1  Create an empty partition exactly as was done for the dataset.
        self.set_dataset('template')
        if isinstance(self.template_path, (str, Path)):
            tpl = self.get_dataset('template')
            if tpl:
                tpl.read_files(
                file_location=str(self.template_path),
                source='xyz',
                    verbose=self.debug
            )
        else:
            self.logger.info("No template file provided, skipping template partition.")

        pre_tpl = self.get_dataset('template_preload')
        if isinstance(pre_tpl, list):
            self.dataset_add(name='template', dataset=pre_tpl)

        # Every *Structure* gets tagged with lineage metadata so that its origin
        # can be traced throughout the workflow.  Stage‑index *0* indicates the
        # very first step; the empty list is the parent set, and the string
        # "initialization" becomes the *event* label.
        if not self.skip_lineage:
            self.lineage.assign_lineage_info_par_partition(
                self.get_dataset('dataset'), 0, [], "initialization"
            )
            
        self.logger.info("Partitions loaded successfully.")

    
    def export_structures(
        self,
        dataset,
        file_path: str,
        sort: bool = False,
        key: int = 0,
        default_label: str = "config",
    ):
        """
        Export structures in the given dataset.

        Parameters
        ----------
        dataset : dataset
            Partition containing the structures to export.
        file_path : str
            Either:
              • a directory (no extension) → writes to <dir>/<default_label>.xyz, or
              • a full file path with extension (e.g., *.xyz) → writes exactly there.
        sort : bool
            If True, sort by objective 'key' (ascending) and store objectives in metadata.
        key : int
            Objective column to sort by when sort=True.
        default_label : str
            Used as filename when file_path is a directory.
        """
        p = Path(file_path)

        # Determine output directory and final file path
        if p.suffix:  # looks like a file (has extension)
            out_dir   = p.parent
            out_file  = p
        else:         # looks like a directory
            out_dir   = p
            out_file  = out_dir / f"{default_label}.xyz"

        # Create all needed directories
        out_dir.mkdir(parents=True, exist_ok=True)

        # Optional sorting by objective
        if sort:
            objectives = evaluate_objectives( list(dataset.containers), self.objective_funcs)
            for container, objective in zip(dataset.containers, objectives):
                apm = container.AtomPositionManager
                if not isinstance(getattr(apm, "metadata", None), dict):
                    apm.metadata = {}
                apm.metadata["objectives"] = list(objective)

            numeric = np.array([float(obj[key]) for obj in objectives])
            order = np.argsort(numeric, kind="mergesort")
            dataset.containers = [dataset.containers[i] for i in order]

        # Write using a full path (what your exporter expects)
        dataset.export_files(
            file_location=str(out_file),  # full path including filename + extension
            source="xyz",
        )
        return True

    def agent_sync(
        self, 
        dataset: Partition,
        objectives: np.ndarray,
        features: np.ndarray,
    ) -> List[IIndividual]:
        """
        Synchronize local population with the shared multi-agent backend.

        This method facilitates global information exchange by:
          1. Calculating local diversity and selection pressure metrics.
          2. Updating the adaptive Nash-inspired communication policy.
          3. Publishing newly discovered local candidates to the shared pool.
          4. Fetching unseen candidates discovered by other agents.
          5. Performing background maintenance (pruning) of the shared filesystem.

        Args:
            dataset: Partition containing the latest local batch of individuals.
            objectives: Matrix of objective values (N x num_objectives) for the batch.
            features: Matrix of feature descriptors (N x num_features) for the batch.

        Returns:
            A list of new structure containers received from the synchronization agent.
        """
        # 1. Early exit if agent is not configured or inactive
        if not self.agent or not hasattr(self.agent, "active") or not self.agent.active():
            return []

        try:
            # Extract containers to be published
            local_individuals = list(dataset.containers)
            
            # --- 1) ADAPTIVE POLICY UPDATE ---
            # We provide local state feedback to the agent so it can throttle 
            # communication based on global contention and local progress.
            diversity, pressure = self._calculate_sync_metrics(features, objectives)
            
            try:
                if hasattr(self.agent, "update_policy_metrics"):
                    self.agent.update_policy_metrics(diversity=diversity, pressure=pressure)
            except Exception as e:
                self.logger.warning(f"[Agentic_Sync] Policy update failed: {e}")

            # --- 2) PUBLISH CANDIDATES ---
            if local_individuals:
                try:
                    self.agent.append(local_individuals)
                except Exception as e:
                    self.logger.error(f"[Agentic_Sync] Failed to publish local candidates: {e}")

            # --- 3) FETCH GLOBAL DISCOVERIES ---
            discovered_individuals: List[IIndividual] = []
            try:
                # fetch_limit determines how many structures we grab at once
                sync_batches = self.agent.get_batch(prune=False, fetch_limit=self.fetch_limit)
                for batch_id, batch_partition in sync_batches:
                    try:
                        # Defensive extraction of containers from the batch partition
                        batch_items = getattr(batch_partition, "containers", [])
                        discovered_individuals.extend(batch_items)
                    except Exception as e:
                        self.logger.warning(f"[Agentic_Sync] Failed to extract from batch {batch_id}: {e}")
            except Exception as e:
                self.logger.error(f"[Agentic_Sync] Failed to fetch remote batches: {e}")

            # --- 4) MAINTENANCE ---
            # Perform periodic pruning of old batches to prevent storage bloat.
            try:
                if hasattr(self.agent, "attempt_prune"):
                    pruned = self.agent.attempt_prune(interval=60.0, budget=8, min_age=10.0)
                    if pruned and self.logger.isEnabledFor(logging.DEBUG):
                        self.logger.debug(f"[Agentic_Sync] Pruned {pruned} stale batches.")
            except Exception as e:
                self.logger.debug(f"[Agentic_Sync] Maintenance prune failed: {e}")

            return discovered_individuals

        except Exception as e:

            self.logger.error(f"[Agentic_Sync] Critical failure during synchronization: {e}", exc_info=True)
            return []

    def _calculate_sync_metrics(self, features: np.ndarray, objectives: np.ndarray) -> Tuple[float, float]:
        """
        Compute high-level summary metrics for the current batch.
        (Nash-response logic)
        Diversity: Proxy for exploration volume (Sum of feature variances).
        Pressure: Proxy for exploitation intensity (Mean magnitude of objectives).
        """
        diversity = 0.0
        pressure = 0.0

        try:
            # --- Fast diversity estimate (O(NF)) ---
            # (N, F) features
            if features.size > 0 and features.ndim == 2 and features.shape[0] > 1:
                diversity = float(np.var(features, axis=0).sum())
        except Exception:
            pass

        try:
            # --- Fast pressure estimate (O(N)) ---
            # (N, M) objectives
            if objectives.size > 0:
                # Handle both vector and scalar objectives securely
                objs = objectives if objectives.ndim == 2 else objectives.reshape(-1, 1)
                pressure = float(np.linalg.norm(objs, axis=1).mean())
        except Exception:
            pressure = 0.0

        return diversity, pressure

    def validate_candidates(self, individuals: list, features: np.array, remove:Optional[bool] = None) -> list:
        """
        Validate candidate individuals 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.DoE.validate). The validator returns True if the candidate's features
        satisfy the design requirements and False otherwise.

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

        Returns
        -------
        list
            A list of candidate individuals that meet the feature space requirements.
        """

        # 2) Local bindings for speed
        validate = self.DoE.validate if self.DoE else lambda *args, **kwargs: True

        # 3) Single pass: build valid list and penalise invalid
        valid_candidates = []
        physically_invalid_candidates = []
        out_of_doe_candidates = []
        blacklist_invalid_candidates = []

        for struct, feat in zip(individuals, features):
            # 1) DOE‐validation check
            if not self.DoE.validate(features=feat):
                out_of_doe_candidates.append(struct)
                continue  # skip to next structure

            # 2) Physical‐sanity check
            if self.ef_bounds is not None:
                apm = struct.AtomPositionManager
                if apm.E is not None and apm.atomCount != 0:
                    ef_per_atom = apm.E / apm.atomCount
                    if not self.ef_bounds[0] <= ef_per_atom <= self.ef_bounds[1]:
                        physically_invalid_candidates.append(struct)
                        continue

            # 2) blacklist_detector check
            try:
                present, tag = self.blacklist_detector.contains(struct, remove=remove)
            except Exception as ex:
                # Fail-soft: proceed with other checks; optionally log when debug
                if True or self.debug:
                    print(f"[validate_candidates] blacklist check error: {ex}")
                present, tag = (False, None)

            if present:
                # Optionally annotate offending motif for auditability
                apm = struct.AtomPositionManager
                md = getattr(apm, "metadata", None)
                if isinstance(md, dict):
                    md["blacklist_hit"] = tag
                elif hasattr(apm, "metadata"):
                    try:
                        apm.metadata = {"blacklist_hit": tag}
                    except Exception:
                        pass
                blacklist_invalid_candidates.append(struct)
                continue

            # 4) Survives all checks
            valid_candidates.append(struct)

        # --- Export invalid groups (mirroring your existing behavior) ---

        # 3a) Export any out_of_doe invalid individuals
        if out_of_doe_candidates:
            invalid_dataset = Partition()
            invalid_dataset.add_container( out_of_doe_candidates )
            self.export_structures(
                invalid_dataset,
                f'{self.output_path}/out_of_doe.xyz'
            )

        # 3b) Export any physically invalid individuals
        if physically_invalid_candidates:
            invalid_dataset = Partition()
            invalid_dataset.add_container( physically_invalid_candidates )
            self.export_structures(
                invalid_dataset,
                f'{self.output_path}/physically_invalid.xyz'
            )

        if blacklist_invalid_candidates:
            invalid_dataset = Partition()
            invalid_dataset.add_container(blacklist_invalid_candidates)
            self.export_structures(invalid_dataset, f"{self.output_path}/blacklisted.xyz")

        return valid_candidates, out_of_doe_candidates+physically_invalid_candidates+blacklist_invalid_candidates 

    def filter_hash_collisions(self, individuals:list, force_rehash:bool=True, export:Optional[bool]=None):
        """
        Filters out individuals that collide in the hash map (i.e., duplicates).

        Parameters
        ----------
        individuals : list
            List of candidate individuals after the physical model.
        export : bool, optional
            Whether to export basin configurations. If None, uses self.export_hash_collisions.

        Returns
        -------
        unique_structures : list
            Structures that are unique (not in the hash map).
        """

        if export is None:
            export = self.export_hash_collisions

        def _composition_key(comp: dict) -> str:
            """
            Return a deterministic, reduced-formula key such as 'C1-H1-N1-O2'
            for any integer multiple - e.g. both {'C':2,'H':2,'N':2,'O':4}
            and {'C':6,'H':6,'N':6,'O':12} collapse onto the same key.

            • Counts **must** be non-negative integers.  If you store floats,
              cast/round them beforehand.
            • Element order is alphabetical to keep the key canonical.
            """
            counts = [int(comp[el]) for el in comp]
            g = 1 #reduce(gcd, counts) or 1                       # robust gcd
            return "-".join(
                f"{el}{int(comp[el] // g)}"                    # always include ‘1’
                for el in sorted(comp)
            )

        if not self._filter_duplicates:
            return individuals, []

        unique_structures = []
        not_unique_structures = []
        for individual in individuals:
            hash_ = self.hash_map.get_hash(individual, force_rehash=force_rehash)

            '''
            # deduplicate a batch of candidate payload-hashes in one shot
            seen = composite_store.has_hashes(batch_hashes)
            new_candidates = [cand for cand in candidates if not seen[cand.hash]]
            '''

            if not self.get_dataset('dataset').has_hash(hash_):
                unique_structures.append(individual)
                continue

            elif force_rehash:
                dataset = Partition()
                dataset.add_container(individual)
                
                comp_key = _composition_key(individual.AtomPositionManager.atomCountDict)
                if export:
                    self.export_structures(
                        dataset=dataset,
                        file_path=f'{self._output_path}/basin/{comp_key}/{individual.AtomPositionManager.metadata["hash"]}/config_basin.xyz',
                    )

            not_unique_structures.append(individual)

        return unique_structures, not_unique_structures


    def filter_selfcollision(self, individuals:list, remove:bool=False):
        """
        Filters out individuals that have self-collision above the collision factor threshold.

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

        Returns
        -------
        non_colliding : list
            Filtered list of individuals without self-collisions.
        """
        if self.collision_factor < 1e-6:
            return individuals, []

        # 1) Compute collision mask in one pass (C-optimized list comp)
        mask = []
        for s in individuals:
            try:
                m = s.AtomPositionManager.self_collision(factor=self.collision_factor, remove=remove)
            except Exception:
                # Treat corrupted structures as colliding (to be filtered out)
                m = True
            mask.append(m)

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

        # 3) Export invalid individuals, if any
        if colliding:
            dataset = Partition()
            dataset.add_container( colliding )
        return non_colliding, colliding


    def update_main_dataset(
        self,
        dataset: Partition,
        generation: int,
        features_prev: np.ndarray,
        objectives_prev: np.ndarray,
    ) -> int:
        """
        Efficiently updates the main population dataset, performing size pruning if necessary.
        
        This method manages the global population limit (cap). It identifies the worst-performing
        structures from the current population using the provided objective scores and removes 
        them from the storage backend using their persistent IDs to avoid memory overhead.

        Parameters
        ----------
        dataset : Partition
            The new generation of structures to be added.
        generation : int
            Current generation index.
        features_prev : np.ndarray
            Features associated with the current population.
        objectives_prev : np.ndarray
            Multi-objective matrix (or 1D array) for the current population survivors.

        Returns
        -------
        int
            The number of structures discarded during the pruning process.
        """
        cap = self.size_limit
        ds = self.get_dataset("dataset")
        n_discarded = 0
        
        if cap is not None:
            # Check if adding the new generation exceeds the size limit
            ds_size = ds.size
            new_size = dataset.size
            needed_size = ds_size + new_size
            
            if needed_size > cap:
                # Calculate how many structures must be removed
                excess = int(needed_size - cap)
                
                # Fetch IDs from the backend without materializing heavy objects
                all_ids = ds.list_ids()
                n_old = len(all_ids)

                # --- Ranking & Selection ---
                if objectives_prev.ndim == 1:
                    # Single objective: simple sort
                    scores = objectives_prev[:n_old]
                    idx_bad = np.argsort(scores)[-excess:]
                else:
                    # Multi-objective: Use Non-Dominated Sorting for Pareto-aware pruning
                    # We want to keep all fronts and prune from the worst rank downwards.
                    obj_crop = objectives_prev[:n_old]
                    
                    # 1. High-Performance Pareto + Niche Pruning
                    def fast_nondominated_sort_with_niche_preservation(objs, sizes, n_cap):
                        N, K = objs.shape
                        
                        # --- Objective Normalization ---
                        # Crucial for stable Pareto and Crowding calculations
                        obj_min, obj_max = objs.min(axis=0), objs.max(axis=0)
                        span = obj_max - obj_min + 1e-9
                        objs_norm = (objs - obj_min) / span

                        domination_counts = np.zeros(N, dtype=int)
                        dominated_sets = [[] for _ in range(N)]
                        
                        # Vectorized Domination Check (Optimized for thousands)
                        for i in range(N):
                            # j dominates i? 
                            is_dominated = np.all(objs_norm <= objs_norm[i], axis=1) & np.any(objs_norm < objs_norm[i], axis=1)
                            domination_counts[i] = np.sum(is_dominated)
                            dominated_sets[i] = np.where(np.all(objs_norm[i] <= objs_norm, axis=1) & np.any(objs_norm[i] < objs_norm, axis=1))[0]
                        
                        # Standard Non-Dominated Sorting peel
                        fronts = []
                        current_front = np.where(domination_counts == 0)[0].tolist()
                        while current_front:
                            fronts.append(current_front)
                            next_front = []
                            for p in current_front:
                                for q in dominated_sets[p]:
                                    domination_counts[q] -= 1
                                    if domination_counts[q] == 0:
                                        next_front.append(q)
                            current_front = next_front
                        
                        # --- Niche Preservation (Size-grouping) ---
                        # We MUST keep the best individual of EACH size to prevent constants 
                        # from killing structural innovation.
                        must_keep = set()
                        unique_sizes = np.unique(sizes)
                        for s in unique_sizes:
                            idx_s = np.where(sizes == s)[0]
                            # Best MSE in this size niche wins a seat
                            best_mse_idx = idx_s[np.argmin(objs[idx_s, 0])]
                            must_keep.add(best_mse_idx)
                        
                        # 2. Select survivors (starting with must_keep)
                        survivors = list(must_keep)
                        
                        # Fill remaining capacity using Pareto Fronts
                        for front in fronts:
                            if len(survivors) >= n_cap: break
                            
                            # Filter out those already kept by niche logic
                            front_unique = [i for i in front if i not in must_keep]
                            if not front_unique: continue
                            
                            if len(survivors) + len(front_unique) <= n_cap:
                                survivors.extend(front_unique)
                            else:
                                # Prune front using Crowding Distance in normalized space
                                rem = n_cap - len(survivors)
                                idx_f = np.array(front_unique)
                                f_objs = objs_norm[idx_f]
                                dists = np.zeros(len(idx_f))
                                for k in range(K):
                                    idx_sorted = np.argsort(f_objs[:, k])
                                    dists[idx_sorted[0]], dists[idx_sorted[-1]] = np.inf, np.inf
                                    o_range = f_objs[idx_sorted[-1], k] - f_objs[idx_sorted[0], k]
                                    if o_range > 0:
                                        for j in range(1, len(idx_f) - 1):
                                            dists[idx_sorted[j]] += (f_objs[idx_sorted[j+1], k] - f_objs[idx_sorted[j-1], k]) / o_range
                                
                                best_in_front = np.argsort(dists)[-rem:]
                                survivors.extend(idx_f[best_in_front].tolist())
                                break
                        
                        all_indices = set(range(N))
                        survivor_set = set(survivors)
                        return list(all_indices - survivor_set)

                    # For SR, Objective 1 is Size.
                    sizes = obj_crop[:, 1] if obj_crop.shape[1] > 1 else np.zeros(n_old)
                    idx_bad = fast_nondominated_sort_with_niche_preservation(obj_crop, sizes, n_old - excess)


                
                # Map indices to persistent backend IDs for safe removal
                ids_to_remove = [all_ids[i] for i in idx_bad if i < n_old]
                
                if ids_to_remove:
                    ds.remove(ids_to_remove)
                    n_discarded = len(ids_to_remove)


        # --- Integration ---
        # Add the new generation (already in memory) to the survivors
        ds.add(list(dataset.containers))
        
        return n_discarded