class APPROXDIR(Enum):
    """Approximation direction tag for a TLT node.

    Values:
        - INVALID: semantics cannot be guaranteed (construction error)
        - UNDER: under-approximation (subset of the exact set)
        - EXACT: exact semantics
        - OVER: over-approximation (superset of the exact set)

    The enum supports + and * with either another APPROXDIR or an int with the
    conventional meaning used by the primitive combination rules.
    """
    INVALID = None
    UNDER = -1
    EXACT = 0
    OVER = +1

    def __str__(self):
        return f'{self.name}'

    def __radd__(self, other): return self.__add__(other)

    def __add__(self, other: 'APPROXDIR | int') -> 'APPROXDIR':
        """Combine approximation directions additively, propagating INVALID."""
        lhs = self.value
        rhs = other.value if isinstance(other, APPROXDIR) else other
        return APPROXDIR(None if None in (lhs, rhs) else
                         lhs + rhs)

    def __rmul__(self, other): return self.__mul__(other)

    def __mul__(self, other: 'APPROXDIR | int') -> 'APPROXDIR':
        """Combine approximation directions multiplicatively, propagating INVALID."""
        lhs = self.value
        rhs = other.value if isinstance(other, APPROXDIR) else other
        return APPROXDIR(None if None in (lhs, rhs) else
                         lhs * rhs)

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 Impl[R]:
    """Marker base class for concrete implementation backends."""

class ImplClient[R](metaclass=ImplClientMeta):
    """Client-side mixin for declaring and checking required operations.

    Clients refer to objects that use an implementation backend. ImplClient
    ensures to list the names of operations they need the implementation to
    provide via the class attribute __require__. This mixin offers methods to:
        - report which operations are missing on a given Impl
        - test if an implementation satisfies all requirements
        - extend or inherit requirement sets dynamically
    """

    @classmethod
    def missing_ops(cls, impl: Impl[R]) -> list[str]:
        """Return the names of required operations not found on impl."""
        return [
            field
            for field in cls.__require__
            if not hasattr(impl, field)
        ]

    @classmethod
    def is_supported(cls, impl: Impl[R]) -> bool:
        """Return True if impl provides all operations listed in __require__."""
        return not cls.missing_ops(impl)

    def add_requirements(self, funcnames: Sequence[str]) -> None:
        """Add function names to the set of required operations."""
        _require = set(self.__require__)
        _require = _require.union(funcnames)
        self.__require__ = tuple(_require)

    def inherit_requirements(self, *others: ImplClient[R]) -> None:
        """Merge requirement sets from other ImplClient-like classes."""
        for other in others:
            self.add_requirements(other.__require__)

