featherstone 0.1.0

//! Articulated body tree data structure
//!
//! Represents a kinematic tree of rigid bodies connected by joints using
//! generalized (minimal) coordinates. This is the input to all Featherstone
//! algorithms (ABA, RNEA, CRBA).
//!
//! Parent-indexed tree: body i's parent is `parent[i]`, root has parent = -1.
//! Supports branching topologies (humanoid, quadruped) not just serial chains.
//!
//! Position state (q) may have different dimension than velocity state (qd) when
//! quaternion-based joints are used (Spherical: 4 pos params / 3 vel DOF,
//! Floating: 7 pos params / 6 vel DOF).

use nalgebra::{DVector, Vector3};

use super::joint::GenJoint;
use super::spatial::{SpatialInertia, SpatialTransform, SpatialVector};

/// Backward-compatible alias for the joint type enum.
pub type GenJointType = GenJoint;

// ============================================================================
// Body definition
// ============================================================================

/// A single body in the articulated body tree.
#[derive(Clone, Debug)]
pub struct BodyDef {
    /// Body name (from URDF link name)
    pub name: String,
    /// Parent body index (-1 for root)
    pub parent: i32,
    /// Joint connecting this body to parent
    pub joint_type: GenJoint,
    /// Spatial inertia of this body
    pub inertia: SpatialInertia,
    /// Fixed transform from parent body frame to this joint frame
    pub x_tree: SpatialTransform,
    /// Starting index into the generalized position vector q
    pub q_index: usize,
    /// Starting index into the generalized velocity vector qd
    pub v_index: usize,
}

// ============================================================================
// ArticulatedBody
// ============================================================================

/// A kinematic tree of rigid bodies connected by joints.
///
/// Uses generalized (minimal) coordinates — no constraint drift by construction.
/// Supports branching topologies: serial chains, humanoids, quadrupeds.
///
/// # State
/// - `q`: generalized positions (dimension = total_nq, may differ from total_nv for quaternion joints)
/// - `qd`: generalized velocities (dimension = total_nv)
/// - `qdd`: generalized accelerations (dimension = total_nv)
/// - `tau`: applied joint torques/forces (dimension = total_nv)
#[derive(Clone, Debug)]
pub struct ArticulatedBody {
    /// Body definitions (index 0 = first body, not root/world)
    pub bodies: Vec<BodyDef>,
    /// Children of each body (index -> list of child indices)
    children: Vec<Vec<usize>>,
    /// Total position parameters (sum of n_pos for all joints)
    total_nq: usize,
    /// Total velocity DOF (sum of dof for all joints)
    total_nv: usize,

    // === State vectors ===
    /// Generalized positions (dimension = total_nq)
    pub q: DVector<f32>,
    /// Generalized velocities (dimension = total_nv)
    pub qd: DVector<f32>,
    /// Generalized accelerations (dimension = total_nv, output of ABA)
    pub qdd: DVector<f32>,
    /// Applied joint torques/forces (dimension = total_nv)
    pub tau: DVector<f32>,
    /// External spatial forces on each body (in body frame)
    pub f_ext: Vec<SpatialVector>,

    /// Kahan compensation vector for position integration (dimension = total_nq).
    /// Tracks the low-order bits lost during `q += qd * dt` accumulation,
    /// reducing f32 round-off from ~0.15% to ~0.01% over long simulations.
    pub q_compensation: DVector<f32>,

    /// Kahan compensation vector for velocity integration (dimension = total_nv).
    /// Reduces f32 round-off during `qd += qdd * dt` accumulation.
    pub qd_compensation: DVector<f32>,

    /// Gravity vector (spatial: [0; -g] in Featherstone convention)
    pub gravity: SpatialVector,
}

impl Default for ArticulatedBody {
    fn default() -> Self {
        Self::new()
    }
}

impl ArticulatedBody {
    /// Create an empty articulated body
    pub fn new() -> Self {
        Self {
            bodies: Vec::new(),
            children: Vec::new(),
            total_nq: 0,
            total_nv: 0,
            q: DVector::zeros(0),
            qd: DVector::zeros(0),
            qdd: DVector::zeros(0),
            tau: DVector::zeros(0),
            f_ext: Vec::new(),
            q_compensation: DVector::zeros(0),
            qd_compensation: DVector::zeros(0),
            // Featherstone convention: a_0 = -g (base accelerates upward)
            // Earth gravity g = (0, -9.81, 0), so a_0 = (0, 9.81, 0)
            gravity: SpatialVector::new(Vector3::zeros(), Vector3::new(0.0, 9.81, 0.0)),
        }
    }

    /// Set gravity vector (linear acceleration, e.g. [0, -9.81, 0])
    pub fn set_gravity(&mut self, g: Vector3<f32>) {
        // In Featherstone convention, gravity is modeled as base acceleration
        // a_0 = -g (the world accelerates upward)
        self.gravity = SpatialVector::new(Vector3::zeros(), -g);
    }

