use crate::alloc_prelude::*;
use crate::dynamics::solver::contact_constraint::{
    ContactWithCoulombFriction, ContactWithCoulombFrictionBuilder, GenericContactConstraint,
    GenericContactConstraintBuilder,
};
use crate::dynamics::solver::interaction_groups::InteractionGroups;
use crate::dynamics::{ImpulseJoint, JointAxesMask};
use crate::math::DVector;
use parry::math::SimdReal;

#[cfg(feature = "dim3")]
use crate::dynamics::solver::contact_constraint::{
    ContactWithTwistFriction, ContactWithTwistFrictionBuilder,
};

// NOTE: the constraints count from contacts is always 1 since the max number of solver contacts
//       matches the max number of contacts per constraint.
/// The number of velocity-constraint rows generated by `joint`.
pub fn joint_num_constraints(joint: &ImpulseJoint) -> usize {
    joint_data_num_constraints(&joint.data)
}

/// The number of velocity-constraint rows generated by a joint with data `joint`.
pub fn joint_data_num_constraints(joint: &crate::dynamics::GenericJoint) -> usize {
    let locked_axes = joint.locked_axes.bits();
    let motor_axes = joint.motor_axes.bits() & !locked_axes;
    let limit_axes = joint.limit_axes.bits() & !locked_axes;
    let coupled_axes = joint.coupled_axes.bits();

    (motor_axes & !coupled_axes).count_ones() as usize
        + ((motor_axes & coupled_axes) & JointAxesMask::ANG_AXES.bits() != 0) as usize
        + ((motor_axes & coupled_axes) & JointAxesMask::LIN_AXES.bits() != 0) as usize
        + locked_axes.count_ones() as usize
        + (limit_axes & !coupled_axes).count_ones() as usize
        + ((limit_axes & coupled_axes) & JointAxesMask::ANG_AXES.bits() != 0) as usize
        + ((limit_axes & coupled_axes) & JointAxesMask::LIN_AXES.bits() != 0) as usize
}

pub(crate) struct ContactConstraintsSet {
    pub generic_jacobians: DVector,
    /// Greedy body-mask grouper, only used for the overflow-color manifolds
    /// (they may share bodies, so the color-run chunking can't cover them).
    pub interaction_groups: InteractionGroups,

    pub generic_velocity_constraints: Vec<GenericContactConstraint>,
    pub simd_velocity_coulomb_constraints: Vec<ContactWithCoulombFriction<SimdReal>>,
    #[cfg(feature = "dim3")]
    pub simd_velocity_twist_constraints: Vec<ContactWithTwistFriction<SimdReal>>,

    pub generic_velocity_constraints_builder: Vec<GenericContactConstraintBuilder>,
    pub simd_velocity_coulomb_constraints_builder: Vec<ContactWithCoulombFrictionBuilder>,
    #[cfg(feature = "dim3")]
    pub simd_velocity_twist_constraints_builder: Vec<ContactWithTwistFrictionBuilder<SimdReal>>,
}

impl ContactConstraintsSet {
    pub fn new() -> Self {
        Self {
            generic_jacobians: DVector::zeros(0),
            interaction_groups: InteractionGroups::new(),
            generic_velocity_constraints: vec![],
            simd_velocity_coulomb_constraints: vec![],
            generic_velocity_constraints_builder: vec![],
            simd_velocity_coulomb_constraints_builder: vec![],
            #[cfg(feature = "dim3")]
            simd_velocity_twist_constraints: vec![],
            #[cfg(feature = "dim3")]
            simd_velocity_twist_constraints_builder: vec![],
        }
    }
}
