phyz 0.3.0

//! Contact detection and force computation.

use super::material::ContactMaterial;
use crate::collision::{AABB, Collision, gjk_distance_rot, sweep_and_prune};
use crate::math::{SpatialVec, Vec3};
use crate::model::{Geometry as ModelGeometry, Model, State};

/// Convert crate::model::Geometry to crate::collision::Geometry.
fn convert_geometry(g: &ModelGeometry) -> crate::collision::Geometry {
    match g {
        ModelGeometry::Sphere { radius } => crate::collision::Geometry::Sphere { radius: *radius },
        ModelGeometry::Capsule { radius, length } => crate::collision::Geometry::Capsule {
            radius: *radius,
            length: *length,
        },
        ModelGeometry::Box { half_extents } => crate::collision::Geometry::Box {
            half_extents: *half_extents,
        },
        ModelGeometry::Cylinder { radius, height } => crate::collision::Geometry::Cylinder {
            radius: *radius,
            height: *height,
        },
        ModelGeometry::Mesh { vertices, faces } => crate::collision::Geometry::Mesh {
            vertices: vertices.clone(),
            faces: faces.clone(),
        },
        ModelGeometry::Plane { normal } => crate::collision::Geometry::Plane { normal: *normal },
    }
}

/// Find all contacts in the current state.
pub fn find_contacts(
    _model: &Model,
    state: &State,
    geometries: &[Option<ModelGeometry>],
) -> Vec<Collision> {
    let mut contacts = Vec::new();

    // Build AABBs for broad phase. If a body's transform contains any
    // non-finite component (NaN/inf — e.g. left over from an upstream blowup),
    // emit an empty AABB at the origin. The broad phase will skip it via its
    // own finiteness filter, but emitting an empty AABB also keeps the index
    // alignment with `geometries`/`state.body_xform`.
    let mut aabbs = Vec::new();
    for (i, geom_opt) in geometries.iter().enumerate() {
        if let Some(geom) = geom_opt {
            let xform = &state.body_xform[i];
            let pos = xform.pos;
            let rot = xform.rot;
            if !pos_is_finite(&pos) || !rot_is_finite(&rot) {
                // Poisoned transform; degrade gracefully and skip this body
                // by giving it a degenerate, non-finite-tagged AABB.
                aabbs.push(AABB::new(
                    Vec3::new(f64::NAN, f64::NAN, f64::NAN),
                    Vec3::new(f64::NAN, f64::NAN, f64::NAN),
                ));
                continue;
            }
            let collision_geom = convert_geometry(geom);
            let aabb = AABB::from_geometry(&collision_geom, &pos, &rot);
            aabbs.push(aabb);
        } else {
            // No geometry for this body
            aabbs.push(AABB::new(Vec3::zeros(), Vec3::zeros()));
        }
    }

    // Broad phase: find potentially colliding pairs
    let pairs = sweep_and_prune(&aabbs);

    // Narrow phase: GJK/EPA for each pair
    for (i, j) in pairs {
        if let (Some(geom_i), Some(geom_j)) = (&geometries[i], &geometries[j]) {
            let xform_i = &state.body_xform[i];
            let xform_j = &state.body_xform[j];
            let pos_i = xform_i.pos;
            let pos_j = xform_j.pos;
            let rot_i = xform_i.rot;
            let rot_j = xform_j.rot;

            // Defensive: even if the broad phase let through a body with a
            // NaN transform we refuse to produce a contact with a NaN normal.
            if !pos_is_finite(&pos_i) || !pos_is_finite(&pos_j) {
                continue;
            }

            let collision_geom_i = convert_geometry(geom_i);
            let collision_geom_j = convert_geometry(geom_j);
            let dist = gjk_distance_rot(
                &collision_geom_i,
                &collision_geom_j,
                &pos_i,
                &pos_j,
                &rot_i,
                &rot_j,
            );

            if dist < 0.0 {
                // Default: take the normal from the body centres. When the
                // centres coincide (or are within numerical noise) that
                // division gives NaN, so we fall back to an EPA-derived
                // normal, then to +Z. If even EPA can't agree on a
                // direction we drop the contact rather than seeding NaN.
                let center_offset = pos_j - pos_i;
                let normal = if center_offset.norm() > 1e-9 {
                    center_offset.normalize()
                } else if let Some((_, epa_normal)) = crate::collision::epa_penetration_rot(
                    &collision_geom_i,
                    &collision_geom_j,
                    &pos_i,
                    &pos_j,
                    &rot_i,
                    &rot_j,
                ) {
                    if pos_is_finite(&epa_normal) && epa_normal.norm() > 1e-9 {
                        epa_normal.normalize()
                    } else {
                        Vec3::z()
                    }
                } else {
                    Vec3::z()
                };

                if !pos_is_finite(&normal) {
                    continue;
                }

                let contact_point = (pos_i + pos_j) * 0.5;
                let penetration_depth = -dist;

                contacts.push(Collision {
                    body_i: i,
                    body_j: j,
                    contact_point,
                    contact_normal: normal,
                    penetration_depth,
                });
            }
        }
    }

    contacts
}

