sim.py

"""Headless simulation core: the whole model with no rendering and no pygame.

This is the seam the rest of the rebuild hangs off. Everything that is not
drawing lives here, so the simulation can be driven by a UI, a batch sweep, a
replay, or a learning agent without any of them knowing about the others.

``Simulation`` owns the physics, contrast transport, and assessment. Rendering
and input devices sit above it in :mod:`cathsim.ui` and :mod:`cathsim.control`.
Importing this module must never pull in pygame; ``test_simulation.py`` asserts
that in a subprocess.
"""

from __future__ import annotations

import math
from dataclasses import dataclass

import numpy as np

from cathsim.device import Device
from cathsim.physics.dye import DyeField
from cathsim.physics.hemodynamics import WomersleyField
from cathsim.physics.rod import Rod
from cathsim.physics.runner import RodRunner
from cathsim.scenario import Scenario, ScenarioEngine, build_vessel
from cathsim.vessel.geometry import VesselGeometry

__all__ = [
    "Action",
    "ControlInput",
    "Observation",
    "Simulation",
    "bounds_of",
]

FPS = 30
"""Simulation frames per second. One ``Action`` is consumed per frame."""

SUBSTEPS = 16
"""Physics substeps per frame.

Measured as the requirement of the most demanding shipped anatomy: the tortuous
carotid needs 16, a straight vessel is stable at 2. It is set by discretization
and geometric nonlinearity, *not* by device stiffness -- a 17x range of EI does
not move it. Scenarios may override it.
"""

NODES = 30
"""Rod discretization. 30-60 is the validated range."""

INSERTION_LENGTH = 0.06
"""Device arclength inside the anatomy at the start of a run (m)."""

RESERVOIR_ML = 20.0
SHOT_ML = 1.5
DYE_PER_SHOT = 25.0

INJECTION_RATE_ML_S = 4.0
"""Contrast delivery rate while injection is held (mL/s).

A neuro carotid run is around 6-8 mL at 4-5 mL/s, so this is the middle of the
band for a diagnostic run. Hand injection is neither this steady nor this fast,
which is a distinction worth modelling once the injector becomes a control.
"""

INJECTION_JET_SPEED = 0.35
"""Exit velocity of contrast leaving the catheter tip (m/s).

**Estimate.** Order-of-magnitude right for a diagnostic catheter at a few mL/s,
and it is what decides whether an injection crosses the lumen or streams along
one wall. A calibration knob, not a measurement.
"""

MARGIN = 0.015
"""World-space padding around the anatomy for grid and camera bounds (m)."""


def bounds_of(vessel: VesselGeometry) -> tuple[float, float, float, float]:
    """World bounds covering every wall vertex plus a margin.

    Args:
        vessel: Geometry to bound.

    Returns:
        ``(xmin, ymin, xmax, ymax)`` in metres.
    """
    points = [p for wall in vessel.walls for p in wall]
    xs = [p.x for p in points]
    ys = [p.y for p in points]
    return (min(xs) - MARGIN, min(ys) - MARGIN, max(xs) + MARGIN, max(ys) + MARGIN)


@dataclass(frozen=True)
class Action:
    """One frame of operator intent, in simulation units.

    This is the input side of the simulation's interface: a keyboard driver, a
    serial rig, a replayed recording, and a policy all produce this same value.

    Attributes:
        advance: Commanded insertion rate (m/s); positive inserts. Clinical
            range is 5-30 mm/s. Rate, not force: an operator advances a device
            at a controlled speed and applies whatever force that takes, and
            force-controlling it put the whole push on one ~4 ug node.
        torque: Torque applied to the shaft (N*m); positive is counter-clockwise.
        tip_deflection: Active tip deflection command, -1..1.
        inject: Whether contrast injection is requested this frame.
    """

    advance: float = 0.0
    torque: float = 0.0
    tip_deflection: float = 0.0
    inject: bool = False
    fluoro: bool = False
    """Whether the fluoroscopy pedal is down this frame.

    The image is not free and is not continuous. It exists only while the pedal
    is held; the rest of the time the operator is working from the last frame
    acquired, which is the habit the dose model is there to teach. Without a
    pedal there is nothing to be disciplined about.
    """

    fluoro_rate_pps: float = 7.5
    """Pulse rate while the pedal is down (pulses/s).

    A real control with a real trade-off: 7.5 is choppy against 15 but costs
    half the dose, and :class:`cathsim.assessment.DoseLedger` already meters it.
    Physics still runs at the full frame rate -- only the *image* is pulsed, so
    a fast push genuinely can happen between two frames and not be seen.
    """


ControlInput = Action
"""Legacy name for :class:`Action`, kept so existing callers keep working."""


