from __future__ import annotations
import copy
import numpy as np
from typing import Any, List, Callable, Optional

from .base import TransitionKernel


class HardConstraintKernel(TransitionKernel):
    """
    Hard-wall acceptance kernel for Nested Sampling.
    Accepts candidates satisfying constraint(x),
    otherwise retains parent walker.

    Samples from:
        uniform distribution over states satisfying constraint(x).
    """

    def __init__(
        self,
        cfg: Optional[Any] = None,
        logger: Optional[Any] = None,
        ctx: Optional[Any] = None,
        constraint_func: Optional[Callable[[Any, dict], bool]] = None,
        **kwargs
    ):
        self.cfg = cfg
        self.logger = logger
        self.ctx = ctx
        self.constraint = constraint_func or (lambda x, s: True)

    def step(self, parents: Any, candidates: List[Any], **kwargs) -> List[Any]:
        """
        Applies constraint acceptance criteria matching children to parents.
        """
        if not candidates:
            return []

        if hasattr(parents, "containers"):
            parents_list = list(parents.containers)
        else:
            parents_list = list(parents) if not isinstance(parents, list) else parents

        candidates_list = list(candidates) if not isinstance(candidates, list) else candidates

        # shared_state is passed by CompositeKernel allowing dynamic threshold access
        shared_state = kwargs.get("shared_state", {})

        # 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 = []
        accepted = 0

        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

            # Evaluate constraint
            if self.constraint(child, shared_state):
                survivors.append(child)
                accepted += 1
            else:
                # Reject candidate, revert to parent
                survivors.append(copy.deepcopy(parent))

        if self.logger and len(parents_list) > 0:
            self.logger.debug(
                f"[HardConstraintKernel] Accepted {accepted}/{len(parents_list)} candidates."
            )

        return survivors
