# src/util/robot_dynamics.py
# Lightweight wrapper around the URDF-native C++ dynamics implementation.

from __future__ import annotations

import os
import numpy as np
from typing import Any, Callable, List, Sequence, cast

from reforge_core.util.urdf_to_dh.generate_dh import GenerateDhParams
import reforge_core.util.urdf_to_dh.kinematics_helpers as kh
from reforge_core.util.utility import rotation_matrix_to_quaternion

# Uses
try:
    from reforge_core.util import _robot_dynamics_native as _rd_native  # type: ignore[attr-defined]
except Exception as exc:  # pragma: no cover - extension must be built
    _rd_native = None
    _rd_native_import_error: Exception | None = exc
else:
    _rd_native_import_error = None


class Dynamics:
    """Provide kinematics and dynamics utilities derived from a URDF model.

    This class wraps the URDF-native dynamics backend and exposes kinematics
    (e.g., forward kinematics, Jacobians) and dynamics (e.g., mass/inertia
    matrices) needed by calibration and analysis tooling.

    Args:
        urdf_file: Path to the Unified Robot Description Format file.
        initialize: Retained for backward compatibility; no-op.
        initial_angles: Optional initial joint angles used for validation.
        robot_object: Optional robot handle for downstream integrations.

    Side Effects:
        Parses the URDF, builds Denavit-Hartenberg parameters, and constructs
        the native dynamics backend.

    Raises:
        ImportError: If the native dynamics extension is unavailable.
        FileNotFoundError: If the URDF file path does not exist.
        RuntimeError: If the native backend cannot be constructed.
        ValueError: If provided initial joint angles do not match the URDF.

    Preconditions:
        The native dynamics extension is built and importable, and the URDF is
        valid for parameter extraction.
    """

    def __init__(
        self,
        urdf_file: str,
        initialize: bool = False,  # retained for backward compatibility; no-op
        initial_angles: Sequence[float] | None = None,
        robot_object: object | None = None,
        tcp_payload: float | None = None,
        tcp_payload_com: Sequence[float] | None = None,
    ) -> None:
        """Initialize dynamics from a URDF file.

        Args:
            urdf_file: Path to the Unified Robot Description Format file.
            initialize: Retained for backward compatibility; no-op.
            initial_angles: Optional initial joint angles for validation.
            robot_object: Optional robot handle for downstream integrations.
            tcp_payload: Optional payload of the robot for NN prediction of
                payload changes.
            tcp_payload_com: Optional 3x1 center of mass location of the
                payload, defined relative to the origin of the TCP [meters].

        Side Effects:
            Parses the URDF and constructs the native dynamics backend.

        Raises:
            ImportError: If the native dynamics extension is unavailable.
            FileNotFoundError: If the URDF file path does not exist.
            RuntimeError: If the native backend cannot be constructed.
            ValueError: If provided initial joint angles do not match the URDF.

        Preconditions:
            The URDF file exists and the native extension is built.
        """
        if _rd_native is None:
            raise ImportError(
                "reforge_core.util._robot_dynamics_native failed to load; ensure the extension is built and importable"
            ) from _rd_native_import_error

        # Initialize URDF location
        self.urdf_file = urdf_file
        self.robot_object = robot_object
        self.tcp_payload = 0.0 if tcp_payload is None else float(tcp_payload)
        self.tcp_payload_com = self._validate_tcp_payload_com(tcp_payload_com)
        if self.tcp_payload < 0.0:
            raise ValueError("tcp_payload must be greater than or equal to 0")

        # Generate DH parameters from the URDF file
        if not os.path.exists(self.urdf_file):
            raise FileNotFoundError(f"URDF file not found at {self.urdf_file}")

        self.params_node = GenerateDhParams(self.urdf_file)
        self.base_prefix_tf = np.eye(4)
        self.tool_offset_tf = np.eye(4)
        self._urdf_joint_frames: list[dict] = []
        self._native_model: _rd_native.RobotDynamics | None = None

        self.compute_params()
        # Use only moving joints (revolute/prismatic). Fallback to all joints if not provided.
        self.num_joints = len(
            getattr(self.params_node, "moving_joints", self.params_node.urdf_joints)
        )
        self.base_prefix_tf = getattr(self.params_node, "base_prefix_tf", np.eye(4))
        self.tool_offset_tf = getattr(self.params_node, "tool_offset_tf", np.eye(4))
        for name in getattr(self.params_node, "moving_joints", []):
            jd = self.params_node.urdf_joints[name]
            self._urdf_joint_frames.append(
                {
                    "xyz": np.asarray(jd["xyz"], dtype=float),
                    "rpy": np.asarray(jd["rpy"], dtype=float),
                    "axis": np.asarray(jd["axis"], dtype=float),
                }
            )

        if self._urdf_joint_frames:
            origins = [f["xyz"] for f in self._urdf_joint_frames]
            rpy = [f["rpy"] for f in self._urdf_joint_frames]
            axes = [f["axis"] for f in self._urdf_joint_frames]
            masses = [
                self.params_node.inertia_dict["mass"][i] for i in range(self.num_joints)
            ]
            coms = [
                np.asarray(self.params_node.inertia_dict["com"][i], dtype=float)
                for i in range(self.num_joints)
            ]
            inertias = [
                np.asarray(
                    self.params_node.inertia_dict["inertia_tensor"][i], dtype=float
                )
                for i in range(self.num_joints)
            ]
            self._apply_tcp_payload_to_terminal_link(masses, coms, inertias)
            self._native_model = _rd_native.RobotDynamics(
                origins,
                rpy,
                axes,
                masses,
                coms,
                inertias,
                self.base_prefix_tf,
                self.tool_offset_tf,
            )

        if self._native_model is None:
            raise RuntimeError("Failed to construct Dynamics from URDF data.")

        if initialize:
            initial_angles = (
                list(initial_angles)
                if initial_angles is not None
                else [0.0] * self.num_joints
            )
            if len(initial_angles) != self.num_joints:
                raise ValueError(
                    "Number of initial joint angles does not match URDF specification: "
                    f"{len(initial_angles)} vs {self.num_joints}"
                )
            # _robot_dynamics_ext does not require an explicit initialize step.

    def compute_params(self) -> None:
        """Parse the URDF and compute kinematic and inertial parameters.

        Returns:
            `None`.

        Side Effects:
            Updates internal parameter structures derived from the URDF.

        Raises:
            FileNotFoundError: If the URDF file is missing.
            RuntimeError: If URDF parsing fails.

        Preconditions:
            `self.urdf_file` points to a readable URDF file.
        """
        self.params_node.parse_urdf()
        self.params_node.calculate_tfs_in_world_frame()
        self.params_node.calculate_params()

    def _validate_tcp_payload_com(
        self, tcp_payload_com: Sequence[float] | None
    ) -> np.ndarray:
        """Return payload COM as a validated 3-vector in TCP coordinates."""
        if tcp_payload_com is None:
            return np.zeros(3, dtype=float)
        com = np.asarray(tcp_payload_com, dtype=float)
        if com.shape != (3,):
            raise ValueError(
                f"tcp_payload_com must be a 3-vector, received shape {com.shape}"
            )
        return com

    def _apply_tcp_payload_to_terminal_link(
        self,
        masses: list[float],
        coms: list[np.ndarray],
        inertias: list[np.ndarray],
    ) -> None:
        """Merge TCP payload as a point mass into the terminal moving link."""
        if self.tcp_payload <= 0.0:
            return
        if not masses or not coms or not inertias:
            raise RuntimeError("Cannot apply payload: inertial arrays are empty")

        terminal_idx = len(masses) - 1
        mass_link = float(masses[terminal_idx])
        com_link = np.asarray(coms[terminal_idx], dtype=float)
        inertia_link = np.asarray(inertias[terminal_idx], dtype=float)

        tool_tf = np.asarray(self.tool_offset_tf, dtype=float)
        payload_com_terminal = tool_tf[:3, 3] + tool_tf[:3, :3] @ self.tcp_payload_com

        mass_payload = self.tcp_payload
        mass_total = mass_link + mass_payload
        com_total = (
            mass_link * com_link + mass_payload * payload_com_terminal
        ) / mass_total

        delta_link = com_link - com_total
        delta_payload = payload_com_terminal - com_total
        identity3 = np.eye(3, dtype=float)
        inertia_total = (
            inertia_link
            + mass_link
            * (
                np.dot(delta_link, delta_link) * identity3
                - np.outer(delta_link, delta_link)
            )
            + mass_payload
            * (
                np.dot(delta_payload, delta_payload) * identity3
                - np.outer(delta_payload, delta_payload)
            )
        )

        masses[terminal_idx] = mass_total
        coms[terminal_idx] = com_total
        inertias[terminal_idx] = inertia_total
        print(
            "Applied TCP payload to terminal link: "
            f"mass={mass_payload:.6f}, "
            f"tcp_com={self.tcp_payload_com.tolist()}, "
            f"terminal_com={payload_com_terminal.tolist()}"
        )

    def initialize_model(self) -> None:
        """Retained for backward compatibility; no-op initializer.

        Returns:
            `None`.

        Side Effects:
            None.

        Raises:
            RuntimeError: If the native backend is unavailable.

        Preconditions:
            The native dynamics backend is constructed.
        """
        self._ensure_ext()

    # ------------------------------------------------------------------
    # Helper accessors
    # ------------------------------------------------------------------
    def _ensure_ext(self) -> None:
        """Validate that the native dynamics backend is available.

        Returns:
            `None`.

        Side Effects:
            None.

        Raises:
            RuntimeError: If the native backend is not constructed.

        Preconditions:
            None.
        """
        if self._native_model is None:
            raise RuntimeError(
                "Robot dynamics backend is unavailable; check the URDF inputs or extension build."
            )

    def _as_eigen_vector(
        self, joint_angles: Sequence[float] | np.ndarray
    ) -> np.ndarray:
        """Validate joint angles and return a NumPy vector.

        Ensures the input is numeric and matches `self.num_joints`.

        Args:
            joint_angles: Sequence of joint angles in radians.

        Returns:
            `np.ndarray` of shape `(num_joints,)`.

        Side Effects:
            None.

        Raises:
            ValueError: If the input shape does not match `num_joints`.

        Preconditions:
            Input can be converted to a numeric NumPy array and has length
            `self.num_joints`.
        """
        arr = np.asarray(joint_angles, dtype=float)
        if arr.shape != (self.num_joints,):
            raise ValueError(
                f"Expected {self.num_joints} joint angles, received shape {arr.shape}"
            )
        return arr

    def _native(self) -> Any:
        """Return the initialized native backend model."""
        self._ensure_ext()
        return cast(Any, self._native_model)

    # ------------------------------------------------------------------
    # Public API: mass/inertia properties
    # ------------------------------------------------------------------
    def get_mass_matrix(self, joint_angles: Sequence[float] | np.ndarray) -> np.ndarray:
        """Return the joint-space mass matrix for the given joint angles.

        Args:
            joint_angles: Joint angles in radians.

        Returns:
            `np.ndarray` mass matrix.

        Side Effects:
            None.

        Raises:
            RuntimeError: If the native backend is unavailable.
            ValueError: If `joint_angles` has the wrong shape.

        Preconditions:
            `joint_angles` length matches `self.num_joints`.
        """
        native = self._native()
        return np.asarray(native.mass_matrix(self._as_eigen_vector(joint_angles)))

    def get_jacobian_matrix(
        self, joint_angles: Sequence[float] | np.ndarray
    ) -> np.ndarray:
        """Return the geometric Jacobian for the given joint angles.

        Args:
            joint_angles: Joint angles in radians.

        Returns:
            `np.ndarray` Jacobian matrix.

        Side Effects:
            None.

        Raises:
            RuntimeError: If the native backend is unavailable.
            ValueError: If `joint_angles` has the wrong shape.

        Preconditions:
            `joint_angles` length matches `self.num_joints`.
        """
        native = self._native()
        return np.asarray(native.jacobian_matrix(self._as_eigen_vector(joint_angles)))

    # ------------------------------------------------------------------
    # Public API: kinematics
    # ------------------------------------------------------------------
    def _axis_angle_rot(self, axis: np.ndarray, theta: float) -> np.ndarray:
        """Compute a rotation matrix from an axis-angle representation.

        Args:
            axis: Rotation axis vector.
            theta: Rotation angle in radians.

        Returns:
            `np.ndarray` 3x3 rotation matrix.

        Side Effects:
            None.

        Raises:
            ValueError: If `axis` has zero length.

        Preconditions:
            `axis` is a non-zero 3-vector.
        """
        axis = np.asarray(axis, dtype=float)
        axis = axis / np.linalg.norm(axis)
        x, y, z = axis
        c = np.cos(theta)
        s = np.sin(theta)
        R = np.array(
            [
                [c + x * x * (1 - c), x * y * (1 - c) - z * s, x * z * (1 - c) + y * s],
                [y * x * (1 - c) + z * s, c + y * y * (1 - c), y * z * (1 - c) - x * s],
                [z * x * (1 - c) - y * s, z * y * (1 - c) + x * s, c + z * z * (1 - c)],
            ]
        )
        return R

    def _forward_urdf_state(
        self, joint_angles: Sequence[float]
    ) -> tuple[np.ndarray, list[np.ndarray], list[np.ndarray]]:
        """Compute forward kinematics directly from URDF joint frames.

        Args:
            joint_angles: Joint angles in radians.

        Returns:
            `tuple[np.ndarray, list[np.ndarray], list[np.ndarray]]` containing:
            TCP transform without tool offset, joint origins in world frame,
            and joint axes in world frame.

        Side Effects:
            None.

        Raises:
            ValueError: If `joint_angles` has the wrong shape.

        Preconditions:
            `joint_angles` length matches `self.num_joints`.
        """
        q = self._as_eigen_vector(joint_angles)
        T = np.asarray(self.base_prefix_tf, dtype=float)
        origins_world: list[np.ndarray] = []
        axes_world: list[np.ndarray] = []
        for qi, frame in zip(q, self._urdf_joint_frames):
            R_origin = kh.get_extrinsic_rotation(frame["rpy"])
            T_origin = np.eye(4)
            T_origin[:3, :3] = R_origin
            T_origin[:3, 3] = frame["xyz"]

            R_joint = self._axis_angle_rot(frame["axis"], qi)
            T_joint = np.eye(4)
            T_joint[:3, :3] = R_joint

            T_before_joint = T @ T_origin
            origins_world.append(T_before_joint[:3, 3])
            axes_world.append(T_before_joint[:3, :3] @ frame["axis"])

            T = T_before_joint @ T_joint

        return T, origins_world, axes_world

    def _urdf_fk_transform(self, joint_angles: Sequence[float]) -> np.ndarray:
        """Compute forward-kinematics transform with tool offset applied.

        Args:
            joint_angles: Joint angles in radians.

        Returns:
            `np.ndarray` 4x4 transform of the tool frame in world coordinates.

        Side Effects:
            None.

        Raises:
            ValueError: If `joint_angles` has the wrong shape.

        Preconditions:
            `joint_angles` length matches `self.num_joints`.
        """
        T, _, _ = self._forward_urdf_state(joint_angles)
        return T @ np.asarray(self.tool_offset_tf, dtype=float)

    def get_transformation_matrix(
        self, joint_angles: Sequence[float] | np.ndarray
    ) -> np.ndarray:
        """Return the tool-frame transform for the given joint angles.

        Args:
            joint_angles: Joint angles in radians.

        Returns:
            `np.ndarray` 4x4 transform matrix.

        Side Effects:
            None.

        Raises:
            RuntimeError: If the native backend is unavailable.
            ValueError: If `joint_angles` has the wrong shape.

        Preconditions:
            `joint_angles` length matches `self.num_joints`.
        """
        native = self._native()
        return np.asarray(native.tcp_transform(self._as_eigen_vector(joint_angles)))

    def get_individual_transformation_matrix(
        self, index: int, joint_angles: Sequence[float]
    ) -> np.ndarray:
        """Return the transform for an individual joint frame.

        Args:
            index: Joint index to retrieve.
            joint_angles: Joint angles in radians.

        Returns:
            `np.ndarray` 4x4 transform matrix for the joint frame.

        Side Effects:
            None.

        Raises:
            RuntimeError: If the native backend is unavailable.
            ValueError: If `joint_angles` has the wrong shape.
            IndexError: If `index` is out of range.

        Preconditions:
            `joint_angles` length matches `self.num_joints`.
        """
        native = self._native()
        transforms = native.individual_transforms(self._as_eigen_vector(joint_angles))
        if index < 0 or index >= len(transforms):
            raise IndexError(
                f"Joint index {index} out of range for {len(transforms)} transforms"
            )
        return np.asarray(transforms[index])

    def get_rotation_matrix(self, joint_angles: Sequence[float]) -> np.ndarray:
        """Return the tool-frame rotation matrix for the given joint angles.

        Args:
            joint_angles: Joint angles in radians.

        Returns:
            `np.ndarray` 3x3 rotation matrix.

        Side Effects:
            None.

        Raises:
            RuntimeError: If the native backend is unavailable.
            ValueError: If `joint_angles` has the wrong shape.

        Preconditions:
            `joint_angles` length matches `self.num_joints`.
        """
        T = self.get_transformation_matrix(joint_angles)
        return np.asarray(T[:3, :3])

    def get_individual_rotation_matrix(
        self, index: int, joint_angles: Sequence[float] | np.ndarray
    ) -> np.ndarray:
        """Return the rotation matrix for an individual joint frame.

        Args:
            index: Joint index to retrieve.
            joint_angles: Joint angles in radians.

        Returns:
            `np.ndarray` 3x3 rotation matrix for the joint frame.

        Side Effects:
            None.

        Raises:
            RuntimeError: If the native backend is unavailable.
            ValueError: If `joint_angles` has the wrong shape.
            IndexError: If `index` is out of range.

        Preconditions:
            `joint_angles` length matches `self.num_joints`.
        """
        native = self._native()
        rotations = native.individual_rotations(self._as_eigen_vector(joint_angles))
        if index < 0 or index >= len(rotations):
            raise IndexError(
                f"Joint index {index} out of range for {len(rotations)} rotations"
            )
        return np.asarray(rotations[index])

    def get_position_vectors(
        self, joint_angles: Sequence[float] | None = None
    ) -> List[np.ndarray]:
        """Return joint origin positions in the world frame.

        Args:
            joint_angles: Joint angles in radians. Defaults to zeros.

        Returns:
            `list[np.ndarray]` of joint origin vectors in world coordinates.

        Side Effects:
            None.

        Raises:
            ValueError: If `joint_angles` has the wrong shape.

        Preconditions:
            `joint_angles` length matches `self.num_joints` when provided.
        """
        angles = joint_angles if joint_angles is not None else [0.0] * self.num_joints
        _, origins, _ = self._forward_urdf_state(angles)
        return [np.asarray(vec) for vec in origins]

    def get_forward_kinematics(
        self, joint_angles: Sequence[float] | np.ndarray
    ) -> tuple[np.ndarray, np.ndarray]:
        """Return end-effector position and orientation for joint angles.

        Args:
            joint_angles: Joint angles in radians.

        Returns:
            `tuple[np.ndarray, np.ndarray]` containing position (x, y, z) and
            orientation quaternion (x, y, z, w).

        Side Effects:
            None.

        Raises:
            RuntimeError: If the native backend is unavailable.
            ValueError: If `joint_angles` has the wrong shape.

        Preconditions:
            `joint_angles` length matches `self.num_joints`.
        """
        transform = self.get_transformation_matrix(joint_angles)
        position = np.asarray(transform[:3, 3])
        orientation = rotation_matrix_to_quaternion(np.asarray(transform[:3, :3]))
        return position, orientation

    def get_inverse_kinematics(
        self,
        target_position: np.ndarray[float, np.dtype[np.float32]],
        initial_angles: np.ndarray[float, np.dtype[np.float32]],
        max_iters: int = 200,
        tol: float = 1e-8,
        step_limit: float = 0.1,
        damping: float = 1e-3,
    ) -> np.ndarray:
        """Compute inverse kinematics via iterative Jacobian pseudo-inverse.

        Args:
            target_position: Desired end-effector position (x, y, z).
            initial_angles: Initial joint angles for the solver.
            max_iters: Maximum number of iterations.
            tol: Convergence tolerance on position error.
            step_limit: Maximum joint step per iteration.
            damping: Damping factor for the pseudo-inverse.

        Returns:
            `np.ndarray` joint angles that reach the target within tolerance.

        Side Effects:
            Prints a warning if convergence fails.

        Raises:
            ValueError: If input shapes are invalid.

        Preconditions:
            `target_position` has shape (3,) and `initial_angles` matches
            `self.num_joints`.
        """
        if target_position.shape != (3,):
            raise ValueError(
                f"Target position must be shape (3,), got {target_position.shape}"
            )
        if initial_angles.shape != (self.num_joints,):
            raise ValueError(
                f"Initial angles must be shape ({self.num_joints},), got {initial_angles.shape}"
            )

        def fk_func(q: np.ndarray) -> np.ndarray:
            """Evaluate forward-kinematics position for IK iterations.

            Args:
                q: Joint vector [rad].
            Returns:
                `np.ndarray` tool position `[x, y, z]` in meters.
            Side Effects:
                None.
            Raises:
                ValueError: If `q` has invalid shape.
            Preconditions:
                `q` matches `self.num_joints`.
            """
            return self.get_transformation_matrix(q)[:3, 3]

        def jacobian_func(q: np.ndarray) -> np.ndarray:
            """Evaluate translational Jacobian block for IK iterations.

            Args:
                q: Joint vector [rad].
            Returns:
                `np.ndarray` Jacobian block shaped `(3, num_joints)`.
            Side Effects:
                None.
            Raises:
                ValueError: If `q` has invalid shape.
            Preconditions:
                `q` matches `self.num_joints`.
            """
            return self.get_jacobian_matrix(q)[:3, :]

        q_sol, converged, iters = self._iterative_inverse_kinematics(
            fk_func,
            jacobian_func,
            initial_angles,
            target_position,
            max_iters=max_iters,
            tol=tol,
            step_limit=step_limit,
            damping=damping,
        )

        if not converged:
            print(
                f"Warning: IK did not converge within the maximum iterations ({iters} iters)."
            )

        return q_sol

    # ------------------------------------------------------------------
    # Public API: utilities
    # ------------------------------------------------------------------
    def get_gravity_torque(
        self, joint_angles: Sequence[float], g_world: Sequence[float] | None = None
    ) -> np.ndarray:
        """Return gravity-compensation torque for the given joint angles.

        Args:
            joint_angles: Joint angles in radians.
            g_world: Gravity vector in world coordinates. Defaults to [0, 0, -9.81].

        Returns:
            `np.ndarray` joint torque vector.

        Side Effects:
            None.

        Raises:
            RuntimeError: If the native backend is unavailable.
            ValueError: If `joint_angles` has the wrong shape or `g_world` is invalid.

        Preconditions:
            `joint_angles` length matches `self.num_joints`.
        """
        if g_world is None:
            g = np.array([0.0, 0.0, -9.81], dtype=float)
        else:
            g = np.asarray(g_world, dtype=float)
            if g.shape != (3,):
                raise ValueError(f"gravity vector must be shape (3,), got {g.shape}")
        native = self._native()
        return np.asarray(native.gravity_torque(self._as_eigen_vector(joint_angles), g))

    def get_joint_axes(self, joint_angles: Sequence[float]) -> List[np.ndarray]:
        """Return joint axis vectors in the world frame.

        Args:
            joint_angles: Joint angles in radians.

        Returns:
            `list[np.ndarray]` of joint axis vectors.

        Side Effects:
            None.

        Raises:
            ValueError: If `joint_angles` has the wrong shape.

        Preconditions:
            `joint_angles` length matches `self.num_joints`.
        """
        _, _, axes = self._forward_urdf_state(joint_angles)
        return [np.asarray(vec) for vec in axes]

    # ------------------------------------------------------------------
    # Private API: utilities
    # ------------------------------------------------------------------
    def _damped_pseudoinverse(self, J: np.ndarray, lam: float) -> np.ndarray:
        """Compute a damped least-squares pseudoinverse.

        Args:
            J: Jacobian matrix.
            lam: Damping factor.

        Returns:
            `np.ndarray` pseudoinverse matrix.

        Side Effects:
            None.

        Raises:
            None.

        Preconditions:
            `J` is a 2D array.
        """
        m, n = J.shape
        A = J @ J.T + (lam**2) * np.eye(m)
        # Solve A X = I instead of explicitly inverting A
        X = np.linalg.solve(A, np.eye(m))
        return J.T @ X  # shape (n, m)

    def _iterative_inverse_kinematics(
        self,
        fk_func: Callable[[np.ndarray], np.ndarray],
        jacobian_func: Callable[[np.ndarray], np.ndarray],
        q0: np.ndarray,
        x_target: np.ndarray,
        *,
        max_iters: int = 200,
        tol: float = 1e-4,
        step_limit: float = 0.1,  # α on the slide (max ||Δq|| per iteration)
        damping: float = 1e-3,  # λ, set to 0 to use plain pseudoinverse
    ) -> tuple[np.ndarray, bool, int]:
        """Solve inverse kinematics using iterative Jacobian updates.

        Args:
            fk_func: Function mapping q -> x (end-effector position).
            jacobian_func: Function mapping q -> J(q).
            q0: Initial joint configuration.
            x_target: Desired end-effector position.
            max_iters: Maximum number of iterations.
            tol: Stop when ||x_target - x_current|| < tol.
            step_limit: Maximum joint update norm per iteration.
            damping: Damping factor for the pseudo-inverse.

        Returns:
            `tuple[np.ndarray, bool, int]` containing solution `q`, convergence
            flag, and iteration count.

        Side Effects:
            None.

        Raises:
            None.

        Preconditions:
            `fk_func` and `jacobian_func` are defined for all iterates.
        """
        q = np.asarray(q0, dtype=float).copy()
        x_target = np.asarray(x_target, dtype=float)

        for k in range(max_iters):
            # 1. Forward kinematics
            x_current = np.asarray(fk_func(q), dtype=float)

            # 2. Task-space error
            e = x_target - x_current
            err_norm = np.linalg.norm(e)

            if err_norm < tol:
                return q, True, k

            # 3. Jacobian at current configuration
            J = np.asarray(jacobian_func(q), dtype=float)

            # 4. Compute pseudo-inverse (damped or plain)
            if damping > 0.0:
                J_pinv = self._damped_pseudoinverse(J, damping)
            else:
                J_pinv = np.linalg.pinv(J)

            # 5. Joint update (treat e as desired task-space velocity)
            dq = J_pinv @ e

            # 6. Limit step size as on the slide: if ||dq|| > α, scale it
            dq_norm = np.linalg.norm(dq)
            if dq_norm > step_limit:
                dq = (step_limit / dq_norm) * dq

            # 7. Integrate
            q = q + dq

        # If we reach here, we did not hit the error threshold
        return q, False, max_iters