class ImplClientMeta(type):
    """Metaclass that aggregates required operation names across inheritance.

    Classes using this metaclass can declare a tuple[str, ...] in __require__.
    During class creation, this metaclass unions all __require__ tuples found
    in the base classes (that also use ImplClientMeta) with the one declared
    on the subclass, storing the result back into __require__.
    """

    __require__: ClassVar[tuple[str, ...]]

    def __new__(mcs, name: str, bases: tuple[type, ...], namespace: dict[str, Any]) -> None:
        _require = namespace.setdefault('__require__', ())

        require = set(_require)
        for base in bases:
            if isinstance(base, ImplClientMeta):
                require = require.union(base.__require__)

        namespace['__require__'] = tuple(require)

        return super().__new__(mcs, name, bases, namespace)                

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 TLT[R](ImplClient[R]):
    """Temporal Logic Tree node parameterized by a concrete set type R.

    A TLT encapsulates a tuple-based logic formula, an implementation-agnostic
    set builder realizing its semantics, an approximation direction, and a
    mapping from free proposition names to bound SetBuilders.

    Construction:
        - TLT.select(primitives) must be called to choose available operators.
        - TLT(arg, **where) dispatches based on the type of `arg`:
            str          -> proposition
            SetBuilder   -> constant set, assigned a unique proposition name
            tuple(TLExpr)-> formula; recursively constructs children via primitives
            TLT          -> existing tree, optionally updated with new bindings

    Realization:
        - Call .realize(impl) to obtain an R from an implementation.
        - .assert_realizable() checks missing props/ops and approximation validity.
    """

    __primitives__ = idict({}) # Null set of primitives

    @classmethod
    def select(cls, primitives):
        """Select the set of primitive operator implementations for this class."""
        cls.__primitives__ = primitives

    @classmethod
    def construct(cls, arg: TLTLike[R], **kwds: TLTLike[R]) -> TLT[R]:
        """Explicit constructor alias to apply usual dispatch rules."""
        return cls(arg, **kwds)

    def __new__(cls, arg: TLTLike[R], **kwds: TLTLike[R]) -> TLT[R]:
        """Construct a TLT by dispatching on the kind of `arg`."""
        if 'where' in kwds: kwds |= kwds.pop('where')
        return (cls.__new_from_tlt__(arg, **kwds)      if isinstance(arg, TLT) else
                cls.__new_from_prop__(arg, **kwds)     if isinstance(arg, str) else
                cls.__new_from_builder__(arg, **kwds)  if isinstance(arg, Callable) else
                cls.__new_from_formula__(arg, **kwds))

    @classmethod
    @TLTDebugger.wrap('tlt')
    def __new_from_tlt__(cls, tlt: TLT[R], **kwds: TLTLike[R]) -> TLT[R]:
        """Rebuild a TLT from an existing one, optionally rebinding free props."""
        # Create new TLT with updated information. With this, we can bind new 
        # TLT-like objects to free variables/propositions. We do not allow 
        # updates to existing propositions. First collect all updatable props.
        kwds = {prop: kwds[prop] if prop in kwds else sb
                for prop, sb in tlt._setmap.items()
                if prop in kwds or sb is not None}

        # It is best to reconstruct builders in the "top-level" TLT when there
        # is updates to the tree since we have modified how 
        return cls.__new_from_formula__(tlt._formula, **kwds)

    @classmethod
    @TLTDebugger.wrap('prop')
    def __new_from_prop__(cls, prop: str, **kwds: TLTLike[R]) -> TLT[R]:
        """Create a TLT from a proposition name, optionally binding it from kwds."""
        return (cls(kwds.pop(prop), **kwds) if prop in kwds else
                cls.__new_init__(prop, ReferredSet(prop), setmap={prop: None}))

    @classmethod
    @TLTDebugger.wrap('builder')
    def __new_from_builder__(cls, sb: SetBuilder, **kwds: TLTLike[R]) -> TLT[R]:
        """Create a TLT from a constant SetBuilder, assigning it a unique name."""
        # Assume a formula "_0" where the set exactly represent the prop "_0".
        # In reality, we define a unique ID `uid` instead of "_0". We add one
        # extra level of indirection with a ReferredSet for `uid` and 
        # letting `setmap` hold the actual set builder `sb`. This way, the root
        # TLT will hold a full formula that refers to even constant sets
        # (information is not lost as when internally binding constant sets to
        # builder functions).
        name = f'_{sb.uid}'
        formula = name # UID is a proposition
        self = cls.__new_init__(formula, ReferredSet(name), setmap={name: sb})

        self.inherit_requirements(sb)

        return self

    @classmethod
    @TLTDebugger.wrap('formula')
    def __new_from_formula__(cls, formula: 'TLExpr', **kwds: TLTLike[R]) -> TLT[R]:
        """Create a TLT from a TLExpr by invoking registered primitives."""
        if isinstance(formula, str):
            # If formula is a string, it is a proposition
            return cls.__new_from_prop__(formula, **kwds)

        # Otherwise, it is a non-trivial formula
        head, *tail = formula

        assert head in cls.__primitives__, \
            f'Unknown operator `{head}` in formula `{formula}`. ' \
            f'Available operators: {cls.__primitives__.keys()}'

        args = [cls(arg, **kwds) for arg in tail]  # make TLTs of formula args
        return cls.__primitives__[head](*args)

    @classmethod
    def __new_init__(cls, formula=..., builder=..., approx=..., setmap=..., times=...):
        """Initialize a bare TLT instance with the provided internal fields."""
        self = super(TLT, cls).__new__(cls)

        # Lock the selected language for this instance
        self.__primitives__ = self.__primitives__

        self._formula = formula if formula is not ... else '_0'

        # If constructed with the absurd set, then the TLT is also absurd, i.e. cannot be realized.
        self._builder = builder if builder is not ... else ABSURD

        self._approx = approx if approx is not ... else APPROXDIR.EXACT

        # Sets are associated with names using ReferredSets.
        self._setmap = setmap if setmap is not ... else idict()

        self.inherit_requirements(self._builder)

        return self

    _formula: TLExpr
    _builder: SetBuilder[R]
    _approx: APPROXDIR
    _setmap: SetMap[R]

    def __repr__(self) -> str:
        cls = type(self).__name__
        approx = str(self._approx)
        formula = str(self._formula)
        return f'{cls}({approx}, {formula})'

    def realize(self, impl: Impl[R], memoize: bool = False) -> R:
        """Realize this TLT into a concrete set R using `impl`.

        Raises if:
            - some proposition is unbound
            - required Impl operations are missing
            - approximation is INVALID

        memoize=True is reserved for a future optimization where realized sets
        may be cached on the node.
        """
        self.assert_realizable(impl)
        out = self._builder(impl, **self._setmap)
        if memoize:
            raise NotImplementedError() # TODO: builder = Set(out)
        return out

    def assert_realizable(self, impl: Optional[Impl[R]] = None) -> None:
        """Validate that this TLT can be realized, optionally against `impl`."""
        for name, sb in self._setmap.items():
            if sb is None: 
                raise Exception(f'Missing proposition `{name}`')
        if impl is not None and bool(missing := self.missing_ops(impl)):
            raise Exception(f'Missing from implementation: {", ".join(missing)}')
        if self._approx is APPROXDIR.INVALID:
            raise Exception('Invalid approximation. TLT operational semantics of formula does not hold.')

    def is_realizable(self, impl: Optional[Impl[R]] = None) -> bool:
        """Return True if assert_realizable would succeed."""
        try:
            self.assert_realizable(impl)
        except Exception:
            return False
        else:
            return True

    def iter_frml(self, formula: Optional[TLExpr] = None, **kwds):
        """Yield sub-formulas in post-order; optionally only terminals.

        Parameters:
            - formula: root to traverse (defaults to self._formula)
            - only_terminals: if True, yield only terminal sub-formulas
        """
        only_terminals = kwds.get('only_terminals', False)
        if formula is None:
            formula = self._formula
        _, *args = formula
        for arg in args:
            yield from self.iter_frml(arg, **kwds)
        if not (only_terminals and args):
            yield formula

    def iter_free(self):
        """Iterate names of free (unbound) propositions in this TLT."""
        yield from filter(lambda p: self._setmap[p] is None, self._setmap)

