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

_EPS = 1e-12

class MetropolisKernel(TransitionKernel):
    """
    Standard Metropolis-Hastings acceptance kernel.
    Handles NVT and GCMC acceptance criteria.
    """

    def __init__(self, sampling_temperature: float, rng=None):
        self.sampling_temperature = sampling_temperature
        self._rng = rng or random.Random()

    def get_delta_phi_total(self, delta_phi: float, child: Any, parent: Any, **kwargs) -> float:
        """
        Hook for subclasses to modify the objective delta (e.g. GCMC).
        """
        return delta_phi

    def get_temperature(self, idx: int, walker_id: int, **kwargs) -> float:
        """
        Hook for subclasses to determine the temperature (e.g. Replica Exchange).
        """
        return self.sampling_temperature

    def step(
        self,
        parents: Any,
        candidates: List[Any],
        **kwargs
    ) -> List[Any]:
        """
        Applies Metropolis-Hastings acceptance criteria with detailed ensemble support.
        """
        if not candidates:
            return []

        # Ensure candidates list
        candidates_list = list(candidates) if not isinstance(candidates, list) else candidates
        parents_list = list(parents) if not isinstance(parents, list) else parents

        # 1. Initialize walker_ids if missing (first pass)
        self._ensure_walker_id(parents_list)

        # 2. Match parents and children by walker_id
        children_by_walker = self._match_parents_and_children(parents_list, candidates_list)

        survivors = []

        for idx, parent in enumerate(parents_list):
            wid = self._get_walker_id(parent, default=idx)
            child = children_by_walker.get(wid)

            if child is None:
                # Child filtered out (invalid, collided, etc.); implicitly rejected
                survivors.append(copy.deepcopy(parent))
                continue

            phi_child = self._get_phi(child)
            phi_parent = self._get_phi(parent)

            # If any energy is missing (e.g. first step), always accept
            if phi_child is None or phi_parent is None:
                accept = True
            else:
                delta_phi = phi_child - phi_parent
                delta_phi_total = self.get_delta_phi_total(delta_phi, child, parent, **kwargs)
                
                # Compute detailed balance Q-correction
                log_q_delta = 0.0
                apm = getattr(child, "AtomPositionManager", None)
                if apm is not None:
                    meta = getattr(apm, "metadata", {})
                    db_list = meta.get("detailed_balance", [])
                    
                    for db in db_list:
                        log_q_fwd_step = db.get("log_p_o_given_x", 0.0) + db.get("log_q_fwd_internal", 0.0)
                        rev_func = db.get("log_q_rev_func")
                        trace = db.get("trace")
                        
                        if rev_func and trace:
                            try:
                                log_q_rev_internal = rev_func(child, trace)
                            except Exception:
                                log_q_rev_internal = -math.inf
                        else:
                            log_q_rev_internal = 0.0
                            
                        log_q_rev_step = db.get("log_p_o_given_x", 0.0) + log_q_rev_internal
                        log_q_delta += (log_q_rev_step - log_q_fwd_step)

                T_curr = self.get_temperature(idx, wid, **kwargs) + _EPS
                total_delta = (-delta_phi_total / T_curr) + log_q_delta
                
                if total_delta >= 0:
                    accept = True
                else:
                    prob = math.exp(total_delta)
                    accept = self._rng.random() < prob

            # Clean up ephemeral unserializable functions from child
            if apm is not None and "detailed_balance" in getattr(apm, "metadata", {}):
                del apm.metadata["detailed_balance"]

            if accept:
                survivors.append(child)
            else:
                survivors.append(copy.deepcopy(parent))

        # Sweep all candidates/survivors to ensure none leak unserializable functions
        for struct in candidates_list + survivors:
            apm = getattr(struct, "AtomPositionManager", None)
            if apm is not None and "detailed_balance" in getattr(apm, "metadata", {}):
                del apm.metadata["detailed_balance"]

        return survivors
