spatial-math 0.4.0-beta.1

// Copyright (C) 2020-2025 spatial-math authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::ops::{Add, AddAssign, Mul, Sub};

use super::exts::Vec3Ext;
use super::matrix::SymmetricMat3;
use super::{Mat3, SpatialForceVector, SpatialMotionVector, Vec3};
use crate::{Real, SMatrix, VEC3_ZERO};

/// Rigid body spatial inertia matrix implementing Featherstone's spatial inertia formulation.
///
/// This represents the complete 6×6 inertia matrix of a rigid body in spatial algebra,
/// combining both rotational and translational inertia properties. The spatial inertia
/// relates spatial motion to spatial force: f = I·v
///
/// # Matrix Structure
///
/// The spatial inertia matrix has the form:
/// ```text
///     | I_3×3   h× |
/// I = |            |
///     | -h×     m   |
/// ```
///
/// where:
/// - **`I_3×3`**: 3×3 rotational inertia tensor about the body-fixed frame origin
/// - **`h`**: mass × center of mass position (first moment of mass)
/// - **`m`**: scalar mass
/// - **`h×`**: cross-product matrix of h (skew-symmetric matrix)
///
/// # Physical Meaning
///
/// The spatial inertia captures both:
/// 1. **Rotational inertia**: resistance to angular acceleration
/// 2. **Linear inertia**: resistance to linear acceleration
/// 3. **Coupling effects**: how linear motion creates moments and vice versa
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[derive(Clone, Copy, Debug)]
pub struct RigidBodyInertia {
    /// Total mass of the rigid body.
    pub mass: Real,

    /// First moment of mass: `h = mass × center_of_mass`
    ///
    /// This vector couples linear and angular motion in the spatial inertia matrix.
    /// When expressed about the center of mass, this becomes zero.
    pub h: Vec3,

    /// Rotational inertia tensor expressed as a symmetric 3×3 matrix.
    ///
    /// This is the lower triangular representation of the 3×3 rotational
    /// inertia tensor about the origin of the body-fixed coordinate frame.
    pub inertia: SymmetricMat3,
}

impl Default for RigidBodyInertia {
    fn default() -> Self {
        Self::new(1.0, Vec3::zeros(), SymmetricMat3::ONE)
    }
}

impl RigidBodyInertia {
    /// Zero rigid body inertia (massless body).
    ///
    /// Useful for initialization and as the identity element
    /// for certain inertia operations.
    pub const ZERO: RigidBodyInertia = RigidBodyInertia {
        mass: 0.0,
        h: VEC3_ZERO,
        inertia: SymmetricMat3::ZERO,
    };

    /// Create a new rigid body spatial inertia from physical properties.
    ///
    /// This method constructs the spatial inertia matrix using the parallel axis theorem
    /// to shift the inertia from the center of mass to the specified reference point.
    ///
    /// # Arguments
    ///
    /// * `mass` - Total mass of the rigid body (must be positive)
    /// * `center_of_mass` - Center of mass position relative to the reference frame origin
    /// * `inertia_com` - 3×3 rotational inertia tensor about the center of mass
    ///
    /// # Mathematical Transformation
    ///
    /// The parallel axis theorem transforms inertia from center of mass (`CoM`) to reference point:
    /// ```text
    /// I_origin = I_com + m[(c·c)I - ccᵀ]
    /// ```
    /// where c is the center of mass vector and I is the 3×3 identity matrix.
    ///
    /// # Physical Interpretation
    ///
    /// - If `center_of_mass = [0, 0, 0]`, the inertia is expressed about the `CoM`
    /// - The resulting inertia relates spatial motion to spatial force about the origin
    /// - Positive mass ensures physically meaningful inertia properties
    ///
    /// # Panics
    ///
    /// Panics if mass is not positive, as negative or zero mass is non-physical.
    #[inline]
    pub fn new(mass: Real, center_of_mass: Vec3, inertia_com: SymmetricMat3) -> Self {
        assert!(mass > 0.0, "mass must be positive:{mass}");
        //TODO: check inertia_com is positive definite
        let h = mass * center_of_mass;
        let cx = center_of_mass.into_cross_matrix();
        // parallel axis theorem
        let top_left = inertia_com - SymmetricMat3::from(mass * cx * cx);
        Self {
            mass,
            h,
            inertia: top_left,
        }
    }

    /// Scale the rigid body inertia by a scalar factor.
    ///
    /// This operation multiplies all inertia components by the given scalar,
    /// effectively scaling the mass and all inertia properties proportionally.
    ///
    /// # Mathematical Meaning
    ///
    /// Given spatial inertia I and scalar s:
    /// ```text
    /// I_scaled = s · I
    /// ```
    ///
    /// This is useful for scaling body properties or creating
    /// proportional inertia variations.
    #[inline]
    #[must_use]
    pub fn scale(&self, scalar: Real) -> Self {
        Self {
            mass: self.mass * scalar,
            h: self.h * scalar,
            inertia: self.inertia.scale(scalar),
        }
    }

