Blame for runtime/backends/actioncodec2/physical_codec.py

"""Configuration-driven physical quantization."""

from __future__ import annotations

import itertools
import os
from typing import Protocol

import numpy as np

from ._so3 import (
    conjugate_quaternion as _conjugate,
)
from ._so3 import (
    euler_xyz_to_quaternion as _euler_xyz_to_quaternion,
)
from ._so3 import (
    multiply_quaternions as _multiply,
)
from ._so3 import (
    quaternion_to_euler_xyz as _quaternion_to_euler_xyz,
)
from ._so3 import (
    quaternion_to_rotvec as _quaternion_to_rotvec,
)
from ._so3 import (
    rotvec_to_quaternion as _rotvec_to_quaternion,
)
from ._validation import broadcast_vector
from .configuration import ActionCodec2Config, ComponentConfig, ConfigError
from .native_backend import native_kernel, require_native, resolve_backend
from .primitives import PrimitiveStore
from .semantic import difference_additive as _difference_additive
from .semantic import integrate_additive as _integrate_additive

_FLOAT32_SLACK = np.float32(8 * np.finfo(np.float32).eps)
_IDENTITY = np.asarray((0.0, 0.0, 0.0, 1.0), dtype=np.float64)

_SELECTION_CODES = {"nearest": 0, "smooth": 1, "momentum_lazy": 2}


def _constraint(component: ComponentConfig, order: int) -> tuple[np.ndarray, np.ndarray] | None:
    value = component.constraints.get(f"order_{order}")
    if value is None:
        return None
    return (
        np.asarray(value.minimum, dtype=np.float32),
        np.asarray(value.maximum, dtype=np.float32),
    )


def _scale(component: ComponentConfig) -> np.ndarray:
    return broadcast_vector(
        component.input_transform.scale,
        len(component.dims),
        dtype=np.float32,
        name="input_transform.scale",
    )


def _transform(component: ComponentConfig, actions: np.ndarray) -> np.ndarray:
    values = np.asarray(actions, dtype=np.float32) * _scale(component)
    if component.input_transform.clip:
        bounds = _constraint(component, component.input_order)
        if bounds is None:
            raise ConfigError(f"clipping requires constraints.order_{component.input_order}")
        values = np.clip(values, bounds[0], bounds[1])
    return np.asarray(values, dtype=np.float32)


class _ComponentRuntime(Protocol):
    def encode_batch(
        self, actions: np.ndarray, previous_first_order: np.ndarray | None = None
    ) -> np.ndarray: ...

    def decode_batch(
        self, cells: np.ndarray, previous_first_order: np.ndarray | None = None
    ) -> np.ndarray: ...


