from __future__ import annotations
from typing import Any
import numpy as np
from .metropolis import MetropolisKernel
from .nested_sampling import NestedSamplingKernel
from .identity import IdentityKernel
from .nvt import NVTKernel
from .gcmc import GCMCKernel
from .replica_acceptance import ReplicaAcceptanceKernel
from .replica_exchange import ReplicaExchangeKernel
from .gibbs_transfer import GibbsTransferKernel
from .composite import CompositeKernel
from .hard_constraint import HardConstraintKernel

def make_transition_kernel(cfg: Any, selector: Any, logger: Any, evaluator: Any = None, ctx: Any = None) -> Any:
    """
    Factory function to create the appropriate transition kernel(s) 
    based on the configuration.
    """
    ens = getattr(cfg, 'ensemble', None)
    kind = getattr(ens, 'kind', 'none')
    
    sel_method = getattr(selector, 'selection_method', 'trajectory')
    sampling_T = getattr(selector, 'sampling_temperature', 1.0)
    rng = getattr(selector, '_rng', None)

    # 1. Base Acceptance Kernel
    if kind == "identity" or (kind == "none" and sel_method != "trajectory"):
        kernels = [IdentityKernel()]
    elif kind == "nvt":
        T = getattr(ens.nvt, 'temperature', sampling_T)
        kernels = [NVTKernel(temperature=T, rng=rng)]
    elif kind == "gcmc":
        T = getattr(ens.gcmc, 'temperature', sampling_T)
        mu = getattr(ens.gcmc, 'chemical_potential', {})
        kernels = [GCMCKernel(temperature=T, mu=mu, rng=rng)]
    elif kind == "replica_exchange":
        T_ladder = getattr(ens.replica_exchange, 'T_ladder', [sampling_T])
        kernels = [ReplicaAcceptanceKernel(T_ladder=T_ladder, rng=rng)]
    elif kind == "nested_sampling":
        ns_kernel = NestedSamplingKernel(cfg=cfg, logger=logger, ctx=ctx)
        
        # NS implicitly requires candidates to be bounded by E_worst (ns_phi_star)
        def ns_constraint(state, shared):
            val = float("inf")
            apm = getattr(state, "AtomPositionManager", None)
            if apm is not None and getattr(apm, "E", None) is not None:
                val = apm.E
            elif hasattr(state, "info") and "energy" in state.info:
                val = state.info["energy"]

            return val < shared.get("ns_phi_star", float("inf"))
            
        hc_kernel = HardConstraintKernel(
            constraint_func=ns_constraint,
            cfg=cfg,
            logger=logger,
            ctx=ctx,
            rng=rng
        )
        kernels = [ns_kernel, hc_kernel]
    elif kind == "hard_constraint":
        kernels = [HardConstraintKernel(cfg=cfg, logger=logger, ctx=ctx, rng=rng)]
    else:
        # Default Metropolis Acceptance (must be trajectory mode here)
        kernels = [MetropolisKernel(
            sampling_temperature=sampling_T,
            rng=rng
        )]
    
    # 2. Add ensemble moves if configured
    if ens and kind not in ("none", "identity", "nvt", "gcmc"):
        if kind == "replica_exchange":
            kernels.append(ReplicaExchangeKernel(cfg=cfg, logger=logger, ctx=ctx))
        elif kind == "gibbs":
            kernels.append(GibbsTransferKernel(cfg=cfg, logger=logger, evaluator=evaluator, ctx=ctx))
            
    # 3. If multiple kernels, wrap in Composite
    if len(kernels) > 1:
        return CompositeKernel(kernels=kernels)
    
    return kernels[0]