@dataclass(frozen=True)
class Observation:
    """One frame of simulation state, as seen by a consumer.

    This is the output side of the interface and the unit of the recorded
    trace: after-action review plots it, and the skill classifier draws its
    features from it. Everything here is a plain number so a run serialises
    without the physics objects.

    Attributes:
        elapsed_s: Simulated time since the run started (s).
        tip_position: Catheter tip world position (m).
        tip_heading: Tip segment heading (rad).
        base_position: Proximal (handle) end world position (m).
        contact_force: Wall contact force this frame (N).
        peak_force: Largest contact force seen so far (N).
        fit: Force Impact over Time, the integral of contact force (N*s).
            Literature bands: expert 1.71, novice 9.11 (Rafii-Tari et al. 2017).
        contrast_ml: Contrast injected so far (mL).
        reservoir_ml: Contrast remaining (mL).
        segments_outside_lumen: Device segments outside every branch lumen.
            Non-zero means the device has left the anatomy.
        status: Assessment status, e.g. ``IN_PROGRESS`` or ``SUCCESS``.
    """

    elapsed_s: float
    tip_position: tuple[float, float]
    tip_heading: float
    base_position: tuple[float, float]
    contact_force: float
    peak_force: float
    fit: float
    contrast_ml: float
    reservoir_ml: float
    segments_outside_lumen: int
    status: str