    /// Add two rigid body inertias.
    ///
    /// This operation combines two rigid bodies into a composite body,
    /// assuming they are expressed in the same coordinate frame and about
    /// the same reference point.
    ///
    /// # Physical Meaning
    ///
    /// The resulting inertia represents the combined properties of both bodies:
    /// - Total mass is the sum of individual masses
    /// - First moment of mass (h) is additive
    /// - Rotational inertia tensors are additive
    ///
    /// # Use Case
    ///
    /// Commonly used when constructing composite bodies or combining
    /// multiple rigid components into a single equivalent body.
    #[inline]
    #[must_use]
    pub fn add(&self, rhs: Self) -> Self {
        Self {
            mass: self.mass + rhs.mass,
            h: self.h + rhs.h,
            inertia: self.inertia.add(rhs.inertia),
        }
    }

    /// Multiply spatial inertia with a spatial motion vector to get spatial force.
    ///
    /// This implements the fundamental dynamics equation: f = I·v
    /// where f is spatial force, I is spatial inertia, and v is spatial motion.
    ///
    /// # Mathematical Formulation
    ///
    /// Given motion vector v = [ω, v] and spatial inertia I:
    /// ```text
    /// f = I · v = [n]
    ///           [f]
    ///
    /// where:
    /// n = I_3×3 · ω + h × v  (moment/torque)
    /// f = m·v - h × ω        (force)
    /// ```
    ///
    /// # Physical Interpretation
    ///
    /// - The moment n includes rotational inertia effects and coupling from linear motion
    /// - The force f includes translational inertia effects and coupling from angular motion
    /// - Cross terms (h × v and h × ω) represent the coupling between rotation and translation
    ///
    /// This is the core operation in forward dynamics for calculating forces from accelerations.
    #[inline]
    pub fn mul_vec(&self, motion_vec: SpatialMotionVector) -> SpatialForceVector {
        let motion_w = motion_vec.top;
        let motion_v = motion_vec.bottom;
        let h = self.h;
        let m = self.mass;
        let inertia = self.inertia;
        // n = I * w + h x v
        let force_n = inertia.mul_vec(motion_w) + h.cross(&motion_v);
        let force_f = m * motion_v - h.cross(&motion_w);
        SpatialForceVector::from_pair(force_n, force_f)
    }

    pub fn matrix(&self) -> SMatrix<6, 6> {
        let mut result = SMatrix::<6, 6>::zeros();

        // top left
        result
            .fixed_view_mut::<3, 3>(0, 0)
            .copy_from(&self.inertia.mat3());
        let h_crm = self.h.into_cross_matrix();

        // top right
        result.fixed_view_mut::<3, 3>(0, 3).copy_from(&h_crm);

        // bottom left
        result
            .fixed_view_mut::<3, 3>(3, 0)
            .copy_from(&h_crm.transpose());

        // bootom right
        let mass = Mat3::identity().scale(self.mass);
        result.fixed_view_mut::<3, 3>(3, 3).copy_from(&mass);

        result
    }

    /// Convert the `RigidBodyInertia` to an array of 10 elements. The data layout is:
    ///
    /// ```text
    /// | mass | h.x | h.y | h.z | inertia[0]|...|inertia[9]
    /// ```
    #[inline]
    pub fn into_array(self) -> [Real; 10] {
        let mut array = [0.0; 10];
        array[0] = self.mass;
        array[1] = self.h.x;
        array[2] = self.h.y;
        array[3] = self.h.z;
        let inertia_array = self.inertia.into_array();
        array[4] = inertia_array[0];
        array[5] = inertia_array[1];
        array[6] = inertia_array[2];
        array[7] = inertia_array[3];
        array[8] = inertia_array[4];
        array[9] = inertia_array[5];
        array
    }
}

impl From<RigidBodyInertia> for [Real; 10] {
    #[inline]
    fn from(val: RigidBodyInertia) -> Self {
        val.into_array()
    }
}

impl Mul<Real> for RigidBodyInertia {
    type Output = Self;

    #[inline]
    fn mul(self, rhs: Real) -> Self::Output {
        self.scale(rhs)
    }
}

impl Add<Self> for RigidBodyInertia {
    type Output = Self;

    #[inline]
    fn add(self, rhs: Self) -> Self::Output {
        RigidBodyInertia::add(&self, rhs)
    }
}

impl Mul<SpatialMotionVector> for RigidBodyInertia {
    type Output = SpatialForceVector;

