from __future__ import annotations
from typing import Protocol, List, Any, Iterator, Optional, Dict, Union, Sequence, runtime_checkable
import numpy as np

# --- Data Representation Protocols ---

class IAtomPositionManager(Protocol):
    """Protocol for managing atomic coordinates, symbols, and metadata."""
    @property
    def atomPositions(self) -> np.ndarray: ...
    @property
    def atomLabels(self) -> List[str]: ...
    @property
    def latticeVectors(self) -> np.ndarray: ...
    @property
    def atomCount(self) -> int: ...
    @property
    def E(self) -> Optional[float]: ...
    @property
    def metadata(self) -> Dict[str, Any]: ...
    @property
    def atomCountDict(self) -> Dict[str, int]: ...
    
    def set_atomPositions(self, pos: np.ndarray) -> None: ...
    def set_atomLabels(self, labels: List[str]) -> None: ...
    def set_latticeVectors(self, cell: np.ndarray) -> None: ...
    def set_E(self, E: float) -> None: ...

class IIndividual(Protocol):
    """Protocol for an evolutionary individual (structure container)."""
    @property
    def AtomPositionManager(self) -> IAtomPositionManager: ...

# --- Engine Component Protocols ---

class IHash(Protocol):
    def get_hash(self, individual: IIndividual, force_rehash: bool = False) -> str: ...
    def add_structure(self, individual: IIndividual, force_rehash: bool = False) -> bool: ...
    def already_visited(self, individual: IIndividual) -> bool: ...

class IAgent(Protocol):
    def active(self) -> bool: ...
    def append(self, obj: Union[IIndividual, List[IIndividual]]) -> bool: ...
    def flush(self) -> List[IIndividual]: ...
    def get_batch(self, prune: bool = False, fetch_limit: Optional[int] = None) -> Iterator[Any]: ...
    def update_policy_metrics(self, diversity: float, pressure: float) -> bool: ...

class ILineage(Protocol):
    def assign_lineage_info_par_partition(self, partition: Any, generation: int, parents: List[str], event: str) -> None: ...

class IDoE(Protocol):
    """Protocol for Design of Experiments (DoE) and constraint management."""
    def validate(self, features: np.ndarray) -> bool: ...
    def set_name_mapping(self, mapping: Dict[str, int]) -> Any: ...
    def add_constraint(self, constraint: Callable) -> None: ...

class IPopulation(Protocol):
    def get_dataset(self, name: str) -> Any: ...
    def set_population(self, name: str, population: List[IIndividual]) -> bool: ...
    def size(self, name: Optional[str] = None) -> int: ...

class IThermostat(Protocol):
    def update(self, generation: int, stall: int = 0, signals: Any = None) -> float: ...

class ISelector(Protocol):
    def select(self, objectives: np.ndarray, features: np.ndarray, size: int) -> np.ndarray: ...
    def set_temperature(self, temperature: float) -> bool: ...

class ITransitionKernel(Protocol):
    def step(self, parents: Any, candidates: List[IIndividual], population: IPopulation, simulator: ISimulator, ctx: Any) -> List[IIndividual]: ...
    def allows_agent_sync(self) -> bool: ...

class IMutation(Protocol):
    def mutate(self, individual: IIndividual, ctx: Any) -> IIndividual: ...

class ICrossover(Protocol):
    def crossover(self, parent1: IIndividual, parent2: IIndividual, ctx: Any) -> IIndividual: ...

class IVariation(Protocol):
    def mutate(self, individual: IIndividual, ctx: Any) -> IIndividual: ...

class IForeignerGenerator(Protocol):
    def generate(self, ctx: Any) -> List[IIndividual]: ...

class IGenerative(Protocol):
    def is_active(self, generation: int) -> bool: ...
    def generate_targets(
        self,
        *,
        generation: int,
        features: np.ndarray,
        objectives: np.ndarray,
        temperature: float,
    ) -> np.ndarray | None: ...

class IEvaluator(Protocol):
    def evaluate(self, individuals: List[IIndividual], ctx: Any) -> List[IIndividual]: ...
    def get_features(self, individuals: List[IIndividual]) -> np.ndarray: ...
    def get_objectives(self, individuals: List[IIndividual]) -> np.ndarray: ...

class ISimulator(Protocol):
    def run(self, individuals: List[IIndividual], **kwargs: Any) -> List[IIndividual]: ...

@runtime_checkable
class IConvergence(Protocol):
    def check(
        self,
        generation: int,
        objectives: np.ndarray,
        features: np.ndarray,
        accepted_structures: int = 0,
        proposed_structures: int = 0,
    ) -> Dict[str, Any]: ...

    def is_converge(self) -> bool: ...
    def improvement_found(self) -> bool: ...
    def get_stagnation(self) -> int: ...

class ILogger(Protocol):
    def gen_start(self, gen: int, population: IPopulation) -> None: ...
    def gen_end(self, gen: int, population: IPopulation) -> None: ...
    def info(self, msg: str, *args: Any, **kwargs: Any) -> None: ...
    def debug(self, msg: str, *args: Any, **kwargs: Any) -> None: ...
    def error(self, msg: str, *args: Any, **kwargs: Any) -> None: ...

class IPlotter(Protocol):
    def gen_start(self, gen: int, population: IPopulation) -> None: ...
    def gen_end(self, gen: int, population: IPopulation) -> None: ...