    /// Add a body to the tree.
    ///
    /// # Arguments
    /// - `name`: Body name (from URDF link)
    /// - `parent`: Parent body index, or -1 for root-level body
    /// - `joint_type`: Joint connecting this body to parent
    /// - `inertia`: Spatial inertia of this body
    /// - `x_tree`: Fixed transform from parent body to this joint
    ///
    /// # Returns
    /// Index of the newly added body
    pub fn add_body(
        &mut self,
        name: impl Into<String>,
        parent: i32,
        joint_type: GenJoint,
        inertia: SpatialInertia,
        x_tree: SpatialTransform,
    ) -> usize {
        let body_index = self.bodies.len();
        let q_index = self.total_nq;
        let v_index = self.total_nv;
        let npos = joint_type.n_pos();
        let dof = joint_type.dof();

        // Get default position values for this joint (e.g. unit quaternion)
        let default_q = joint_type.default_q();

        self.bodies.push(BodyDef {
            name: name.into(),
            parent,
            joint_type,
            inertia,
            x_tree,
            q_index,
            v_index,
        });

        // Update children list
        self.children.push(Vec::new());
        if parent >= 0 {
            let pidx = parent as usize;
            assert!(
                pidx < body_index,
                "add_body: parent index {} out of range (0..{})",
                pidx, body_index
            );
            self.children[pidx].push(body_index);
        }

        // Expand state vectors
        self.total_nq += npos;
        self.total_nv += dof;

        // Resize q and copy default values
        let mut new_q = DVector::zeros(self.total_nq);
        for i in 0..q_index {
            new_q[i] = self.q[i];
        }
        for (i, &val) in default_q.iter().enumerate() {
            new_q[q_index + i] = val;
        }
        self.q = new_q;

        // Resize Kahan compensation vector (zero-initialized for new elements)
        let mut new_comp = DVector::zeros(self.total_nq);
        for i in 0..q_index {
            new_comp[i] = self.q_compensation[i];
        }
        self.q_compensation = new_comp;

        // Resize velocity vectors (zero-initialized is correct)
        let mut new_qd = DVector::zeros(self.total_nv);
        let mut new_qdd = DVector::zeros(self.total_nv);
        let mut new_tau = DVector::zeros(self.total_nv);
        for i in 0..v_index {
            new_qd[i] = self.qd[i];
            new_qdd[i] = self.qdd[i];
            new_tau[i] = self.tau[i];
        }
        self.qd = new_qd;
        self.qdd = new_qdd;
        self.tau = new_tau;

        // Resize Kahan velocity compensation
        let mut new_qd_comp = DVector::zeros(self.total_nv);
        for i in 0..v_index {
            new_qd_comp[i] = self.qd_compensation[i];
        }
        self.qd_compensation = new_qd_comp;

        self.f_ext.push(SpatialVector::zero());

        body_index
    }

    /// Number of bodies in the tree
    pub fn body_count(&self) -> usize {
        self.bodies.len()
    }

    /// Total velocity degrees of freedom (dimension of qd, qdd, tau)
    pub fn dof_count(&self) -> usize {
        self.total_nv
    }

    /// Total position parameters (dimension of q)
    pub fn nq(&self) -> usize {
        self.total_nq
    }

    /// Get parent index for body i (-1 for root-level bodies)
    pub fn parent(&self, i: usize) -> i32 {
        self.bodies[i].parent
    }

    /// Get children of body i
    pub fn children(&self, i: usize) -> &[usize] {
        &self.children[i]
    }

    /// Get body by name
    pub fn body_by_name(&self, name: &str) -> Option<usize> {
        self.bodies.iter().position(|b| b.name == name)
    }

    /// Get the q slice for body i's joint (position parameters)
    pub fn joint_q(&self, i: usize) -> &[f32] {
        let body = &self.bodies[i];
        let npos = body.joint_type.n_pos();
        if npos == 0 {
            return &[];
        }
        &self.q.as_slice()[body.q_index..body.q_index + npos]
    }

    /// Get the qd slice for body i's joint (velocity DOF)
    pub fn joint_qd(&self, i: usize) -> &[f32] {
        let body = &self.bodies[i];
        let dof = body.joint_type.dof();
        if dof == 0 {
            return &[];
        }
        &self.qd.as_slice()[body.v_index..body.v_index + dof]
    }

    /// Set joint position for body i
    pub fn set_joint_q(&mut self, i: usize, values: &[f32]) {
        let body = &self.bodies[i];
        let npos = body.joint_type.n_pos();
        let start = body.q_index;
        for (j, &val) in values.iter().take(npos).enumerate() {
            self.q[start + j] = val;
            // Reset Kahan compensation — fresh value, not accumulated
            self.q_compensation[start + j] = 0.0;
        }
    }

    /// Set joint velocity for body i
    pub fn set_joint_qd(&mut self, i: usize, values: &[f32]) {
        let body = &self.bodies[i];
        let dof = body.joint_type.dof();
        let start = body.v_index;
        for (j, &val) in values.iter().take(dof).enumerate() {
            self.qd[start + j] = val;
            // Reset Kahan compensation — fresh value, not accumulated
            self.qd_compensation[start + j] = 0.0;
        }
    }

    /// Set joint torque for body i
    pub fn set_joint_tau(&mut self, i: usize, values: &[f32]) {
        let body = &self.bodies[i];
        let dof = body.joint_type.dof();
        let start = body.v_index;
        for (j, &val) in values.iter().take(dof).enumerate() {
            self.tau[start + j] = val;
        }
    }

