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

class GibbsTransferKernel(TransitionKernel):
    """
    Kernel for Gibbs Ensemble moves (particle/volume transfer).
    """

    def __init__(self, cfg, logger, evaluator: Any = None, ctx: Any = None):
        self.cfg = cfg
        self.logger = logger
        # Evaluator no longer needed for step, but keeping for factory compatibility
        self.evaluator = evaluator
        self.ctx = ctx

    def step(
        self,
        parents: Any,
        candidates: List[Any],
        **kwargs
    ) -> List[Any]:
        """
        Processes acceptance for Gibbs ensemble moves tagging in candidates.
        """
        ens = getattr(self.cfg, 'ensemble', None)
        if not ens or ens.kind != "gibbs":
            return candidates

        # 1. Match candidates and parents
        parents_list = list(parents) if not isinstance(parents, list) else parents
        candidates_list = list(candidates) if not isinstance(candidates, list) else candidates

        self._ensure_walker_id(parents_list)
        parent_by_walker = {self._get_walker_id(p): p for p in parents_list}
        
        # 2. Identify Gibbs pairs
        gibbs_pairs = {}
        processed_candidates = []
        
        for c in candidates_list:
            apm = getattr(c, "AtomPositionManager", None)
            meta = apm.metadata if apm is not None else getattr(c, "info", {})
            if meta.get("transformation") == "gibbs_exchange":
                pair_id = meta.get("gibbs_pair_id")
                if pair_id not in gibbs_pairs:
                    gibbs_pairs[pair_id] = []
                gibbs_pairs[pair_id].append(c)
            else:
                processed_candidates.append(c)

        T = self.ctx.temperature if self.ctx else 1.0
        swaps_v = 0
        swaps_p = 0
        
        # 3. Process each pair
        for pair_id, pair_candidates in gibbs_pairs.items():
            if len(pair_candidates) != 2:
                # Malformed pair, reject (restore parents)
                for c in pair_candidates:
                    wid = self._get_walker_id(c)
                    if wid in parent_by_walker:
                        processed_candidates.append(copy.deepcopy(parent_by_walker[wid]))
                continue

            c1, c2 = pair_candidates
            wid1, wid2 = self._get_walker_id(c1), self._get_walker_id(c2)
            p1, p2 = parent_by_walker.get(wid1), parent_by_walker.get(wid2)
            
            if not p1 or not p2:
                # Missing parent, keep candidates as-is or reject? Reject is safer.
                processed_candidates.extend(pair_candidates)
                continue

            apm_c1 = getattr(c1, "AtomPositionManager", None)
            meta_c1 = apm_c1.metadata if apm_c1 is not None else getattr(c1, "info", {})
            move_type = meta_c1.get("gibbs_move_type")
            
            apm_c2 = getattr(c2, "AtomPositionManager", None)
            apm_p1 = getattr(p1, "AtomPositionManager", None)
            apm_p2 = getattr(p2, "AtomPositionManager", None)
            
            if not apm_c1 or not apm_c2 or not apm_p1 or not apm_p2:
                processed_candidates.extend([c1, c2])
                continue

            U_c1, U_c2 = self._get_phi(c1), self._get_phi(c2)
            U_p1, U_p2 = self._get_phi(p1), self._get_phi(p2)
            
            accept = False
            
            # If any energy is missing (inf), always accept
            if any(u > 1e100 for u in [U_c1, U_c2, U_p1, U_p2]):
                accept = True
            else:
                dU = (U_c1 + U_c2) - (U_p1 + U_p2)
                
                if move_type == "volume":
                    V1_new, V2_new = apm_c1.get_volume(), apm_c2.get_volume()
                    V1_old, V2_old = apm_p1.get_volume(), apm_p2.get_volume()
                    N1, N2 = len(apm_c1.atomPositions), len(apm_c2.atomPositions)
                    
                    delta_phi = - (dU / T) + N1 * math.log(V1_new / V1_old) + N2 * math.log(V2_new / V2_old)
                    if delta_phi >= 0 or random.random() < math.exp(delta_phi):
                        accept = True
                        swaps_v += 1
                elif move_type == "particle":
                    V1, V2 = apm_c1.get_volume(), apm_c2.get_volume()
                    
                    # Species-aware particle counts
                    species = meta_c1.get("gibbs_species")
                    if species:
                        N1_old = apm_p1.atomCountDict.get(species, 0)
                        N2_old = apm_p2.atomCountDict.get(species, 0)
                        N1_new = apm_c1.atomCountDict.get(species, 0)
                        N2_new = apm_c2.atomCountDict.get(species, 0)
                    else:
                        # Fallback to total counts if species not tagged
                        N1_old, N2_old = len(apm_p1.atomPositions), len(apm_p2.atomPositions)
                        N1_new, N2_new = len(apm_c1.atomPositions), len(apm_c2.atomPositions)
                    
                    # We need to know which direction the swap went (1->2 or 2->1)
                    if N1_new < N1_old: # 1 -> 2
                        delta_phi = - (dU / T) + math.log((N1_old * V2) / ((N2_old + 1) * V1))
                    else: # 2 -> 1
                        delta_phi = - (dU / T) + math.log((N2_old * V1) / ((N1_old + 1) * V2))
                    
                    if delta_phi >= 0 or random.random() < math.exp(delta_phi):
                        accept = True
                        swaps_p += 1
            
            if accept:
                processed_candidates.extend([c1, c2])
            else:
                processed_candidates.extend([copy.deepcopy(p1), copy.deepcopy(p2)])

        if (swaps_v > 0 or swaps_p > 0) and self.logger:
            self.logger.info(f"[Transition] Gibbs: Accepted {swaps_v} volume and {swaps_p} particle swaps.")
            
        return processed_candidates

    def allows_agent_sync(self) -> bool:
        # Gibbs ensemble moves typically require preserving individual Markov chains 
        # without external immigration to maintain detailed balance.
        return False
