SystemSimulationApp/app/simulation/solvers/mechanical.py

from __future__ import annotations

from dataclasses import dataclass
from math import isfinite
import os
from typing import Callable, Literal, Mapping, Sequence

from app.simulation.components.amesim.mechanical.translational import (
    AmesimLstp00a,
    AmesimMecmas21,
)
from app.simulation.core.base import DynamicComponent
from app.simulation.solvers.solver import StateTransition
from app.simulation.systems.network import SimulationNetwork


ConstraintMode = Literal["uninitialized", "free", "lower", "upper"]
MechanicalAbsoluteToleranceMode = Literal["legacy", "contact-aware-v1"]
DenseState = Callable[[float], Sequence[float]]


MECHANICAL_ATOL_MODE_ENVIRONMENT_VARIABLE = (
    "SIMULATION_MECHANICAL_ATOL_MODE"
)


def _requested_mechanical_absolute_tolerance_mode(
) -> MechanicalAbsoluteToleranceMode:
    value = os.environ.get(
        MECHANICAL_ATOL_MODE_ENVIRONMENT_VARIABLE,
        "legacy",
    ).strip().lower()
    if value == "legacy":
        return "legacy"
    if value in {"contact-aware-v1", "contact_aware_v1", "contact-aware"}:
        return "contact-aware-v1"
    raise ValueError(
        f"{MECHANICAL_ATOL_MODE_ENVIRONMENT_VARIABLE} must be "
        "'legacy' or 'contact-aware-v1'."
    )


@dataclass(frozen=True)
class MechanicalToleranceGroupPlan:
    """One rigid-coordinate group's state tolerances and proof result."""

    components: tuple[str, ...]
    contacts: tuple[str, ...]
    eligible: bool
    reason: str
    velocity_atol: float
    position_atol: float
    minimum_dvel: float | None
    minimum_contact_damping_length: float | None
    minimum_damping_strength_ratio: float | None
    minimum_force_limited_velocity_atol: float | None

    def as_dict(self) -> dict[str, object]:
        return {
            "components": list(self.components),
            "contacts": list(self.contacts),
            "eligible": self.eligible,
            "reason": self.reason,
            "velocityAtol": self.velocity_atol,
            "positionAtol": self.position_atol,
            "minimumDvel": self.minimum_dvel,
            "minimumContactDampingLength": (
                self.minimum_contact_damping_length
            ),
            "minimumDampingStrengthRatio": (
                self.minimum_damping_strength_ratio
            ),
            "minimumForceLimitedVelocityAtol": (
                self.minimum_force_limited_velocity_atol
            ),
        }


@dataclass(frozen=True)
class MechanicalAbsoluteTolerancePlan:
    """State-aligned absolute tolerances with auditable group proofs."""

    mode: MechanicalAbsoluteToleranceMode
    default_atol: float
    legacy_mechanical_atol: float
    values: tuple[float, ...]
    groups: tuple[MechanicalToleranceGroupPlan, ...]

    def as_dict(self) -> dict[str, object]:
        legacy_value = min(
            self.default_atol,
            self.legacy_mechanical_atol,
        )
        relaxed_groups = tuple(
            group
            for group in self.groups
            if group.velocity_atol > legacy_value
        )
        return {
            "mode": self.mode,
            "defaultAtol": self.default_atol,
            "legacyMechanicalAtol": self.legacy_mechanical_atol,
            "stateCount": len(self.values),
            "groupCount": len(self.groups),
            "eligibleGroupCount": sum(group.eligible for group in self.groups),
            "relaxedVelocityGroupCount": len(relaxed_groups),
            "relaxedVelocityStateCount": len(relaxed_groups),
            "relaxedPositionStateCount": 0,
            "minimumEffectiveAtol": (
                min(self.values) if self.values else None
            ),
            "maximumEffectiveAtol": (
                max(self.values) if self.values else None
            ),
            "groups": [group.as_dict() for group in self.groups],
        }