    /// Set external force on body i (in body frame)
    pub fn set_external_force(&mut self, i: usize, force: SpatialVector) {
        self.f_ext[i] = force;
    }

    /// Clear all external forces
    pub fn clear_external_forces(&mut self) {
        for f in &mut self.f_ext {
            *f = SpatialVector::zero();
        }
    }

    /// Reset all state to zero (positions get default values for quaternion joints)
    pub fn reset(&mut self) {
        // Reset positions to defaults (unit quaternion for quat joints)
        for body in &self.bodies {
            let default_q = body.joint_type.default_q();
            for (j, &val) in default_q.iter().enumerate() {
                self.q[body.q_index + j] = val;
            }
        }
        self.qd.fill(0.0);
        self.qdd.fill(0.0);
        self.tau.fill(0.0);
        self.q_compensation.fill(0.0);
        self.qd_compensation.fill(0.0);
        self.clear_external_forces();
    }

    /// Normalize all quaternion-based joint positions (call after integration)
    pub fn normalize_quaternions(&mut self) {
        for body in &self.bodies {
            let npos = body.joint_type.n_pos();
            if npos > 0 {
                let start = body.q_index;
                let q_slice = &mut self.q.as_mut_slice()[start..start + npos];
                body.joint_type.normalize_q(q_slice);
            }
        }
    }

    /// Iterate bodies in base-to-tip order (forward pass: 0, 1, 2, ...)
    pub fn iter_base_to_tip(&self) -> impl Iterator<Item = usize> {
        0..self.bodies.len()
    }

    /// Iterate bodies in tip-to-base order (backward pass: n-1, n-2, ..., 0)
    pub fn iter_tip_to_base(&self) -> impl Iterator<Item = usize> {
        (0..self.bodies.len()).rev()
    }

    /// Check if body i is a root-level body (parent = -1)
    pub fn is_root(&self, i: usize) -> bool {
        self.bodies[i].parent < 0
    }

    /// Check if body i is a leaf (no children)
    pub fn is_leaf(&self, i: usize) -> bool {
        self.children[i].is_empty()
    }

    /// Get the joint transform for body i at current q
    pub fn joint_transform(&self, i: usize) -> SpatialTransform {
        let body = &self.bodies[i];
        body.joint_type.transform(self.joint_q(i))
    }

    /// Get the motion subspace for body i
    pub fn joint_motion_subspace(&self, i: usize) -> Vec<SpatialVector> {
        let body = &self.bodies[i];
        body.joint_type.motion_subspace(self.joint_q(i))
    }

    /// Get the total transform from world to body i (compute on the fly)
    pub fn world_transform(&self, i: usize) -> SpatialTransform {
        let body = &self.bodies[i];
        let x_j = body.joint_type.transform(self.joint_q(i));
        let x_local = x_j.compose(&body.x_tree);

        if body.parent < 0 {
            x_local
        } else {
            let x_parent = self.world_transform(body.parent as usize);
            x_parent.compose(&x_local)
        }
    }
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;
    use nalgebra::{Matrix3, Vector3};

    fn make_simple_inertia(mass: f32) -> SpatialInertia {
        SpatialInertia::from_mass_inertia_at_com(mass, Matrix3::identity() * 0.01 * mass)
    }

    #[test]
    fn test_empty_body() {
        let ab = ArticulatedBody::new();
        assert_eq!(ab.body_count(), 0);
        assert_eq!(ab.dof_count(), 0);
        assert_eq!(ab.nq(), 0);
    }

    #[test]
    fn test_single_revolute() {
        let mut ab = ArticulatedBody::new();
        let idx = ab.add_body(
            "link1",
            -1,
            GenJoint::Revolute {
                axis: Vector3::z(),
            },
            make_simple_inertia(1.0),
            SpatialTransform::identity(),
        );
        assert_eq!(idx, 0);
        assert_eq!(ab.body_count(), 1);
        assert_eq!(ab.dof_count(), 1);
        assert_eq!(ab.nq(), 1);
        assert!(ab.is_root(0));
        assert!(ab.is_leaf(0));
    }

    #[test]
    fn test_serial_chain_6dof() {
        let mut ab = ArticulatedBody::new();
        for i in 0..6 {
            let parent = if i == 0 { -1 } else { (i - 1) as i32 };
            ab.add_body(
                format!("link{i}"),
                parent,
                GenJoint::Revolute {
                    axis: Vector3::z(),
                },
                make_simple_inertia(1.0),
                SpatialTransform::from_translation(Vector3::new(0.1, 0.0, 0.0)),
            );
        }
        assert_eq!(ab.body_count(), 6);
        assert_eq!(ab.dof_count(), 6);
        assert_eq!(ab.nq(), 6);
        assert!(ab.is_root(0));
        assert!(!ab.is_root(1));
        assert!(ab.is_leaf(5));
        assert!(!ab.is_leaf(0));
        assert_eq!(ab.children(0), &[1]);
        assert_eq!(ab.children(4), &[5]);
    }

