exoskeleton/code/core/energy_audit.py

"""Independent H3/H4 metrics reconstructed from immutable raw logs."""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np


@dataclass(frozen=True)
class EnergyAuditResult:
    energy_before: np.ndarray
    energy_preclip: np.ndarray
    energy_after: np.ndarray
    floor_deficit: np.ndarray
    projected_shadow_energy: np.ndarray
    candidate_shadow_energy: np.ndarray
    candidate_power: np.ndarray
    accepted_power: np.ndarray
    max_floor_deficit: float
    projected_floor_deficit: float
    candidate_floor_deficit: float
    delta_B: float
    projection_distortion: float
    released_energy: float
    preclip_log_max_error: float | None
    downstream_modification_max: float

    @property
    def passed_floor_gate(self) -> bool:
        return self.max_floor_deficit <= 1e-12


def _sample_matrix(value, name: str) -> np.ndarray:
    array = np.asarray(value, dtype=float)
    if array.ndim != 2 or array.shape[0] == 0 or array.shape[1] == 0:
        raise ValueError(f"{name} must be a non-empty 2-D array")
    if not np.all(np.isfinite(array)):
        raise ValueError(f"{name} must contain only finite values")
    return array


def audit_haptic_energy(
    *,
    tau_candidate: np.ndarray,
    tau_projected: np.ndarray,
    tau_accepted: np.ndarray,
    qd_master: np.ndarray,
    dt: float | np.ndarray,
    energy_initial: float,
    energy_min: float,
    energy_max: float,
    epsilon_tau: float = 1e-12,
    logged_preclip: np.ndarray | None = None,
) -> EnergyAuditResult:
    """Recompute the H4 budget and transparency endpoints.

    The deterministic budget gate uses ``tau_accepted`` because this is the
    actual actuator-port increment.  The same-run shadow uses the logged
    pre-projection candidate and never feeds the counterfactual back into the
    simulated trajectory.
    """
    candidate = _sample_matrix(tau_candidate, "tau_candidate")
    projected = _sample_matrix(tau_projected, "tau_projected")
    accepted = _sample_matrix(tau_accepted, "tau_accepted")
    velocity = _sample_matrix(qd_master, "qd_master")
    if not (
        candidate.shape == projected.shape == accepted.shape == velocity.shape
    ):
        raise ValueError("all torque and velocity arrays must have equal shapes")
    count = candidate.shape[0]
    delta = np.broadcast_to(np.asarray(dt, dtype=float), (count,)).copy()
    if np.any(delta <= 0.0) or not np.all(np.isfinite(delta)):
        raise ValueError("dt must contain finite positive values")
    if not (
        np.isfinite(energy_initial)
        and np.isfinite(energy_min)
        and np.isfinite(energy_max)
        and energy_min <= energy_initial <= energy_max
    ):
        raise ValueError("energy values must satisfy min <= initial <= max")
    if epsilon_tau <= 0.0:
        raise ValueError("epsilon_tau must be positive")

    candidate_power = np.einsum("ij,ij->i", candidate, velocity)
    accepted_power = np.einsum("ij,ij->i", accepted, velocity)
    energy_before = np.empty(count, dtype=float)
    energy_preclip = np.empty(count, dtype=float)
    energy_after = np.empty(count, dtype=float)
    floor_deficit = np.empty(count, dtype=float)
    energy = float(energy_initial)
    for index in range(count):
        energy_before[index] = energy
        preclip = energy - accepted_power[index] * delta[index]
        energy_preclip[index] = preclip
        floor_deficit[index] = max(0.0, energy_min - preclip)
        energy = float(np.clip(preclip, energy_min, energy_max))
        energy_after[index] = energy

    candidate_shadow = np.empty(count + 1, dtype=float)
    projected_shadow = np.empty(count + 1, dtype=float)
    candidate_shadow[0] = energy_initial
    projected_shadow[0] = energy_initial
    for index in range(count):
        candidate_shadow[index + 1] = min(
            energy_max,
            candidate_shadow[index] - candidate_power[index] * delta[index],
        )
        projected_shadow[index + 1] = min(
            energy_max,
            projected_shadow[index] - accepted_power[index] * delta[index],
        )
    candidate_B = float(
        np.max(np.maximum(0.0, energy_min - candidate_shadow))
    )
    projected_B = float(
        np.max(np.maximum(0.0, energy_min - projected_shadow))
    )
    numerator = float(
        np.sum(np.linalg.norm(accepted - candidate, axis=1) * delta)
    )
    denominator = float(
        np.sum(np.linalg.norm(candidate, axis=1) * delta) + epsilon_tau
    )

    preclip_error = None
    if logged_preclip is not None:
        logged = np.asarray(logged_preclip, dtype=float).reshape(-1)
        if logged.shape != (count,) or not np.all(np.isfinite(logged)):
            raise ValueError("logged_preclip must be finite with one value per sample")
        preclip_error = float(np.max(np.abs(logged - energy_preclip)))

    return EnergyAuditResult(
        energy_before=energy_before,
        energy_preclip=energy_preclip,
        energy_after=energy_after,
        floor_deficit=floor_deficit,
        projected_shadow_energy=projected_shadow,
        candidate_shadow_energy=candidate_shadow,
        candidate_power=candidate_power,
        accepted_power=accepted_power,
        max_floor_deficit=float(np.max(floor_deficit)),
        projected_floor_deficit=projected_B,
        candidate_floor_deficit=candidate_B,
        delta_B=candidate_B - projected_B,
        projection_distortion=numerator / denominator,
        released_energy=float(np.sum(accepted_power * delta)),
        preclip_log_max_error=preclip_error,
        downstream_modification_max=float(
            np.max(np.linalg.norm(accepted - projected, axis=1))
        ),
    )