featherstone 0.1.0

//! Joint limit constraints using MuJoCo-style spring-damper barriers
//!
//! Produces smooth constraint forces that prevent joints from exceeding their
//! position, velocity, and effort limits. Uses a spring-damper barrier at limits
//! rather than hard clamping, ensuring differentiable dynamics.
//!
//! Compatible with ABA: limit forces are treated as applied joint torques (added to tau).

/// Per-joint limit specification and barrier parameters.
///
/// The barrier model applies force when a joint enters the margin zone near a limit:
/// ```text
/// F_limit = -stiffness * penetration - damping * velocity
/// ```
/// where penetration = how far past the margin the joint has gone.
/// This produces zero force away from limits and smooth force near limits.
#[derive(Clone, Debug)]
pub struct JointLimits {
    /// Lower position limit per DOF (radians or meters)
    pub lower: Vec<f32>,
    /// Upper position limit per DOF
    pub upper: Vec<f32>,
    /// Maximum velocity per DOF (rad/s or m/s). 0 = unlimited.
    pub max_velocity: Vec<f32>,
    /// Maximum effort per DOF (Nm or N). 0 = unlimited.
    pub max_effort: Vec<f32>,
    /// Barrier stiffness (N/m or Nm/rad) — how hard the barrier pushes back
    pub stiffness: f32,
    /// Barrier damping (Ns/m or Nms/rad) — prevents oscillation at limits
    pub damping: f32,
    /// Margin: start applying force this far before the hard limit (rad or m)
    pub margin: f32,
    /// Whether limits are active
    pub enabled: bool,
}

impl Default for JointLimits {
    fn default() -> Self {
        Self {
            lower: Vec::new(),
            upper: Vec::new(),
            max_velocity: Vec::new(),
            max_effort: Vec::new(),
            stiffness: 1000.0,
            damping: 100.0,
            margin: 0.01,
            enabled: true,
        }
    }
}

impl JointLimits {
    /// Create joint limits for a single-DOF joint (revolute or prismatic).
    pub fn single(lower: f32, upper: f32) -> Self {
        Self {
            lower: vec![lower],
            upper: vec![upper],
            max_velocity: vec![0.0],
            max_effort: vec![0.0],
            ..Default::default()
        }
    }

    /// Create joint limits with full URDF parameters.
    pub fn from_urdf_params(
        lower: f32,
        upper: f32,
        max_velocity: f32,
        max_effort: f32,
    ) -> Self {
        Self {
            lower: vec![lower],
            upper: vec![upper],
            max_velocity: vec![max_velocity],
            max_effort: vec![max_effort],
            ..Default::default()
        }
    }

    /// Create unlimited (no limits enforced).
    pub fn unlimited(dof: usize) -> Self {
        Self {
            lower: vec![f32::NEG_INFINITY; dof],
            upper: vec![f32::INFINITY; dof],
            max_velocity: vec![0.0; dof],
            max_effort: vec![0.0; dof],
            enabled: false,
            ..Default::default()
        }
    }

    /// Number of DOFs these limits cover
    pub fn dof(&self) -> usize {
        self.lower.len()
    }

    /// Set barrier parameters.
    pub fn with_barrier(mut self, stiffness: f32, damping: f32, margin: f32) -> Self {
        self.stiffness = stiffness;
        self.damping = damping;
        self.margin = margin;
        self
    }