    #[test]
    fn test_branching_tree() {
        let mut ab = ArticulatedBody::new();
        let torso = ab.add_body(
            "torso",
            -1,
            GenJoint::Fixed,
            make_simple_inertia(10.0),
            SpatialTransform::identity(),
        );
        let left = ab.add_body(
            "left_arm",
            torso as i32,
            GenJoint::Revolute {
                axis: Vector3::z(),
            },
            make_simple_inertia(1.0),
            SpatialTransform::from_translation(Vector3::new(0.0, 0.3, 0.0)),
        );
        let right = ab.add_body(
            "right_arm",
            torso as i32,
            GenJoint::Revolute {
                axis: Vector3::z(),
            },
            make_simple_inertia(1.0),
            SpatialTransform::from_translation(Vector3::new(0.0, -0.3, 0.0)),
        );

        assert_eq!(ab.body_count(), 3);
        assert_eq!(ab.dof_count(), 2);
        assert_eq!(ab.children(torso), &[left, right]);
    }

    #[test]
    fn test_floating_base() {
        let mut ab = ArticulatedBody::new();
        ab.add_body(
            "base",
            -1,
            GenJoint::Floating,
            make_simple_inertia(5.0),
            SpatialTransform::identity(),
        );
        ab.add_body(
            "arm",
            0,
            GenJoint::Revolute {
                axis: Vector3::z(),
            },
            make_simple_inertia(1.0),
            SpatialTransform::from_translation(Vector3::new(0.2, 0.0, 0.0)),
        );

        assert_eq!(ab.dof_count(), 7); // 6 (floating) + 1 (revolute)
        assert_eq!(ab.nq(), 8); // 7 (floating: 3+4) + 1 (revolute)
    }

    #[test]
    fn test_spherical_dimensions() {
        let mut ab = ArticulatedBody::new();
        ab.add_body(
            "ball",
            -1,
            GenJoint::Spherical,
            make_simple_inertia(1.0),
            SpatialTransform::identity(),
        );
        assert_eq!(ab.dof_count(), 3);
        assert_eq!(ab.nq(), 4);

        // Default q should be unit quaternion
        let q = ab.joint_q(0);
        assert_eq!(q.len(), 4);
        assert!((q[0] - 1.0).abs() < 1e-10); // w=1
    }

    #[test]
    fn test_joint_dof_types() {
        assert_eq!(GenJoint::Fixed.dof(), 0);
        assert_eq!(
            GenJoint::Revolute {
                axis: Vector3::z()
            }
            .dof(),
            1
        );
        assert_eq!(
            GenJoint::Prismatic {
                axis: Vector3::x()
            }
            .dof(),
            1
        );
        assert_eq!(GenJoint::Spherical.dof(), 3);
        assert_eq!(GenJoint::Floating.dof(), 6);
        assert_eq!(
            GenJoint::Planar {
                normal: Vector3::z()
            }
            .dof(),
            3
        );
    }

    #[test]
    fn test_set_get_joint_state() {
        let mut ab = ArticulatedBody::new();
        ab.add_body(
            "link1",
            -1,
            GenJoint::Revolute {
                axis: Vector3::z(),
            },
            make_simple_inertia(1.0),
            SpatialTransform::identity(),
        );
        ab.add_body(
            "link2",
            0,
            GenJoint::Revolute {
                axis: Vector3::z(),
            },
            make_simple_inertia(1.0),
            SpatialTransform::from_translation(Vector3::new(0.5, 0.0, 0.0)),
        );

        ab.set_joint_q(0, &[1.0]);
        ab.set_joint_q(1, &[2.0]);
        assert_eq!(ab.joint_q(0), &[1.0]);
        assert_eq!(ab.joint_q(1), &[2.0]);

        ab.set_joint_qd(0, &[0.5]);
        assert_eq!(ab.joint_qd(0), &[0.5]);
    }

    #[test]
    fn test_body_by_name() {
        let mut ab = ArticulatedBody::new();
        ab.add_body(
            "link_a",
            -1,
            GenJoint::Revolute {
                axis: Vector3::z(),
            },
            make_simple_inertia(1.0),
            SpatialTransform::identity(),
        );
        ab.add_body(
            "link_b",
            0,
            GenJoint::Revolute {
                axis: Vector3::x(),
            },
            make_simple_inertia(1.0),
            SpatialTransform::identity(),
        );

        assert_eq!(ab.body_by_name("link_a"), Some(0));
        assert_eq!(ab.body_by_name("link_b"), Some(1));
        assert_eq!(ab.body_by_name("nonexistent"), None);
    }

    #[test]
    fn test_reset() {
        let mut ab = ArticulatedBody::new();
        ab.add_body(
            "link1",
            -1,
            GenJoint::Revolute {
                axis: Vector3::z(),
            },
            make_simple_inertia(1.0),
            SpatialTransform::identity(),
        );
        ab.set_joint_q(0, &[1.5]);
        ab.set_joint_qd(0, &[0.5]);
        ab.reset();
        assert_eq!(ab.joint_q(0), &[0.0]);
        assert_eq!(ab.joint_qd(0), &[0.0]);
    }

