SystemSimulationApp/app/simulation/solvers/mechanical.py

from __future__ import annotations

from dataclasses import dataclass
from typing import Callable, Literal, Mapping, Sequence

from app.simulation.components.amesim.mechanical.translational import (
    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"]
DenseState = Callable[[float], Sequence[float]]


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

    components: tuple[AmesimMecmas21, ...]
    mode: ConstraintMode = "uninitialized"

    @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)

    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
        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 <= 0.0
            and total_force <= 0.0
        ):
            return "lower"
        if (
            upper is not None
            and position >= upper - self._boundary_tolerance(upper)
            and velocity >= 0.0
            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.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.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.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)

        return tuple(
            MechanicalConstraintGroup(tuple(components))
            for components in masses_by_roots.values()
        )

    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 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_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 < 0.0
            ):
                candidates.append((previous_time, group, "lower", lower))
            elif (
                lower is not None
                and previous_position > 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 > 0.0
                and current_velocity < 0.0
                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 > 0.0
            ):
                candidates.append((previous_time, group, "upper", upper))
            elif (
                upper is not None
                and previous_position < 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 < 0.0
                and current_velocity > 0.0
                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)