    #[inline]
    fn mul(self, rhs: SpatialMotionVector) -> Self::Output {
        self.mul_vec(rhs)
    }
}

/// Articulated body spatial inertia matrix for complex multi-body systems.
///
/// This represents the effective inertia of an articulated body, which accounts for
/// the inertia of all bodies in the subtree beyond a given joint. Unlike rigid body
/// inertia, articulated body inertia captures the dynamic coupling between joints
/// in a kinematic chain.
///
/// # Matrix Structure
///
/// The articulated body inertia has the form:
/// ```text
///      | I   H |
/// I_A = |      |
///      | Hᵀ  M |
/// ```
///
/// where:
/// - **I**: 3×3 matrix representing rotational inertia effects
/// - **H**: 3×3 matrix representing coupling between rotation and translation
/// - **M**: 3×3 matrix representing translational inertia effects
///
/// # Physical Meaning
///
/// Articulated body inertia represents the effective inertia seen at a joint
/// when all bodies beyond that joint move together. It accounts for:
///
/// 1. **Propagated inertia**: Inertia of all bodies in the subtree
/// 2. **Dynamic coupling**: How motion at one joint affects other joints
/// 3. **Constraint forces**: Forces transmitted through the kinematic chain
///
/// # Use in Dynamics
///
/// This is a key component in Featherstone's Articulated Body Algorithm:
/// - Used in forward dynamics to compute joint accelerations
/// - Captures the influence of the entire subtree on each joint
/// - Enables O(n) complexity dynamics for tree-structured mechanisms
///
/// # Storage
///
/// The matrix stores 21 independent elements (6 for I, 9 for H, 6 for M),
/// optimized for the symmetric structure of the top-left and bottom-right blocks.
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[derive(Clone, Debug, Default)]
pub struct ArticulatedBodyInertia {
    /// Top-left 3×3 block (I) of the articulated body inertia matrix.
    ///
    /// This symmetric matrix represents rotational inertia effects and
    /// couples angular accelerations to resulting moments.
    pub top_left: SymmetricMat3,

    /// Top-right 3×3 block (H) of the articulated body inertia matrix.
    ///
    /// This general (non-symmetric) matrix represents cross-coupling
    /// between rotational and translational motion in the articulated system.
    pub h: SMatrix<3, 3>,

    /// Bottom-right 3×3 block (M) of the articulated body inertia matrix.
    ///
    /// This symmetric matrix represents translational inertia effects and
    /// couples linear accelerations to resulting forces.
    pub mass: SymmetricMat3,
}

impl ArticulatedBodyInertia {
    #[inline]
    #[must_use]
    pub fn add_arti_inertia(&self, rhs: &Self) -> Self {
        Self {
            top_left: self.top_left.add(rhs.top_left),
            h: self.h + rhs.h,
            mass: self.mass.add(rhs.mass),
        }
    }

    #[inline]
    #[must_use]
    pub fn sub_arti_inertia(&self, rhs: &Self) -> Self {
        Self {
            top_left: self.top_left - rhs.top_left,
            h: self.h - rhs.h,
            mass: self.mass - rhs.mass,
        }
    }

    /// Add a rigid body inertia to this articulated body inertia.
    ///
    /// This operation is used in the articulated body algorithm when
    /// propagating inertia from child bodies to parent bodies. It combines
    /// the effective articulated body inertia of the subtree with the
    /// rigid body inertia of the current body.
    ///
    /// # Physical Meaning
    ///
    /// This represents adding the inertia of a rigid body to an existing
    /// articulated body, effectively extending the articulated system by one body.
    /// The operation accounts for:
    /// - Adding the rigid body's mass to the translational inertia
    /// - Including cross-coupling effects from the rigid body's first moment
    /// - Adding the rigid body's rotational inertia to the system
    ///
    /// # Algorithm Context
    ///
    /// This is a crucial step in Featherstone's articulated body algorithm
    /// for computing the effective inertia at each joint in a kinematic tree.
    #[inline]
    #[must_use]
    pub fn add_spat_inertia(&self, rhs: &RigidBodyInertia) -> Self {
        let mass = self.mass.add(SymmetricMat3::ONE.scale(rhs.mass));
        let h = self.h + rhs.h.into_cross_matrix();
        let inertia = self.top_left + rhs.inertia;

        Self {
            top_left: inertia,
            h,
            mass,
        }
    }

    #[inline]
    pub fn mul_motion_vec(&self, motion_vec: SpatialMotionVector) -> SpatialForceVector {
        let motion_w = motion_vec.top;
        let motion_v = motion_vec.bottom;
        let h = self.h;
        let n = self.top_left.mul_vec(motion_w) + h * motion_v;
        //TODO: implement transpose_mul_vec for mat3
        let f = self.mass.mul_vec(motion_v) + h.transpose() * motion_w;
        SpatialForceVector::from_pair(n, f)
    }

