from __future__ import annotations
import random
import math
from typing import Any, List
from sage_lib.partition.Partition import Partition
from .base import TransitionKernel

class ReplicaExchangeKernel(TransitionKernel):
    """
    Kernel for Replica Exchange (Parallel Tempering) moves.
    """

    def __init__(self, cfg, logger, ctx: Any = None):
        self.cfg = cfg
        self.logger = logger
        self.ctx = ctx

    def step(
        self,
        parents: Any,
        candidates: List[Any],
        **kwargs
    ) -> List[Any]:
        """
        Executes replica exchange swaps.
        """
        ens = getattr(self.cfg, 'ensemble', None)
        if not ens or ens.kind != "replica_exchange":
            return candidates

        replica_cfg = ens.replica_exchange
        generation = self.ctx.generation if self.ctx else 0
        T_ladder = replica_cfg.T_ladder

        if generation % replica_cfg.exchange_every != 0:
            return candidates

        # 1. Ensure walker IDs and get consistent list
        candidates_list = list(candidates) if not isinstance(candidates, list) else candidates
        self._ensure_walker_id(candidates_list)
        
        n_replicas = len(candidates_list)
        n_ladder = len(T_ladder)
        
        # Standard even-odd swapping pattern for 1D temperature ladder
        offset = 0 if (generation // replica_cfg.exchange_every) % 2 == 0 else 1
        swaps = 0
        
        for i in range(offset, n_replicas - 1, 2):
            idx1, idx2 = i, i + 1
            w1, w2 = candidates_list[idx1], candidates_list[idx2]
            T1 = T_ladder[idx1 % n_ladder]
            T2 = T_ladder[idx2 % n_ladder]
            
            E1, E2 = self._get_phi(w1), self._get_phi(w2)
            
            # If any energy is missing, skip/accept? For exchange, skipping or accepting 
            # might depend, but per user request, we should consider it "accepted" 
            # or at least not fail. Let's force acceptance if missing.
            if E1 is None or E2 is None:
                accept = True
            else:
                delta = (1.0/T1 - 1.0/T2) * (E1 - E2)
                
                if delta >= 0.0:
                    accept = True
                else:
                    accept = random.random() < math.exp(delta)
                
            if accept:
                candidates_list[idx1], candidates_list[idx2] = w2, w1
                swaps += 1
                
        if swaps > 0 and self.logger:
            self.logger.info(f"[Exchange] Replica: Accepted {swaps} temperature cross-swaps.")

        return candidates_list

    def allows_agent_sync(self) -> bool:
        return True
