# Copyright (C) Istituto Italiano di Tecnologia (IIT). All rights reserved.

import numpy.typing as npt

from adam.core.constants import Representations
from adam.core.spatial_math import ArrayLike, SpatialMath
from adam.model import Model, Node, Joint


class RBDAlgorithms:
    """This is a small class that implements Rigid body algorithms retrieving robot quantities, for Floating Base systems - as humanoid robots."""

    def __init__(self, model: Model, math: SpatialMath) -> None:
        """
        Args:
            model (Model): the adam.model representing the robot
            math (SpatialMath): the spatial math.
        """

        self.model = model
        self.NDoF = model.NDoF
        self.root_link = self.model.tree.root
        self.math = math
        self.frame_velocity_representation = (
            Representations.MIXED_REPRESENTATION
        )  # default
        # Cache root quantities that are reused at every call
        self._root_spatial_transform = self.math.spatial_transform(
            self.math.factory.eye(3),
            self.math.factory.zeros((3, 1)),
        )
        self._root_motion_subspace = self.math.factory.eye(6)
        self._prepare_tree_cache()

    def set_frame_velocity_representation(self, representation: Representations):
        """Sets the frame velocity representation

        Args:
            representation (str): The representation of the frame velocity
        """
        self.frame_velocity_representation = representation

    def crba(self, base_transform: npt.ArrayLike, joint_positions: npt.ArrayLike):
        """
        Batched Composite Rigid Body Algorithm (CRBA) + Orin's Centroidal Momentum Matrix.

        Args:
            base_transform (npt.ArrayLike): The homogenous transform from base to world frame
            joint_positions (npt.ArrayLike): The joints position

        Returns:
            M, Jcm (npt.ArrayLike, npt.ArrayLike): The mass matrix and the centroidal momentum matrix
        """
        base_transform, joint_positions = self._convert_to_arraylike(
            base_transform, joint_positions
        )

        math = self.math
        n = self.NDoF
        node_count = self._node_count
        node_indices = self._node_indices
        rev_indices = self._rev_node_indices
        parent_indices = self._parent_indices
        joint_indices = self._joint_indices_per_node
        motion_subspaces = self._motion_subspaces
        inertias = self._spatial_inertias
        joints = self._joints_per_node
        root_idx = self._root_index
        joint_to_node = self._joint_index_to_node

        if len(joint_positions.shape) >= 2:
            batch_shape = joint_positions.shape[:-1]
        elif base_transform.ndim > 2:
            batch_shape = base_transform.shape[:-2]
        else:
            batch_shape = ()

        if joint_positions.shape[-1] > 0:
            zero_q = math.zeros_like(joint_positions[..., 0])
        else:
            zero_q = math.zeros_like(base_transform[..., 0, 0])

        def tile_batch(arr):
            if not batch_shape:
                return arr
            reps = batch_shape + (1,) * len(arr.shape)
            return math.tile(arr, reps)

        Xup = [None] * node_count
        Phi = [None] * node_count
        Ic = [None] * node_count

        for idx in node_indices:
            Ic[idx] = tile_batch(inertias[idx])
            if idx == root_idx:
                Xup[idx] = tile_batch(self._root_spatial_transform)
                Phi[idx] = tile_batch(self._root_motion_subspace)
            else:
                joint = joints[idx]
                joint_idx = joint_indices[idx]
                q_i = (
                    joint_positions[..., joint_idx] if joint_idx is not None else zero_q
                )
                X_current = (
                    joint.spatial_transform(q=q_i)
                    if joint is not None
                    else self._root_spatial_transform
                )
                Xup[idx] = X_current
                Phi[idx] = tile_batch(motion_subspaces[idx])

        Xup_T = [math.swapaxes(Xup[i], -2, -1) for i in range(node_count)]
        Ic_comp = Ic[:]

        for idx in rev_indices:
            parent = parent_indices[idx]
            if parent >= 0:
                Ic_comp[parent] = Ic_comp[parent] + math.mtimes(
                    math.mtimes(Xup_T[idx], Ic_comp[idx]),
                    Xup[idx],
                )

        def block_index(node_idx: int) -> int | None:
            if node_idx == root_idx:
                return 0
            joint_idx = joint_indices[node_idx]
            if joint_idx is None:
                return None
            return 1 + int(joint_idx)

        sizes = [6] + [1] * n
        blocks: list[list[npt.ArrayLike | None]] = [
            [None for _ in range(n + 1)] for _ in range(n + 1)
        ]

        for idx in rev_indices:
            Phi_i = Phi[idx]
            Phi_T = math.swapaxes(Phi_i, -2, -1)
            F = math.mtimes(Ic_comp[idx], Phi_i)
            ri = block_index(idx)

            if idx == root_idx:
                blocks[0][0] = math.mtimes(Phi_T, F)
                continue

            if ri is not None:
                blocks[ri][ri] = math.mtimes(Phi_T, F)

            current = idx
            F_path = F
            while current != root_idx:
                parent = parent_indices[current]
                F_path = math.mtimes(Xup_T[current], F_path)
                rj = block_index(parent)
                if rj is None:
                    current = parent
                    continue
                Bij = math.mtimes(math.swapaxes(F_path, -2, -1), Phi[parent])
                if ri is None:
                    current = parent
                    continue
                blocks[ri][rj] = Bij
                blocks[rj][ri] = math.swapaxes(Bij, -2, -1)
                current = parent

        row_tensors = []
        for r in range(n + 1):
            row_blocks = []
            for c in range(n + 1):
                block = blocks[r][c]
                if block is None:
                    block = math.factory.zeros(batch_shape + (sizes[r], sizes[c]))
                row_blocks.append(block)
            row_tensors.append(math.concatenate(row_blocks, axis=-1))
        M = math.concatenate(row_tensors, axis=-2)

        A = math.swapaxes(M[..., :3, 3:6], -2, -1) / M[..., 0, 0][..., None, None]
        I3 = math.factory.eye(batch_shape + (3,))
        Z3 = math.factory.zeros(batch_shape + (3, 3))
        top = math.concatenate([I3, A], axis=-1)
        bot = math.concatenate([Z3, I3], axis=-1)
        O_X_G = math.concatenate([top, bot], axis=-2)

        X_G = [None] * node_count
        for idx in node_indices:
            if idx == root_idx:
                X_G[idx] = O_X_G
            else:
                parent = parent_indices[idx]
                X_G[idx] = math.mtimes(Xup[idx], X_G[parent])

        J_base = math.mtimes(
            math.swapaxes(X_G[root_idx], -2, -1),
            math.mtimes(Ic_comp[root_idx], Phi[root_idx]),
        )

        zero_col = math.factory.zeros(batch_shape + (6, 1))
        joint_cols = []
        for jidx in range(n):
            node_idx = joint_to_node.get(jidx)
            if node_idx is not None:
                col = math.mtimes(
                    math.swapaxes(X_G[node_idx], -2, -1),
                    math.mtimes(Ic_comp[node_idx], Phi[node_idx]),
                )
            else:
                col = zero_col
            joint_cols.append(col)

        Jcm = math.concatenate([J_base] + joint_cols, axis=-1)

        if (
            self.frame_velocity_representation
            == Representations.BODY_FIXED_REPRESENTATION
        ):
            return M, Jcm

        if self.frame_velocity_representation == Representations.MIXED_REPRESENTATION:
            Xm = math.adjoint_mixed_inverse(base_transform)
            In = math.factory.eye(batch_shape + (n,))
            Z6n = math.factory.zeros(batch_shape + (6, n))
            Zn6 = math.factory.zeros(batch_shape + (n, 6))

            top = math.concatenate([Xm, Z6n], axis=-1)
            bot = math.concatenate([Zn6, In], axis=-1)
            X_to_mixed = math.concatenate([top, bot], axis=-2)

            M_mixed = math.mtimes(
                math.swapaxes(X_to_mixed, -2, -1), math.mtimes(M, X_to_mixed)
            )
            Jcm_mixed = math.mtimes(
                math.swapaxes(Xm, -2, -1), math.mtimes(Jcm, X_to_mixed)
            )
            return M_mixed, Jcm_mixed

        if (
            self.frame_velocity_representation
            == Representations.INERTIAL_FIXED_REPRESENTATION
        ):
            Xi = math.adjoint_inverse(base_transform)
            In = math.factory.eye(batch_shape + (n,))
            Z6n = math.factory.zeros(batch_shape + (6, n))
            Zn6 = math.factory.zeros(batch_shape + (n, 6))

            top = math.concatenate([Xi, Z6n], axis=-1)
            bot = math.concatenate([Zn6, In], axis=-1)
            X_to_inertial = math.concatenate([top, bot], axis=-2)

            M_inertial = math.mtimes(
                math.swapaxes(X_to_inertial, -2, -1), math.mtimes(M, X_to_inertial)
            )
            Xm = math.adjoint_mixed_inverse(base_transform)  # rotation-only, for cmm
            wrench_transform = math.swapaxes(Xm, -2, -1)
            Jcm_inertial = math.mtimes(
                wrench_transform, math.mtimes(Jcm, X_to_inertial)
            )
            return M_inertial, Jcm_inertial

        raise ValueError(
            f"Unknown frame velocity representation: {self.frame_velocity_representation}"
        )

    def forward_kinematics(
        self, frame, base_transform: npt.ArrayLike, joint_positions: npt.ArrayLike
    ) -> npt.ArrayLike:
        """Computes the forward kinematics relative to the specified `frame`.

        Args:
            frame (str): The frame to which the fk will be computed
            base_transform (npt.ArrayLike): The homogenous transform from base to world frame
            joint_positions (npt.ArrayLike): The joints position

        Returns:
            I_H_L (npt.ArrayLike): The fk represented as Homogenous transformation matrix
        """
        base_transform, joint_positions = self._convert_to_arraylike(
            base_transform, joint_positions
        )
        chain = self.model.get_joints_chain(self.root_link, frame)
        I_H_L = base_transform
        for joint in chain:
            q_ = (
                joint_positions[..., joint.idx]
                if joint.idx is not None
                else self.math.zeros_like(joint_positions[..., 0])
            )
            H_joint = joint.homogeneous(q=q_)
            I_H_L = I_H_L @ H_joint
        return I_H_L

    def joints_jacobian(
        self, frame: str, joint_positions: npt.ArrayLike
    ) -> npt.ArrayLike:
        """Returns the Jacobian relative to the specified `frame`.

        Args:
            frame (str): The frame to which the jacobian will be computed
            base_transform (npt.ArrayLike): The homogenous transform from base to world frame
            joint_positions (npt.ArrayLike): The joints position

        Returns:
            J (npt.ArrayLike): The Joints Jacobian relative to the frame
        """
        joint_positions = self._convert_to_arraylike(joint_positions)

        batch_size = (
            joint_positions.shape[:-1] if len(joint_positions.shape) >= 2 else ()
        )

        chain = self.model.get_joints_chain(self.root_link, frame)
        eye = self.math.factory.eye(batch_size + (4,))
        B_H_j = eye
        B_H_L = self.forward_kinematics(frame, eye, joint_positions)
        L_H_B = self.math.homogeneous_inverse(B_H_L)
        cols = [None] * self.NDoF
        for joint in chain:
            q = (
                joint_positions[..., joint.idx]
                if joint.idx is not None
                else self.math.zeros_like(joint_positions[..., 0])
            )
            H_j = joint.homogeneous(q=q)
            B_H_j = B_H_j @ H_j
            L_H_j = L_H_B @ B_H_j
            if joint.idx is not None:
                cols[joint.idx] = self.math.adjoint(L_H_j) @ joint.motion_subspace()

        zero_col = self.math.zeros(batch_size + (6, 1))
        cols = [zero_col if col is None else col for col in cols]
        # Use concatenate instead of stack to get (batch, 6, joints) directly
        return self.math.concatenate(cols, axis=-1)

    def jacobian(
        self, frame: str, base_transform: npt.ArrayLike, joint_positions: npt.ArrayLike
    ) -> npt.ArrayLike:
        """Returns the Jacobian for `frame`.

        Args:
            frame (str): The frame to which the jacobian will be computed
            base_transform (npt.ArrayLike): The homogenous transform from base to world frame
            joint_positions (npt.ArrayLike): The joints position

        Returns:
            npt.ArrayLike: The Jacobian for the specified frame
        """

        base_transform, joint_positions = self._convert_to_arraylike(
            base_transform, joint_positions
        )

        eye = self.math.factory.eye(4)
        J = self.joints_jacobian(frame, joint_positions)
        B_H_L = self.forward_kinematics(frame, eye, joint_positions)
        L_X_B = self.math.adjoint_inverse(B_H_L)
        J_tot = self.math.concatenate([L_X_B, J], axis=-1)
        w_H_B = base_transform
        if (
            self.frame_velocity_representation
            == Representations.BODY_FIXED_REPRESENTATION
        ):
            return J_tot
        if self.frame_velocity_representation == Representations.MIXED_REPRESENTATION:
            w_H_L = w_H_B @ B_H_L
            LI_X_L = self.math.adjoint_mixed(w_H_L)

            top_left = self.math.adjoint_mixed_inverse(base_transform)
            top_right = self.math.factory.zeros(
                joint_positions.shape[:-1] + (6, self.NDoF)
            )
            bottom_left = self.math.factory.zeros(
                joint_positions.shape[:-1] + (self.NDoF, 6)
            )
            bottom_right = self.math.factory.eye(
                joint_positions.shape[:-1] + (self.NDoF,)
            )
            top = self.math.concatenate([top_left, top_right], axis=-1)
            bottom = self.math.concatenate([bottom_left, bottom_right], axis=-1)
            X = self.math.concatenate([top, bottom], axis=-2)

            J_tot = LI_X_L @ J_tot @ X
            return J_tot

        if (
            self.frame_velocity_representation
            == Representations.INERTIAL_FIXED_REPRESENTATION
        ):
            w_H_L = w_H_B @ B_H_L
            I_X_L = self.math.adjoint(w_H_L)

            top_left = self.math.adjoint_inverse(base_transform)
            top_right = self.math.factory.zeros(
                joint_positions.shape[:-1] + (6, self.NDoF)
            )
            bottom_left = self.math.factory.zeros(
                joint_positions.shape[:-1] + (self.NDoF, 6)
            )
            bottom_right = self.math.factory.eye(
                joint_positions.shape[:-1] + (self.NDoF,)
            )
            top = self.math.concatenate([top_left, top_right], axis=-1)
            bottom = self.math.concatenate([bottom_left, bottom_right], axis=-1)
            X = self.math.concatenate([top, bottom], axis=-2)

            return I_X_L @ J_tot @ X

        raise NotImplementedError(
            "Only BODY_FIXED_REPRESENTATION, MIXED_REPRESENTATION and INERTIAL_FIXED_REPRESENTATION are implemented"
        )

    def relative_jacobian(
        self, frame: str, joint_positions: npt.ArrayLike
    ) -> npt.ArrayLike:
        """Returns the Jacobian between the root link and a specified `frame`.

        Args:
            frame (str): The tip of the chain
            joint_positions (npt.ArrayLike): The joints position

        Returns:
            J (npt.ArrayLike): The 6 x NDoF Jacobian between the root and the `frame`
        """
        joint_positions = self._convert_to_arraylike(joint_positions)
        if (
            self.frame_velocity_representation
            == Representations.BODY_FIXED_REPRESENTATION
        ):
            return self.joints_jacobian(frame, joint_positions)
        elif self.frame_velocity_representation == Representations.MIXED_REPRESENTATION:
            eye = self.math.factory.eye(4)
            B_H_L = self.forward_kinematics(frame, eye, joint_positions)
            LI_X_L = self.math.adjoint_mixed(B_H_L)
            return LI_X_L @ self.joints_jacobian(frame, joint_positions)

        elif (
            self.frame_velocity_representation
            == Representations.INERTIAL_FIXED_REPRESENTATION
        ):
            eye = self.math.factory.eye(4)
            B_H_L = self.forward_kinematics(frame, eye, joint_positions)
            I_X_L = self.math.adjoint(B_H_L)
            return I_X_L @ self.joints_jacobian(frame, joint_positions)

        raise NotImplementedError(
            "Only BODY_FIXED_REPRESENTATION, MIXED_REPRESENTATION and INERTIAL_FIXED_REPRESENTATION are implemented"
        )

    def jacobian_dot(
        self,
        frame: str,
        base_transform: npt.ArrayLike,
        joint_positions: npt.ArrayLike,
        base_velocity: npt.ArrayLike,
        joint_velocities: npt.ArrayLike,
    ) -> npt.ArrayLike:
        """Returns the Jacobian time derivative for `frame`.

        Args:
            frame (str): The frame to which the jacobian will be computed
            base_transform (npt.ArrayLike): The homogenous transform from base to world frame
            joint_positions (npt.ArrayLike): The joints position
            base_velocity (npt.ArrayLike): The spatial velocity of the base
            joint_velocities (npt.ArrayLike): The joints velocities

        Returns:
            J_dot (npt.ArrayLike): The Jacobian derivative relative to the frame

        """

        base_transform, joint_positions, base_velocity, joint_velocities = (
            self._convert_to_arraylike(
                base_transform, joint_positions, base_velocity, joint_velocities
            )
        )

        batch_size = base_transform.shape[:-2] if len(base_transform.shape) > 2 else ()

        I4 = self.math.factory.eye(batch_size + (4,))

        B_H_L = self.forward_kinematics(frame, I4, joint_positions)
        L_H_B = self.math.homogeneous_inverse(B_H_L)

        if self.frame_velocity_representation == Representations.MIXED_REPRESENTATION:
            # frame C = B[I] for mixed representation
            B_v_C = self.math.mxv(
                self.math.adjoint_mixed_inverse(base_transform), base_velocity
            )
        elif (
            self.frame_velocity_representation
            == Representations.INERTIAL_FIXED_REPRESENTATION
        ):
            # frame C = I for inertial-fixed representation
            B_v_C = self.math.mxv(
                self.math.adjoint_inverse(base_transform), base_velocity
            )
        elif (
            self.frame_velocity_representation
            == Representations.BODY_FIXED_REPRESENTATION
        ):
            # frame C = B for body-fixed representation
            B_v_C = base_velocity
        else:
            raise NotImplementedError(
                "Only BODY_FIXED_REPRESENTATION, MIXED_REPRESENTATION, and INERTIAL_FIXED_REPRESENTATION are implemented"
            )

        v = self.math.mxv(self.math.adjoint(L_H_B), B_v_C)
        a = self.math.mxv(self.math.adjoint_derivative(L_H_B, v), B_v_C)

        J_base_full = self.math.adjoint_inverse(B_H_L)
        J_base_cols = [J_base_full[..., :, i : i + 1] for i in range(6)]
        J_dot_base_full = self.math.adjoint_derivative(L_H_B, v)
        J_dot_base_cols = [J_dot_base_full[..., :, i : i + 1] for i in range(6)]

        cols: list = [None] * self.NDoF
        cols_dot: list = [None] * self.NDoF

        B_H_j = I4
        chain = self.model.get_joints_chain(self.root_link, frame)

        for joint in chain:
            if joint.idx is not None:
                q = joint_positions[..., joint.idx]
                q_dot = joint_velocities[..., joint.idx]
            else:
                q = self.math.factory.zeros(batch_size + ())
                q_dot = self.math.factory.zeros(batch_size + ())

            H_j = joint.homogeneous(q=q)
            B_H_j = B_H_j @ H_j
            L_H_j = L_H_B @ B_H_j
            S = joint.motion_subspace()
            J_j = self.math.adjoint(L_H_j) @ S

            v = v + self.math.vxs(J_j, q_dot)
            J_dot_j = self.math.adjoint_derivative(L_H_j, v) @ S
            a = a + self.math.vxs(J_dot_j, q_dot)

            if joint.idx is not None:
                cols[joint.idx] = J_j
                cols_dot[joint.idx] = J_dot_j

        zero_col = self.math.factory.zeros(batch_size + (6, 1))
        cols = [zero_col if c is None else c for c in cols]
        cols_dot = [zero_col if c is None else c for c in cols_dot]

        J = self.math.concatenate([*J_base_cols, *cols], axis=-1)
        J_dot = self.math.concatenate([*J_dot_base_cols, *cols_dot], axis=-1)

        if (
            self.frame_velocity_representation
            == Representations.BODY_FIXED_REPRESENTATION
        ):
            return J_dot

        elif self.frame_velocity_representation == Representations.MIXED_REPRESENTATION:
            adj = self.math.adjoint_mixed
            adj_derivative = self.math.adjoint_mixed_derivative
            adj_inverse = self.math.adjoint_mixed_inverse
        elif (
            self.frame_velocity_representation
            == Representations.INERTIAL_FIXED_REPRESENTATION
        ):
            adj = self.math.adjoint
            adj_derivative = self.math.adjoint_derivative
            adj_inverse = self.math.adjoint_inverse
        else:
            raise NotImplementedError(
                "Only BODY_FIXED_REPRESENTATION, MIXED_REPRESENTATION and INERTIAL_FIXED_REPRESENTATION are implemented"
            )
        I_H_L = self.forward_kinematics(frame, base_transform, joint_positions)
        I_X_L = adj(I_H_L)
        I_v_L = self.math.mxv(I_X_L, v)
        I_X_L_dot = adj_derivative(I_H_L, I_v_L)

        adj_inv = adj_inverse(base_transform)

        Z_6xN = self.math.factory.zeros(batch_size + (6, self.NDoF))
        Z_Nx6 = self.math.factory.zeros(batch_size + (self.NDoF, 6))
        I_N = self.math.factory.eye(batch_size + (self.NDoF,))

        top = self.math.concatenate([adj_inv, Z_6xN], axis=-1)
        bottom = self.math.concatenate([Z_Nx6, I_N], axis=-1)
        X = self.math.concatenate([top, bottom], axis=-2)

        B_H_I = self.math.homogeneous_inverse(base_transform)
        B_H_I_deriv = adj_derivative(B_H_I, -B_v_C)

        Z_NxN = self.math.factory.zeros(batch_size + (self.NDoF, self.NDoF))
        topd = self.math.concatenate([B_H_I_deriv, Z_6xN], axis=-1)
        bottomd = self.math.concatenate([Z_Nx6, Z_NxN], axis=-1)
        X_dot = self.math.concatenate([topd, bottomd], axis=-2)

        return (I_X_L_dot @ J @ X) + (I_X_L @ J_dot @ X) + (I_X_L @ J @ X_dot)

    def CoM_position(
        self, base_transform: npt.ArrayLike, joint_positions: npt.ArrayLike
    ) -> npt.ArrayLike:
        """Returns the CoM position

        Args:
            base_transform (npt.ArrayLike): The homogenous transform from base to world frame
            joint_positions (npt.ArrayLike): The joints position

        Returns:
            com (npt.ArrayLike): The CoM position
        """
        base_transform, joint_positions = self._convert_to_arraylike(
            base_transform, joint_positions
        )
        batch_size = base_transform.shape[:-2] if len(base_transform.shape) > 2 else ()

        com_pos = self.math.factory.zeros(batch_size + (3,))
        for item in self.model.tree:
            link = item.link
            I_H_l = self.forward_kinematics(link.name, base_transform, joint_positions)
            H_link = link.homogeneous()
            I_H_l = I_H_l @ H_link
            link_pos = I_H_l[..., :3, 3:4]
            com_pos += self.math.vxs(link_pos, link.inertial.mass)
        com_pos /= self._convert_to_arraylike(self.get_total_mass())
        return com_pos

    def CoM_jacobian(
        self, base_transform: npt.ArrayLike, joint_positions: npt.ArrayLike
    ) -> npt.ArrayLike:
        """Returns the center of mass (CoM) Jacobian using the centroidal momentum matrix.

        Args:
            base_transform (npt.ArrayLike): The homogenous transform from base to world frame
            joint_positions (npt.ArrayLike): The joints position

        Returns:
            J_com (npt.ArrayLike): The CoM Jacobian
        """
        base_transform, joint_positions = self._convert_to_arraylike(
            base_transform, joint_positions
        )
        batch_size = base_transform.shape[:-2] if len(base_transform.shape) > 2 else ()
        ori_frame_velocity_representation = self.frame_velocity_representation
        self.frame_velocity_representation = Representations.MIXED_REPRESENTATION
        _, Jcm = self.crba(base_transform, joint_positions)
        Xm = self.math.adjoint_mixed(base_transform)
        In = self.math.factory.eye(batch_size + (self.NDoF,))
        Z6n = self.math.factory.zeros(batch_size + (6, self.NDoF))
        Zn6 = self.math.factory.zeros(batch_size + (self.NDoF, 6))
        if (
            ori_frame_velocity_representation
            == Representations.BODY_FIXED_REPRESENTATION
        ):

            top = self.math.concatenate([Xm, Z6n], axis=-1)
            bot = self.math.concatenate([Zn6, In], axis=-1)
            X = self.math.concatenate([top, bot], axis=-2)
            Jcm = Jcm @ X
        elif (
            ori_frame_velocity_representation
            == Representations.INERTIAL_FIXED_REPRESENTATION
        ):
            Xi = self.math.adjoint_inverse(base_transform)
            A = self.math.mtimes(Xm, Xi)
            top = self.math.concatenate([A, Z6n], axis=-1)
            bot = self.math.concatenate([Zn6, In], axis=-1)
            X = self.math.concatenate([top, bot], axis=-2)
            Jcm = Jcm @ X
        self.frame_velocity_representation = ori_frame_velocity_representation
        return Jcm[..., :3, :] / self._convert_to_arraylike(self.get_total_mass())

    def get_total_mass(self):
        """Returns the total mass of the robot

        Returns:
            mass: The total mass
        """
        return self.model.get_total_mass()

    def rnea(
        self,
        base_transform: npt.ArrayLike,
        joint_positions: npt.ArrayLike,
        base_velocity: npt.ArrayLike,
        joint_velocities: npt.ArrayLike,
        g: npt.ArrayLike,
    ) -> npt.ArrayLike:
        """
        Batched Recursive Newton-Euler (reduced: no joint/base accelerations, no external forces).

        Args:
            base_transform (npt.ArrayLike): The homogenous transform from base to world frame
            joint_positions (npt.ArrayLike): The joints position
            base_velocity (npt.ArrayLike): The base spatial velocity
            joint_velocities (npt.ArrayLike): The joints velocities
            g (npt.ArrayLike): The gravity vector

        Returns:
            tau (npt.ArrayLike): The vector of generalized forces
        """
        base_transform, joint_positions, base_velocity, joint_velocities, g = (
            self._convert_to_arraylike(
                base_transform, joint_positions, base_velocity, joint_velocities, g
            )
        )

        math = self.math
        n = self.NDoF
        node_count = self._node_count
        parent_indices = self._parent_indices
        joint_indices = self._joint_indices_per_node
        motion_subspaces = self._motion_subspaces
        inertias = self._spatial_inertias
        joints = self._joints_per_node
        root_idx = self._root_index
        node_indices = self._node_indices
        rev_indices = self._rev_node_indices

        batch_shape = (
            tuple(base_transform.shape[:-2]) if base_transform.ndim > 2 else ()
        )

        if (
            self.frame_velocity_representation
            == Representations.INERTIAL_FIXED_REPRESENTATION
        ):
            gravity_X = math.adjoint_inverse(base_transform)
        else:
            gravity_X = math.adjoint_mixed_inverse(base_transform)  # (...,6,6)

        if (
            self.frame_velocity_representation
            == Representations.BODY_FIXED_REPRESENTATION
        ):
            # frame C = B for body-fixed representation
            B_X_C = math.factory.eye(batch_shape + (6,))
            transformed_acc = math.factory.zeros(batch_shape + (6,))
        elif self.frame_velocity_representation == Representations.MIXED_REPRESENTATION:
            # frame C = B[I] for mixed representation
            B_X_C = math.adjoint_mixed_inverse(base_transform)
            omega = base_velocity[..., 3:]
            vlin = base_velocity[..., :3]
            skew_omega_times_vlin = math.mxv(math.skew(omega), vlin)
            top3 = -math.mxv(B_X_C[..., :3, :3], skew_omega_times_vlin)
            bot3 = math.factory.zeros(batch_shape + (3,))
            transformed_acc = math.concatenate([top3, bot3], axis=-1)
        elif (
            self.frame_velocity_representation
            == Representations.INERTIAL_FIXED_REPRESENTATION
        ):
            # frame C = I for inertial-fixed representation
            B_X_C = math.adjoint_inverse(base_transform)
            transformed_acc = math.factory.zeros(batch_shape + (6,))
        else:
            raise NotImplementedError(
                "Only BODY_FIXED_REPRESENTATION, MIXED_REPRESENTATION and INERTIAL_FIXED_REPRESENTATION are implemented"
            )

        a0 = -(math.mxv(gravity_X, g)) + transformed_acc

        if n > 0:
            zero_q = math.zeros_like(joint_positions[..., 0])
            zero_qd = math.zeros_like(joint_velocities[..., 0])
        else:
            zero_q = math.zeros_like(base_velocity[..., 0])
            zero_qd = zero_q

        Ic = [None] * node_count
        Xup = [None] * node_count
        v = [None] * node_count
        a = [None] * node_count
        f = [None] * node_count

        for idx in node_indices:
            Ic[idx] = inertias[idx]

            if idx == root_idx:
                Xup[idx] = self._root_spatial_transform
                v[idx] = math.mxv(B_X_C, base_velocity)
                a[idx] = math.mxv(Xup[idx], a0)
                continue

            joint = joints[idx]
            parent = parent_indices[idx]
            joint_idx = joint_indices[idx]

            q = joint_positions[..., joint_idx] if joint_idx is not None else zero_q
            qd = joint_velocities[..., joint_idx] if joint_idx is not None else zero_qd

            X = joint.spatial_transform(q=q)
            Xup[idx] = X

            Phi_i = motion_subspaces[idx]
            phi_qd = math.vxs(Phi_i, qd)
            v[idx] = math.mxv(X, v[parent]) + phi_qd
            a[idx] = math.mxv(X, a[parent]) + math.mxv(
                math.spatial_skew(v[idx]), phi_qd
            )

            f[idx] = math.mxv(Ic[idx], a[idx]) + math.mxv(
                math.spatial_skew_star(v[idx]), math.mxv(Ic[idx], v[idx])
            )

        # Root wrench contribution (skipped in loop above)
        f[root_idx] = math.mxv(Ic[root_idx], a[root_idx]) + math.mxv(
            math.spatial_skew_star(v[root_idx]), math.mxv(Ic[root_idx], v[root_idx])
        )

        tau_base = None
        tau_joint_cols = [None] * n if n > 0 else []

        for idx in rev_indices:
            Phi_i = motion_subspaces[idx]
            Fi = f[idx]
            Phi_T = math.swapaxes(Phi_i, -2, -1)

            if idx == root_idx:
                tau_base = math.mxv(Phi_T, Fi)
            else:
                joint_idx = joint_indices[idx]
                if joint_idx is not None:
                    tau_joint_cols[joint_idx] = math.mxv(Phi_T, Fi)
                parent = parent_indices[idx]
                if parent >= 0:
                    f[parent] = f[parent] + math.mxv(
                        math.swapaxes(Xup[idx], -2, -1), Fi
                    )

        tau_base = math.mxv(math.swapaxes(B_X_C, -2, -1), tau_base)

        if n > 0:
            zero_tau = math.factory.zeros(batch_shape + (1,))
            tau_joints = [
                col if col is not None else zero_tau for col in tau_joint_cols
            ]
            tau_joints_vec = math.concatenate(tau_joints, axis=-1)
        else:
            tau_joints_vec = math.factory.zeros(batch_shape + (0,))

        return math.concatenate([tau_base, tau_joints_vec], axis=-1)

    def aba(
        self,
        base_transform: npt.ArrayLike,
        joint_positions: npt.ArrayLike,
        base_velocity: npt.ArrayLike,
        joint_velocities: npt.ArrayLike,
        joint_torques: npt.ArrayLike,
        g: npt.ArrayLike,
        external_wrenches: dict[str, npt.ArrayLike] | None = None,
    ) -> npt.ArrayLike:
        """Featherstone Articulated Body Algorithm for floating-base forward dynamics.

        Args:
            base_transform (npt.ArrayLike): The homogenous transform from base to world frame
            joint_positions (npt.ArrayLike): The joints position
            base_velocity (npt.ArrayLike): The spatial velocity of the base
            joint_velocities (npt.ArrayLike): The joints velocities
            joint_torques (npt.ArrayLike): The joints torques/forces
            g (npt.ArrayLike | None, optional): The gravity vector
            external_wrenches (dict[str, npt.ArrayLike] | None, optional): A dictionary of external wrenches applied to specific links. Keys are link names, and values are 6D wrench vectors. Defaults to None.

        Returns:
            accelerations (npt.ArrayLike): The spatial acceleration of the base and joints accelerations
        """
        math = self.math
        n = self.NDoF
        node_count = self._node_count
        parent_indices = self._parent_indices
        joint_indices = self._joint_indices_per_node
        motion_subspaces = self._motion_subspaces
        inertias = self._spatial_inertias
        joints = self._joints_per_node
        root_idx = self._root_index
        node_indices = self._node_indices
        rev_indices = self._rev_node_indices

        (
            base_transform,
            joint_positions,
            base_velocity,
            joint_velocities,
            joint_torques,
            g,
        ) = self._convert_to_arraylike(
            base_transform,
            joint_positions,
            base_velocity,
            joint_velocities,
            joint_torques,
            g,
        )

        batch_shape = (
            tuple(base_transform.shape[:-2]) if base_transform.ndim > 2 else ()
        )

        if external_wrenches is not None:
            generalized_ext = math.factory.zeros(batch_shape + (6 + n,))
            for frame, wrench in external_wrenches.items():
                wrench_arr = self._convert_to_arraylike(wrench)
                J = self.jacobian(frame, base_transform, joint_positions)
                generalized_ext = generalized_ext + math.mxv(
                    math.swapaxes(J, -2, -1), wrench_arr
                )

            base_ext = generalized_ext[..., :6]
            joint_ext = (
                generalized_ext[..., 6:] if n > 0 else math.zeros_like(joint_torques)
            )
        else:
            base_ext = math.factory.zeros(batch_shape + (6,))
            joint_ext = math.zeros_like(joint_torques)

        joint_torques_eff = joint_torques + joint_ext

        if self.frame_velocity_representation == Representations.MIXED_REPRESENTATION:
            B_X_I = math.adjoint_mixed_inverse(base_transform)
            I_X_B = math.adjoint_mixed(base_transform)
        elif (
            self.frame_velocity_representation
            == Representations.INERTIAL_FIXED_REPRESENTATION
        ):
            B_X_I = math.adjoint_inverse(base_transform)
            I_X_B = math.adjoint(base_transform)
        elif (
            self.frame_velocity_representation
            == Representations.BODY_FIXED_REPRESENTATION
        ):
            B_X_I = I_X_B = math.factory.eye(batch_shape + (6,))
        else:
            raise NotImplementedError(
                "Only BODY_FIXED_REPRESENTATION, MIXED_REPRESENTATION and INERTIAL_FIXED_REPRESENTATION are implemented"
            )

        base_velocity_body = math.mxv(B_X_I, base_velocity)

        B_star_BI = math.swapaxes(I_X_B, -2, -1)
        base_ext_body = math.mxv(B_star_BI, base_ext)

        a0_input = math.mxv(math.adjoint_mixed_inverse(base_transform), g)

        def zeros6():
            return math.factory.zeros(batch_shape + (6,))

        def eye6():
            return math.factory.eye(batch_shape + (6,))

        def expand_to_match(vec, reference):
            expanded = math.expand_dims(vec, axis=-1)
            expanded_ndim = expanded.ndim
            reference_ndim = reference.ndim
            if expanded_ndim != reference_ndim:
                expanded = math.expand_dims(expanded, axis=-1)
            return expanded

        if n > 0:
            zero_q = math.zeros_like(joint_positions[..., 0])
        else:
            zero_q = math.zeros_like(base_velocity[..., 0])

        def tile_batch(arr):
            if not batch_shape:
                return arr
            reps = batch_shape + (1,) * len(arr.shape)
            return math.tile(arr, reps)

        Xup = [None] * node_count
        Scols: list[ArrayLike | None] = [None] * node_count
        v = [None] * node_count
        c = [None] * node_count
        IA = [None] * node_count
        pA = [None] * node_count
        g_acc = [None] * node_count

        for idx in node_indices:
            IA[idx] = tile_batch(inertias[idx])

            if idx == root_idx:
                Xup[idx] = eye6()
                v[idx] = base_velocity_body
                c[idx] = zeros6()
                g_acc[idx] = a0_input
            else:
                parent = parent_indices[idx]
                joint = joints[idx]
                joint_idx = joint_indices[idx]
                q_i = (
                    joint_positions[..., joint_idx] if joint_idx is not None else zero_q
                )
                X_current = joint.spatial_transform(q=q_i)
                Xup[idx] = X_current

                g_acc[idx] = math.mxv(X_current, g_acc[parent])

                if joint_idx is not None:
                    Si = tile_batch(motion_subspaces[idx])
                    Scols[idx] = Si
                    qd_i = joint_velocities[..., joint_idx]
                    vJ = math.vxs(Si, qd_i)
                else:
                    Scols[idx] = None
                    vJ = zeros6()

                v[idx] = math.mxv(X_current, v[parent]) + vJ
                c[idx] = math.mxv(math.spatial_skew(v[idx]), vJ)

            pA[idx] = math.mxv(
                math.spatial_skew_star(v[idx]), math.mxv(IA[idx], v[idx])
            )

        d_list: list[ArrayLike | None] = [None] * node_count
        inv_d_list: list[ArrayLike | None] = [None] * node_count
        u_list: list[ArrayLike | None] = [None] * node_count
        U_list: list[ArrayLike | None] = [None] * node_count

        for idx in rev_indices:
            if idx == root_idx:
                continue

            parent = parent_indices[idx]
            Xpt = math.swapaxes(Xup[idx], -2, -1)
            Si = Scols[idx]
            if Scols[idx] is not None:
                U_i = math.mtimes(IA[idx], Si)
                d_i = math.mtimes(math.swapaxes(Si, -2, -1), U_i)
                joint_idx = joint_indices[idx]
                tau_vec = joint_torques_eff[..., joint_idx]
                Si_T_pA = math.mxv(math.swapaxes(Si, -2, -1), pA[idx])[..., 0]
                u_i = tau_vec - Si_T_pA

                d_list[idx] = d_i
                u_list[idx] = u_i
                U_list[idx] = U_i

                inv_d = math.inv(d_i)
                inv_d_list[idx] = inv_d
                Ia = IA[idx] - math.mtimes(
                    U_i, math.mtimes(inv_d, math.swapaxes(U_i, -2, -1))
                )
                gain = math.mtimes(inv_d, expand_to_match(u_i, inv_d))
                gain_vec = gain[..., 0]
                pa = pA[idx] + math.mxv(Ia, c[idx]) + math.mxv(U_i, gain_vec)
            else:
                Ia = IA[idx]
                pa = pA[idx] + math.mxv(Ia, c[idx])

            IA[parent] = IA[parent] + math.mtimes(math.mtimes(Xpt, Ia), Xup[idx])
            pA[parent] = pA[parent] + math.mxv(Xpt, pa)

        rhs_root = base_ext_body - pA[root_idx] + math.mxv(IA[root_idx], a0_input)
        a_base = math.solve(IA[root_idx], rhs_root)

        a = [None] * node_count
        a[root_idx] = a_base

        qdd_entries: list[ArrayLike | None] = [None] * n if n > 0 else []

        for idx in node_indices:
            if idx == root_idx:
                continue

            parent = parent_indices[idx]
            a_pre = math.mxv(Xup[idx], a[parent]) + c[idx]
            free_acc = g_acc[idx]
            rel_acc = a_pre - free_acc if free_acc is not None else a_pre

            Si = Scols[idx]
            joint_idx = joint_indices[idx]

            if Si is not None and joint_idx is not None:
                U_i = U_list[idx]
                U_T_rel_acc = math.mxv(math.swapaxes(U_i, -2, -1), rel_acc)[..., 0]
                num = u_list[idx] - U_T_rel_acc
                inv_d = inv_d_list[idx]
                num_expanded = expand_to_match(num, inv_d)
                gain_qdd = math.mtimes(inv_d, num_expanded)
                qdd_col = gain_qdd[..., 0]
                if joint_idx < n:
                    qdd_entries[joint_idx] = qdd_col
                a_correction_vec = math.mxv(Si, qdd_col)
                a[idx] = a_pre + a_correction_vec
            else:
                a[idx] = a_pre

        if n > 0:
            zero_col = math.factory.zeros(batch_shape + (1,))
            qdd_cols = [
                entry if entry is not None else zero_col for entry in qdd_entries
            ]
            joint_qdd = math.concatenate(qdd_cols, axis=-1)
        else:
            joint_qdd = math.factory.zeros(batch_shape + (0,))

        if self.frame_velocity_representation == Representations.MIXED_REPRESENTATION:
            Xm = math.adjoint_mixed(base_transform)
            base_vel_mixed = math.mxv(Xm, base_velocity_body)
            Xm_dot = math.adjoint_mixed_derivative(base_transform, base_vel_mixed)
            base_acc = math.mxv(Xm, a_base) + math.mxv(Xm_dot, base_velocity_body)
        elif (
            self.frame_velocity_representation
            == Representations.INERTIAL_FIXED_REPRESENTATION
        ):
            X = math.adjoint(base_transform)
            base_vel_inertial = math.mxv(X, base_velocity_body)
            X_dot = math.adjoint_derivative(base_transform, base_vel_inertial)
            base_acc = math.mxv(X, a_base) + math.mxv(X_dot, base_velocity_body)
        else:
            base_acc = a_base

        return math.concatenate([base_acc, joint_qdd], axis=-1)

    def _convert_to_arraylike(self, *args):
        """Convert inputs to ArrayLike if they are not already.
        Args:
            *args: Input arguments to be converted.
        Returns:
            Converted arguments as ArrayLike.
        """
        if not args:
            raise ValueError("At least one argument is required")

        converted = []
        for arg in args:
            if isinstance(arg, ArrayLike):
                converted.append(arg)
            else:
                converted.append(self.math.asarray(arg))
        return converted[0] if len(converted) == 1 else converted

    def _prepare_tree_cache(self) -> None:
        """Pre-compute static tree data so the dynamic algorithms avoid repeated Python work."""
        nodes = list(self.model.tree)
        node_count = len(nodes)
        self._node_indices = tuple(range(node_count))
        self._rev_node_indices = tuple(reversed(self._node_indices))
        self._parent_indices = [-1] * node_count
        self._joint_indices_per_node: list[int | None] = [None] * node_count
        self._motion_subspaces: list[npt.ArrayLike] = [None] * node_count
        self._spatial_inertias: list[npt.ArrayLike] = [None] * node_count
        self._joints_per_node: list[Joint | None] = [None] * node_count
        self._joint_index_to_node: dict[int, int] = {}

        for idx, node in enumerate(nodes):
            link, joint, parent_link = node.get_elements()
            self._joints_per_node[idx] = joint
            self._spatial_inertias[idx] = link.spatial_inertia()
            parent_idx = (
                self.model.tree.get_idx_from_name(parent_link.name)
                if parent_link is not None
                else -1
            )
            self._parent_indices[idx] = parent_idx
            if joint is None:
                self._motion_subspaces[idx] = self._root_motion_subspace
                self._joint_indices_per_node[idx] = None
            else:
                self._motion_subspaces[idx] = joint.motion_subspace()
                self._joint_indices_per_node[idx] = joint.idx
                if joint.idx is not None:
                    self._joint_index_to_node[int(joint.idx)] = idx

        self._root_index = self.model.tree.get_idx_from_name(self.root_link)
        self._node_count = node_count
