exoskeleton/code/core/command_allocator.py

"""Total-command allocation and accepted haptic-increment reconstruction."""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np


def _vector(value, size: int | None, name: str) -> np.ndarray:
    array = np.asarray(value, dtype=float).reshape(-1)
    if size is not None and array.shape != (size,):
        raise ValueError(f"{name} must have shape ({size},), got {array.shape}")
    if array.size == 0 or not np.all(np.isfinite(array)):
        raise ValueError(f"{name} must be non-empty and finite")
    return array


@dataclass(frozen=True)
class PreparedAllocation:
    compensation: np.ndarray
    haptic_raw: np.ndarray
    haptic_candidate: np.ndarray
    lower_haptic_bound: np.ndarray
    upper_haptic_bound: np.ndarray
    compensation_limited: bool
    haptic_limited: bool


@dataclass(frozen=True)
class AcceptedAllocation:
    compensation_requested: np.ndarray
    haptic_projected: np.ndarray
    total_requested: np.ndarray
    total_quantized: np.ndarray
    total_accepted: np.ndarray
    compensation_accepted: np.ndarray
    haptic_accepted: np.ndarray
    downstream_modified: bool
    quantization_active: bool
    derating_active: bool
    readback_used: bool


class CommandAllocator:
    """Reserve total actuator headroom before the final energy projection.

    Call :meth:`prepare` before energy supervision.  Pass its
    ``haptic_candidate`` through the selected energy supervisor, then call
    :meth:`finalize`.  The latter never silently claims the projected torque was
    accepted: quantization, derating, or drive readback are exposed and the
    accepted haptic increment is reconstructed explicitly.
    """

    def __init__(
        self,
        total_torque_limit: float | np.ndarray,
        *,
        quantization_step: float | np.ndarray | None = None,
    ):
        limits = np.asarray(total_torque_limit, dtype=float)
        if limits.ndim == 0:
            limits = limits.reshape(1)
        else:
            limits = limits.reshape(-1)
        if limits.size == 0 or np.any(~np.isfinite(limits)) or np.any(limits <= 0.0):
            raise ValueError("total_torque_limit must contain finite positive values")
        self.total_torque_limit = limits
        if quantization_step is None:
            self.quantization_step = np.zeros_like(limits)
        else:
            steps = np.asarray(quantization_step, dtype=float)
            if steps.ndim == 0:
                steps = np.full(limits.size, float(steps))
            else:
                steps = steps.reshape(-1)
            if steps.shape != limits.shape or np.any(~np.isfinite(steps)) or np.any(steps < 0.0):
                raise ValueError(
                    "quantization_step must be scalar or match torque limits"
                )
            self.quantization_step = steps

    def prepare(
        self,
        compensation: np.ndarray,
        haptic_raw: np.ndarray,
    ) -> PreparedAllocation:
        size = self.total_torque_limit.size
        comp_raw = _vector(compensation, size, "compensation")
        haptic = _vector(haptic_raw, size, "haptic_raw")
        comp = np.clip(
            comp_raw, -self.total_torque_limit, self.total_torque_limit
        )
        lower = -self.total_torque_limit - comp
        upper = self.total_torque_limit - comp
        candidate = np.clip(haptic, lower, upper)
        return PreparedAllocation(
            compensation=comp.copy(),
            haptic_raw=haptic.copy(),
            haptic_candidate=candidate,
            lower_haptic_bound=lower,
            upper_haptic_bound=upper,
            compensation_limited=bool(np.any(np.abs(comp - comp_raw) > 1e-12)),
            haptic_limited=bool(np.any(np.abs(candidate - haptic) > 1e-12)),
        )

    def finalize(
        self,
        prepared: PreparedAllocation,
        haptic_projected: np.ndarray,
        *,
        derating: float | np.ndarray = 1.0,
        accepted_total: np.ndarray | None = None,
        accepted_compensation: np.ndarray | None = None,
    ) -> AcceptedAllocation:
        size = self.total_torque_limit.size
        projected = _vector(haptic_projected, size, "haptic_projected")
        tolerance = 1e-12
        if np.any(projected < prepared.lower_haptic_bound - tolerance) or np.any(
            projected > prepared.upper_haptic_bound + tolerance
        ):
            raise ValueError(
                "haptic_projected exceeds reserved headroom; projection must "
                "not enlarge the prepared candidate"
            )
        total_requested = prepared.compensation + projected

        quantized = total_requested.copy()
        active_steps = self.quantization_step > 0.0
        quantized[active_steps] = (
            np.round(quantized[active_steps] / self.quantization_step[active_steps])
            * self.quantization_step[active_steps]
        )

        derating_vector = np.asarray(derating, dtype=float)
        if derating_vector.ndim == 0:
            derating_vector = np.full(size, float(derating_vector))
        else:
            derating_vector = derating_vector.reshape(-1)
        if (
            derating_vector.shape != (size,)
            or np.any(~np.isfinite(derating_vector))
            or np.any(derating_vector < 0.0)
            or np.any(derating_vector > 1.0)
        ):
            raise ValueError("derating must be scalar/vector in [0, 1]")
        accepted_limits = derating_vector * self.total_torque_limit
        expected_accepted = np.clip(quantized, -accepted_limits, accepted_limits)
        readback_used = accepted_total is not None
        accepted = (
            expected_accepted
            if accepted_total is None
            else _vector(accepted_total, size, "accepted_total")
        )
        if np.any(np.abs(accepted) > accepted_limits + tolerance):
            raise ValueError("accepted_total exceeds the declared derated limit")

        comp_accepted = (
            prepared.compensation
            if accepted_compensation is None
            else _vector(
                accepted_compensation, size, "accepted_compensation"
            )
        )
        haptic_accepted = accepted - comp_accepted
        return AcceptedAllocation(
            compensation_requested=prepared.compensation.copy(),
            haptic_projected=projected.copy(),
            total_requested=total_requested,
            total_quantized=quantized,
            total_accepted=accepted.copy(),
            compensation_accepted=comp_accepted.copy(),
            haptic_accepted=haptic_accepted,
            downstream_modified=bool(
                np.any(np.abs(haptic_accepted - projected) > tolerance)
            ),
            quantization_active=bool(
                np.any(np.abs(quantized - total_requested) > tolerance)
            ),
            derating_active=bool(np.any(derating_vector < 1.0 - tolerance)),
            readback_used=readback_used,
        )