"""Python interface that mirrors the covalent shaping pipeline used in C++."""

import math
import os
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np

from control.python.base_shaper import BaseShaper
from control.python.map_loader import ModelLoader
from util.robot_dynamics import Dynamics

# TODO: Remove hard-coded values and read from CSV/config.
ROBOT_NUM_JOINTS = 6
ROBOT_NUM_SHAPED_AXES = 3
ROBOT_SIDE_LENGTH = 0.0007
ROBOT_BASE_HEIGHT = 0.364
DEFAULT_PROB_THRESH = 0.5


@dataclass
class RobotState:
    """Represent robot state required for shaper inference.

    Args:
        joint_angles: Joint angles `[num_joints]` in radians.
        tcp_position: Optional tool center point position `[x, y, z]` in meters.
    Side Effects:
        None.
    Raises:
        None.
    Preconditions:
        Arrays are finite and ordered to match model joint/axis conventions.
    """

    joint_angles: np.ndarray
    tcp_position: np.ndarray | None = None


@dataclass
class ShapedSample:
    """Hold shaped command values for a single control sample.

    Args:
        positions: Shaped joint positions `[num_joints]` in radians.
        velocities: Shaped joint velocities `[num_joints]` in radians/s.
        accelerations: Shaped joint accelerations `[num_joints]` in radians/s^2.
    Side Effects:
        None.
    Raises:
        None.
    Preconditions:
        Arrays align with the same joint ordering.
    """

    positions: np.ndarray
    velocities: np.ndarray
    accelerations: np.ndarray


@dataclass
class ShapedTrajectory:
    """Store shaped trajectory arrays and time stamps.

    Args:
        positions: Shaped positions `[N, num_joints]` in radians.
        velocities: Shaped velocities `[N, num_joints]` in radians/s.
        accelerations: Shaped accelerations `[N, num_joints]` in radians/s^2.
        time: Time vector `[N]` in seconds.
    Side Effects:
        None.
    Raises:
        None.
    Preconditions:
        All arrays share the same sample count `N`.
    """

    positions: np.ndarray
    velocities: np.ndarray
    accelerations: np.ndarray
    time: np.ndarray


def _import_torch_cpu_first() -> Any:
    """Import torch from the active environment before any system libtorch.

    Args:
        None.
    Returns:
        `module` torch module loaded from the active Python environment.
    Side Effects:
        May prepend to `LD_LIBRARY_PATH` in the current process.
    Raises:
        ImportError: If torch is not installed.
    Preconditions:
        None.
    """

    torch_lib_dir = (
        Path(sys.prefix)
        / "lib"
        / f"python{sys.version_info.major}.{sys.version_info.minor}"
        / "site-packages"
        / "torch"
        / "lib"
    )
    if torch_lib_dir.exists():
        ld_library_path = os.environ.get("LD_LIBRARY_PATH", "")
        if str(torch_lib_dir) not in ld_library_path.split(os.pathsep):
            os.environ["LD_LIBRARY_PATH"] = str(torch_lib_dir) + (
                os.pathsep + ld_library_path if ld_library_path else ""
            )

    import torch  # type: ignore

    return torch


torch = _import_torch_cpu_first()


def _gradient(samples: np.ndarray, sample_time: float) -> np.ndarray:
    """Compute first derivative of one 1D sample stream.

    Args:
        samples: Signal values `[N]`.
        sample_time: Sample period [s].
    Returns:
        `np.ndarray` derivative `[N]` in units per second.
    Side Effects:
        None.
    Raises:
        ValueError: If `sample_time <= 0`.
    Preconditions:
        `samples` is a one-dimensional numeric array.
    """

    if sample_time <= 0.0:
        raise ValueError("sample_time must be positive")

    n_samples = samples.size
    gradient = np.zeros_like(samples)

    if n_samples <= 1:
        return gradient
    if n_samples == 2:
        diff = (samples[1] - samples[0]) / sample_time
        gradient[:] = diff
        return gradient

    gradient[0] = (samples[1] - samples[0]) / sample_time
    for idx in range(1, n_samples - 1):
        gradient[idx] = (samples[idx + 1] - samples[idx - 1]) / (2.0 * sample_time)
    gradient[-1] = (samples[-1] - samples[-2]) / sample_time
    return gradient