    /// Compute limit forces for each DOF.
    ///
    /// Returns a Vec of forces (dimension = dof). These should be added to `tau`
    /// before running ABA forward dynamics.
    ///
    /// The barrier force model:
    /// - When q < lower + margin: F = -stiffness * (q - (lower + margin)) - damping * qd
    /// - When q > upper - margin: F = -stiffness * (q - (upper - margin)) - damping * qd
    /// - Otherwise: F = 0
    ///
    /// This gives smooth, differentiable force near limits with zero force away from them.
    pub fn compute_limit_force(&self, q: &[f32], qd: &[f32]) -> Vec<f32> {
        if !self.enabled {
            return vec![0.0; self.dof()];
        }

        let n = self.dof();
        let mut forces = vec![0.0; n];

        for (i, force) in forces.iter_mut().enumerate().take(n) {
            let qi = q.get(i).copied().unwrap_or(0.0);
            let qdi = qd.get(i).copied().unwrap_or(0.0);
            let lo = self.lower[i];
            let hi = self.upper[i];

            // Skip if limits are infinite (no constraint)
            if lo == f32::NEG_INFINITY && hi == f32::INFINITY {
                continue;
            }

            let barrier_lo = lo + self.margin;
            let barrier_hi = hi - self.margin;

            if qi < barrier_lo {
                // Penetrating lower limit barrier
                let penetration = barrier_lo - qi;
                // Only damp velocity moving INTO the limit (negative direction)
                let damp_vel = qdi.min(0.0);
                *force = self.stiffness * penetration - self.damping * damp_vel;
            } else if qi > barrier_hi {
                // Penetrating upper limit barrier
                let penetration = qi - barrier_hi;
                // Only damp velocity moving INTO the limit (positive direction)
                let damp_vel = qdi.max(0.0);
                *force = -self.stiffness * penetration - self.damping * damp_vel;
            }
        }

        forces
    }

    /// Clamp applied torques/forces to effort limits.
    ///
    /// Returns clamped values. Pass-through if max_effort is 0 (unlimited).
    pub fn clamp_effort(&self, tau: &[f32]) -> Vec<f32> {
        let n = self.dof();
        let mut clamped = vec![0.0; n];

        for (i, val) in clamped.iter_mut().enumerate().take(n) {
            let t = tau.get(i).copied().unwrap_or(0.0);
            let max_e = self.max_effort.get(i).copied().unwrap_or(0.0);
            *val = if max_e > 0.0 {
                t.clamp(-max_e, max_e)
            } else {
                t
            };
        }

        clamped
    }

    /// Clamp velocities to velocity limits.
    ///
    /// Returns clamped values. Pass-through if max_velocity is 0 (unlimited).
    pub fn clamp_velocity(&self, qd: &[f32]) -> Vec<f32> {
        let n = self.dof();
        let mut clamped = vec![0.0; n];

        for (i, val) in clamped.iter_mut().enumerate().take(n) {
            let v = qd.get(i).copied().unwrap_or(0.0);
            let max_v = self.max_velocity.get(i).copied().unwrap_or(0.0);
            *val = if max_v > 0.0 {
                v.clamp(-max_v, max_v)
            } else {
                v
            };
        }

        clamped
    }

    /// Check if any position limit is currently violated (past hard limit).
    pub fn is_violated(&self, q: &[f32]) -> bool {
        for i in 0..self.dof() {
            let qi = q.get(i).copied().unwrap_or(0.0);
            if qi < self.lower[i] || qi > self.upper[i] {
                return true;
            }
        }
        false
    }

