src/fpgm/motion/trajectory.py

"""Plans an :class:`~fpgm.motion.types.ActionSpec` into a joint trajectory.

Realises each waypoint with :class:`fpgm.robot.ik.IKSolver`, one 15 Hz grid
tick at a time, producing a :class:`~fpgm.motion.types.JointTrajectory`.

:class:`TrajectoryPlanner` walks each step of an
:class:`~fpgm.motion.types.ActionSpec` and, for a
:class:`~fpgm.motion.types.CartesianWaypoint`, drives
:class:`fpgm.robot.ik.IKSolver` along a
:func:`fpgm.motion.profile.trapezoidal_profile` built on the shared SE3
interpolation parameter ``s``. Two design choices matter enough to call out
here rather than leave implicit in the code:

**The segment's start pose is the FK of the actual previous ``q``, never the
previous segment's *commanded* target.** IK never converges to exact-zero
residual (``position_tol_m=1e-4`` by default) -- if segment N+1 interpolated
from segment N's *commanded* pose, that last few tenths of a millimetre of
residual would be silently dropped on the floor every segment, and over a
long plan those drops either accumulate into a visible discontinuity or (worse)
mask a real, growing tracking problem. Starting from ``arm.fk(q_prev, g_prev)``
instead means whatever residual segment N left behind becomes part of segment
N+1's actual path -- it gets *absorbed and corrected*, not hidden.

**IK is warm-started from the previous sample throughout**, including across
gripper/dwell steps (where ``q`` does not change, so the warm start is exact).
:mod:`fpgm.robot.ik`'s own measurements show this is the regime where damped
least squares converges in single-digit iterations at ~100% success -- a cold
restart is only needed if warm-starting itself is failing, which would be a
sign something upstream is wrong (e.g. a target outside the reachable
workspace), not something to route around by feeding IK a fresh ready-pose
guess every segment.
"""

from __future__ import annotations

from dataclasses import dataclass, field

import numpy as np

from fpgm.geometry.transforms import matrix_to_rotvec, rotvec_to_matrix
from fpgm.motion.profile import ProfileConfig, ScalarProfile, trapezoidal_profile
from fpgm.motion.types import (
    ActionSpec,
    CartesianWaypoint,
    Dwell,
    GripperAction,
    IssueKind,
    JointTrajectory,
    PlanningError,
    TrajectoryIssue,
)
from fpgm.robot.ik import IKSolver
from fpgm.robot.kinematics import ArmKinematics
from fpgm.robot.urdf import RobotModel
from fpgm.utils.logging import get_logger

logger = get_logger(__name__)

#: Labels used in segment_index bookkeeping / summarize_plan, one per ActionStep type.
_LABEL_CARTESIAN = "cartesian"
_LABEL_GRIPPER = "gripper"
_LABEL_DWELL = "dwell"


@dataclass
class PlannerConfig:
    """Tunables for :class:`TrajectoryPlanner`.

    Attributes:
        profile: Acceleration/rate limits and grid rate for
            :func:`~fpgm.motion.profile.trapezoidal_profile`.
        on_failure: ``"raise"`` (default, matching this repo's eager-validation
            posture elsewhere -- see :mod:`fpgm.robot.ik`'s own docstring)
            raises :class:`~fpgm.motion.types.PlanningError` the moment a
            step cannot be realised. ``"best_effort"`` instead records a
            :class:`~fpgm.motion.types.TrajectoryIssue`, keeps the
            closest-achieved ``q``, sets
            :attr:`~fpgm.motion.types.JointTrajectory.success` to False, and
            keeps planning the remaining steps (warm-started from that
            closest ``q`` -- still the best available seed, exactly as
            :meth:`fpgm.robot.ik.IKSolver.solve_sequence` already treats a
            failed solve's ``q``).
        max_retime_iterations: How many times a
            :class:`~fpgm.motion.types.CartesianWaypoint` segment is re-timed
            (slowed down, more grid ticks) before its peak per-joint velocity
            ratio is accepted as unfixable and reported as a
            :attr:`~fpgm.motion.types.IssueKind.SINGULARITY` instead.
        joint_limit_atol: Distance (radians) from a joint's bound within which
            :attr:`~fpgm.motion.types.JointTrajectory.joints_at_limit` reports True.
        velocity_clamp_tol: Slack on the ``ratio <= 1.0`` velocity-clamp check,
            purely to absorb float noise in the ratio computation itself --
            not a physical margin.
    """

    profile: ProfileConfig = field(default_factory=ProfileConfig)
    on_failure: str = "raise"
    max_retime_iterations: int = 3
    joint_limit_atol: float = 1e-6
    velocity_clamp_tol: float = 1e-6

    def __post_init__(self) -> None:
        if self.on_failure not in ("raise", "best_effort"):
            raise ValueError(
                f'on_failure must be "raise" or "best_effort", got {self.on_failure!r}'
            )