class _AdditiveRuntime:
    def __init__(self, component: ComponentConfig, primitives: tuple[np.ndarray, ...]) -> None:
        self.config = component
        self.primitives = tuple(np.asarray(value, dtype=np.float32) for value in primitives)
        self.width = len(component.dims)
        self.order = component.primitive_order
        self._second_order_axes = (
            tuple(axis.astype(np.float64) for axis in self.primitives) if self.order == 2 else ()
        )
        self._second_order_sorted = tuple(
            bool(np.all(axis[1:] > axis[:-1])) for axis in self._second_order_axes
        )
        self._second_order_zero_bounds = tuple(
            (int(np.searchsorted(axis, 0)), int(np.searchsorted(axis, 0, side="right")) - 1)
            if ordered
            else (0, 0)
            for axis, ordered in zip(
                self._second_order_axes, self._second_order_sorted, strict=True
            )
        )
        raw_position_bounds = _constraint(component, 0)
        self.position_bounds = (
            None
            if raw_position_bounds is None
            else (
                np.broadcast_to(raw_position_bounds[0], (self.width,)).copy(),
                np.broadcast_to(raw_position_bounds[1], (self.width,)).copy(),
            )
        )
        self.limit = broadcast_vector(
            component.distortion_limit,
            self.width,
            dtype=np.float32,
            name="distortion_limit",
        )
        table_width = max(len(values) for values in self.primitives)
        table = np.full((self.width, table_width), np.inf, dtype=np.float32)
        for coordinate, values in enumerate(self.primitives):
            table[coordinate, : len(values)] = values
        self._native_table = np.ascontiguousarray(table)
        self._native_bins = np.asarray([len(values) for values in self.primitives], dtype=np.int32)
        self._native_binary = np.zeros(self.width, dtype=np.uint8)

    def _pick_batch(self, request: np.ndarray, direction: np.ndarray) -> np.ndarray:
        flat_request = np.asarray(request, dtype=np.float32).reshape(-1, self.width)
        flat_direction = np.asarray(direction, dtype=np.float32).reshape(-1, self.width)
        kernel = native_kernel("choose_cells")
        if kernel is not None:
            return kernel(
                self._native_table,
                self._native_bins,
                flat_request,
                flat_direction,
                self.limit,
                self.config.primitive_selection == "smooth",
            ).reshape(request.shape)
        choice = np.empty(flat_request.shape, dtype=np.int64)
        for coordinate, primitives in enumerate(self.primitives):
            distance = np.abs(flat_request[:, coordinate, None] - primitives[None, :])
            nearest = np.argmin(distance, axis=1)
            if self.config.primitive_selection == "nearest":
                choice[:, coordinate] = nearest
                continue
            feasible = distance <= self.limit[coordinate] + _FLOAT32_SLACK
            preferred = feasible & (
                np.sign(primitives)[None, :] * flat_direction[:, coordinate, None] >= 0.0
            )
            candidates = np.where(np.any(preferred, axis=1, keepdims=True), preferred, feasible)
            smooth = np.argmin(np.where(candidates, distance, np.inf), axis=1)
            choice[:, coordinate] = np.where(np.any(feasible, axis=1), smooth, nearest)
        return choice.reshape(request.shape)

    def _advance_batch(self, state: np.ndarray, primitive: np.ndarray) -> None:
        higher = primitive
        for level in range(self.order - 1, -1, -1):
            state[:, level] += higher
            higher = state[:, level]
        if self.order == 1 and self.position_bounds is not None:
            state[:, 0] = np.clip(state[:, 0], self.position_bounds[0], self.position_bounds[1])

    def encode_batch(
        self, actions: np.ndarray, previous_first_order: np.ndarray | None = None
    ) -> np.ndarray:
        """Quantize ``[B,T,d]`` actions, vectorizing all trajectories per step."""
        if self.order == 2:
            return self._encode_second_order(actions, previous_first_order)
        transformed = _transform(self.config, actions)
        target = _integrate_additive(transformed, self.config.input_order)
        batch, steps, _ = target.shape
        if self.order == 0:
            return self._pick_batch(target, np.zeros_like(target))
        if resolve_backend("auto") == "native":
            lower, upper = (
                self.position_bounds
                if self.position_bounds is not None
                else (
                    np.full(self.width, -np.inf, dtype=np.float32),
                    np.full(self.width, np.inf, dtype=np.float32),
                )
            )
            starts = np.zeros((batch, self.width), dtype=np.float32)
            cells = require_native().physical_encode_batch(
                self._native_table,
                self._native_bins,
                self._native_binary,
                np.ascontiguousarray(lower, dtype=np.float32),
                np.ascontiguousarray(upper, dtype=np.float32),
                np.ascontiguousarray(target, dtype=np.float32),
                starts,
                starts,
                self.order,
                min(batch, os.cpu_count() or 1),
                np.ascontiguousarray(self.limit, dtype=np.float32),
                self.config.primitive_selection == "smooth",
            )
            return np.asarray(cells, dtype=np.int64)

        cells = np.empty(target.shape, dtype=np.int64)
        state = np.zeros((batch, self.order, self.width), dtype=np.float32)
        direction = np.zeros((batch, self.width), dtype=np.float32)
        zero = np.zeros((batch, self.width), dtype=np.float32)
        for step in range(steps):
            baseline = state.copy()
            self._advance_batch(baseline, zero)
            request = target[:, step] - baseline[:, 0]
            choice = self._pick_batch(request, direction)
            primitive = np.stack(
                [self.primitives[index][choice[:, index]] for index in range(self.width)],
                axis=1,
            )
            self._advance_batch(state, primitive)
            moved = primitive != 0.0
            direction = np.where(moved, np.sign(primitive), direction)
            cells[:, step] = choice
        return cells

    def decode_batch(
        self, cells: np.ndarray, previous_first_order: np.ndarray | None = None
    ) -> np.ndarray:
        """Decode ``[B,T,d]`` cells to actions in the component's input units."""
        if self.order == 2:
            return self._decode_second_order(cells, previous_first_order)
        indices = np.asarray(cells, dtype=np.int64)
        kernel = native_kernel("additive_decode")
        if kernel is not None:
            lower, upper = (
                self.position_bounds
                if self.position_bounds is not None
                else (
                    np.full(self.width, -np.inf, dtype=np.float32),
                    np.full(self.width, np.inf, dtype=np.float32),
                )
            )
            return kernel(
                self._native_table,
                indices,
                _scale(self.config),
                lower,
                upper,
                self.config.input_order,
                self.order,
            )
        if self.order == 0:
            positions = np.stack(
                [self.primitives[index][indices[:, :, index]] for index in range(self.width)],
                axis=2,
            )
        elif self.position_bounds is not None and resolve_backend("auto") == "native":
            batch = len(indices)
            starts = np.zeros((batch, self.width), dtype=np.float32)
            positions = require_native().physical_decode_batch(
                self._native_table,
                self._native_bins,
                self._native_binary,
                np.ascontiguousarray(self.position_bounds[0], dtype=np.float32),
                np.ascontiguousarray(self.position_bounds[1], dtype=np.float32),
                np.ascontiguousarray(indices, dtype=np.int32),
                starts,
                starts,
                self.order,
                min(batch, os.cpu_count() or 1),
            )
        else:
            picked = np.stack(
                [self.primitives[index][indices[:, :, index]] for index in range(self.width)],
                axis=2,
            )
            batch, steps, _ = indices.shape
            positions = np.empty(indices.shape, dtype=np.float32)
            state = np.zeros((batch, self.order, self.width), dtype=np.float32)
            for step in range(steps):
                self._advance_batch(state, picked[:, step])
                positions[:, step] = state[:, 0]
        actions = _difference_additive(positions, self.config.input_order)
        return np.asarray(actions / _scale(self.config), dtype=np.float32)

    def _encode_second_order(
        self, actions: np.ndarray, previous_first_order: np.ndarray
    ) -> np.ndarray:
        transformed = _transform(self.config, actions).astype(np.float64)
        velocity = previous_first_order.astype(np.float64) * _scale(self.config)
        target = transformed
        if self.config.input_order == 2:
            target = np.cumsum(target, axis=1) + velocity[:, None]
        if self.config.input_order >= 1:
            target = np.cumsum(target, axis=1)
        if resolve_backend("auto") == "native" and hasattr(
            require_native(), "additive_second_order_encode_batch"
        ):
            return np.asarray(
                require_native().additive_second_order_encode_batch(
                    self._native_table,
                    self._native_bins,
                    np.ascontiguousarray(target),
                    np.ascontiguousarray(velocity),
                    self.limit,
                    self.config.primitive_selection == "smooth",
                    min(
                        len(actions),
                        max(1, int(os.environ.get("ACTIONCODEC2_PHYSICAL_THREADS", "16"))),
                    ),
                ),
                dtype=np.int64,
            )
        position = np.zeros_like(velocity)
        direction = np.zeros_like(velocity)
        cells = np.empty(actions.shape, dtype=np.int64)
        for step in range(actions.shape[1]):
            request = target[:, step] - (position + velocity)

            choice = self._pick_second_order(request, direction)
            primitive = np.stack(
                [axis[choice[:, j]] for j, axis in enumerate(self.primitives)], axis=1
            )
            velocity += primitive
            position += velocity
            direction = np.where(primitive != 0, np.sign(primitive), direction)
            cells[:, step] = choice
        return cells

    def _pick_second_order(self, request: np.ndarray, direction: np.ndarray) -> np.ndarray:
        choice = np.empty(request.shape, dtype=np.int64)
        for j, axis in enumerate(self._second_order_axes):
            if self._second_order_sorted[j] and (
                self.config.primitive_selection == "smooth" or len(axis) > 32
            ):
                right = np.searchsorted(axis, request[:, j])
                left = np.clip(right - 1, 0, len(axis) - 1)
                right = np.minimum(right, len(axis) - 1)
                nearest = np.where(
                    np.abs(request[:, j] - axis[left]) <= np.abs(request[:, j] - axis[right]),
                    left,
                    right,
                )

                tied = (nearest > 0) & (
                    np.abs(request[:, j] - axis[np.maximum(nearest - 1, 0)])
                    == np.abs(request[:, j] - axis[nearest])
                )
                if np.any(tied):
                    nearest[tied] = np.argmin(np.abs(request[tied, j, None] - axis), axis=1)
                if self.config.primitive_selection == "smooth":
                    nonnegative, nonpositive = self._second_order_zero_bounds[j]
                    preferred = np.where(
                        direction[:, j] > 0,
                        np.maximum(nearest, min(nonnegative, len(axis) - 1)),
                        np.where(
                            direction[:, j] < 0,
                            np.minimum(nearest, max(nonpositive, 0)),
                            nearest,
                        ),
                    )
                    admissible = (np.sign(axis[preferred]) * direction[:, j] >= 0) & (
                        np.abs(request[:, j] - axis[preferred]) <= self.limit[j]
                    )
                    nearest = np.where(admissible, preferred, nearest)
                choice[:, j] = nearest
                continue

            distance = np.abs(request[:, j, None] - axis)
            nearest = np.argmin(distance, axis=1)
            if self.config.primitive_selection == "nearest":
                choice[:, j] = nearest
                continue
            feasible = distance <= self.limit[j]
            preferred = feasible & (np.sign(axis)[None] * direction[:, j, None] >= 0)
            candidates = np.where(preferred.any(axis=1, keepdims=True), preferred, feasible)
            smooth = np.argmin(np.where(candidates, distance, np.inf), axis=1)
            choice[:, j] = np.where(feasible.any(axis=1), smooth, nearest)
        return choice

    def _decode_second_order(
        self, cells: np.ndarray, previous_first_order: np.ndarray
    ) -> np.ndarray:
        kernel = native_kernel("additive_reconstruct")
        if kernel is not None:
            return kernel(
                self._native_table,
                cells,
                previous_first_order,
                _scale(self.config),
                self.config.input_order,
                self.order,
            )
        picked = np.stack(
            [axis[cells[:, :, j]] for j, axis in enumerate(self.primitives)], axis=2
        ).astype(np.float64)
        if self.config.input_order == 2:
            result = picked
        else:
            result = np.cumsum(picked, axis=1) + previous_first_order[:, None] * _scale(self.config)
            if self.config.input_order == 0:
                result = np.cumsum(result, axis=1)
        return np.asarray(result / _scale(self.config), dtype=np.float32)