fn pos_is_finite(v: &Vec3) -> bool {
    v.x.is_finite() && v.y.is_finite() && v.z.is_finite()
}

fn rot_is_finite(m: &crate::math::Mat3) -> bool {
    for i in 0..3 {
        for j in 0..3 {
            if !m[(i, j)].is_finite() {
                return false;
            }
        }
    }
    true
}

/// Compute contact forces for all contacts using body spatial velocities.
///
/// `body_velocities` should come from forward_kinematics (linear part of spatial velocity).
/// Returns spatial forces for each body in body frame.
pub fn contact_forces(
    contacts: &[Collision],
    state: &State,
    materials: &[ContactMaterial],
    body_velocities: Option<&[SpatialVec]>,
) -> Vec<SpatialVec> {
    let nbodies = state.body_xform.len();
    let mut forces = vec![SpatialVec::zero(); nbodies];

    for contact in contacts {
        let i = contact.body_i;
        let j = contact.body_j;

        // Get material (use first body's material, or default)
        let material = if materials.is_empty() {
            &ContactMaterial::default()
        } else {
            &materials[i.min(materials.len() - 1)]
        };

        // Extract linear velocities from spatial velocities
        let vel_i = body_velocities
            .and_then(|vels| vels.get(i))
            .map(|v| v.linear)
            .unwrap_or(Vec3::zeros());

        let vel_j = if j == usize::MAX {
            // Ground — zero velocity
            Vec3::zeros()
        } else {
            body_velocities
                .and_then(|vels| vels.get(j))
                .map(|v| v.linear)
                .unwrap_or(Vec3::zeros())
        };

        // Compute force
        let force = super::compute_contact_force(contact, material, &vel_i, &vel_j);
        let f_linear = force.linear;

        // Apply equal and opposite forces AT THE CONTACT POINT. Sign
        // convention follows Goal 1; the torque component follows Goal 4
        // (τ = r × F where r is from the body's frame origin to the contact
        // point in world frame). For ground contacts, only body i is updated.
        if j == usize::MAX {
            let r_i = contact.contact_point - state.body_xform[i].pos;
            forces[i] = forces[i] + SpatialVec::new(r_i.cross(&f_linear), f_linear);
        } else {
            let r_i = contact.contact_point - state.body_xform[i].pos;
            let r_j = contact.contact_point - state.body_xform[j].pos;
            forces[i] = forces[i] + SpatialVec::new(r_i.cross(&(-f_linear)), -f_linear);
            forces[j] = forces[j] + SpatialVec::new(r_j.cross(&f_linear), f_linear);
        }
    }

    forces
}

