class AppliedSet[R](SetBuilder[R]):
    """Defer a call to `Impl.<funcname>(*args)` where `args` are realized builders.

    - Accumulates required `Impl` methods from children and adds `funcname`.
    - Propagates and de-duplicates children's free variables.
    - On realization, calls child builders first, then invokes the `Impl` method.
    - Wraps child exceptions to pinpoint which argument failed.
    """

    def __init__(self, funcname: str, *builders: SetBuilder[R]) -> None:
        self.funcname = funcname
        self.builders = builders

        _require = (funcname,)

        for builder in self.builders:
            _require += builder.__require__
            self.free += tuple(name for name in builder.free if name not in self.free)

        self.add_requirements(_require)        

    def __call__(self, impl: Impl[R], **m: SetBuilder[R]) -> R:
        try:
            func = getattr(impl, self.funcname)
        except AttributeError as e:
            raise AttributeError(f'Impl {impl.__class__.__name__} does not support "{self.funcname}".') from e

        args = []
        for i, sb in enumerate(self.builders):
            try:
                args.append(sb(impl, **m))
            except Exception as e:
                E = type(e)
                raise E(f'When applying "{self.funcname}" on argument {i}, received: {e!s}') from e

        return func(*args)

class BoundedSet[R](SetBuilder[R]):
    """Axis-aligned box possibly unbounded on one side per axis.

    Bounds mapping: `name -> (vmin, vmax)`. Use Ellipsis to denote an open side
    (e.g., (0, ...) or (..., 1)). For periodic axes where `vmax < vmin`, the
    range wraps around.

    Example:
    ```
    # Left half-circle bounded in y.
    A = BoundedSet(y=(-0.5, 0.5), theta=(+pi/2, -pi/2))
    ```

    Requires:
        - `Impl < AxesImpl`
        - `Impl.complement(inp: R) -> R`
        - `Impl.halfspace(normal, offset, axes, ...) -> R`
        - `Impl.intersect(inp1: R, inp2: R) -> R`
    """

    __require__ = ('complement', 'halfspace', 'intersect')

    def __init__(self, **bounds: list[int]) -> None:
        self.bounds = bounds

    def __call__(self, impl: Impl[R], **m: SetBuilder[R]) -> R:
        s = impl.complement(impl.empty())
        for name, (vmin, vmax) in self.bounds.items():
            i = impl.axis(name)
            if vmin is Ellipsis:
                assert vmax is not Ellipsis, f'Invalid bounds for axis {impl.axis_name(i)}, there must be either an upper or lower bound.'
                upper_bound = impl.halfspace(normal=[0 if i != j else -1 for j in range(impl.ndim)],
                                             offset=[0 if i != j else vmax for j in range(impl.ndim)])
                axis_range = upper_bound
            elif vmax is Ellipsis:
                assert vmin is not Ellipsis, f'Invalid bounds for axis {impl.axis_name(i)}, there must be either an upper or lower bound.'
                lower_bound = impl.halfspace(normal=[0 if i != j else +1 for j in range(impl.ndim)],
                                             offset=[0 if i != j else vmin for j in range(impl.ndim)])
                axis_range = lower_bound
            elif impl.axis_is_periodic(i) and vmax < vmin:
                upper_bound = impl.halfspace(normal=[0 if i != j else -1 for j in range(impl.ndim)],
                                             offset=[0 if i != j else vmin for j in range(impl.ndim)])
                lower_bound = impl.halfspace(normal=[0 if i != j else +1 for j in range(impl.ndim)],
                                             offset=[0 if i != j else vmax for j in range(impl.ndim)])
                axis_range = impl.complement(impl.intersect(upper_bound, lower_bound))
            else:
                # NOTE: See similar assertion in TVHJImpl's halfspace
                amin, amax = impl.axis_bounds(i)
                assert amin < vmin < amax, f'For dimension "{name}", {amin} < {vmin=} < {amax}. Use Ellipsis (...) to indicate subset stretching to the space boundary.'
                assert amin < vmax < amax, f'For dimension "{name}", {amin} < {vmax=} < {amax}. Use Ellipsis (...) to indicate subset stretching to the space boundary.'
                upper_bound = impl.halfspace(normal=[0 if i != j else -1 for j in range(impl.ndim)],
                                             offset=[0 if i != j else vmax for j in range(impl.ndim)])
                lower_bound = impl.halfspace(normal=[0 if i != j else +1 for j in range(impl.ndim)],
                                             offset=[0 if i != j else vmin for j in range(impl.ndim)])
                axis_range = impl.intersect(upper_bound, lower_bound)
            s = impl.intersect(s, axis_range)
        return s

class HalfSpaceSet[R](SetBuilder[R]):
    """Half-space described by the normal and offset of a hyperplane.

    Note: The set is in the direction of the normal.

    Parameters:
        normal: coefficients along each axis
        offset: offsets along each axis
        axes: axis indices (or str if using `AxesImpl`) in the `Impl`'s coordinate system
        kwds: forwarded to `Impl.halfspace`

    Requires:
        - `Impl.halfspace(normal, offset, axes, ...) -> R`
    """

    __require__ = ('halfspace',)

    def __init__(
        self,
        normal: list[float],
        offset: list[float],
        axes: list[Axis],
        **kwds: Any,
    ) -> None:
        assert len(axes) == len(normal) == len(offset)
        self.normal = normal
        self.offset = offset
        self.axes = axes
        self.kwds = kwds

    def __call__(self, impl: Impl[R], **m: SetBuilder[R]) -> R:
        return impl.halfspace(normal=self.normal, 
                              offset=self.offset, 
                              axes=[impl.axis(ax) for ax in self.axes],
                              **self.kwds)