    #[test]
    fn test_reset_spherical() {
        let mut ab = ArticulatedBody::new();
        ab.add_body(
            "ball",
            -1,
            GenJoint::Spherical,
            make_simple_inertia(1.0),
            SpatialTransform::identity(),
        );
        ab.set_joint_q(0, &[0.5, 0.5, 0.5, 0.5]);
        ab.reset();
        let q = ab.joint_q(0);
        assert!((q[0] - 1.0).abs() < 1e-10); // w=1 (identity quat)
        assert!(q[1].abs() < 1e-10);
    }

    #[test]
    fn test_revolute_transform() {
        let jt = GenJoint::Revolute {
            axis: Vector3::z(),
        };
        let x = jt.transform(&[std::f32::consts::FRAC_PI_2]);
        let p = x.rotation * Vector3::x();
        assert!((p.y - 1.0).abs() < 1e-5);
        assert!(p.x.abs() < 1e-5);
    }

    #[test]
    fn test_prismatic_transform() {
        let jt = GenJoint::Prismatic {
            axis: Vector3::x(),
        };
        let x = jt.transform(&[0.5]);
        assert!((x.translation.x - 0.5).abs() < 1e-10);
        assert!(x.translation.y.abs() < 1e-10);
    }

    #[test]
    fn test_fixed_transform_identity() {
        let jt = GenJoint::Fixed;
        let x = jt.transform(&[]);
        assert!((x.rotation - nalgebra::Matrix3::identity()).norm() < 1e-10);
    }

    #[test]
    fn test_motion_subspace_revolute() {
        let jt = GenJoint::Revolute {
            axis: Vector3::z(),
        };
        let s = jt.motion_subspace(&[0.0]);
        assert_eq!(s.len(), 1);
        assert!((s[0].angular().z - 1.0).abs() < 1e-10);
        assert!(s[0].linear().norm() < 1e-10);
    }

    #[test]
    fn test_motion_subspace_prismatic() {
        let jt = GenJoint::Prismatic {
            axis: Vector3::y(),
        };
        let s = jt.motion_subspace(&[0.0]);
        assert!(s[0].angular().norm() < 1e-10);
        assert!((s[0].linear().y - 1.0).abs() < 1e-10);
    }

    #[test]
    fn test_motion_subspace_floating() {
        let jt = GenJoint::Floating;
        let s = jt.motion_subspace(&[0.0; 7]);
        assert_eq!(s.len(), 6);
    }

    #[test]
    fn test_iteration_order() {
        let mut ab = ArticulatedBody::new();
        for i in 0..5 {
            let parent = if i == 0 { -1 } else { (i - 1) as i32 };
            ab.add_body(
                format!("link{i}"),
                parent,
                GenJoint::Revolute {
                    axis: Vector3::z(),
                },
                make_simple_inertia(1.0),
                SpatialTransform::identity(),
            );
        }

        let fwd: Vec<usize> = ab.iter_base_to_tip().collect();
        assert_eq!(fwd, vec![0, 1, 2, 3, 4]);

        let bwd: Vec<usize> = ab.iter_tip_to_base().collect();
        assert_eq!(bwd, vec![4, 3, 2, 1, 0]);
    }

    #[test]
    fn test_quadruped_topology() {
        let mut ab = ArticulatedBody::new();
        let base = ab.add_body(
            "base",
            -1,
            GenJoint::Floating,
            make_simple_inertia(10.0),
            SpatialTransform::identity(),
        );
        for leg_name in &["fl", "fr", "rl", "rr"] {
            let hip = ab.add_body(
                format!("{leg_name}_hip"),
                base as i32,
                GenJoint::Revolute {
                    axis: Vector3::y(),
                },
                make_simple_inertia(0.5),
                SpatialTransform::identity(),
            );
            ab.add_body(
                format!("{leg_name}_knee"),
                hip as i32,
                GenJoint::Revolute {
                    axis: Vector3::y(),
                },
                make_simple_inertia(0.3),
                SpatialTransform::from_translation(Vector3::new(0.0, 0.0, -0.2)),
            );
        }
        assert_eq!(ab.body_count(), 9);
        assert_eq!(ab.dof_count(), 14); // 6 (floating) + 8 (4 legs * 2 revolute)
        assert_eq!(ab.nq(), 15); // 7 (floating) + 8 (revolute)
        assert_eq!(ab.children(base).len(), 4);
    }

    #[test]
    fn test_normalize_quaternions() {
        let mut ab = ArticulatedBody::new();
        ab.add_body(
            "ball",
            -1,
            GenJoint::Spherical,
            make_simple_inertia(1.0),
            SpatialTransform::identity(),
        );
        // Set non-normalized quaternion
        ab.set_joint_q(0, &[2.0, 0.0, 0.0, 0.0]);
        ab.normalize_quaternions();
        let q = ab.joint_q(0);
        let norm = (q[0] * q[0] + q[1] * q[1] + q[2] * q[2] + q[3] * q[3]).sqrt();
        assert!((norm - 1.0).abs() < 1e-7);
    }

