from __future__ import annotations
from typing import Any, List
from .base import TransitionKernel

class CompositeKernel(TransitionKernel):
    """
    Chains multiple transition kernels sequentially.
    Useful for applying acceptance criteria followed by ensemble moves.
    Kernels can share information by reading/writing to kwargs['shared_state'].
    """

    def __init__(self, kernels: List[TransitionKernel]):
        self.kernels = kernels

    def step(
        self,
        parents: Any,
        candidates: List[Any],
        **kwargs
    ) -> List[Any]:
        """
        Sequentially applies each kernel in the composite.
        """
        # Initialize a shared state dictionary to allow kernels to communicate
        kwargs.setdefault('shared_state', {})
        
        current_state = candidates
        for kernel in self.kernels:
            current_state = kernel.step(
                parents=parents,
                candidates=current_state,
                **kwargs
            )
        return current_state

    def allows_agent_sync(self) -> bool:
        # If any kernel disallows sync, the composite disallows it.
        return all(k.allows_agent_sync() for k in self.kernels)