class ReferredSet[R](SetBuilder[R]):
    """Reference a named free variable resolved from the realization mapping.

    `ReferredSet('X')(impl, X=some_builder)` realizes to `some_builder(impl, ...)`.
    This is useful in two ways:
        1. We can be lazy when constructing the call tree, i.e. we allow users to
           define which builder to use at a later stage.
        2. This essentially allow variables to exist within the call tree which avoids
           having to reconstruct an entire tree in some cases.
    """

    def __init__(self, name: str) -> None:
        self.free += (name,)

    def __call__(self, impl: Impl[R], **m: SetBuilder[R]) -> R:
        name, = self.free
        sb = m.pop(name)
        return sb(impl, **m)

class Set[R](SetBuilder[R]):
    """Wrap a concrete set value R and return it unchanged on realization."""

    def __init__(self, arg: R) -> None:
        self.arg = arg

    def __call__(self, impl: Impl[R], **m: SetBuilder[R]) -> R:
        return self.arg

class SetBuilder[R](ImplClient[R]):
    """Abstract base for all set builders.

    Responsibilities:
        - Be callable with an implementation `Impl[R]` to produce a concrete set `R`.
        - Track required `Impl` operations through `ImplClient`.
        - Track free variable names (see `ReferredSet`).

    Subclasses should implement `__call__`, which is called to realize the sets.
    """

    free: tuple[str, ...] = ()

    def __call__(self, impl: Impl[R], **m: SetBuilder[R]) -> R:
        raise NotImplementedError(f"{type(self).__name__}.__call__ is not implemented. SetBuilders must implement __call__.")

    def __repr__(self) -> str:
        """Return a compact identifier for the builder instance."""
        cls = type(self)
        ptr = hash(self)
        return f'<{cls.__name__} at {ptr:#0{18}x}>'

    @property
    def uid(self) -> str:
        """Stable hexadecimal id derived from the object hash."""
        # Simple way to create a unique id from a python function.
        # - hash(sb) returns the function pointer (I think)
        # - Convert to bytes to get capture full 64-bit value (incl. zeroes)
        # - Convert to hex-string
        return hash(self).to_bytes(8,"big").hex()

class primitive[R, **P]:
    """Decorator class for declaring TLT primitives.

    Usage:
    - Specify the operator form and argument names:
      ```
      @primitive(AND('_1', '_2'))
      def And(_1: TLTLike[R], _2: TLTLike[R]) -> tuple[SetBuilder[R], APPROXDIR]:
          ...
      ```
      The decorator registers a callable under the operator head ("AND") that,
      when invoked with TLTLike arguments, returns a constructed TLT[R]. The
      arguments to the decorated function are mapped by name from the provided
      operands by order as they appear in the formula. The formula must consist
      of only a single connective.

    Notes:
    - The decorated function receives named arguments matching the formula's
      parameter names (strings in the TLExpr tail) and should return the pair
      (SetBuilder[R], APPROXDIR) describing the primitive's semantics and its
      approximation contribution.
    """

    type PrimitiveFunc = Callable[P, TLT[R]]
    type PrimitiveFactory = Callable[P, tuple[SetBuilder[R], APPROXDIR]]

    formula: TLExpr
    head: str
    tail: tuple[TLExpr, ...]

    def __init__(self, formula: TLExpr):
        """Bind this primitive to a formula shape.

        Args:
          formula: a TLExpr operator tuple whose head is the operator name and
                   whose tail elements are unique proposition names (str)
        """
        assert isinstance(formula, tuple), "Formula must be a tuple"
        assert isinstance(formula[0], str), "Operator name of formula must be a string"
        assert all(isinstance(arg, str) for arg in formula[1:]), "Arguments of formula must be propositions (string) for primitives"
        head, *tail = formula
        assert len(tail) == len(set(tail)), "Argument names must be unique"
        self.formula = formula
        self.head = head
        self.tail = tail

    def __call__(self, func: PrimitiveFactory) -> idict[str, PrimitiveFunc]:
        """Decorator entry: wrap `func` to build a TLT node for this primitive.

        The wrapped function:
        - Validates/assembles the operator formula with the child formulas
        - Calls `func` with named arguments mapped from the provided children
        - Constructs a TLT using the returned (SetBuilder, APPROXDIR)
        - Merges children's set maps and inherits requirements

        Returns:
            idict mapping {connective: function}
        """

        @wraps(func)
        def wrapper(*args: P.args) -> TLT[R]:
            formula = (self.head, *[arg._formula if isinstance(arg, TLT) else arg for arg in args])
            builder, approx = func(**{name: arg for name, arg in zip(self.tail, args)})
            setmap = idict(sum([list(arg._setmap.items()) for arg in args if isinstance(arg, TLT)], []))
            tree = TLT.__new_init__(formula, builder, approx, setmap)
            tree.inherit_requirements(*args)
            return tree

        return idict({self.head: wrapper})

    def define_as(self, equiv: TLExpr) -> idict[str, Callable[..., TLT]]:
        """Define this operator via a pure logical equivalence.

        This produces a thin wrapper that expands calls to this operator into
        the provided equivalent TLExpr, without supplying a new SetBuilder.

        Constraints:
        - `equiv` must be well-formed
        - `equiv` must reference exactly the same set of proposition names as
          the original operator formula

        Returns:
            idict mapping {connective: expansion_function}
        """
        assert (x := get_malformed(equiv)) is None, f"Equivalent formula must be well-formed, not {x}"
        assert get_props(self.formula) == get_props(equiv), "Formula and equivalent must have the same set of propositions"

        func = lambda *args: TLT(equiv, **{name: arg for name, arg in zip(self.tail, args)})

        return idict({self.head: func})