    #[test]
    fn test_mixed_nq_nv() {
        let mut ab = ArticulatedBody::new();
        // Spherical (nq=4, nv=3) + Revolute (nq=1, nv=1) + Floating (nq=7, nv=6)
        ab.add_body(
            "ball",
            -1,
            GenJoint::Spherical,
            make_simple_inertia(1.0),
            SpatialTransform::identity(),
        );
        ab.add_body(
            "arm",
            0,
            GenJoint::Revolute {
                axis: Vector3::z(),
            },
            make_simple_inertia(0.5),
            SpatialTransform::from_translation(Vector3::new(0.3, 0.0, 0.0)),
        );

        assert_eq!(ab.nq(), 5);  // 4 + 1
        assert_eq!(ab.dof_count(), 4); // 3 + 1

        // Verify indexing
        assert_eq!(ab.bodies[0].q_index, 0);
        assert_eq!(ab.bodies[0].v_index, 0);
        assert_eq!(ab.bodies[1].q_index, 4);
        assert_eq!(ab.bodies[1].v_index, 3);

        // Set and get joint states
        ab.set_joint_q(0, &[0.707, 0.0, 0.0, 0.707]);
        ab.set_joint_q(1, &[1.5]);
        assert_eq!(ab.joint_q(0).len(), 4);
        assert_eq!(ab.joint_q(1), &[1.5]);

        ab.set_joint_qd(0, &[0.1, 0.2, 0.3]);
        ab.set_joint_qd(1, &[0.5]);
        assert_eq!(ab.joint_qd(0), &[0.1, 0.2, 0.3]);
        assert_eq!(ab.joint_qd(1), &[0.5]);
    }

    // ======================================================================
    // External force and tree operation tests
    // ======================================================================

    #[test]
    fn test_set_and_clear_external_forces() {
        let mut ab = ArticulatedBody::new();
        ab.add_body(
            "link1", -1,
            GenJoint::Revolute { axis: Vector3::z() },
            make_simple_inertia(1.0),
            SpatialTransform::identity(),
        );
        ab.add_body(
            "link2", 0,
            GenJoint::Revolute { axis: Vector3::z() },
            make_simple_inertia(1.0),
            SpatialTransform::from_translation(Vector3::new(1.0, 0.0, 0.0)),
        );

        // Apply force to link 1
        let force = SpatialVector::new(Vector3::zeros(), Vector3::new(0.0, -9.81, 0.0));
        ab.set_external_force(1, force);

        // Verify it's set
        assert!(ab.f_ext[1].data.norm() > 0.0);

        // Clear all
        ab.clear_external_forces();
        for f in &ab.f_ext {
            assert!(f.data.norm() < 1e-10, "External force should be zero after clear");
        }
    }

    #[test]
    fn test_body_count_matches_add() {
        let mut ab = ArticulatedBody::new();
        assert_eq!(ab.body_count(), 0);

        ab.add_body("a", -1, GenJoint::Fixed, make_simple_inertia(1.0), SpatialTransform::identity());
        assert_eq!(ab.body_count(), 1);

        ab.add_body("b", 0, GenJoint::Fixed, make_simple_inertia(1.0), SpatialTransform::identity());
        assert_eq!(ab.body_count(), 2);
    }

    #[test]
    fn test_set_gravity() {
        let mut ab = ArticulatedBody::new();
        // Featherstone convention: stored as base accel = -g
        ab.set_gravity(Vector3::new(0.0, -9.81, 0.0));
        // Stored value should be (0, 9.81, 0) (negated for base acceleration)
        assert!(ab.gravity.linear().y > 9.0, "Gravity stored as base accel: {:?}", ab.gravity.linear());

        // Change to moon gravity
        ab.set_gravity(Vector3::new(0.0, -1.62, 0.0));
        assert!((ab.gravity.linear().y - 1.62).abs() < 0.01);
    }

    // ── SLAM Cycle 6: Body intent + error path tests ─────────────────

    #[test]
    fn intent_body_nq_equals_sum_of_joint_npos() {
        // Intent: position vector length = sum of each joint's n_pos()
        let mut body = ArticulatedBody::new();
        body.add_body("l1", -1, GenJointType::Revolute { axis: Vector3::z() },
            SpatialInertia::sphere(1.0, 0.1), SpatialTransform::identity());
        body.add_body("l2", 0, GenJointType::Spherical,
            SpatialInertia::sphere(1.0, 0.1),
            SpatialTransform::from_translation(Vector3::new(0.0, -1.0, 0.0)));

        // Revolute = 1 pos param, Spherical = 4 pos params (quaternion)
        assert_eq!(body.q.len(), 5, "nq should be 1 + 4 = 5");
        // Revolute = 1 DOF, Spherical = 3 DOF
        assert_eq!(body.qd.len(), 4, "nv should be 1 + 3 = 4");
    }

    #[test]
    fn intent_body_empty_has_zero_dimensions() {
        let body = ArticulatedBody::new();
        assert_eq!(body.q.len(), 0);
        assert_eq!(body.qd.len(), 0);
        assert_eq!(body.qdd.len(), 0);
        assert_eq!(body.tau.len(), 0);
        assert_eq!(body.body_count(), 0);
        assert_eq!(body.dof_count(), 0);
    }