class _SO3Runtime:
    def __init__(self, component: ComponentConfig, primitives: tuple[np.ndarray, ...]) -> None:
        self.config = component
        self.primitives = tuple(np.asarray(value, dtype=np.float64) for value in primitives)
        self.left = component.multiplication == "left"
        self.limit = float(component.distortion_limit)
        if component.primitive_order == 2:
            bounds = component.constraints["order_2"]
            corner = np.maximum(np.abs(bounds.minimum), np.abs(bounds.maximum))
            corner = np.broadcast_to(corner, (3,))
            radius = (
                np.linalg.norm(corner)
                if component.representation == "rotation_vector"
                else np.sum(corner)
            )
            self.second_order_radius = min(np.pi, float(radius) + 3 * self.limit)
        table_width = max(len(values) for values in self.primitives)
        table = np.full((3, table_width), np.inf, dtype=np.float64)
        for coordinate, values in enumerate(self.primitives):
            table[coordinate, : len(values)] = values
        self._native_table = table
        self._native_bins = np.asarray([len(values) for values in self.primitives], dtype=np.int32)
        self.lazy = component.primitive_selection == "momentum_lazy"
        self._zero = np.asarray([int(np.argmin(np.abs(values))) for values in self.primitives])
        if self.lazy and any(
            values[index] != 0.0 for values, index in zip(self.primitives, self._zero, strict=True)
        ):
            raise ValueError("momentum_lazy selection requires an exact zero primitive per axis")

    def _from_representation(self, values: np.ndarray) -> np.ndarray:
        shape = values.shape[:-1]
        flat = np.asarray(values).reshape(-1, 3)
        converted = (
            _rotvec_to_quaternion(flat)
            if self.config.representation == "rotation_vector"
            else _euler_xyz_to_quaternion(flat)
        )
        return converted.reshape(*shape, 4)

    def _to_representation(self, values: np.ndarray) -> np.ndarray:
        shape = values.shape[:-1]
        flat = np.asarray(values).reshape(-1, 4)
        converted = (
            _quaternion_to_rotvec(flat)
            if self.config.representation == "rotation_vector"
            else _quaternion_to_euler_xyz(flat)
        )
        return converted.reshape(*shape, 3)

    def _residual(self, target: np.ndarray, reconstruction: np.ndarray) -> np.ndarray:
        return (
            _multiply(target, _conjugate(reconstruction))
            if self.left
            else _multiply(_conjugate(reconstruction), target)
        )

    def _apply(self, state: np.ndarray, increment: np.ndarray) -> np.ndarray:
        return _multiply(increment, state) if self.left else _multiply(state, increment)

    def _nearest_batch(self, request: np.ndarray) -> np.ndarray:
        return np.stack(
            [
                np.argmin(np.abs(request[:, index, None] - values[None, :]), axis=1)
                for index, values in enumerate(self.primitives)
            ],
            axis=1,
        ).astype(np.int64, copy=False)

    def _smooth(
        self,
        request: np.ndarray,
        reconstruction: np.ndarray,
        target: np.ndarray,
        direction: np.ndarray,
        nearest: np.ndarray,
        previous: np.ndarray | None = None,
    ) -> np.ndarray:
        axes = []
        for coordinate, values in enumerate(self.primitives):
            candidates = np.flatnonzero(
                np.abs(values - request[coordinate]) <= 2.0 * self.limit + 1e-12
            )
            if nearest[coordinate] not in candidates:
                candidates = np.append(candidates, nearest[coordinate])
            axes.append(candidates)

        best_any: tuple | None = None
        best_preferred: tuple | None = None
        repeat = previous is not None and tuple(int(v) for v in previous)
        for raw_choice in itertools.product(*axes):
            choice = tuple(int(value) for value in raw_choice)
            rotvec = np.asarray([self.primitives[index][choice[index]] for index in range(3)])
            rebuilt = self._apply(reconstruction, _rotvec_to_quaternion(rotvec))
            error = float(np.linalg.norm(_quaternion_to_rotvec(self._residual(target, rebuilt))))
            if error > self.limit + 1e-12:
                continue
            if self.lazy and choice == repeat:
                return np.asarray(choice, dtype=np.int64)
            offset = np.abs(np.asarray(choice) - self._zero)
            candidate = (
                (int(np.count_nonzero(offset)), int(offset.sum()), error, choice)
                if self.lazy
                else (error, choice)
            )
            if best_any is None or candidate < best_any:
                best_any = candidate

            if np.dot(rotvec, direction) >= -1e-15 and (
                best_preferred is None or candidate < best_preferred
            ):
                best_preferred = candidate
        selected = best_preferred or best_any
        return nearest if selected is None else np.asarray(selected[-1], dtype=np.int64)

    def encode_batch(
        self, actions: np.ndarray, previous_first_order: np.ndarray | None = None
    ) -> np.ndarray:
        """Quantize ``[B,T,3]`` rotations with a vectorized batch state."""
        if self.config.primitive_order == 2:
            return self._encode_second_order(actions, previous_first_order)
        transformed = _transform(self.config, actions).astype(np.float64)
        if resolve_backend("auto") == "native":
            cells = require_native().so3_encode_batch(
                self._native_table,
                self._native_bins,
                np.ascontiguousarray(transformed, dtype=np.float32),
                self.config.input_order,
                self.config.representation == "rotation_vector",
                self.left,
                self.limit,
                _SELECTION_CODES[self.config.primitive_selection],
                min(len(transformed), os.cpu_count() or 1),
            )
            return np.asarray(cells, dtype=np.int64)
        converted = self._from_representation(transformed)
        batch, steps = transformed.shape[:2]
        true = np.broadcast_to(_IDENTITY, (batch, 4)).copy()
        reconstruction = true.copy()
        direction = np.zeros((batch, 3), dtype=np.float64)
        cells = np.empty((batch, steps, 3), dtype=np.int64)
        previous = np.broadcast_to(self._zero, (batch, 3))
        for step in range(steps):
            value = converted[:, step]
            true = value if self.config.input_order == 0 else self._apply(true, value)
            residual = _quaternion_to_rotvec(self._residual(true, reconstruction))
            choice = self._nearest_batch(residual)
            if self.config.primitive_selection != "nearest":
                choice = np.stack(
                    [
                        self._smooth(
                            residual[index],
                            reconstruction[index],
                            true[index],
                            direction[index],
                            choice[index],
                            previous[index] if self.lazy else None,
                        )
                        for index in range(batch)
                    ]
                )
            previous = choice
            primitive = np.stack(
                [self.primitives[index][choice[:, index]] for index in range(3)], axis=1
            )
            reconstruction = self._apply(reconstruction, _rotvec_to_quaternion(primitive))
            moved = np.linalg.norm(primitive, axis=1) > 0.0
            direction[moved] = primitive[moved]
            cells[:, step] = choice
        return cells

    def decode_batch(
        self, cells: np.ndarray, previous_first_order: np.ndarray | None = None
    ) -> np.ndarray:
        """Decode ``[B,T,3]`` cells in the configured boundary representation."""
        kernel = native_kernel("so3_reconstruct")
        if kernel is not None:
            previous = previous_first_order
            if previous is None:
                previous = np.zeros((len(cells), 3), dtype=np.float32)
            return kernel(
                self._native_table,
                cells,
                previous,
                _scale(self.config),
                self.config.input_order,
                self.config.primitive_order,
                self.left,
                self.config.representation == "rotation_vector",
                getattr(self, "second_order_radius", np.inf),
            )
        if self.config.primitive_order == 2:
            return self._decode_second_order(cells, previous_first_order)
        indices = np.asarray(cells, dtype=np.int64)
        rotvec = np.stack(
            [self.primitives[index][indices[:, :, index]] for index in range(3)], axis=2
        )
        if self.config.input_order == 1:
            represented = (
                rotvec
                if self.config.representation == "rotation_vector"
                else self._to_representation(_rotvec_to_quaternion(rotvec.astype(np.float32)))
            )
            actions = np.asarray(represented / _scale(self.config), dtype=np.float32)

            return _difference_additive(np.cumsum(actions, axis=1, dtype=np.float32), 1)

        batch, steps = indices.shape[:2]
        reconstruction = np.broadcast_to(_IDENTITY, (batch, 4)).copy()
        quaternion = np.empty((batch, steps, 4), dtype=np.float64)
        primitive_quaternion = _rotvec_to_quaternion(rotvec)
        for step in range(steps):
            reconstruction = self._apply(reconstruction, primitive_quaternion[:, step])
            quaternion[:, step] = reconstruction
        represented = self._to_representation(quaternion)
        return np.asarray(represented / _scale(self.config), dtype=np.float32)

    def _encode_second_order(
        self, actions: np.ndarray, previous_first_order: np.ndarray
    ) -> np.ndarray:
        converted = self._from_representation(_transform(self.config, actions).astype(np.float64))
        increment = self._from_representation(
            previous_first_order.astype(np.float64) * _scale(self.config)
        )
        if resolve_backend("auto") == "native" and hasattr(
            require_native(), "so3_second_order_encode_batch"
        ):
            return np.asarray(
                require_native().so3_second_order_encode_batch(
                    self._native_table,
                    self._native_bins,
                    np.ascontiguousarray(converted),
                    np.ascontiguousarray(increment),
                    self.config.input_order,
                    self.left,
                    self.limit,
                    self.second_order_radius,
                    self.config.primitive_selection == "smooth",
                    min(
                        len(actions),
                        max(1, int(os.environ.get("ACTIONCODEC2_PHYSICAL_THREADS", "16"))),
                    ),
                ),
                dtype=np.int64,
            )
        batch, steps = actions.shape[:2]
        true = np.broadcast_to(_IDENTITY, (batch, 4)).copy()
        reconstruction = true.copy()
        direction = np.zeros((batch, 3), dtype=np.float64)
        cells = np.empty((batch, steps, 3), dtype=np.int64)
        for step in range(steps):
            true = (
                converted[:, step]
                if self.config.input_order == 0
                else self._apply(true, converted[:, step])
            )
            desired = self._residual(true, reconstruction)
            correction = self._residual(desired, increment)
            request = _quaternion_to_rotvec(correction)
            choice = self._nearest_batch(request)
            if self.config.primitive_selection == "smooth":
                choice = self._smooth_second_order_batch(request, correction, direction, choice)
            primitive = np.stack(
                [axis[choice[:, j]] for j, axis in enumerate(self.primitives)], axis=1
            )
            primitive = self._second_order_primitive(primitive)
            increment = self._apply(increment, _rotvec_to_quaternion(primitive))
            reconstruction = self._apply(reconstruction, increment)
            moved = np.linalg.norm(primitive, axis=1) > 0
            direction[moved] = primitive[moved]
            cells[:, step] = choice
        return cells

    def _second_order_primitive(self, rotvec: np.ndarray) -> np.ndarray:
        norm = np.linalg.norm(rotvec, axis=-1, keepdims=True)
        return rotvec * np.minimum(
            1.0, self.second_order_radius / np.maximum(norm, np.finfo(float).tiny)
        )

    def _smooth_second_order_batch(
        self,
        request: np.ndarray,
        correction: np.ndarray,
        direction: np.ndarray,
        nearest: np.ndarray,
    ) -> np.ndarray:
        if len(request) > 64:
            return np.concatenate(
                [
                    self._smooth_second_order_batch(
                        request[i : i + 64],
                        correction[i : i + 64],
                        direction[i : i + 64],
                        nearest[i : i + 64],
                    )
                    for i in range(0, len(request), 64)
                ]
            )
        axes, counts = [], []
        for j, axis in enumerate(self.primitives):
            mask = np.abs(axis[None] - request[:, j, None]) <= 2 * self.limit + 1e-12
            mask[np.arange(len(request)), nearest[:, j]] = True
            count = mask.sum(axis=1)
            rows, columns = np.nonzero(mask)
            slots = np.arange(len(rows)) - np.repeat(np.cumsum(count) - count, count)
            candidates = np.zeros((len(request), int(count.max())), dtype=np.int64)
            candidates[rows, slots] = columns
            axes.append(candidates)
            counts.append(count)
        counts = np.stack(counts, axis=1)
        widths = counts.max(axis=0)
        total = int(np.prod(widths))
        best = nearest.copy()
        best_error = np.full(len(request), np.inf)
        best_preferred = np.zeros(len(request), dtype=bool)
        for offset in range(0, total, 4096):
            lattice = np.stack(
                np.unravel_index(np.arange(offset, min(total, offset + 4096)), tuple(widths)),
                axis=1,
            )
            valid = np.all(lattice[None] < counts[:, None], axis=2)
            choices = np.stack(
                [
                    axis[
                        np.arange(len(request))[:, None],
                        np.minimum(lattice[None, :, j], counts[:, j, None] - 1),
                    ]
                    for j, axis in enumerate(axes)
                ],
                axis=2,
            )
            rotvec = np.stack(
                [axis[choices[:, :, j]] for j, axis in enumerate(self.primitives)], axis=2
            )
            rotvec = self._second_order_primitive(rotvec)
            primitive = _rotvec_to_quaternion(rotvec)
            residual = _multiply(_conjugate(primitive), correction[:, None])
            error = np.linalg.norm(_quaternion_to_rotvec(residual), axis=2)
            valid &= error <= self.limit + 1e-12
            preferred = valid & (np.sum(rotvec * direction[:, None], axis=2) >= -1e-15)
            has_preferred = preferred.any(axis=1)
            eligible = np.where(has_preferred[:, None], preferred, valid)
            scores = np.where(eligible, error, np.inf)
            index = np.argmin(scores, axis=1)
            score = scores[np.arange(len(request)), index]
            improve = np.isfinite(score) & (
                (has_preferred & ~best_preferred)
                | ((has_preferred == best_preferred) & (score < best_error))
            )
            best[improve] = choices[np.arange(len(request)), index][improve]
            best_error[improve] = score[improve]
            best_preferred[improve] = has_preferred[improve]
        return best

    def _decode_second_order(
        self, cells: np.ndarray, previous_first_order: np.ndarray
    ) -> np.ndarray:
        rotvec = np.stack([axis[cells[:, :, j]] for j, axis in enumerate(self.primitives)], axis=2)
        rotvec = self._second_order_primitive(rotvec)
        changes = _rotvec_to_quaternion(rotvec)
        increment = self._from_representation(
            previous_first_order.astype(np.float64) * _scale(self.config)
        )
        reconstruction = np.broadcast_to(_IDENTITY, increment.shape).copy()
        output = np.empty((*cells.shape[:2], 4), dtype=np.float64)
        for step in range(cells.shape[1]):
            increment = self._apply(increment, changes[:, step])
            reconstruction = self._apply(reconstruction, increment)
            output[:, step] = increment if self.config.input_order == 1 else reconstruction
        return np.asarray(self._to_representation(output) / _scale(self.config), dtype=np.float32)


