from __future__ import annotations
from typing import Any, List, Optional
from abc import ABC, abstractmethod

class TransitionKernel(ABC):
    """
    Interface for transition kernels in the GA engine.
    A kernel defines how to move from a current state to a next state.
    """

    @abstractmethod
    def step(
        self,
        parents: Any,
        candidates: List[Any],
        **kwargs
    ) -> List[Any]:
        """
        Applies the transition logic.
        """
        raise NotImplementedError

    def _get_walker_id(self, s: Any, default: Optional[int] = None) -> Optional[int]:
        """
        Retrieves the persistent walker_id from a structure, favoring apm.metadata.
        """
        apm = getattr(s, "AtomPositionManager", None)
        if apm is not None:
            if getattr(apm, "metadata", None) is None:
                apm.metadata = {}
            
            # Check apm.metadata first
            wid = apm.metadata.get("walker_id")
            if wid is not None:
                return wid
            
            # Fallback to info but migrate to apm if found
            if hasattr(s, "info"):
                wid = s.info.get("walker_id")
                if wid is not None:
                    apm.metadata["walker_id"] = wid
                    return wid
        
        # Structure-level info fallback
        if hasattr(s, "info"):
            return s.info.get("walker_id", default)
            
        return default

    def _ensure_walker_id(self, structures: List[Any]):
        """
        Assigns a persistent walker_id to apm.metadata if missing.
        """
        for idx, s in enumerate(structures):
            apm = getattr(s, "AtomPositionManager", None)
            if apm is not None:
                if getattr(apm, "metadata", None) is None:
                    apm.metadata = {}
                
                if "walker_id" not in apm.metadata:
                    # Check info first before using index as default
                    info_wid = getattr(s, "info", {}).get("walker_id")
                    apm.metadata["walker_id"] = info_wid if info_wid is not None else idx
            elif hasattr(s, "info"):
                if "walker_id" not in s.info:
                    s.info["walker_id"] = idx

    def _get_phi(self, struct: Any) -> float:
        """
        Extracts the energy/objective value from a structure, favoring apm.metadata.
        """
        apm = getattr(struct, "AtomPositionManager", None)
        
        # 2. Try the volatile energy from APM
        val = None
        if apm is not None and getattr(apm, "E", None) is not None:
            val = apm.E

        return float(val) if val is not None else float("inf")

    def _match_parents_and_children(self, parents: List[Any], candidates: List[Any]) -> dict:
        """
        Matches candidates to parents using walker_id, falling back to index.
        """
        children_by_walker = {}
        for idx, child in enumerate(candidates):
            wid = self._get_walker_id(child)
            # Fallback to index if no walker_id yet (first pass)
            if wid is None:
                wid = idx
            children_by_walker[wid] = child
        return children_by_walker

    def allows_agent_sync(self) -> bool:
        """
        Whether this kernel allows population-wide synchronization (agent_sync).
        Default is True. Use False for kernels that require strict Markov Chain integrity (e.g. Gibbs).
        """
        return True