    #[test]
    fn intent_body_branching_tree_correct_count() {
        // Humanoid-like branching: root → left_arm + right_arm
        let mut body = ArticulatedBody::new();
        body.add_body("torso", -1, GenJointType::Fixed,
            SpatialInertia::sphere(5.0, 0.2), SpatialTransform::identity());
        body.add_body("left_arm", 0, GenJointType::Revolute { axis: Vector3::z() },
            SpatialInertia::sphere(1.0, 0.1),
            SpatialTransform::from_translation(Vector3::new(-0.5, 0.0, 0.0)));
        body.add_body("right_arm", 0, GenJointType::Revolute { axis: Vector3::z() },
            SpatialInertia::sphere(1.0, 0.1),
            SpatialTransform::from_translation(Vector3::new(0.5, 0.0, 0.0)));

        assert_eq!(body.body_count(), 3);
        assert_eq!(body.dof_count(), 2); // Fixed=0 + Revolute=1 + Revolute=1
    }

    // ── SLAM Cycle 1: Body tree intent/property tests ─────────────────

    #[test]
    fn intent_parent_index_always_less_than_child() {
        // Intent: Tree invariant — parent index must be less than child index
        let mut body = ArticulatedBody::new();
        body.add_body("root", -1, GenJointType::Fixed, SpatialInertia::zero(), SpatialTransform::identity());
        body.add_body("l1", 0, GenJointType::Revolute { axis: Vector3::z() },
            SpatialInertia::sphere(1.0, 0.1), SpatialTransform::identity());
        body.add_body("l2", 0, GenJointType::Revolute { axis: Vector3::z() },
            SpatialInertia::sphere(1.0, 0.1), SpatialTransform::identity());
        body.add_body("l3", 1, GenJointType::Revolute { axis: Vector3::z() },
            SpatialInertia::sphere(1.0, 0.1), SpatialTransform::identity());

        for i in 0..body.body_count() {
            let p = body.bodies[i].parent;
            if p >= 0 {
                assert!((p as usize) < i, "parent({i})={p} must be < {i}");
            }
        }
    }

    #[test]
    fn intent_state_dimensions_consistent() {
        // Intent: q.len() == total_nq, qd.len() == total_nv
        let mut body = ArticulatedBody::new();
        body.add_body("f", -1, GenJointType::Floating,
            SpatialInertia::sphere(1.0, 0.1), SpatialTransform::identity());
        body.add_body("r", 0, GenJointType::Revolute { axis: Vector3::z() },
            SpatialInertia::sphere(0.5, 0.1), SpatialTransform::identity());

        assert_eq!(body.q.len(), body.nq(), "q.len()={} != nq()={}", body.q.len(), body.nq());
        assert_eq!(body.qd.len(), body.dof_count(), "qd.len()={} != nv()={}", body.qd.len(), body.dof_count());
        assert_eq!(body.qdd.len(), body.dof_count());
        assert_eq!(body.tau.len(), body.dof_count());
        assert_eq!(body.f_ext.len(), body.body_count());
    }

    #[test]
    fn intent_reset_restores_default_state() {
        // Intent: reset() zeros velocities and restores default positions (including quaternions)
        let mut body = ArticulatedBody::new();
        body.add_body("link", -1, GenJointType::Spherical,
            SpatialInertia::sphere(1.0, 0.1), SpatialTransform::identity());
        body.qd[0] = 5.0;
        body.qd[1] = 3.0;
        body.q[0] = 0.5; // corrupt quaternion

        body.reset();

        // Quaternion should be identity [1, 0, 0, 0]
        assert!((body.q[0] - 1.0).abs() < 1e-6, "q[0] (w) should be 1.0 after reset, got {}", body.q[0]);
        for i in 0..body.dof_count() {
            assert!(body.qd[i].abs() < 1e-6, "qd[{i}] should be 0 after reset");
        }
    }

    // ── SLAM Cycle 6: Tests for untested public functions ─────────────

    #[test]
    fn test_set_joint_tau() {
        let mut body = ArticulatedBody::new();
        body.add_body("l1", -1, GenJointType::Revolute { axis: Vector3::z() },
            SpatialInertia::sphere(1.0, 0.1), SpatialTransform::identity());
        body.add_body("l2", 0, GenJointType::Revolute { axis: Vector3::z() },
            SpatialInertia::sphere(0.5, 0.1), SpatialTransform::identity());

        body.set_joint_tau(0, &[5.0]);
        body.set_joint_tau(1, &[-3.0]);

        assert!((body.tau[0] - 5.0).abs() < 1e-6, "tau[0]={}", body.tau[0]);
        assert!((body.tau[1] - (-3.0)).abs() < 1e-6, "tau[1]={}", body.tau[1]);
    }

    #[test]
    fn test_joint_transform_revolute() {
        let mut body = ArticulatedBody::new();
        body.add_body("link", -1, GenJointType::Revolute { axis: Vector3::z() },
            SpatialInertia::sphere(1.0, 0.1), SpatialTransform::identity());
        body.q[0] = 0.0;

        let t = body.joint_transform(0);
        let rot_diff = (t.rotation - nalgebra::Matrix3::identity()).norm();
        assert!(rot_diff < 1e-5, "q=0 should give identity transform, got diff={rot_diff}");
    }