@dataclass
class MechanicalConstraintGroup:
    """MECMAS21 inertias that share one rigid translational coordinate."""

    components: tuple[AmesimMecmas21, ...]
    mode: ConstraintMode = "uninitialized"
    contact_components: tuple[AmesimLstp00a, ...] = ()

    @property
    def representative(self) -> AmesimMecmas21:
        return self.components[0]

    @property
    def total_mass(self) -> float:
        return sum(component.mass for component in self.components)

    @property
    def ideal_components(self) -> tuple[AmesimMecmas21, ...]:
        return tuple(
            component
            for component in self.components
            if component.uses_ideal_endstops
        )

    @property
    def discrete_endstop_components(self) -> tuple[AmesimMecmas21, ...]:
        return tuple(
            component
            for component in self.components
            if int(component.stoptype) in {1, 3}
        )

    @property
    def lower_bound(self) -> float | None:
        components = self.discrete_endstop_components
        return max((component.xmin for component in components), default=None)

    @property
    def upper_bound(self) -> float | None:
        components = self.discrete_endstop_components
        return min((component.xmax for component in components), default=None)

    @staticmethod
    def _boundary_tolerance(bound: float) -> float:
        return 1.0e-12 * max(abs(bound), 1.0)

    @staticmethod
    def _velocity_tolerance(velocity: float) -> float:
        """Treat only floating-point-scale motion as stationary at a stop.

        Implicit solvers perturb every state while constructing a numerical
        Jacobian.  Around an ideal endstop those perturbations must not switch
        the unilateral constraint on and off; doing so turns a zero constrained
        acceleration into the full outward-force acceleration across a
        machine-scale velocity delta.  The tolerance is deliberately far below
        MECMAS21's physical ``dvel`` threshold so real release motion is kept.
        """

        return 1.0e-12 * max(abs(velocity), 1.0)

    def reset_mode(self) -> None:
        self.mode = "uninitialized"

    def release(self) -> None:
        self.mode = "free"

    def synchronize_state(self) -> list[float]:
        reference = self.representative
        velocity_scale = max(
            [abs(component.v) for component in self.components] + [1.0]
        )
        position_scale = max(
            [abs(component.x) for component in self.components] + [1.0]
        )
        if any(
            abs(component.v - reference.v) > 1.0e-10 * velocity_scale
            or abs(component.x - reference.x) > 1.0e-10 * position_scale
            for component in self.components[1:]
        ):
            names = ", ".join(component.name for component in self.components)
            raise ValueError(
                "Rigidly connected MECMAS21 components must have consistent "
                f"initial x/v states: {names}."
            )

        lower = self.lower_bound
        upper = self.upper_bound
        names = ", ".join(component.name for component in self.components)
        if lower is not None and upper is not None and lower > upper:
            raise ValueError(
                "Rigidly connected MECMAS21 components have incompatible discrete "
                f"endstop limits: {names}."
            )

        position = reference.x
        below_lower = (
            lower is not None
            and position < lower - self._boundary_tolerance(lower)
        )
        above_upper = (
            upper is not None
            and position > upper + self._boundary_tolerance(upper)
        )
        if below_lower or above_upper:
            raise ValueError(
                f"Initial MECMAS21 position {position:g} is outside the discrete "
                f"endstop limits for: {names}."
            )
        if lower is not None and position < lower:
            position = lower
        if upper is not None and position > upper:
            position = upper

        state = [reference.v, position]
        self.set_state_vector(state)
        return state

    def set_state_vector(self, values: Sequence[float]) -> None:
        state = [float(value) for value in values]
        for component in self.components:
            component.set_state_vector(state)

    def total_unconstrained_force(self) -> float:
        return sum(
            component.mass * component.unconstrained_acceleration()
            for component in self.components
        )

    def _static_endstop_side(self, total_force: float) -> str | None:
        position = self.representative.x
        velocity = self.representative.v
        velocity_tolerance = self._velocity_tolerance(velocity)
        lower = self.lower_bound
        upper = self.upper_bound
        # MECMAS21's dvel is the friction stick threshold.  Its discrete
        # endstops release by motion direction; velocity away from a stop is free.
        if (
            lower is not None
            and position <= lower + self._boundary_tolerance(lower)
            and velocity <= velocity_tolerance
            and total_force <= 0.0
        ):
            return "lower"
        if (
            upper is not None
            and position >= upper - self._boundary_tolerance(upper)
            and velocity >= -velocity_tolerance
            and total_force >= 0.0
        ):
            return "upper"
        return None

    def lock(self, side: Literal["lower", "upper"]) -> None:
        self.mode = side

    def impact_velocity(
        self,
        side: Literal["lower", "upper"],
        incoming_velocity: float,
    ) -> float:
        """Return the post-impact velocity for the active group boundary."""

        bound = self.lower_bound if side == "lower" else self.upper_bound
        if bound is None:
            return float(incoming_velocity)
        parameter_name = "xmin" if side == "lower" else "xmax"
        active_components = tuple(
            component
            for component in self.discrete_endstop_components
            if abs(float(getattr(component, parameter_name)) - bound)
            <= self._boundary_tolerance(bound)
        )
        if any(int(component.stoptype) == 1 for component in active_components):
            return 0.0

        restitution_components = tuple(
            component
            for component in active_components
            if int(component.stoptype) == 3
        )
        speed = abs(float(incoming_velocity))
        threshold = max(
            (component.restdvel for component in restitution_components),
            default=0.0,
        )
        if speed <= threshold:
            return 0.0

        # A rigid group cannot satisfy two different simultaneous rebounds;
        # use the most dissipative active stop after plastic priority.
        restitution = min(
            (component.restcoeff for component in restitution_components),
            default=0.0,
        )
        outgoing_speed = restitution * speed
        return outgoing_speed if side == "lower" else -outgoing_speed

    def update_acceleration(self) -> float:
        """Resolve the current ideal constraint without committing event mode.

        ODE solvers may evaluate rejected or out-of-order trial states.  The
        derivative calculation therefore cannot change ``mode``; only an
        accepted state transition may commit a discrete impact mode.
        """

        total_force = self.total_unconstrained_force()
        if self._static_endstop_side(total_force) is not None:
            for component in self.components:
                component.set_constraint_motion(
                    0.0,
                    velocity=0.0,
                )
            return 0.0

        acceleration = total_force / self.total_mass
        for component in self.components:
            component.set_constraint_motion(acceleration)
        return acceleration