class ShaperInterface:
    """Run map-based modal inference and apply per-axis input shaping.

    Args:
        sample_time: Controller sample period [s].
        model_directory: Directory containing per-axis map models.
        python_src_root: Root source path injected into `sys.path` for runtime imports.
        urdf_filepath: URDF file used by dynamics calculations.
        num_axes: Number of axes that are shaped.
        side_length: Kinematic side-length parameter [m].
        base_height: Kinematic base-height parameter [m].
        num_joints: Number of robot joints.
        prob_thresh: Probability threshold for selecting two-mode shaping.
    Side Effects:
        Loads model files and URDF dynamics model into memory.
    Raises:
        ValueError: If scalar constructor inputs are out of range.
        FileNotFoundError: If model files or URDF are missing.
    Preconditions:
        `model_directory` contains model files matching configured axes.
    """

    def __init__(
        self,
        sample_time: float,
        model_directory: str,
        python_src_root: str,
        urdf_filepath: str,
        num_axes: int = ROBOT_NUM_SHAPED_AXES,
        side_length: float = ROBOT_SIDE_LENGTH,
        base_height: float = ROBOT_BASE_HEIGHT,
        num_joints: int = ROBOT_NUM_JOINTS,
        prob_thresh: float = DEFAULT_PROB_THRESH,
    ) -> None:
        """Initialize model inference and per-axis shapers.

        Args:
            sample_time: Controller sample period [s].
            model_directory: Directory containing `axis*_model.pt` files.
            python_src_root: Source root path to add to import path.
            urdf_filepath: URDF path for robot dynamics.
            num_axes: Number of shaped axes.
            side_length: Side-length geometry parameter [m].
            base_height: Base-height geometry parameter [m].
            num_joints: Number of robot joints.
            prob_thresh: Threshold for two-mode shaper selection.
        Returns:
            `None`.
        Side Effects:
            Mutates `sys.path` and loads model/dynamics resources.
        Raises:
            ValueError: If dimensions or sample time are invalid.
        Preconditions:
            `urdf_filepath` and model files exist and are readable.
        """

        if sample_time <= 0.0:
            raise ValueError("Sample time must be positive")
        if num_joints <= 0:
            raise ValueError("Number of joints must be positive")
        if num_axes <= 0 or num_axes > num_joints:
            raise ValueError("Number of shaped axes must be in (0, num_joints]")

        if python_src_root not in sys.path:
            sys.path.insert(0, python_src_root)

        self.Ts = sample_time
        self.num_axes = num_axes
        self.num_joints = num_joints
        self.side_length = side_length
        self.base_height = base_height
        self.prob_thresh = prob_thresh
        self.model_directory = model_directory
        self.store_buffer: list[list[int]] = [[] for _ in range(num_axes)]

        self._torch = torch
        self.robot_dynamics = Dynamics(urdf_file=urdf_filepath, initialize=True)
        self.map_fitter = ModelLoader.load(model_directory, num_axes, 3, [64, 64])

        self.shapers = [BaseShaper(self.Ts) for _ in range(num_axes)]
        self.last_input_position: np.ndarray | None = None
        self.last_input_velocity: np.ndarray | None = None
        self.last_input_acceleration: np.ndarray | None = None
        self._logged_once: list[bool] = [False] * max(1, self.num_axes)
        self._last_vibration_control_enabled: bool | None = None

    def reset(self) -> None:
        """Reset internal shaper states and stream continuity markers.

        Args:
            None.
        Returns:
            `None`.
        Side Effects:
            Replaces per-axis shaper objects and clears stream cache.
        Raises:
            None.
        Preconditions:
            None.
        """

        self.shapers = [BaseShaper(self.Ts) for _ in range(self.num_axes)]
        self.last_input_position = None
        self.last_input_velocity = None
        self.last_input_acceleration = None

    def _handle_vibration_mode_transition(self, enable_vibration_control: bool) -> bool:
        """Reset stale shaper history on `disabled -> enabled` transitions.

        Args:
            enable_vibration_control: Requested mode for the next processing call.
        Returns:
            `bool` true if this call triggered an internal reset.
        Side Effects:
            May reset shaper state and updates mode-tracking metadata.
        Raises:
            None.
        Preconditions:
            None.
        """

        toggled_disabled_to_enabled = (
            enable_vibration_control and self._last_vibration_control_enabled is False
        )
        if toggled_disabled_to_enabled:
            # While pass-through mode is active we do not advance shaper buffers.
            # Reset before re-enabling to avoid shaping against stale history.
            self.reset()
        self._last_vibration_control_enabled = enable_vibration_control
        return toggled_disabled_to_enabled

    def compute_forward_kinematics(self, joint_angles: np.ndarray) -> np.ndarray:
        """Compute TCP position from joint angles.

        Args:
            joint_angles: Joint angles `[num_joints]` in radians.
        Returns:
            `np.ndarray` TCP position `[x, y, z]` in meters.
        Side Effects:
            None.
        Raises:
            RuntimeError: If dynamics backend forward kinematics fails.
        Preconditions:
            `joint_angles` matches model joint count.
        """

        position, _ = self.robot_dynamics.get_forward_kinematics(joint_angles)
        return np.asarray(position)

    def compute_inertia(self, joint_angles: np.ndarray) -> np.ndarray:
        """Compute diagonal joint inertia terms from mass matrix.

        Args:
            joint_angles: Joint angles `[num_joints]` in radians.
        Returns:
            `np.ndarray` diagonal inertia terms `[num_joints]`.
        Side Effects:
            None.
        Raises:
            RuntimeError: If dynamics backend mass matrix calculation fails.
        Preconditions:
            `joint_angles` matches model joint count.
        """

        mass_matrix = self.robot_dynamics.get_mass_matrix(joint_angles)
        return np.asarray(np.diag(mass_matrix))

    def cartesian_to_polar(self, xyz: np.ndarray) -> tuple[float, float]:
        """Convert Cartesian TCP position into `(v, r)` map features.

        Args:
            xyz: Cartesian TCP coordinates `[x, y, z]` in meters.
        Returns:
            `tuple[float, float]` as `(v, r)` where `v` is angle [rad] and `r` is radius [m].
        Side Effects:
            None.
        Raises:
            ValueError: If `xyz` is not a 3-vector.
        Preconditions:
            `xyz` is finite.
        """

        xyz = np.asarray(xyz, dtype=float)
        if xyz.shape != (3,):
            raise ValueError("xyz must have shape (3,)")

        abs_xyz = np.abs(xyz)
        reach_axis = int(np.argmax(abs_xyz))
        height_axis = 2
        height = xyz[height_axis]

        max_reach = float(abs_xyz[reach_axis])
        tol = 1e-9
        depth_axis = -1
        for axis in range(3):
            matches_height = abs(abs_xyz[axis] - abs(height)) < tol
            matches_reach = abs(abs_xyz[axis] - max_reach) < tol
            if not matches_height and not matches_reach:
                depth_axis = axis
                break

        if depth_axis == -1:
            for axis in range(3):
                if axis != height_axis:
                    depth_axis = axis
                    break

        long_axis_value = float(xyz[reach_axis])
        width_value = float(xyz[depth_axis])
        height_value = float(xyz[height_axis])

        dx = math.sqrt(
            long_axis_value * long_axis_value + (self.side_length - width_value) ** 2
        )
        dz = height_value - self.base_height
        radius = math.hypot(dx, dz)
        angle = math.atan2(dz, dx)
        return angle, radius

    def _finite_difference_axes(self, samples: np.ndarray) -> np.ndarray:
        """Compute per-axis finite differences for an `(N, D)` sample matrix.

        Args:
            samples: Input samples shaped `(N, D)`.
        Returns:
            `np.ndarray` finite differences shaped `(N, D)`.
        Side Effects:
            None.
        Raises:
            ValueError: If `samples` is not 2D.
        Preconditions:
            `samples` is numeric and finite.
        """

        if samples.ndim != 2:
            raise ValueError("samples must be a 2D array")

        _, num_cols = samples.shape
        diff = np.zeros_like(samples)
        for col in range(num_cols):
            diff[:, col] = _gradient(samples[:, col], self.Ts)
        return diff

    def finite_difference_first(self, samples: np.ndarray) -> np.ndarray:
        """Compute first derivative for each trajectory axis.

        Args:
            samples: Position samples `[N, D]`.
        Returns:
            `np.ndarray` velocity samples `[N, D]`.
        Side Effects:
            None.
        Raises:
            ValueError: If `samples` is not 2D.
        Preconditions:
            `self.Ts > 0`.
        """

        return self._finite_difference_axes(samples)

    def finite_difference_second(self, first: np.ndarray) -> np.ndarray:
        """Compute second derivative for each trajectory axis.

        Args:
            first: Velocity samples `[N, D]`.
        Returns:
            `np.ndarray` acceleration samples `[N, D]`.
        Side Effects:
            None.
        Raises:
            ValueError: If `first` is not 2D.
        Preconditions:
            `self.Ts > 0`.
        """

        return self._finite_difference_axes(first)

    def _rotate_xyz_to_zero_base(
        self, joint_angles: np.ndarray, xyz: np.ndarray
    ) -> tuple[np.ndarray, bool]:
        """Rotate TCP position into a base-zero frame.

        Args:
            joint_angles: Joint angles `[num_joints]` in radians.
            xyz: TCP coordinates `[x, y, z]` in meters.
        Returns:
            `tuple[np.ndarray, bool]` with rotated vector and fallback-used flag.
        Side Effects:
            None.
        Raises:
            None.
        Preconditions:
            `xyz` is shape `(3,)`.
        """

        if joint_angles.size == 0:
            return xyz, False

        try:
            current_rot = self.robot_dynamics.get_individual_rotation_matrix(
                0, joint_angles
            )
            joint_angles_zero = np.asarray(joint_angles, dtype=float).copy()
            joint_angles_zero[0] = 0.0
            zero_rot = self.robot_dynamics.get_individual_rotation_matrix(
                0, joint_angles_zero
            )
            base_rotation = current_rot @ zero_rot.T
            rotated = base_rotation.T @ xyz
            return rotated, False
        except Exception:
            # Developer: Codex (2026-02-10)
            # Keep a deterministic fallback so calibration can proceed when detailed
            # per-link rotation matrices are unavailable from the dynamics backend.
            base_angle = float(joint_angles[0])
            c = math.cos(-base_angle)
            s = math.sin(-base_angle)
            rot = np.array([[c, -s, 0.0], [s, c, 0.0], [0.0, 0.0, 1.0]], dtype=float)
            return rot @ xyz, True

    def compute_nn_inputs(self, state: RobotState) -> tuple[float, float, np.ndarray]:
        """Build map-network inputs from current robot state.

        Args:
            state: Current robot state.
        Returns:
            `tuple[float, float, np.ndarray]` as `(v_rad, r_m, inertia_vector)`.
        Side Effects:
            Prints one warning line when fallback base-rotation logic is used.
        Raises:
            RuntimeError: If dynamics calls fail and no fallback path can recover.
        Preconditions:
            `state.joint_angles` matches configured `num_joints`.
        """

        inertia = self.compute_inertia(state.joint_angles)
        xyz = (
            state.tcp_position
            if state.tcp_position is not None
            else self.compute_forward_kinematics(state.joint_angles)
        )

        rotated_xyz, used_fallback = self._rotate_xyz_to_zero_base(
            state.joint_angles,
            np.asarray(xyz, dtype=float),
        )
        if used_fallback:
            print(
                "Warning: Robot dynamics model returned no joint rotations; "
                "falling back to Z-axis base rotation assumption."
            )

        v_rad, r_m = self.cartesian_to_polar(rotated_xyz)
        return v_rad, r_m, inertia

    def process_sample(
        self,
        command: np.ndarray,
        state: RobotState,
        command_dot: np.ndarray | None = None,
        command_ddot: np.ndarray | None = None,
        enable_vibration_control: bool = True,
    ) -> ShapedSample:
        """Process one command sample with optional vibration shaping.

        Args:
            command: Joint command `[num_joints]` in radians.
            state: Robot state for feature calculation.
            command_dot: Optional command velocity `[num_joints]` in radians/s.
            command_ddot: Optional command acceleration `[num_joints]` in radians/s^2.
            enable_vibration_control: Whether to apply shaping.
        Returns:
            `ShapedSample` containing position/velocity/acceleration vectors.
        Side Effects:
            Updates per-axis shaper state and stream cache.
        Raises:
            ValueError: If command vector shapes are invalid.
        Preconditions:
            Model inference backend is initialized.
        """
        self._handle_vibration_mode_transition(enable_vibration_control)
        dot_from_fd = command_dot is None
        ddot_from_fd = command_ddot is None

        command = np.asarray(command, dtype=float)
        if command.shape != (self.num_joints,):
            raise ValueError("Command length mismatch")

        if command_dot is None:
            command_dot = np.zeros(self.num_joints, dtype=float)
        else:
            command_dot = np.asarray(command_dot, dtype=float)
            if command_dot.shape != (self.num_joints,):
                raise ValueError("command_dot length mismatch")

        if command_ddot is None:
            command_ddot = np.zeros(self.num_joints, dtype=float)
        else:
            command_ddot = np.asarray(command_ddot, dtype=float)
            if command_ddot.shape != (self.num_joints,):
                raise ValueError("command_ddot length mismatch")

        if not enable_vibration_control:
            out_pos = command.copy()
            if self.last_input_position is not None and dot_from_fd:
                out_vel = (out_pos - self.last_input_position) / self.Ts
            else:
                out_vel = command_dot.copy()

            if self.last_input_velocity is not None and ddot_from_fd:
                out_acc = (out_vel - self.last_input_velocity) / self.Ts
            else:
                out_acc = command_ddot.copy()

            self.last_input_position = out_pos
            self.last_input_velocity = out_vel
            self.last_input_acceleration = out_acc
            return ShapedSample(out_pos, out_vel, out_acc)

        state_copy = RobotState(
            joint_angles=np.asarray(state.joint_angles, dtype=float),
            tcp_position=(
                None
                if state.tcp_position is None
                else np.asarray(state.tcp_position, dtype=float)
            ),
        )
        if state_copy.joint_angles.shape != (self.num_joints,):
            state_copy.joint_angles = command
        if state_copy.tcp_position is None:
            state_copy.tcp_position = self.compute_forward_kinematics(
                state_copy.joint_angles
            )

        v_rad, r_m, inertia = self.compute_nn_inputs(state_copy)

        out_pos = command.copy()
        out_vel = command_dot.copy()
        out_acc = command_ddot.copy()

        with torch.no_grad():
            features = torch.zeros((self.num_axes, 3), dtype=torch.float32)
            v_deg = float(v_rad * 180.0 / math.pi)
            r_mm = float(r_m * 1000.0)
            for axis in range(self.num_axes):
                features[axis, 0] = v_deg
                features[axis, 1] = r_mm
                features[axis, 2] = float(inertia[axis])

            order_probs, mode_params = self.map_fitter.infer_batch(features)
            order_probs = order_probs.squeeze(-1).double()
            mode_params = mode_params.double()

        for axis in range(self.num_axes):
            prob_second = float(order_probs[axis].item())
            wn1 = float(mode_params[axis, 0].item())
            z1_log = float(mode_params[axis, 1].item())
            wn2 = float(mode_params[axis, 2].item())
            z2_log = float(mode_params[axis, 3].item())

            z1 = math.exp(z1_log)
            z2 = math.exp(z2_log)
            two_mode = prob_second > self.prob_thresh

            if not self._logged_once[axis]:
                print(
                    f"[ShaperInterfacePy] Axis {axis} inference: wn1={wn1} z1={z1} prob_second={prob_second}"
                )
                if two_mode:
                    print(
                        f"[ShaperInterfacePy] Axis {axis} second mode wn2={wn2} z2={z2}"
                    )
                self._logged_once[axis] = True

            params = (
                np.array([[wn1, z1], [wn2, z2]], dtype=float)
                if two_mode
                else np.array([[wn1, z1]], dtype=float)
            )

            try:
                y, v, a = self.shapers[axis].shape_sample(float(command[axis]), params)
                self.store_buffer[axis].append(len(self.shapers[axis]._x_buf))
                out_pos[axis] = y
                out_vel[axis] = v
                out_acc[axis] = a
            except ValueError as exc:
                print(f"[ShaperInterfacePy] axis {axis} shaping skipped: {exc}")

        self.last_input_position = out_pos
        self.last_input_velocity = out_vel
        self.last_input_acceleration = out_acc
        return ShapedSample(out_pos, out_vel, out_acc)

    def shape_sample(
        self,
        command: np.ndarray,
        state: RobotState,
        command_dot: np.ndarray | None = None,
        command_ddot: np.ndarray | None = None,
    ) -> ShapedSample:
        """Backward-compatible alias for `process_sample(..., enable_vibration_control=True)`.

        Args:
            command: Joint command `[num_joints]` in radians.
            state: Robot state for feature calculation.
            command_dot: Optional command velocity `[num_joints]` in radians/s.
            command_ddot: Optional command acceleration `[num_joints]` in radians/s^2.
        Returns:
            `ShapedSample` with vibration control enabled.
        Side Effects:
            Same as `process_sample`.
        Raises:
            ValueError: If command vector shapes are invalid.
        Preconditions:
            Model inference backend is initialized.
        """

        return self.process_sample(
            command=command,
            state=state,
            command_dot=command_dot,
            command_ddot=command_ddot,
            enable_vibration_control=True,
        )

    def process_trajectory(
        self,
        command: np.ndarray,
        states: list[RobotState],
        command_dot: np.ndarray | None = None,
        command_ddot: np.ndarray | None = None,
        time_vector: list[float] | None = None,
        enable_vibration_control: bool = True,
    ) -> ShapedTrajectory:
        """Process a full trajectory with optional vibration shaping.

        Args:
            command: Command trajectory `[N, num_joints]` in radians.
            states: State list length `N` or empty list for command-only shaping.
            command_dot: Optional velocity trajectory `[N, num_joints]` in radians/s.
            command_ddot: Optional acceleration trajectory `[N, num_joints]` in radians/s^2.
            time_vector: Optional sample times `[N]` in seconds.
            enable_vibration_control: Whether to apply shaping.
        Returns:
            `ShapedTrajectory` containing processed trajectory fields.
        Side Effects:
            Mutates shaper internal history and stream continuity cache.
        Raises:
            ValueError: If command/state/derivative shapes are inconsistent.
        Preconditions:
            Dynamics backend and map models are initialized.
        """
        reset_for_mode_change = self._handle_vibration_mode_transition(
            enable_vibration_control
        )

        command = np.asarray(command, dtype=float)
        if command.ndim != 2 or command.shape[1] != self.num_joints:
            raise ValueError("command must be (N, num_joints)")

        samples = command.shape[0]
        if states and len(states) != samples:
            raise ValueError("states length must match samples")

        if command_dot is None:
            command_dot = self.finite_difference_first(command)
        else:
            command_dot = np.asarray(command_dot, dtype=float)
            if command_dot.shape != command.shape:
                raise ValueError("command_dot shape mismatch")

        if command_ddot is None:
            command_ddot = self.finite_difference_second(command_dot)
        else:
            command_ddot = np.asarray(command_ddot, dtype=float)
            if command_ddot.shape != command.shape:
                raise ValueError("command_ddot shape mismatch")

        times = (
            list(time_vector)
            if time_vector is not None
            else [float(i) * self.Ts for i in range(samples)]
        )
        if len(times) != samples:
            raise ValueError("time_vector length mismatch")

        first_position = command[0]
        first_velocity = command_dot[0]
        first_acceleration = command_ddot[0]

        if not enable_vibration_control:
            self.last_input_position = command[-1].copy()
            self.last_input_velocity = command_dot[-1].copy()
            self.last_input_acceleration = command_ddot[-1].copy()
            return ShapedTrajectory(
                command.copy(),
                command_dot.copy(),
                command_ddot.copy(),
                np.asarray(times, dtype=float),
            )

        continue_stream = False
        if (
            self.last_input_position is not None
            and self.last_input_velocity is not None
            and self.last_input_acceleration is not None
        ):
            position_match = np.allclose(
                self.last_input_position, first_position, rtol=1e-8, atol=1e-10
            )
            velocity_match = np.allclose(
                self.last_input_velocity, first_velocity, rtol=1e-8, atol=1e-10
            )
            acceleration_match = np.allclose(
                self.last_input_acceleration,
                first_acceleration,
                rtol=1e-8,
                atol=1e-10,
            )
            continue_stream = position_match and velocity_match and acceleration_match

        if not continue_stream and not reset_for_mode_change:
            self.reset()

        traj_pos = np.zeros((samples, self.num_joints), dtype=float)
        traj_vel = np.zeros_like(traj_pos)
        traj_acc = np.zeros_like(traj_pos)
        traj_time = np.asarray(times, dtype=float)

        for sample_index in range(samples):
            state_i = (
                states[sample_index]
                if states
                else RobotState(joint_angles=command[sample_index])
            )
            if state_i.tcp_position is None:
                joint_angles = np.asarray(state_i.joint_angles, dtype=float)
                state_i = RobotState(
                    joint_angles=joint_angles,
                    tcp_position=self.compute_forward_kinematics(joint_angles),
                )

            shaped = self.process_sample(
                command[sample_index],
                state_i,
                command_dot[sample_index],
                command_ddot[sample_index],
                enable_vibration_control=True,
            )
            traj_pos[sample_index, :] = shaped.positions
            traj_vel[sample_index, :] = shaped.velocities
            traj_acc[sample_index, :] = shaped.accelerations

        tails: list[list[tuple[float, float, float]]] = []
        max_tail = 0
        for axis in range(self.num_axes):
            tail = self.shapers[axis].finalize()
            tails.append(tail)
            max_tail = max(max_tail, len(tail))

        if max_tail > 0:
            extra = max_tail
            last_pos = traj_pos[-1].copy()
            traj_pos = np.vstack([traj_pos, np.zeros((extra, self.num_joints))])
            traj_vel = np.vstack([traj_vel, np.zeros((extra, self.num_joints))])
            traj_acc = np.vstack([traj_acc, np.zeros((extra, self.num_joints))])
            last_time = (
                traj_time[-1] if traj_time.size else float(samples - 1) * self.Ts
            )
            traj_time = np.concatenate(
                [
                    traj_time,
                    np.array(
                        [last_time + (step + 1) * self.Ts for step in range(extra)],
                        dtype=float,
                    ),
                ]
            )

            for tail_index in range(max_tail):
                pos = last_pos.copy()
                vel = np.zeros(self.num_joints, dtype=float)
                acc = np.zeros(self.num_joints, dtype=float)
                for axis in range(self.num_axes):
                    if tail_index < len(tails[axis]):
                        y, v, a = tails[axis][tail_index]
                        pos[axis] = y
                        vel[axis] = v
                        acc[axis] = a

                row = samples + tail_index
                traj_pos[row, :] = pos
                traj_vel[row, :] = vel
                traj_acc[row, :] = acc

        self.last_input_position = command[-1]
        self.last_input_velocity = command_dot[-1]
        self.last_input_acceleration = command_ddot[-1]

        return ShapedTrajectory(traj_pos, traj_vel, traj_acc, traj_time)

    def shape_trajectory(
        self,
        command: np.ndarray,
        states: list[RobotState],
        command_dot: np.ndarray | None = None,
        command_ddot: np.ndarray | None = None,
        time_vector: list[float] | None = None,
    ) -> ShapedTrajectory:
        """Backward-compatible alias for `process_trajectory(..., enable_vibration_control=True)`.

        Args:
            command: Command trajectory `[N, num_joints]` in radians.
            states: State list length `N` or empty list for command-only shaping.
            command_dot: Optional velocity trajectory `[N, num_joints]` in radians/s.
            command_ddot: Optional acceleration trajectory `[N, num_joints]` in radians/s^2.
            time_vector: Optional sample times `[N]` in seconds.
        Returns:
            `ShapedTrajectory` with vibration control enabled.
        Side Effects:
            Same as `process_trajectory`.
        Raises:
            ValueError: If command/state/derivative shapes are inconsistent.
        Preconditions:
            Dynamics backend and map models are initialized.
        """

        return self.process_trajectory(
            command=command,
            states=states,
            command_dot=command_dot,
            command_ddot=command_ddot,
            time_vector=time_vector,
            enable_vibration_control=True,
        )
