from typing import Any, List, Dict
import math
from .metropolis import MetropolisKernel

class GCMCKernel(MetropolisKernel):
    """
    Metropolis-Hastings kernel for the Grand Canonical ensemble (GCMC).
    Handles chemical potential terms in the acceptance criteria.
    """
    def __init__(self, temperature: float, mu: Dict[str, float], rng=None):
        super().__init__(
            sampling_temperature=temperature,
            rng=rng
        )
        self.mu = mu

    def get_delta_phi_total(self, delta_phi: float, child: Any, parent: Any, **kwargs) -> float:
        """
        Includes chemical potential in the energy delta.
        """
        delta_mu_N = 0.0
        apm_child = getattr(child, "AtomPositionManager", None)
        apm_parent = getattr(parent, "AtomPositionManager", None)
        
        if apm_child is not None and apm_parent is not None:
            counts_child = apm_child.atomCountDict
            counts_parent = apm_parent.atomCountDict
            for species, mu_val in self.mu.items():
                delta_n = counts_child.get(species, 0) - counts_parent.get(species, 0)
                delta_mu_N += mu_val * delta_n

        return delta_phi - delta_mu_N