    /// Get the maximum penetration depth (how far past a hard limit).
    /// Returns 0 if within limits.
    pub fn max_penetration(&self, q: &[f32]) -> f32 {
        let mut max_pen = 0.0_f32;
        for i in 0..self.dof() {
            let qi = q.get(i).copied().unwrap_or(0.0);
            let pen_lo = (self.lower[i] - qi).max(0.0);
            let pen_hi = (qi - self.upper[i]).max(0.0);
            max_pen = max_pen.max(pen_lo).max(pen_hi);
        }
        max_pen
    }
}

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

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

    #[test]
    fn test_single_limit_creation() {
        let lim = JointLimits::single(-1.57, 1.57);
        assert_eq!(lim.dof(), 1);
        assert_relative_eq!(lim.lower[0], -1.57);
        assert_relative_eq!(lim.upper[0], 1.57);
    }

    #[test]
    fn test_zero_force_within_limits() {
        let lim = JointLimits::single(-1.57, 1.57);
        let f = lim.compute_limit_force(&[0.0], &[0.0]);
        assert_relative_eq!(f[0], 0.0);

        let f = lim.compute_limit_force(&[1.0], &[0.0]);
        assert_relative_eq!(f[0], 0.0);
    }

    #[test]
    fn test_lower_limit_barrier() {
        let lim = JointLimits::single(-1.57, 1.57)
            .with_barrier(1000.0, 100.0, 0.01);

        // At lower limit boundary (within margin)
        let q = [-1.57]; // at the hard limit
        let f = lim.compute_limit_force(&q, &[0.0]);
        assert!(f[0] > 0.0, "Should push away from lower limit: {}", f[0]);

        // Force should increase with deeper penetration
        let f_deep = lim.compute_limit_force(&[-1.6], &[0.0]);
        assert!(f_deep[0] > f[0], "Deeper penetration → larger force");
    }

    #[test]
    fn test_upper_limit_barrier() {
        let lim = JointLimits::single(-1.57, 1.57)
            .with_barrier(1000.0, 100.0, 0.01);

        let f = lim.compute_limit_force(&[1.57], &[0.0]);
        assert!(f[0] < 0.0, "Should push away from upper limit: {}", f[0]);
    }

    #[test]
    fn test_damping_at_limit() {
        let lim = JointLimits::single(-1.57, 1.57)
            .with_barrier(1000.0, 100.0, 0.01);

        // Moving toward limit: damping should resist
        let f_moving = lim.compute_limit_force(&[-1.57], &[-1.0])[0];
        let f_still = lim.compute_limit_force(&[-1.57], &[0.0])[0];

        // With negative velocity into limit, damping increases restoring force
        assert!(
            f_moving > f_still,
            "Moving into limit should produce more force: moving={f_moving} still={f_still}"
        );
    }

    #[test]
    fn test_effort_clamping() {
        let lim = JointLimits::from_urdf_params(-1.57, 1.57, 10.0, 50.0);
        let clamped = lim.clamp_effort(&[100.0]);
        assert_relative_eq!(clamped[0], 50.0);

        let clamped = lim.clamp_effort(&[-100.0]);
        assert_relative_eq!(clamped[0], -50.0);

        let clamped = lim.clamp_effort(&[30.0]);
        assert_relative_eq!(clamped[0], 30.0);
    }

    #[test]
    fn test_velocity_clamping() {
        let lim = JointLimits::from_urdf_params(-1.57, 1.57, 10.0, 50.0);
        let clamped = lim.clamp_velocity(&[20.0]);
        assert_relative_eq!(clamped[0], 10.0);
    }

    #[test]
    fn test_unlimited() {
        let lim = JointLimits::unlimited(3);
        assert_eq!(lim.dof(), 3);
        assert!(!lim.enabled);
        let f = lim.compute_limit_force(&[100.0, 200.0, 300.0], &[0.0, 0.0, 0.0]);
        assert!(f.iter().all(|&v| v == 0.0));
    }

    #[test]
    fn test_is_violated() {
        let lim = JointLimits::single(-1.57, 1.57);
        assert!(!lim.is_violated(&[0.0]));
        assert!(!lim.is_violated(&[1.57]));
        assert!(lim.is_violated(&[1.58]));
        assert!(lim.is_violated(&[-1.58]));
    }

    #[test]
    fn test_max_penetration() {
        let lim = JointLimits::single(-1.57, 1.57);
        assert_relative_eq!(lim.max_penetration(&[0.0]), 0.0);
        assert_relative_eq!(lim.max_penetration(&[1.67]), 0.1, epsilon = 1e-5);
        assert_relative_eq!(lim.max_penetration(&[-1.67]), 0.1, epsilon = 1e-5);
    }

    #[test]
    fn test_pendulum_with_limits() {
        // Simulate a pendulum hitting a +-90° limit
        use super::super::body::ArticulatedBody;
        use super::super::joint::GenJoint;
        use super::super::spatial::{SpatialInertia, SpatialTransform};
        use super::super::aba::aba_forward_dynamics;
        use nalgebra::{Matrix3, Vector3};

        let mut body = ArticulatedBody::new();
        body.set_gravity(Vector3::new(0.0, -9.81, 0.0));

        let inertia = SpatialInertia::from_mass_inertia(
            1.0,
            Vector3::new(0.3, 0.0, 0.0),
            Matrix3::from_diagonal(&Vector3::new(0.01, 0.01, 0.01)),
        );

        body.add_body(
            "pendulum",
            -1,
            GenJoint::Revolute { axis: Vector3::z() },
            inertia,
            SpatialTransform::identity(),
        );

        let limits = JointLimits::single(
            -std::f32::consts::FRAC_PI_2,
            std::f32::consts::FRAC_PI_2,
        )
        .with_barrier(5000.0, 500.0, 0.05);

        // Start near the upper limit
        body.set_joint_q(0, &[1.5]);

        let dt = 0.001;
        let mut max_q = 0.0_f32;

        // Simulate for 500 steps
        for _ in 0..500 {
            // Compute limit forces
            let limit_force = limits.compute_limit_force(body.joint_q(0), body.joint_qd(0));
            body.set_joint_tau(0, &[limit_force[0]]);

            aba_forward_dynamics(&mut body);

            // Semi-implicit Euler integration
            let qdd = body.qdd[0];
            let new_qd = body.joint_qd(0)[0] + qdd * dt;
            let new_q = body.joint_q(0)[0] + new_qd * dt;

            body.set_joint_qd(0, &[new_qd]);
            body.set_joint_q(0, &[new_q]);

            max_q = max_q.max(new_q.abs());
        }

        // Joint should never overshoot by more than ~0.001 rad
        let overshoot = (max_q - std::f32::consts::FRAC_PI_2).max(0.0);
        assert!(
            overshoot < 0.05,
            "Joint should stay near limits. Max overshoot: {overshoot} rad, max_q: {max_q}"
        );
    }

    #[test]
    fn test_force_continuity() {
        // Force should be continuous (no discontinuity at margin boundary)
        let lim = JointLimits::single(-1.57, 1.57)
            .with_barrier(1000.0, 0.0, 0.1); // No damping, large margin

        let margin_boundary = 1.57 - 0.1; // = 1.47
        let f_inside = lim.compute_limit_force(&[margin_boundary - 0.001], &[0.0]);
        let f_boundary = lim.compute_limit_force(&[margin_boundary], &[0.0]);
        let f_outside = lim.compute_limit_force(&[margin_boundary + 0.001], &[0.0]);

        // At boundary: force should be exactly 0
        assert_relative_eq!(f_boundary[0], 0.0, epsilon = 1e-6);
        // Just inside margin: force should be tiny
        assert!(f_inside[0].abs() < 0.01);
        // Just outside margin: force should be tiny
        assert!(f_outside[0].abs() < 2.0);
        // Force should transition smoothly
        assert!(
            (f_outside[0] - f_boundary[0]).abs() < 2.0,
            "Force should be continuous"
        );
    }

    // ======================================================================
    // Edge case tests for limit checking
    // ======================================================================

    #[test]
    fn test_is_violated_within_bounds() {
        let lim = JointLimits::single(-1.0, 1.0);
        assert!(!lim.is_violated(&[0.0]));
        assert!(!lim.is_violated(&[-1.0])); // at lower bound
        assert!(!lim.is_violated(&[1.0])); // at upper bound
    }

    #[test]
    fn test_is_violated_outside_bounds() {
        let lim = JointLimits::single(-1.0, 1.0);
        assert!(lim.is_violated(&[-1.001]));
        assert!(lim.is_violated(&[1.001]));
    }

    #[test]
    fn test_max_penetration_no_violation() {
        let lim = JointLimits::single(-1.0, 1.0);
        assert_eq!(lim.max_penetration(&[0.0]), 0.0);
        assert_eq!(lim.max_penetration(&[-1.0]), 0.0);
        assert_eq!(lim.max_penetration(&[1.0]), 0.0);
    }

    #[test]
    fn test_max_penetration_with_violation() {
        let lim = JointLimits::single(-1.0, 1.0);
        let pen = lim.max_penetration(&[1.5]);
        assert!((pen - 0.5).abs() < 1e-5, "Penetration should be 0.5, got {pen}");

        let pen_lo = lim.max_penetration(&[-1.3]);
        assert!((pen_lo - 0.3).abs() < 1e-5, "Lower pen should be 0.3, got {pen_lo}");
    }

    #[test]
    fn test_clamp_effort_within_limit() {
        let lim = JointLimits::from_urdf_params(-1.0, 1.0, 0.0, 10.0);
        let clamped = lim.clamp_effort(&[5.0]);
        assert_eq!(clamped[0], 5.0);
    }

    #[test]
    fn test_clamp_effort_exceeds_limit() {
        let lim = JointLimits::from_urdf_params(-1.0, 1.0, 0.0, 10.0);
        let clamped = lim.clamp_effort(&[15.0]);
        assert_eq!(clamped[0], 10.0);

        let clamped_neg = lim.clamp_effort(&[-15.0]);
        assert_eq!(clamped_neg[0], -10.0);
    }

    #[test]
    fn test_clamp_effort_zero_max() {
        // Zero max_effort means no clamping
        let lim = JointLimits::single(-1.0, 1.0); // defaults to max_effort=0
        let clamped = lim.clamp_effort(&[100.0]);
        assert_eq!(clamped[0], 100.0, "Zero max_effort should not clamp");
    }

    #[test]
    fn intent_limit_force_zero_away_from_limits() {
        let lim = JointLimits::single(-1.0, 1.0);
        let force = lim.compute_limit_force(&[0.0], &[0.0]);
        assert!(force[0].abs() < 1e-6, "center force should be zero: {}", force[0]);
    }

    #[test]
    fn intent_violation_detected_past_hard_limit() {
        let lim = JointLimits::single(-1.0, 1.0);
        assert!(!lim.is_violated(&[0.0]));
        assert!(lim.is_violated(&[1.5]));
        assert!(lim.is_violated(&[-1.5]));
    }

    #[test]
    fn intent_limit_force_opposes_violation() {
        // Physics intent: limit force should push back toward valid range
        let lim = JointLimits::single(-1.0, 1.0);

        // At upper limit (q=1.2, beyond +1.0)
        let force_upper = lim.compute_limit_force(&[1.2], &[0.0]);
        assert!(force_upper[0] < 0.0,
            "force at upper violation should push back (negative): {}", force_upper[0]);

        // At lower limit (q=-1.2, beyond -1.0)
        let force_lower = lim.compute_limit_force(&[-1.2], &[0.0]);
        assert!(force_lower[0] > 0.0,
            "force at lower violation should push back (positive): {}", force_lower[0]);
    }

    #[test]
    fn intent_limit_force_increases_with_violation_depth() {
        let lim = JointLimits::single(-1.0, 1.0);

        let force_small = lim.compute_limit_force(&[1.1], &[0.0]);
        let force_large = lim.compute_limit_force(&[1.5], &[0.0]);

        // Deeper violation → larger restoring force (in magnitude)
        assert!(force_large[0].abs() > force_small[0].abs(),
            "deeper violation should produce larger force: small={}, large={}",
            force_small[0], force_large[0]);
    }

    #[test]
    fn intent_clamp_velocity_respects_max() {
        let mut lim = JointLimits::single(-1.0, 1.0);
        lim.max_velocity = vec![5.0];
        let clamped = lim.clamp_velocity(&[10.0]);
        assert!((clamped[0] - 5.0).abs() < 1e-6, "should clamp to max_velocity: {}", clamped[0]);
        let clamped_neg = lim.clamp_velocity(&[-10.0]);
        assert!((clamped_neg[0] - (-5.0)).abs() < 1e-6, "should clamp negative: {}", clamped_neg[0]);
        let not_clamped = lim.clamp_velocity(&[2.0]);
        assert!((not_clamped[0] - 2.0).abs() < 1e-6, "within range should not change: {}", not_clamped[0]);
    }

    // ── SLAM Cycle 3: Limits proptest and intent tests ────────────────

    use proptest::prelude::*;

    proptest! {
        #[test]
        fn prop_zero_force_within_limits(
            q in -0.9f32..0.9,
            qd in -5.0f32..5.0,
        ) {
            let lim = JointLimits::single(-1.0, 1.0);
            let force = lim.compute_limit_force(&[q], &[qd]);
            // q is within [-1.0+margin, 1.0-margin] so force should be zero
            // margin default is 0.01, so q ∈ [-0.99, 0.99] → zero force
            prop_assert!((force[0]).abs() < 1e-3,
                "Force should be ~0 within limits: q={q}, f={}", force[0]);
        }

        #[test]
        fn prop_clamp_effort_bounded(
            tau in -100.0f32..100.0,
            max in 1.0f32..50.0,
        ) {
            let mut lim = JointLimits::single(-1.0, 1.0);
            lim.max_effort = vec![max];
            let clamped = lim.clamp_effort(&[tau]);
            prop_assert!(clamped[0] >= -max - 1e-6 && clamped[0] <= max + 1e-6,
                "Clamped tau={} should be in [-{max}, {max}]", clamped[0]);
        }

        #[test]
        fn prop_violation_agrees_with_penetration(
            q in -3.0f32..3.0,
        ) {
            let lim = JointLimits::single(-1.57, 1.57);
            let violated = lim.is_violated(&[q]);
            let pen = lim.max_penetration(&[q]);
            if violated {
                prop_assert!(pen > 0.0, "violated but penetration={pen}");
            } else {
                prop_assert!(pen <= 1e-6, "not violated but penetration={pen}");
            }
        }
    }

    #[test]
    fn intent_disabled_limits_produce_zero_force() {
        let mut lim = JointLimits::single(-1.0, 1.0);
        lim.enabled = false;
        let force = lim.compute_limit_force(&[5.0], &[10.0]);
        assert!(force[0].abs() < 1e-10,
            "Disabled limits should produce zero force, got {}", force[0]);
    }

    #[test]
    fn intent_barrier_force_opposes_violation() {
        // At upper limit violation, force should be negative (push back)
        let lim = JointLimits::single(-1.0, 1.0);
        let force = lim.compute_limit_force(&[1.5], &[0.0]);
        assert!(force[0] < 0.0,
            "Force at upper violation should be negative (restoring), got {}", force[0]);

        // At lower limit violation, force should be positive (push back)
        let force_low = lim.compute_limit_force(&[-1.5], &[0.0]);
        assert!(force_low[0] > 0.0,
            "Force at lower violation should be positive (restoring), got {}", force_low[0]);
    }

    // ── SLAM Cycle 13: Force continuity and multi-DOF tests ───────────

    #[test]
    fn intent_force_continuous_at_margin_boundary() {
        // Force should transition smoothly from 0 to positive at margin boundary
        let lim = JointLimits::single(-1.0, 1.0).with_barrier(1000.0, 100.0, 0.05);
        let f_inside = lim.compute_limit_force(&[0.94], &[0.0]); // inside margin
        let _f_boundary = lim.compute_limit_force(&[0.95], &[0.0]); // at margin
        let f_outside = lim.compute_limit_force(&[0.96], &[0.0]); // in margin zone

        // Force should increase as we approach upper limit
        assert!(f_inside[0].abs() <= f_outside[0].abs() + 1e-3,
            "force should increase toward limit: inside={}, outside={}",
            f_inside[0], f_outside[0]);
    }

    #[test]
    fn test_multi_dof_limits() {
        // 3-DOF joint with different limits per DOF
        let lim = JointLimits {
            lower: vec![-1.0, -2.0, -0.5],
            upper: vec![1.0, 2.0, 0.5],
            max_velocity: vec![10.0, 10.0, 10.0],
            max_effort: vec![50.0, 50.0, 50.0],
            stiffness: 1000.0,
            damping: 100.0,
            margin: 0.01,
            enabled: true,
        };

        assert_eq!(lim.dof(), 3);

        // DOF 2 is violated (0.6 > 0.5), others are within limits
        let force = lim.compute_limit_force(&[0.0, 0.0, 0.6], &[0.0, 0.0, 0.0]);
        assert!(force[0].abs() < 1e-3, "DOF 0 within limits: f={}", force[0]);
        assert!(force[1].abs() < 1e-3, "DOF 1 within limits: f={}", force[1]);
        assert!(force[2] < 0.0, "DOF 2 violated upper: f={}", force[2]);
    }
}