    #[test]
    fn test_joint_motion_subspace_revolute() {
        let mut body = ArticulatedBody::new();
        body.add_body("link", -1, GenJointType::Revolute { axis: Vector3::z() },
            SpatialInertia::sphere(1.0, 0.1), SpatialTransform::identity());

        let s = body.joint_motion_subspace(0);
        assert_eq!(s.len(), 1, "Revolute joint has 1 DOF");
        // Z-axis revolute: angular part should have z component
        assert!(s[0].angular().z.abs() > 0.5, "revolute Z should have angular Z component");
    }

    #[test]
    fn test_world_transform_chain() {
        let mut body = ArticulatedBody::new();
        body.add_body("l1", -1, GenJointType::Revolute { axis: Vector3::z() },
            SpatialInertia::sphere(1.0, 0.1), SpatialTransform::identity());
        body.add_body("l2", 0, GenJointType::Revolute { axis: Vector3::z() },
            SpatialInertia::sphere(0.5, 0.1),
            SpatialTransform::from_translation(Vector3::new(0.0, -1.0, 0.0)));

        let t0 = body.world_transform(0);
        let t1 = body.world_transform(1);

        // At q=0, body 1 should be 1m below body 0
        assert!(t0.translation.norm() < 1e-4, "root should be at origin");
        assert!((t1.translation.y - (-1.0)).abs() < 0.1,
            "child should be ~1m below root: y={}", t1.translation.y);
    }

    // ── SLAM Cycle 12: Body proptests ─────────────────────────────────

    use proptest::prelude::*;

    proptest! {
        #[test]
        fn prop_body_count_matches_additions(n in 1usize..10) {
            let mut body = ArticulatedBody::new();
            for i in 0..n {
                let parent = if i == 0 { -1 } else { (i - 1) as i32 };
                body.add_body(format!("l{i}"), parent,
                    GenJoint::Revolute { axis: Vector3::z() },
                    SpatialInertia::sphere(1.0, 0.1), SpatialTransform::identity());
            }
            prop_assert_eq!(body.body_count(), n);
            prop_assert_eq!(body.dof_count(), n); // n revolute joints = n DOF
            prop_assert_eq!(body.nq(), n); // revolute: n_pos = 1 each
        }

        #[test]
        fn prop_external_forces_cleared(n in 1usize..5) {
            let mut body = ArticulatedBody::new();
            for i in 0..n {
                let parent = if i == 0 { -1 } else { (i - 1) as i32 };
                body.add_body(format!("l{i}"), parent,
                    GenJoint::Revolute { axis: Vector3::z() },
                    SpatialInertia::sphere(1.0, 0.1), SpatialTransform::identity());
            }
            for i in 0..n {
                body.set_external_force(i, SpatialVector::new(
                    Vector3::new(1.0, 2.0, 3.0), Vector3::new(4.0, 5.0, 6.0)));
            }
            body.clear_external_forces();
            for i in 0..n {
                prop_assert!(body.f_ext[i].data.norm() < 1e-10,
                    "external force[{i}] should be zero after clear");
            }
        }
    }

    #[test]
    fn edge_fixed_joint_contributes_zero_dof() {
        // Fixed joints should contribute 0 DOF but still be a body in the tree
        let mut body = ArticulatedBody::new();
        body.add_body("root", -1, GenJoint::Revolute { axis: Vector3::z() },
            SpatialInertia::sphere(1.0, 0.1), SpatialTransform::identity());
        body.add_body("fixed", 0, GenJoint::Fixed,
            SpatialInertia::sphere(2.0, 0.1),
            SpatialTransform::from_translation(Vector3::new(0.0, -0.5, 0.0)));
        body.add_body("tip", 1, GenJoint::Revolute { axis: Vector3::z() },
            SpatialInertia::sphere(0.5, 0.1),
            SpatialTransform::from_translation(Vector3::new(0.0, -0.5, 0.0)));

        assert_eq!(body.body_count(), 3, "should have 3 bodies");
        assert_eq!(body.dof_count(), 2, "Fixed joint adds 0 DOF: 1+0+1=2");
        assert_eq!(body.nq(), 2, "position params: 1+0+1=2");
    }

    #[test]
    fn intent_body_by_name_returns_correct_index() {
        let mut body = ArticulatedBody::new();
        body.add_body("alpha", -1, GenJoint::Revolute { axis: Vector3::z() },
            SpatialInertia::sphere(1.0, 0.1), SpatialTransform::identity());
        body.add_body("beta", 0, GenJoint::Revolute { axis: Vector3::z() },
            SpatialInertia::sphere(0.5, 0.1), SpatialTransform::identity());
        body.add_body("gamma", 1, GenJoint::Prismatic { axis: Vector3::y() },
            SpatialInertia::sphere(0.3, 0.1), SpatialTransform::identity());

        assert_eq!(body.body_by_name("alpha"), Some(0));
        assert_eq!(body.body_by_name("beta"), Some(1));
        assert_eq!(body.body_by_name("gamma"), Some(2));
        assert_eq!(body.body_by_name("nonexistent"), None);
    }
}