    /// take the lower triangle part.
    #[inline]
    pub fn from_symmertic_matrix(matrix: SMatrix<6, 6>) -> Self {
        let inertia = matrix.fixed_view::<3, 3>(0, 0).into();
        let h = matrix.fixed_view::<3, 3>(0, 3).into();
        let mass = matrix.fixed_view::<3, 3>(3, 3).into();
        Self {
            top_left: inertia,
            h,
            mass,
        }
    }

    pub fn matrix(&self) -> SMatrix<6, 6> {
        let mut result = SMatrix::<6, 6>::zeros();

        // top left
        result
            .fixed_view_mut::<3, 3>(0, 0)
            .copy_from(&self.top_left.mat3());

        // top right
        result.fixed_view_mut::<3, 3>(0, 3).copy_from(&self.h);

        // bottom left
        result
            .fixed_view_mut::<3, 3>(3, 0)
            .copy_from(&self.h.transpose());

        // bootom right
        result
            .fixed_view_mut::<3, 3>(3, 3)
            .copy_from(&self.mass.mat3());

        result
    }
}

impl Add<Self> for ArticulatedBodyInertia {
    type Output = Self;

    #[inline]
    fn add(self, rhs: Self) -> Self::Output {
        self.add_arti_inertia(&rhs)
    }
}

impl Add<Self> for &ArticulatedBodyInertia {
    type Output = ArticulatedBodyInertia;

    #[inline]
    fn add(self, rhs: Self) -> ArticulatedBodyInertia {
        self.add_arti_inertia(rhs)
    }
}

impl AddAssign<&ArticulatedBodyInertia> for ArticulatedBodyInertia {
    #[inline]
    fn add_assign(&mut self, rhs: &ArticulatedBodyInertia) {
        *self = self.add_arti_inertia(rhs);
    }
}

impl Sub<Self> for ArticulatedBodyInertia {
    type Output = Self;

    #[inline]
    fn sub(self, rhs: Self) -> Self::Output {
        self.sub_arti_inertia(&rhs)
    }
}

impl Sub<ArticulatedBodyInertia> for &ArticulatedBodyInertia {
    type Output = ArticulatedBodyInertia;

    #[inline]
    fn sub(self, rhs: ArticulatedBodyInertia) -> Self::Output {
        self.sub_arti_inertia(&rhs)
    }
}

impl Add<RigidBodyInertia> for ArticulatedBodyInertia {
    type Output = Self;

    #[inline]
    fn add(self, rhs: RigidBodyInertia) -> Self::Output {
        self.add_spat_inertia(&rhs)
    }
}

impl From<RigidBodyInertia> for ArticulatedBodyInertia {
    #[inline]
    fn from(value: RigidBodyInertia) -> Self {
        Self {
            top_left: value.inertia,
            h: value.h.into_cross_matrix(),
            mass: SymmetricMat3::ONE.scale(value.mass),
        }
    }
}

impl Mul<SpatialMotionVector> for ArticulatedBodyInertia {
    type Output = SpatialForceVector;

    #[inline]
    fn mul(self, rhs: SpatialMotionVector) -> Self::Output {
        self.mul_motion_vec(rhs)
    }
}

impl Mul<SpatialMotionVector> for &ArticulatedBodyInertia {
    type Output = SpatialForceVector;

    #[inline]
    fn mul(self, rhs: SpatialMotionVector) -> Self::Output {
        self.mul_motion_vec(rhs)
    }
}

#[cfg(test)]
mod tests {
    use approx::assert_relative_eq;

    use super::*;
    use crate::vec3;

    #[test]
    fn test_rbi_construct() {
        let mass = 2.0;
        let com = vec3(1.0, 2.0, 3.0);
        let inertia_com = SymmetricMat3::from_array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
        let rbi = RigidBodyInertia::new(mass, com, inertia_com);

        let expected_inertia =
            inertia_com.mat3() + mass * (com.dot(&com) * Mat3::identity() - com * com.transpose());
        assert_relative_eq!(rbi.inertia.mat3(), expected_inertia, epsilon = 1e-6);
    }

    #[test]
    fn test_inertia_mul_motion_vec() {
        let mass = 2.0;
        let com = vec3(1.0, 2.0, 3.0);
        let inertia_com = SymmetricMat3::from_array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
        let rbi = RigidBodyInertia::new(mass, com, inertia_com);

        let v = SpatialMotionVector::from_array([1., 2., 3., 4., 5., 6.]);
        let f = rbi * v;
        let expected_f = rbi.matrix() * v.into_vec6();
        assert_relative_eq!(f.into_vec6(), expected_f, epsilon = 1e-6);
    }
}