class TLTDebugger:
    """Internal helper to add lightweight debugging and tracing to TLT building.

    Use TLT.debug(print=True, indent=2) to enable, and TLT.nodebug() to reset.
    Arbitrary keyword args starting with '__debug' are collected and forwarded.
    """

    __debug = {}
    __debug_defaults = {'print': False, 'indent': 0}

    @classmethod
    def nodebug(cls):
        """Reset debug options to defaults for subsequent constructions."""
        for k in list(cls.__debug):
            cls.__debug.pop(k, None) # remove to set default

    @classmethod
    def debug(cls, **kwds):
        """Set debug options.

        Supported options (with prefix '__debug_' accepted on wrapper):
            - print: bool, enable/disable printing
            - indent: int, indentation level increment per nested step

        Pass Ellipsis to unset a debug option.
        """
        cls.__debug.update(collect_prefix(kwds, '__debug_'))
        for k, v in cls.__debug.items():
            if v is Ellipsis:
                cls.__debug.pop(k, None) # remove to set default

    @classmethod
    def wrap(cls, msg: str):
        """Decorator to print debug information for TLT construction.

        - Enable with `TLT.debug(print=True)`.
        - Choose indentation size with `TLT.debug(indent=...)`.
        """

        def decorator(func):
            @wraps(func)
            def wrapper(*args, **kwds):

                dbg = cls.__debug_defaults | collect_prefix(kwds, '__debug', remove=True)

                if dbg['print']:
                    print('> ' * dbg['indent'] + msg.ljust(7), kwds)
                    dbg['indent'] = dbg['indent'] + 1

                kwds |= prefix_keys(dbg, '__debug_')
                return func(*args, **kwds)
            return wrapper
        return decorator