class TrajectoryPlanner:
    """Realises an :class:`~fpgm.motion.types.ActionSpec` as a joint trajectory.

    Needs all three of the robot stack's pieces: ``arm`` (fast FK/Jacobian),
    ``ik`` (built against that same ``arm``), and ``robot`` (only consulted
    for the gripper drive joint's velocity limit and travel range -- see
    :meth:`_plan_gripper_segment`).
    """

    def __init__(
        self,
        arm: ArmKinematics,
        ik: IKSolver,
        robot: RobotModel,
        cfg: PlannerConfig | None = None,
    ) -> None:
        self._arm = arm
        self._ik = ik
        self._robot = robot
        self._cfg = cfg if cfg is not None else PlannerConfig()

    @property
    def config(self) -> PlannerConfig:
        return self._cfg

    # ------------------------------------------------------------------ #
    # Public entry point
    # ------------------------------------------------------------------ #

    def plan(self, spec: ActionSpec, q_start: np.ndarray, g_start: float) -> JointTrajectory:
        """Plan a full joint trajectory for ``spec``, starting at ``(q_start, g_start)``.

        Args:
            spec: Ordered Cartesian/gripper/dwell steps.
            q_start: ``(7,)`` starting joint configuration.
            g_start: Starting normalised gripper value in ``[0, 1]``.

        Returns:
            A :class:`~fpgm.motion.types.JointTrajectory` beginning with a
            single ``t=0`` row for ``(q_start, g_start)``, followed by every
            15 Hz-grid sample produced by ``spec``'s steps in order.

        Raises:
            PlanningError: If ``on_failure="raise"`` (the default) and any
                step cannot be realised -- an unreachable Cartesian target, or
                a per-joint velocity demand that re-timing could not fix.
        """
        arm, cfg = self._arm, self._cfg
        q_start = np.asarray(q_start, dtype=np.float64).reshape(arm.n_joints)
        g_start = float(np.clip(g_start, 0.0, 1.0))
        dt = 1.0 / cfg.profile.rate_hz

        times: list[float] = [0.0]
        qs: list[np.ndarray] = [q_start.copy()]
        gs: list[float] = [g_start]
        tcps: list[np.ndarray] = [arm.fk(q_start, g_start)]
        pos_errs: list[float] = [0.0]
        ori_errs: list[float] = [0.0]
        min_svs: list[float] = [_min_singular_value(arm, q_start, g_start)]
        seg_idx: list[int] = [-1]
        labels: dict[int, str] = {}
        issues: list[TrajectoryIssue] = []
        any_hard_failure = False

        q_cur, g_cur = q_start.copy(), g_start
        pose_cur = tcps[0].copy()
        t_cur = 0.0

        for step_i, step in enumerate(spec):
            if isinstance(step, CartesianWaypoint):
                labels[step_i] = _LABEL_CARTESIAN
                q_seq, tcp_seq, pe, oe, msv, step_issues, hard = self._plan_cartesian_segment(
                    step_i, step, q_cur, g_cur, pose_cur
                )
                n = q_seq.shape[0]
                for k in range(n):
                    t_cur += dt
                    times.append(t_cur)
                    qs.append(q_seq[k])
                    gs.append(g_cur)
                    tcps.append(tcp_seq[k])
                    pos_errs.append(float(pe[k]))
                    ori_errs.append(float(oe[k]))
                    min_svs.append(float(msv[k]))
                    seg_idx.append(step_i)
                if n:
                    q_cur = q_seq[-1]
                pose_cur = arm.fk(q_cur, g_cur)  # real q's FK, not the commanded target (see above)

            elif isinstance(step, GripperAction):
                labels[step_i] = _LABEL_GRIPPER
                g_seq, tcp_seq, msv, g_target = self._plan_gripper_segment(step, q_cur, g_cur)
                n = g_seq.shape[0]
                for k in range(n):
                    t_cur += dt
                    times.append(t_cur)
                    qs.append(q_cur.copy())
                    gs.append(float(g_seq[k]))
                    tcps.append(tcp_seq[k])
                    pos_errs.append(0.0)
                    ori_errs.append(0.0)
                    min_svs.append(float(msv[k]))
                    seg_idx.append(step_i)
                g_cur = g_target
                pose_cur = arm.fk(q_cur, g_cur)  # gripper travel moves the TCP with q held fixed
                step_issues, hard = [], False

            elif isinstance(step, Dwell):
                labels[step_i] = _LABEL_DWELL
                n_ticks = max(1, int(np.ceil(step.seconds * cfg.profile.rate_hz - 1e-9)))
                tcp_hold = arm.fk(q_cur, g_cur)
                msv_hold = _min_singular_value(arm, q_cur, g_cur)
                for _ in range(n_ticks):
                    t_cur += dt
                    times.append(t_cur)
                    qs.append(q_cur.copy())
                    gs.append(g_cur)
                    tcps.append(tcp_hold)
                    pos_errs.append(0.0)
                    ori_errs.append(0.0)
                    min_svs.append(msv_hold)
                    seg_idx.append(step_i)
                step_issues, hard = [], False

            else:  # pragma: no cover - ActionStep is a closed union
                raise PlanningError(f"unknown ActionSpec step type: {type(step)!r}")

            issues.extend(step_issues)
            if hard:
                any_hard_failure = True
                if cfg.on_failure == "raise":
                    raise PlanningError(
                        f"trajectory planning failed at step {step_i} "
                        f"({type(step).__name__}): "
                        + "; ".join(i.message for i in step_issues)
                    )

        timestamps = np.array(times, dtype=np.float64)
        joint_positions = np.stack(qs, axis=0)
        gripper = np.array(gs, dtype=np.float64)
        tcp_poses = np.stack(tcps, axis=0)
        position_error_m = np.array(pos_errs, dtype=np.float64)
        orientation_error_rad = np.array(ori_errs, dtype=np.float64)
        min_singular_value = np.array(min_svs, dtype=np.float64)
        segment_index = np.array(seg_idx, dtype=np.int32)

        lo, hi = arm.limits[:, 0], arm.limits[:, 1]
        joints_at_limit = (np.abs(joint_positions - lo) < cfg.joint_limit_atol) | (
            np.abs(joint_positions - hi) < cfg.joint_limit_atol
        )
        for i in np.flatnonzero(joints_at_limit.any(axis=1)):
            seg = int(segment_index[i])
            if seg >= 0:
                issues.append(
                    TrajectoryIssue(
                        IssueKind.JOINT_LIMIT, seg, int(i),
                        f"joint(s) {np.flatnonzero(joints_at_limit[i]).tolist()} at limit "
                        f"on sample {i} (segment {seg})",
                    )
                )

        joint_speed_ratio = np.zeros(timestamps.shape[0], dtype=np.float64)
        if timestamps.shape[0] > 1:
            dq = np.diff(joint_positions, axis=0)
            joint_speed_ratio[1:] = np.max(np.abs(dq) / (dt * arm.vel_lim), axis=1)

        traj = JointTrajectory(
            timestamps=timestamps,
            joint_positions=joint_positions,
            gripper=gripper,
            tcp_poses=tcp_poses,
            position_error_m=position_error_m,
            orientation_error_rad=orientation_error_rad,
            min_singular_value=min_singular_value,
            joints_at_limit=joints_at_limit,
            segment_index=segment_index,
            issues=issues,
            success=not any_hard_failure,
            joint_speed_ratio=joint_speed_ratio,
            segment_labels=labels,
        )
        logger.info(
            "planned %d samples, %.3f s:\n%s", len(traj), timestamps[-1], summarize_plan(traj)
        )
        return traj

    # ------------------------------------------------------------------ #
    # Cartesian segments
    # ------------------------------------------------------------------ #

    def _solve_along_profile(
        self,
        profile: ScalarProfile,
        p0: np.ndarray,
        r0: np.ndarray,
        p1: np.ndarray,
        r1: np.ndarray,
        g: float,
        q_seed: np.ndarray,
    ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
        """IK-solve every new sample of ``profile`` along the shared-``s`` SE3 path.

        The relative rotation ``R0^T R1`` is computed once and scaled by each
        sample's ``s`` -- exact SLERP built from
        :func:`~fpgm.geometry.transforms.rotvec_to_matrix`/
        :func:`~fpgm.geometry.transforms.matrix_to_rotvec`, the two primitives
        already pinned to one rotation convention elsewhere in this repo (see
        :mod:`fpgm.motion.profile`'s module docstring).
        """
        arm, ik = self._arm, self._ik
        relative_rotvec = matrix_to_rotvec(r0.T @ r1)
        n = profile.n_ticks
        q_seq = np.empty((n, arm.n_joints), dtype=np.float64)
        tcp_seq = np.empty((n, 4, 4), dtype=np.float64)
        pos_err = np.empty(n, dtype=np.float64)
        ori_err = np.empty(n, dtype=np.float64)
        min_sv = np.empty(n, dtype=np.float64)
        ok = np.empty(n, dtype=bool)

        q_prev = q_seed
        for i in range(n):
            s = float(profile.s[i + 1])  # skip index 0 (t=0, already the segment's start pose)
            target = np.eye(4, dtype=np.float64)
            target[:3, :3] = r0 @ rotvec_to_matrix(s * relative_rotvec)
            target[:3, 3] = p0 + s * (p1 - p0)

            result = ik.solve(target, gripper=g, seed_q=q_prev)
            q_seq[i] = result.q
            tcp_seq[i] = arm.fk(result.q, g)
            pos_err[i] = result.position_error_m
            ori_err[i] = result.orientation_error_rad
            min_sv[i] = result.min_singular_value
            ok[i] = result.success
            q_prev = result.q

        return q_seq, tcp_seq, pos_err, ori_err, min_sv, ok

    def _plan_cartesian_segment(
        self,
        step_i: int,
        step: CartesianWaypoint,
        q_cur: np.ndarray,
        g_cur: float,
        pose_cur: np.ndarray,
    ) -> tuple[
        np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, list[TrajectoryIssue], bool
    ]:
        """Plan one CartesianWaypoint, re-timing the segment if it is too fast for the joints."""
        arm, cfg = self._arm, self._cfg
        dt = 1.0 / cfg.profile.rate_hz
        target = step.matrix()
        p0, r0 = pose_cur[:3, 3], pose_cur[:3, :3]
        p1, r1 = target[:3, 3], target[:3, :3]
        length_m = float(np.linalg.norm(p1 - p0))
        angle_rad = float(np.linalg.norm(matrix_to_rotvec(r0.T @ r1)))

        issues: list[TrajectoryIssue] = []
        n_ticks = 1
        q_seq = tcp_seq = pos_err = ori_err = min_sv = ok = None  # noqa: F841 - assigned in loop
        for attempt in range(cfg.max_retime_iterations + 1):
            profile = trapezoidal_profile(
                length_m, angle_rad, step.speed_mps, cfg.profile, min_ticks=n_ticks
            )
            q_seq, tcp_seq, pos_err, ori_err, min_sv, ok = self._solve_along_profile(
                profile, p0, r0, p1, r1, g_cur, q_cur
            )
            q_chain = np.vstack([q_cur[None, :], q_seq])
            dq = np.diff(q_chain, axis=0)
            ratio = float(np.max(np.abs(dq) / (dt * arm.vel_lim))) if dq.size else 0.0

            if ratio <= 1.0 + cfg.velocity_clamp_tol:
                if attempt > 0:
                    issues.append(
                        TrajectoryIssue(
                            IssueKind.SPEED_CLAMPED, step_i, -1,
                            f"segment {step_i} re-timed {attempt}x to {profile.n_ticks} ticks "
                            f"({profile.duration:.3f} s) to bring peak per-joint velocity "
                            "ratio back under 1.0",
                        )
                    )
                break
            if attempt == cfg.max_retime_iterations:
                issues.append(
                    TrajectoryIssue(
                        IssueKind.SINGULARITY, step_i, -1,
                        f"segment {step_i} still demands {ratio:.2f}x a joint's velocity "
                        f"limit after {cfg.max_retime_iterations} retime attempts -- this is "
                        "a near-singular direction (unbounded joint rate for bounded "
                        "Cartesian rate), not something slowing down can fix",
                    )
                )
                break
            n_ticks = max(n_ticks + 1, int(np.ceil(ratio * profile.n_ticks)))

        assert q_seq is not None
        if not np.all(ok):
            n_failed = int(ok.size - np.sum(ok))
            issues.append(
                TrajectoryIssue(
                    IssueKind.UNREACHABLE, step_i, -1,
                    f"{n_failed}/{ok.size} IK solve(s) on segment {step_i} failed to converge",
                )
            )

        hard = bool(not np.all(ok)) or any(i.kind is IssueKind.SINGULARITY for i in issues)
        return q_seq, tcp_seq, pos_err, ori_err, min_sv, issues, hard

    # ------------------------------------------------------------------ #
    # Gripper segments
    # ------------------------------------------------------------------ #

    def _plan_gripper_segment(
        self, step: GripperAction, q_cur: np.ndarray, g_cur: float
    ) -> tuple[np.ndarray, np.ndarray, np.ndarray, float]:
        """Ramp the gripper under the drive joint's own URDF velocity limit.

        The arm holds ``q_cur`` throughout -- but the TCP still physically
        moves (see :mod:`fpgm.robot.kinematics`'s module docstring: up to
        11.3 mm between fully open and fully closed, because the Robotiq's
        inner fingers are mimic-driven off ``finger_joint`` and the TCP is
        their midpoint). That is why the caller records
        ``position_error_m = 0.0`` for these samples rather than comparing
        the drifting TCP to some fixed target: there never was a Cartesian
        target here to have "tracking error" against, only ``fk(q, g)`` at a
        deliberately-ramping ``g``.

        The minimum duration is derived from
        ``robot.gripper_drive_velocity_limit`` and ``robot.gripper_limits``
        (the drive joint's own ``<limit>``), never a hard-coded sample count
        -- e.g. the Robotiq 2F-85's ``finger_joint`` (2.0 rad/s over a 0.725
        rad range) works out to 0.363 s / 6 samples at 15 Hz for a full
        open<->close, but that number is a *consequence* of the URDF, not
        assumed here.
        """
        arm, robot, cfg = self._arm, self._robot, self._cfg
        g_target = float(np.clip(step.resolve(arm), 0.0, 1.0))
        lower, upper = robot.gripper_limits
        delta_rad = abs(g_target - g_cur) * (upper - lower)
        vel_lim = robot.gripper_drive_velocity_limit
        min_duration = delta_rad / vel_lim if vel_lim > 0 else 0.0
        n = max(1, int(np.ceil(min_duration * cfg.profile.rate_hz - 1e-9)))

        g_seq = g_cur + (g_target - g_cur) * (np.arange(1, n + 1, dtype=np.float64) / n)
        tcp_seq = np.stack([arm.fk(q_cur, float(g)) for g in g_seq], axis=0)
        min_sv = np.array([_min_singular_value(arm, q_cur, float(g)) for g in g_seq])
        return g_seq, tcp_seq, min_sv, g_target


def _min_singular_value(arm: ArmKinematics, q: np.ndarray, g: float) -> float:
    return float(np.linalg.svd(arm.jacobian(q, g), compute_uv=False)[-1])


def summarize_plan(traj: JointTrajectory) -> str:
    """Render one row per :class:`~fpgm.motion.types.ActionSpec` step -- the
    primary debugging artefact for this module, mirroring
    :func:`fpgm.objects.interaction.summarize`. Eyeballing hundreds of ``(4,
    4)`` TCP matrices is not a debugging strategy; a handful of (type,
    samples, duration, length, peak speed, peak joint-speed ratio, IK
    residual, min sigma) rows is.

    Args:
        traj: Output of :meth:`TrajectoryPlanner.plan`.

    Returns:
        A multi-line human-readable table, one row per originating step
        (skipping the single initial ``t=0`` row, ``segment_index == -1``,
        which precedes any step and is not itself a segment).
    """
    if len(traj) <= 1:
        return "(empty trajectory)"
    dt = float(traj.timestamps[1] - traj.timestamps[0])
    positions = traj.tcp_poses[:, :3, 3]
    header = (
        f"{'segment':>7} {'type':<10} {'samples':>7} {'seconds':>8} "
        f"{'length_m':>9} {'peak_mps':>9} {'peak_qratio':>11} {'max_res_m':>10} {'min_sigma':>9}"
    )
    lines = [header]
    seg_ids = sorted({int(s) for s in traj.segment_index if s >= 0})
    for seg in seg_ids:
        idx = np.flatnonzero(traj.segment_index == seg)
        if idx.size == 0:
            continue
        span = np.concatenate([[idx[0] - 1], idx]) if idx[0] > 0 else idx
        deltas = np.linalg.norm(np.diff(positions[span], axis=0), axis=1)
        length_m = float(np.sum(deltas))
        peak_mps = float(np.max(deltas) / dt) if deltas.size else 0.0
        has_ratio = traj.joint_speed_ratio.size > 0
        peak_qratio = float(np.max(traj.joint_speed_ratio[idx])) if has_ratio else float("nan")
        max_res = float(np.max(traj.position_error_m[idx]))
        min_sigma = float(np.min(traj.min_singular_value[idx]))
        label = traj.segment_labels.get(seg, str(seg))
        lines.append(
            f"{seg:>7d} {label:<10} {idx.size:>7d} {idx.size * dt:>8.3f} "
            f"{length_m:>9.4f} {peak_mps:>9.4f} {peak_qratio:>11.3f} "
            f"{max_res:>10.2e} {min_sigma:>9.4f}"
        )
    return "\n".join(lines)