class PhysicalCodec:
    """Compose independent physical components in global action-dimension order."""

    def __init__(self, config: ActionCodec2Config, primitives: PrimitiveStore) -> None:
        if primitives.dimension != config.physical.input_dim:
            raise ConfigError(
                f"primitive dimension {primitives.dimension} does not match input_dim "
                f"{config.physical.input_dim}"
            )
        self.config = config
        self.primitives = primitives
        runtimes: list[tuple[ComponentConfig, _ComponentRuntime]] = []
        for component in config.physical.components.values():
            component_primitives = tuple(primitives.values[index] for index in component.dims)
            runtime: _ComponentRuntime = (
                _AdditiveRuntime(component, component_primitives)
                if component.space == "additive"
                else _SO3Runtime(component, component_primitives)
            )
            runtimes.append((component, runtime))
        self._runtimes = tuple(runtimes)
        self.bins = primitives.bins
        self.requires_previous_first_order = any(c.primitive_order == 2 for c, _ in runtimes)

    def _check_previous_first_order(
        self, value: np.ndarray | None, batch: int
    ) -> np.ndarray | None:
        if value is None:
            if self.requires_previous_first_order:
                raise ValueError(
                    "second-order tokenization requires previous_first_order; no zero initialization is inferred"
                )
            return None
        array = np.asarray(value, dtype=np.float32)
        shape = (batch, self.config.physical.input_dim)
        if array.shape != shape or not np.all(np.isfinite(array)):
            raise ValueError(f"previous_first_order must be finite with shape {shape}")
        return array

    def _check_actions(self, actions: np.ndarray) -> np.ndarray:
        values = np.asarray(actions, dtype=np.float32)
        expected = self.config.physical.input_dim
        if values.ndim != 3 or values.shape[2] != expected or values.shape[1] < 1:
            raise ValueError(f"actions must have non-empty shape [B, T, {expected}]")
        if not np.all(np.isfinite(values)):
            raise ValueError("actions must contain only finite values")
        return values

    def final_first_order(
        self, cells: np.ndarray, *, previous_first_order: np.ndarray
    ) -> np.ndarray:
        """Return ``[B,D]`` state after these cells, without position differencing."""
        values = self._check_cells(cells)
        previous = self._check_previous_first_order(previous_first_order, len(values))
        result = np.zeros((len(values), self.config.physical.input_dim), dtype=np.float32)
        for component, runtime in self._runtimes:
            if component.primitive_order != 2:
                continue
            local = np.take(values, component.dims, axis=2)
            boundary = np.take(previous, component.dims, axis=1)
            picked = np.stack(
                [axis[local[:, :, j]] for j, axis in enumerate(runtime.primitives)], axis=2
            )
            if component.space == "additive":
                final = boundary.astype(np.float64) + np.sum(
                    picked.astype(np.float64), axis=1
                ) / _scale(component)
            else:
                increment = runtime._from_representation(
                    boundary.astype(np.float64) * _scale(component)
                )
                changes = _rotvec_to_quaternion(runtime._second_order_primitive(picked))
                for step in range(local.shape[1]):
                    increment = runtime._apply(increment, changes[:, step])
                final = runtime._to_representation(increment) / _scale(component)
            result[:, component.dims] = final
        return result

    def _check_cells(self, cells: np.ndarray) -> np.ndarray:
        values = np.asarray(cells, dtype=np.int64)
        expected = self.config.physical.input_dim
        if values.ndim != 3 or values.shape[2] != expected or values.shape[1] < 1:
            raise ValueError(f"cells must have non-empty shape [B, T, {expected}]")
        counts = np.asarray(self.bins, dtype=np.int64)
        if values.size and (np.any(values < 0) or np.any(values >= counts[None, None, :])):
            raise ValueError("a cell index is outside its dimension's primitive set")
        return values

    def encode_batch(
        self, actions: np.ndarray, *, previous_first_order: np.ndarray | None = None
    ) -> np.ndarray:
        """Quantize ``[B,T,D]`` with a required ``[B,D]`` boundary for order 2."""
        values = self._check_actions(actions)
        previous = self._check_previous_first_order(previous_first_order, len(values))
        cells = np.empty(values.shape, dtype=np.int64)
        for component, runtime in self._runtimes:
            component_cells = runtime.encode_batch(
                np.take(values, component.dims, axis=2),
                None if previous is None else np.take(previous, component.dims, axis=1),
            )
            for local, global_dimension in enumerate(component.dims):
                cells[:, :, global_dimension] = component_cells[:, :, local]
        return cells

    def decode_batch(
        self, cells: np.ndarray, *, previous_first_order: np.ndarray | None = None
    ) -> np.ndarray:
        """Decode ``[B,T,D]`` using the same ``[B,D]`` boundary as encoding."""
        values = self._check_cells(cells)
        previous = self._check_previous_first_order(previous_first_order, len(values))
        actions = np.empty(values.shape, dtype=np.float32)
        for component, runtime in self._runtimes:
            component_actions = runtime.decode_batch(
                np.take(values, component.dims, axis=2),
                None if previous is None else np.take(previous, component.dims, axis=1),
            )
            for local, global_dimension in enumerate(component.dims):
                actions[:, :, global_dimension] = component_actions[:, :, local]
        return actions