from typing import List, Optional
from .metropolis import MetropolisKernel

class ReplicaAcceptanceKernel(MetropolisKernel):
    """
    Metropolis-Hastings kernel for Replica Exchange.
    Handles temperature ladders for individual replica acceptance.
    """
    def __init__(self, T_ladder: List[float], rng=None):
        super().__init__(
            sampling_temperature=T_ladder[0] if T_ladder else 1.0, # Default, but hooks will override
            rng=rng
        )
        self.T_ladder = T_ladder

    def get_temperature(self, idx: int, walker_id: int, **kwargs) -> float:
        """
        Returns the temperature for the specific replica using walker_id if available.
        """
        if not self.T_ladder:
            return self.sampling_temperature
        target_idx = walker_id if walker_id is not None else idx
        return self.T_ladder[target_idx % len(self.T_ladder)]