/// Compute contact forces with implicit damping for low-mass-body stability.
///
/// Like [`contact_forces`], but the per-contact wrench is computed by
/// [`super::compute_contact_force_implicit`] which uses the body velocity at
/// the END of the step. This makes the contact solve unconditionally stable
/// for any choice of `dt`, contact stiffness, and damping — preventing the
/// "low-mass cube launches off the plate" failure mode of the explicit form.
///
/// `masses[i]` is the effective contact mass of body `i`. Use
/// `f64::INFINITY` for fixed/world bodies. For ground contacts
/// (`body_j == usize::MAX`) the ground is treated as having infinite mass.
pub fn contact_forces_implicit(
    contacts: &[Collision],
    state: &State,
    materials: &[ContactMaterial],
    body_velocities: Option<&[SpatialVec]>,
    masses: &[f64],
    dt: f64,
) -> Vec<SpatialVec> {
    let nbodies = state.body_xform.len();
    let mut forces = vec![SpatialVec::zero(); nbodies];

    for contact in contacts {
        let i = contact.body_i;
        let j = contact.body_j;

        let material = if materials.is_empty() {
            &ContactMaterial::default()
        } else {
            &materials[i.min(materials.len() - 1)]
        };

        let vel_i = body_velocities
            .and_then(|vels| vels.get(i))
            .map(|v| v.linear)
            .unwrap_or(Vec3::zeros());

        let vel_j = if j == usize::MAX {
            Vec3::zeros()
        } else {
            body_velocities
                .and_then(|vels| vels.get(j))
                .map(|v| v.linear)
                .unwrap_or(Vec3::zeros())
        };

        let mass_i = masses.get(i).copied().unwrap_or(f64::INFINITY);
        let mass_j = if j == usize::MAX {
            f64::INFINITY
        } else {
            masses.get(j).copied().unwrap_or(f64::INFINITY)
        };

        let force = super::compute_contact_force_implicit(
            contact, material, &vel_i, &vel_j, mass_i, mass_j, dt,
        );
        let f_linear = force.linear;

        // Apply at the contact point (Goal 4). Sign convention from Goal 1.
        if j == usize::MAX {
            let r_i = contact.contact_point - state.body_xform[i].pos;
            forces[i] = forces[i] + SpatialVec::new(r_i.cross(&f_linear), f_linear);
        } else {
            let r_i = contact.contact_point - state.body_xform[i].pos;
            let r_j = contact.contact_point - state.body_xform[j].pos;
            forces[i] = forces[i] + SpatialVec::new(r_i.cross(&(-f_linear)), -f_linear);
            forces[j] = forces[j] + SpatialVec::new(r_j.cross(&f_linear), f_linear);
        }
    }

    forces
}

/// Find ground plane contacts.
pub fn find_ground_contacts(
    state: &State,
    geometries: &[Option<ModelGeometry>],
    ground_height: f64,
) -> Vec<Collision> {
    let mut contacts = Vec::new();

    for (i, geom_opt) in geometries.iter().enumerate() {
        if let Some(geom) = geom_opt {
            let xform = &state.body_xform[i];
            let pos = xform.pos;

            // Check if body is below ground
            let min_z = match geom {
                ModelGeometry::Sphere { radius } => pos.z - radius,
                ModelGeometry::Box { half_extents } => pos.z - half_extents.z,
                ModelGeometry::Capsule { radius, length } => pos.z - length * 0.5 - radius,
                ModelGeometry::Cylinder { height, .. } => pos.z - height * 0.5,
                _ => pos.z,
            };

            if min_z < ground_height {
                let penetration = ground_height - min_z;
                contacts.push(Collision {
                    body_i: i,
                    body_j: usize::MAX, // Ground is not a body
                    contact_point: Vec3::new(pos.x, pos.y, ground_height),
                    contact_normal: Vec3::z(),
                    penetration_depth: penetration,
                });
            }
        }
    }

    contacts
}