StateEntry = DynamicComponent | MechanicalConstraintGroup


class MechanicalStateReducer:
    """V1 rigid-inertia reduction and event-driven discrete-endstop handling.

    Rigid mechanical effort relations are causalized into one ``[v, x]`` ODE
    coordinate per connected mass group.  ``MECMAS21 stoptype=1`` applies a
    plastic impact, while ``stoptype=3`` applies its restitution coefficient
    above the configured velocity threshold.
    """

    def __init__(
        self,
        network: SimulationNetwork,
        dynamic_components: list[DynamicComponent],
    ) -> None:
        self.network = network
        self.dynamic_components = dynamic_components
        self.groups = self._build_groups()
        self._group_by_component = {
            component.name: group
            for group in self.groups
            for component in group.components
        }
        self.state_entries = self._build_state_entries()
        self._group_state_offsets = self._build_group_state_offsets()

    @staticmethod
    def _port_key(
        variable: str,
        expected_variable: str,
    ) -> tuple[str, str] | None:
        try:
            component, port, variable_name = variable.rsplit(".", 2)
        except ValueError:
            return None
        if variable_name != expected_variable:
            return None
        return component, port

    def _build_groups(self) -> tuple[MechanicalConstraintGroup, ...]:
        mechanical_ports = {
            (component.name, definition.name)
            for component in self.network.components.values()
            for definition in component.active_port_definitions
            if definition.kind == "physical" and definition.domain == "mechanical"
        }
        parents = {
            variable: {key: key for key in mechanical_ports}
            for variable in ("x", "v")
        }

        def find(variable: str, key: tuple[str, str]) -> tuple[str, str]:
            parent = parents[variable]
            root = key
            while parent[root] != root:
                root = parent[root]
            while parent[key] != key:
                next_key = parent[key]
                parent[key] = root
                key = next_key
            return root

        def union(
            variable: str,
            first: tuple[str, str],
            second: tuple[str, str],
        ) -> None:
            first_root = find(variable, first)
            second_root = find(variable, second)
            if first_root != second_root:
                parents[variable][second_root] = first_root

        for connection in self.network.connections:
            first = connection.endpoint_a.key
            second = connection.endpoint_b.key
            if first in mechanical_ports and second in mechanical_ports:
                for variable in ("x", "v"):
                    union(variable, first, second)

        for component in self.network.components.values():
            for equation in component.pressure_flow_equation_residuals():
                if equation.relation != "equal" or equation.role != "effort":
                    continue
                for variable in ("x", "v"):
                    endpoints = [
                        endpoint
                        for equation_variable in equation.variables
                        if (
                            (endpoint := self._port_key(equation_variable, variable))
                            in mechanical_ports
                        )
                    ]
                    for endpoint in endpoints[1:]:
                        union(variable, endpoints[0], endpoint)

        masses = [
            component
            for component in self.dynamic_components
            if isinstance(component, AmesimMecmas21)
        ]
        for component in masses:
            ports = [
                (component.name, definition.name)
                for definition in component.active_port_definitions
                if definition.kind == "physical"
                and definition.domain == "mechanical"
            ]
            for port in ports[1:]:
                for variable in ("x", "v"):
                    union(variable, ports[0], port)

        masses_by_roots: dict[
            tuple[tuple[str, str], tuple[str, str]],
            list[AmesimMecmas21],
        ] = {}
        for component in masses:
            first_port = next(
                (component.name, definition.name)
                for definition in component.active_port_definitions
                if definition.kind == "physical"
                and definition.domain == "mechanical"
            )
            roots = (find("x", first_port), find("v", first_port))
            masses_by_roots.setdefault(roots, []).append(component)

        contacts_by_roots: dict[
            tuple[tuple[str, str], tuple[str, str]],
            dict[str, AmesimLstp00a],
        ] = {}
        for component in self.network.components.values():
            if not isinstance(component, AmesimLstp00a):
                continue
            contact_roots = tuple(
                (
                    find("x", (component.name, definition.name)),
                    find("v", (component.name, definition.name)),
                )
                for definition in component.active_port_definitions
                if (
                    definition.kind == "physical"
                    and definition.domain == "mechanical"
                )
            )
            if (
                len({roots[0] for roots in contact_roots}) < 2
                or len({roots[1] for roots in contact_roots}) < 2
            ):
                # A compliant contact whose two ports resolve to the same
                # rigid coordinate cannot damp that coordinate.  Treating the
                # self-loop as proof would relax an unrelated velocity state.
                continue
            for definition in component.active_port_definitions:
                if (
                    definition.kind != "physical"
                    or definition.domain != "mechanical"
                ):
                    continue
                endpoint = (component.name, definition.name)
                roots = (find("x", endpoint), find("v", endpoint))
                contacts_by_roots.setdefault(roots, {})[
                    component.name
                ] = component

        return tuple(
            MechanicalConstraintGroup(
                components=tuple(components),
                contact_components=tuple(
                    contacts_by_roots.get(roots, {}).values()
                ),
            )
            for roots, components in masses_by_roots.items()
        )

    def _build_state_entries(self) -> tuple[StateEntry, ...]:
        entries: list[StateEntry] = []
        for component in self.dynamic_components:
            group = self._group_by_component.get(component.name)
            if group is None:
                entries.append(component)
            elif group.representative is component:
                entries.append(group)
        return tuple(entries)

    def _build_group_state_offsets(self) -> dict[int, int]:
        offsets: dict[int, int] = {}
        cursor = 0
        for entry in self.state_entries:
            if isinstance(entry, MechanicalConstraintGroup):
                offsets[id(entry)] = cursor
                cursor += 2
            else:
                cursor += entry.state_size
        return offsets

    @property
    def has_state_events(self) -> bool:
        return any(group.discrete_endstop_components for group in self.groups)

    def absolute_tolerance_plan(
        self,
        default: float,
        *,
        mechanical: float = 1.0e-12,
        mode: MechanicalAbsoluteToleranceMode | None = None,
    ) -> MechanicalAbsoluteTolerancePlan:
        """Compile state tolerances without weakening non-smooth coordinates.

        A scalar ``1e-8`` absolute tolerance makes SciPy perturb a zero-valued
        endstop position across the much smaller unilateral boundary band while
        constructing finite-difference Jacobians. Positions and ideal endstop
        states therefore retain the legacy machine-scale tolerance.

        Strongly damped, compliant LSTP contact can instead drive a *free*
        velocity close to zero for hundreds of accepted steps. Only a
        compile-proven smooth-contact group may use the bounded velocity floor;
        the contact position coordinate remains unchanged.
        """
        default_atol = float(default)
        mechanical_atol = float(mechanical)
        if not isfinite(default_atol) or default_atol <= 0.0:
            raise ValueError(
                "default absolute tolerance must be finite and positive."
            )
        if not isfinite(mechanical_atol) or mechanical_atol <= 0.0:
            raise ValueError(
                "mechanical absolute tolerance must be finite and positive."
            )
        selected_mode = mode or _requested_mechanical_absolute_tolerance_mode()
        if selected_mode not in {"legacy", "contact-aware-v1"}:
            raise ValueError(
                "mechanical absolute tolerance mode must be 'legacy' or "
                "'contact-aware-v1'."
            )

        legacy_atol = min(default_atol, mechanical_atol)
        values: list[float] = []
        group_plans: list[MechanicalToleranceGroupPlan] = []
        for entry in self.state_entries:
            if not isinstance(entry, MechanicalConstraintGroup):
                values.extend([default_atol] * entry.state_size)
                continue

            components = entry.components
            contacts = entry.contact_components
            positive_dvel = tuple(
                float(component.dvel)
                for component in components
                if isfinite(float(component.dvel))
                and float(component.dvel) > 0.0
            )
            positive_pdis = tuple(
                float(contact.Pdis)
                for contact in contacts
                if isfinite(float(contact.Pdis))
                and float(contact.Pdis) > 0.0
            )
            minimum_dvel = min(positive_dvel, default=None)
            minimum_pdis = min(positive_pdis, default=None)
            contact_scale_valid = True
            contact_force_velocity_limits_list: list[float] = []
            damping_strength_ratios_list: list[float] = []
            if selected_mode == "contact-aware-v1" and minimum_dvel is not None:
                for contact in contacts:
                    stiffness = float(contact.kcont)
                    damping_length = float(contact.Pdis)
                    damping = float(contact.rcont)
                    if not (
                        isfinite(stiffness)
                        and stiffness > 0.0
                        and isfinite(damping_length)
                        and damping_length > 0.0
                        and isfinite(damping)
                        and damping > 0.0
                    ):
                        contact_scale_valid = False
                        continue
                    elastic_force_scale = stiffness * damping_length
                    damping_force_scale = damping * minimum_dvel
                    if not (
                        isfinite(elastic_force_scale)
                        and elastic_force_scale > 0.0
                        and isfinite(damping_force_scale)
                        and damping_force_scale > 0.0
                    ):
                        contact_scale_valid = False
                        continue
                    force_velocity_limit = (
                        1.0e-3 * elastic_force_scale / damping
                    )
                    damping_strength_ratio = (
                        damping_force_scale / elastic_force_scale
                    )
                    if not (
                        isfinite(force_velocity_limit)
                        and force_velocity_limit > 0.0
                        and isfinite(damping_strength_ratio)
                        and damping_strength_ratio > 0.0
                    ):
                        contact_scale_valid = False
                        continue
                    contact_force_velocity_limits_list.append(
                        force_velocity_limit
                    )
                    damping_strength_ratios_list.append(
                        damping_strength_ratio
                    )
            contact_force_velocity_limits = tuple(
                contact_force_velocity_limits_list
            )
            minimum_force_velocity_atol = min(
                contact_force_velocity_limits,
                default=None,
            )
            damping_strength_ratios = tuple(
                damping_strength_ratios_list
            )
            minimum_damping_strength_ratio = min(
                damping_strength_ratios,
                default=None,
            )

            if selected_mode == "legacy":
                eligible = False
                reason = "legacyMode"
            elif entry.discrete_endstop_components:
                eligible = False
                reason = "discreteEndstop"
            elif any(int(component.stoptype) != 4 for component in components):
                eligible = False
                reason = "unsupportedStopType"
            elif any(
                component.use_friction and float(component.fcoul) != 0.0
                for component in components
            ):
                eligible = False
                reason = "dryFriction"
            elif not contacts:
                eligible = False
                reason = "noFlexibleContact"
            elif any(
                not isfinite(float(contact.Pdis))
                or float(contact.Pdis) <= 0.0
                for contact in contacts
            ):
                eligible = False
                reason = "nonSmoothContactDampingLength"
            elif any(
                not isfinite(float(contact.rcont))
                or float(contact.rcont) <= 0.0
                for contact in contacts
            ):
                eligible = False
                reason = "undampedContact"
            elif any(
                not isfinite(float(contact.kcont))
                or float(contact.kcont) <= 0.0
                for contact in contacts
            ):
                eligible = False
                reason = "invalidContactStiffness"
            elif not contact_scale_valid:
                eligible = False
                reason = "invalidContactScale"
            elif any(
                int(contact.discContactOption) != 1
                for contact in contacts
            ):
                eligible = False
                reason = "clampedContactForce"
            elif len(positive_dvel) != len(components):
                eligible = False
                reason = "invalidVelocityScale"
            elif (
                len(damping_strength_ratios) != len(contacts)
                or minimum_damping_strength_ratio is None
                or minimum_damping_strength_ratio < 1.0
            ):
                eligible = False
                reason = "weakContactDamping"
            else:
                eligible = True
                reason = "eligibleFlexibleContact"

            velocity_atol = legacy_atol
            if eligible:
                assert minimum_dvel is not None
                assert minimum_force_velocity_atol is not None
                velocity_atol = min(
                    default_atol,
                    max(
                        mechanical_atol,
                        min(
                            1.0e-9,
                            1.0e-3 * minimum_dvel,
                            minimum_force_velocity_atol,
                        ),
                    ),
                )
            position_atol = legacy_atol
            values.extend((velocity_atol, position_atol))
            group_plans.append(
                MechanicalToleranceGroupPlan(
                    components=tuple(
                        component.name for component in components
                    ),
                    contacts=tuple(contact.name for contact in contacts),
                    eligible=eligible,
                    reason=reason,
                    velocity_atol=velocity_atol,
                    position_atol=position_atol,
                    minimum_dvel=minimum_dvel,
                    minimum_contact_damping_length=minimum_pdis,
                    minimum_damping_strength_ratio=(
                        minimum_damping_strength_ratio
                    ),
                    minimum_force_limited_velocity_atol=(
                        minimum_force_velocity_atol
                    ),
                )
            )

        return MechanicalAbsoluteTolerancePlan(
            mode=selected_mode,
            default_atol=default_atol,
            legacy_mechanical_atol=mechanical_atol,
            values=tuple(values),
            groups=tuple(group_plans),
        )

    def absolute_tolerances(
        self,
        default: float,
        *,
        mechanical: float = 1.0e-12,
        mode: MechanicalAbsoluteToleranceMode | None = None,
    ) -> list[float]:
        """Return state-aligned values from the auditable tolerance plan."""

        return list(
            self.absolute_tolerance_plan(
                default,
                mechanical=mechanical,
                mode=mode,
            ).values
        )

    def reset_constraint_modes(self) -> None:
        for group in self.groups:
            group.reset_mode()

    def initial_state_vector(self) -> list[float]:
        self.reset_constraint_modes()
        values: list[float] = []
        for entry in self.state_entries:
            if isinstance(entry, MechanicalConstraintGroup):
                values.extend(entry.synchronize_state())
            else:
                values.extend(entry.get_state_vector())
        return values

    def apply_state_vector(self, values: list[float]) -> None:
        cursor = 0
        for entry in self.state_entries:
            state_size = (
                2 if isinstance(entry, MechanicalConstraintGroup) else entry.state_size
            )
            next_cursor = cursor + state_size
            state = values[cursor:next_cursor]
            if isinstance(entry, MechanicalConstraintGroup):
                entry.set_state_vector(state)
            else:
                entry.set_state_vector(state)
            cursor = next_cursor
        if cursor != len(values):
            raise ValueError("State vector length does not match reduced dynamic components.")

    def update_constraint_accelerations(self) -> None:
        for group in self.groups:
            group.update_acceleration()

    def state_derivatives(
        self,
        connected_h: Mapping[str, Mapping[str, float]],
    ) -> list[float]:
        derivatives: list[float] = []
        for entry in self.state_entries:
            component = (
                entry.representative
                if isinstance(entry, MechanicalConstraintGroup)
                else entry
            )
            derivatives.extend(
                component.state_derivative_from_ports(connected_h[component.name])
            )
        return derivatives

    @staticmethod
    def _locate_crossing(
        dense_state: DenseState,
        state_index: int,
        bound: float,
        side: Literal["lower", "upper"],
        start_time: float,
        end_time: float,
    ) -> float:
        lower_time = float(start_time)
        upper_time = float(end_time)
        for _iteration in range(60):
            middle_time = 0.5 * (lower_time + upper_time)
            position = float(dense_state(middle_time)[state_index])
            crossed = position <= bound if side == "lower" else position >= bound
            if crossed:
                upper_time = middle_time
            else:
                lower_time = middle_time
        return upper_time

    @staticmethod
    def _locate_turnaround(
        dense_state: DenseState,
        velocity_index: int,
        side: Literal["lower", "upper"],
        start_time: float,
        end_time: float,
    ) -> float:
        """Locate the velocity reversal preceding a same-step re-impact."""

        lower_time = float(start_time)
        upper_time = float(end_time)
        for _iteration in range(60):
            middle_time = 0.5 * (lower_time + upper_time)
            velocity = float(dense_state(middle_time)[velocity_index])
            turned = velocity <= 0.0 if side == "lower" else velocity >= 0.0
            if turned:
                upper_time = middle_time
            else:
                lower_time = middle_time
        return upper_time

    def state_transition(
        self,
        previous_time: float,
        previous_state: list[float],
        current_time: float,
        current_state: list[float],
        dense_state: DenseState,
    ) -> StateTransition | None:
        """Return the earliest discrete-endstop impact in one accepted ODE step."""

        candidates: list[
            tuple[float, MechanicalConstraintGroup, Literal["lower", "upper"], float]
        ] = []
        for group in self.groups:
            if not group.discrete_endstop_components:
                continue
            velocity_index = self._group_state_offsets[id(group)]
            position_index = velocity_index + 1
            previous_velocity = float(previous_state[velocity_index])
            current_velocity = float(current_state[velocity_index])
            previous_velocity_tolerance = group._velocity_tolerance(previous_velocity)
            current_velocity_tolerance = group._velocity_tolerance(current_velocity)
            previous_position = float(previous_state[position_index])
            current_position = float(current_state[position_index])
            lower = group.lower_bound
            upper = group.upper_bound
            if (
                lower is not None
                and previous_position <= lower + group._boundary_tolerance(lower)
                and previous_velocity < -previous_velocity_tolerance
            ):
                candidates.append((previous_time, group, "lower", lower))
            elif (
                lower is not None
                and previous_position > lower + group._boundary_tolerance(lower)
                and current_position <= lower
            ):
                candidates.append(
                    (
                        self._locate_crossing(
                            dense_state,
                            position_index,
                            lower,
                            "lower",
                            previous_time,
                            current_time,
                        ),
                        group,
                        "lower",
                        lower,
                    )
                )
            elif (
                lower is not None
                and previous_position <= lower
                and previous_velocity > previous_velocity_tolerance
                and current_velocity < -current_velocity_tolerance
                and current_position <= lower
            ):
                turnaround_time = self._locate_turnaround(
                    dense_state,
                    velocity_index,
                    "lower",
                    previous_time,
                    current_time,
                )
                candidates.append(
                    (
                        self._locate_crossing(
                            dense_state,
                            position_index,
                            lower,
                            "lower",
                            turnaround_time,
                            current_time,
                        ),
                        group,
                        "lower",
                        lower,
                    )
                )
            if (
                upper is not None
                and previous_position >= upper - group._boundary_tolerance(upper)
                and previous_velocity > previous_velocity_tolerance
            ):
                candidates.append((previous_time, group, "upper", upper))
            elif (
                upper is not None
                and previous_position < upper - group._boundary_tolerance(upper)
                and current_position >= upper
            ):
                candidates.append(
                    (
                        self._locate_crossing(
                            dense_state,
                            position_index,
                            upper,
                            "upper",
                            previous_time,
                            current_time,
                        ),
                        group,
                        "upper",
                        upper,
                    )
                )
            elif (
                upper is not None
                and previous_position >= upper
                and previous_velocity < -previous_velocity_tolerance
                and current_velocity > current_velocity_tolerance
                and current_position >= upper
            ):
                turnaround_time = self._locate_turnaround(
                    dense_state,
                    velocity_index,
                    "upper",
                    previous_time,
                    current_time,
                )
                candidates.append(
                    (
                        self._locate_crossing(
                            dense_state,
                            position_index,
                            upper,
                            "upper",
                            turnaround_time,
                            current_time,
                        ),
                        group,
                        "upper",
                        upper,
                    )
                )

        if not candidates:
            return None

        event_time = min(candidate[0] for candidate in candidates)
        event_state = [float(value) for value in dense_state(event_time)]
        simultaneous_tolerance = 1.0e-12 * max(abs(event_time), 1.0)
        for candidate_time, group, side, bound in candidates:
            if abs(candidate_time - event_time) > simultaneous_tolerance:
                continue
            velocity_index = self._group_state_offsets[id(group)]
            event_state[velocity_index] = group.impact_velocity(
                side,
                event_state[velocity_index],
            )
            event_state[velocity_index + 1] = bound
            if event_state[velocity_index] == 0.0:
                group.lock(side)
            else:
                group.release()

        return StateTransition(time=event_time, state=event_state)