class Simulation:
    """The whole model, headless: physics, contrast transport, and assessment.

    Deterministic by construction — the engine sub-steps at a fixed dt — so a
    scenario plus an action sequence replays bit-exactly. That is what makes
    recording cheap (see :mod:`cathsim.recording`).
    """

    def __init__(self, scenario: Scenario, *, fps: int = FPS, substeps: int = SUBSTEPS) -> None:
        """Build the simulation stack for ``scenario``.

        Args:
            scenario: Scenario selecting anatomy, flow, target, and limits.
            fps: Frames per second; one action is consumed per frame.
            substeps: Physics substeps per frame.
        """
        self.scenario = scenario
        self.fps = fps
        self.substeps = scenario.substeps or substeps
        self.reset()

    def reset(self) -> Observation:
        """Rebuild the run from scratch and return the opening observation."""
        scenario = self.scenario
        self.vessel = build_vessel(scenario)
        self.bounds = bounds_of(self.vessel)
        self.device = Device.named(scenario.device)

        # ARCH(vessel): branches[0] is an unwritten contract holding up
        # four separate things -- the flow field below, and RodRunner's
        # arclength, outlet test and sheath mouth. Three vessel
        # generators say so in their docstrings and nothing enforces it,
        # so an SVG import or a hand-authored scenario that happens to
        # list the route second is broken in a way that reads as a
        # physics bug. Worse, flow exists in branch 0 only: the Type III
        # arch's decoy and the driving test's two spurs are stagnant, so
        # contrast never washes out of a cul-de-sac. Name it in
        # VesselGeometry (a `route` index, or an explicit inlet) and
        # build the flow field over the tree.
        branch = self.vessel.branches[0]
        centerline = [(p.x, p.y) for p in branch.points]
        # The scenario's ``v_max`` was the peak centreline speed of a parabolic
        # profile, whose cross-section mean is half of it. Converting preserves
        # the volumetric flow rate across the switch to the Womersley field.
        self.flow = WomersleyField(centerline, branch.radii, mean_velocity=scenario.v_max / 2.0)
        self.runner = RodRunner(
            self.vessel,
            self.device,
            self.flow,
            n_nodes=NODES,
            insertion_length=(
                scenario.insertion_length
                if scenario.insertion_length is not None
                else INSERTION_LENGTH
            ),
            substeps=self.substeps,
            fps=self.fps,
            tip_shape=scenario.tip_shape,
            outer=(
                Device.named(scenario.outer_device) if scenario.outer_device is not None else None
            ),
            outer_lag=scenario.outer_lag,
        )
        # TODO(dye): 192x96 is fixed while ``bounds`` is not, so cell
        # aspect ratio is whatever the anatomy happens to have. The
        # open channel is roughly 2:1 and lands square; the M1 segment
        # is ~50 mm end to end and the aneurysm neck is nearly as tall
        # as it is wide, so both get strongly anisotropic cells. That
        # matters because DyeField picks its back-trace step from
        # min(dx, dy) and its diffusion CFL from the same -- so the
        # narrow axis silently sets the substep count for both. Size the
        # grid from the bounds' aspect at a target cell size instead.
        self.dye = DyeField(bounds=self.bounds, nx=192, ny=96, field=self.flow, vessel=self.vessel)
        self.assessment = ScenarioEngine(scenario)
        # TODO(phase4): wire in cathsim.assessment -- Patient, ContrastLedger and
        # DoseLedger are built, calibrated and tested but nothing constructs them
        # yet. Needs a ``patient`` and ``dose_budget`` block in the scenario
        # schema first.

        self.reservoir_ml = RESERVOIR_ML
        self.elapsed = 0.0
        return self.observe()

    @property
    def rod(self) -> Rod:
        """The device being navigated, as a node polyline."""
        return self.runner.rod

    @property
    def finished(self) -> bool:
        """True once the run has hit a terminal condition."""
        return self.assessment.finished

    def apply_input(self, action: Action) -> None:
        """Route tip deflection and contrast injection into the model.

        Split from :meth:`advance` because the interactive demo keeps accepting
        these while the physics is paused at the end of a run.

        Args:
            action: Operator intent for this frame.
        """
        self.runner.set_tip_deflection(action.tip_deflection)
        if not action.inject or self.reservoir_ml <= 0.0:
            return

        # Injection is a *rate*, held for as long as the operator holds it, not
        # a discrete shot per frame. Shots left the lumen visibly beaded --
        # one blob per frame, marching downstream with gaps between them -- and
        # made contrast usage quantised to the frame rate rather than to the
        # time spent on the syringe.
        dt = 1.0 / self.runner.fps
        # FIXME(sim): ``millilitres`` is billed in full whenever the jet
        # reached any cell at all, but inject_jet returns a *cell count*,
        # not a delivered volume. Injecting with the tip against a wall,
        # or half out of the grid, charges the reservoir and the contrast
        # ledger for agent that never entered the lumen. Contrast volume
        # is a graded quantity here, so the discrepancy costs marks.
        millilitres = min(INJECTION_RATE_ML_S * dt, self.reservoir_ml)
        delivered = self.dye.inject_jet(
            self.runner.tip_position(),
            self.runner.rod.tip_heading(),
            self._contrast_fraction(),
            INJECTION_JET_SPEED,
            dt,
        )
        if delivered > 0.0:
            self.reservoir_ml -= millilitres
            self.assessment.record_contrast(millilitres)

    def _contrast_fraction(self) -> float:
        """Share of the flow at the tip that is contrast, 0 to 1.

        What a vessel actually carries is a *mixture*: the injected rate against
        the blood already going past. Injecting 4 mL/s into a vessel carrying
        2.7 mL/s gives ``4 / (4 + 2.7) = 0.60``, and no injection can exceed 1.0
        because that is neat agent.

        Computing it -- rather than depositing a fixed mass -- is what makes the
        same injection behave correctly at any calibre. A mass-based deposit
        read 0.59 in a 1.85 mm neuro vessel and **9.8** in a 4 mm femoral, and
        980% contrast is not a thing. It also makes the clinical trade real: the
        same syringe barely tints a big vessel and saturates a small one, which
        is why injection rates are chosen per target.
        """
        tip = self.runner.tip_position()
        vx, vy = self.flow.velocity_at_points(np.array([tip[0]]), np.array([tip[1]]))
        speed = float(math.hypot(float(vx[0]), float(vy[0])))
        radius = self.runner.lumen_radius_at(tip)
        blood_ml_s = speed * math.pi * radius * radius * 1e6
        return float(np.clip(INJECTION_RATE_ML_S / (INJECTION_RATE_ML_S + blood_ml_s), 0.02, 1.0))

    def advance(self, action: Action) -> None:
        """Advance one frame of physics, contrast transport, and assessment.

        Args:
            action: Operator intent for this frame; supplies the insertion rate
                and hub torque.
        """
        # Advance the cardiac cycle before anything samples the flow, so drag
        # and contrast transport see the same instant of systole.
        self.flow.set_time(self.elapsed)
        cycle = self.flow.waveform(self.elapsed)
        self.dye.set_flow_scale(cycle)
        # Contrast transport and vessel distension ride the same beat.
        self.runner.cycle_scale = cycle

        self.runner.step(advance=action.advance, torque=action.torque)
        self.dye.step(1.0 / self.fps)
        self.elapsed += 1.0 / self.fps

        if self.runner.perforated:
            self.dye.extravasate(self.runner.tip_position(), DYE_PER_SHOT * 0.2)

        self.assessment.update(
            tip_position=self.runner.tip_position(),
            elapsed_s=self.elapsed,
            peak_force=self.runner.peak_force,
            wti=self.runner.fit,
            perforated=self.runner.perforated,
        )

    def step(self, action: Action) -> Observation:
        """Consume one action and return the resulting observation.

        The interface a batch run, a replay, or a policy drives. The interactive
        demo calls :meth:`apply_input` and :meth:`advance` separately instead,
        because it keeps taking input after the run ends.

        Args:
            action: Operator intent for this frame.

        Returns:
            The observation after the frame.
        """
        self.apply_input(action)
        self.advance(action)
        return self.observe()

    def segments_outside_lumen(self) -> int:
        """Count device nodes lying outside every branch lumen.

        Non-zero means the device has left the anatomy. The distance field caps
        the ends of the modelled vessel, so this should stay zero.
        """
        return self.runner.nodes_outside_lumen()

    def observe(self) -> Observation:
        """Snapshot the current state as a plain-number :class:`Observation`."""
        tip = self.runner.tip_position()
        base = self.rod.positions[0]
        return Observation(
            elapsed_s=self.elapsed,
            tip_position=tip,
            tip_heading=self.rod.tip_heading(),
            base_position=(float(base[0]), float(base[1])),
            contact_force=self.runner.contact.force,
            peak_force=self.runner.peak_force,
            fit=self.runner.fit,
            contrast_ml=self.assessment.contrast_ml,
            reservoir_ml=self.reservoir_ml,
            segments_outside_lumen=self.segments_outside_lumen(),
            status=self.assessment.outcome().status,
        )