imstkPbdCollisionHandling.h

/*
** This file is part of the Interactive Medical Simulation Toolkit (iMSTK)
** iMSTK is distributed under the Apache License, Version 2.0.
** See accompanying NOTICE for details.
*/

#pragma once

#include "imstkCollisionHandling.h"
#include "imstkPbdConstraint.h"

#include <unordered_map>

namespace imstk
{
enum PbdContactCase
{
    Vertex,    // Mesh vertex either from Pbd or CollidingObject
    Edge,      // Mesh edge either from Pbd or CollidingObject
    Triangle,  // Mesh triangle either from Pbd or CollidingObject
    Body,      // Body
    Primitive, // Not a mesh at all still could be any CD data type
    None
};
static std::string
getContactCaseStr(PbdContactCase contactCase)
{
    switch (contactCase)
    {
    case PbdContactCase::Body:
        return "Body";
    case PbdContactCase::Edge:
        return "Edge";
    case PbdContactCase::Primitive:
        return "Primitive";
    case PbdContactCase::Triangle:
        return "Triangle";
    case PbdContactCase::Vertex:
        return "Vertex";
    default:
        return "None";
    }
    ;
}

struct PbdCHTableKey
{
    PbdContactCase elemAType;
    PbdContactCase elemBType;
    bool ccd;

    friend std::ostream& operator<<(std::ostream& os, const PbdCHTableKey& dt);

    bool operator==(const PbdCHTableKey& other) const
    {
        return
            (elemAType == other.elemAType)
            && (elemBType == other.elemBType)
            && (ccd == other.ccd);
    }
};
} // namespace imstk

namespace std
{
template<>
struct hash<imstk::PbdCHTableKey>
{
    std::size_t operator()(const imstk::PbdCHTableKey& k) const
    {
        using std::size_t;
        using std::hash;

        // Base on the bit width of each value
        std::size_t v0 = static_cast<std::size_t>(k.elemAType); // first 2 bits
        std::size_t v1 = static_cast<std::size_t>(k.elemBType); // Next 2 bits
        std::size_t v2 = static_cast<std::size_t>(k.ccd);       // Next bit
        return v0 ^ (v1 << 3) ^ (v2 << 5);
    }
};
} // namespace std

namespace imstk
{
class PbdObject;
class PbdModel;
class PointSet;
class PointwiseMap;

class PbdCollisionHandling : public CollisionHandling
{
public:
    enum class ObjType
    {
        PbdDeformable,
        PbdRigid,
        Colliding
    };
    struct CollisionSideData
    {
        CollisionSideData() = default;

        // Objects
        PbdObject* pbdObj       = nullptr;
        CollidingObject* colObj = nullptr;
        ObjType objType = ObjType::Colliding;

        PbdModel* model    = nullptr;
        double compliance  = 0.0;
        double stiffness   = 0.0;
        Geometry* geometry = nullptr;
        PointSet* pointSet = nullptr;
        VecDataArray<double, 3>* vertices = nullptr;
        PointwiseMap* mapPtr = nullptr;
        AbstractDataArray* indicesPtr = nullptr;
        int bodyId = 0;

        Geometry* prevGeometry = nullptr;
    };
    struct ColElemSide
    {
        const CollisionElement* elem  = nullptr;
        const CollisionSideData* data = nullptr;
    };

    PbdCollisionHandling();
    ~PbdCollisionHandling() override;

    IMSTK_TYPE_NAME(PbdCollisionHandling)

    
    double getRestitution() const { return m_restitution; }
    void setRestitution(const double restitution) { m_restitution = restitution; }

    double getFriction() const { return m_friction; }
    void setFriction(const double friction) { m_friction = friction; }

    bool getUseCorrectVelocity() const { return m_useCorrectVelocity; }
    void setUseCorrectVelocity(const bool useCorrectVelocity) { m_useCorrectVelocity = useCorrectVelocity; }

    int getCCDSubsteps() const { return m_ccdSubsteps; }
    void setCCDSubsteps(const int ccdSubsteps) { m_ccdSubsteps = ccdSubsteps; }

    void setEnableBoundaryCollisions(const bool enableBoundaryCollisions) { m_enableBoundaryCollisions = enableBoundaryCollisions; }
    const bool getEnableBoundaryCollisions() const { return m_enableBoundaryCollisions; }

    void setRigidBodyCompliance(const double compliance) { m_compliance = compliance; }
    double getRigidBodyCompliance() const { return m_compliance; }

    void setDeformableStiffnessA(const double stiffness) { m_stiffness[0] = stiffness; }
    double getDeformableStiffnessA() const { return m_stiffness[0]; }
    void setDeformableStiffnessB(const double stiffness) { m_stiffness[1] = stiffness; }
    double getDeformableStiffnessB() const { return m_stiffness[1]; }

    const std::vector<PbdConstraint*>& getConstraints() const { return m_collisionConstraints; }

    std::pair<PbdParticleId, Vec3d> getBodyAndContactPoint(
        const CollisionElement&  elem,
        const CollisionSideData& data);

protected:
    std::array<PbdParticleId, 2> getEdge(
        const CollisionElement&  elem,
        const CollisionSideData& side);
    std::array<PbdParticleId, 3> getTriangle(
        const CollisionElement&  elem,
        const CollisionSideData& side);
    std::array<PbdParticleId, 1> getVertex(
        const CollisionElement&  elem,
        const CollisionSideData& side);

    CollisionSideData getDataFromObject(std::shared_ptr<CollidingObject> obj);

    PbdContactCase getCaseFromElement(const ColElemSide& elem);

    void handle(
        const std::vector<CollisionElement>& elementsA,
        const std::vector<CollisionElement>& elementsB) override;

    void handleElementPair(ColElemSide sideA, ColElemSide sideB);

    // -----------------One-Way Rigid on X Cases-----------------
    virtual void addConstraint_Body_V(
        const ColElemSide& sideA,
        const ColElemSide& sideB);
    virtual void addConstraint_Body_E(
        const ColElemSide& sideA,
        const ColElemSide& sideB);
    virtual void addConstraint_Body_T(
        const ColElemSide& sideA,
        const ColElemSide& sideB);
    // ---------------Two-Way Rigid on Rigid Cases---------------
    virtual void addConstraint_Body_Body(
        const ColElemSide& sideA,
        const ColElemSide& sideB);

    // ----------DeformableMesh on DeformableMesh Cases----------
    virtual void addConstraint_V_T(
        const ColElemSide& sideA,
        const ColElemSide& sideB);
    virtual void addConstraint_E_E(
        const ColElemSide& sideA,
        const ColElemSide& sideB);
    virtual void addConstraint_E_E_CCD(
        const ColElemSide& sideA,
        const ColElemSide& sideB);
    virtual void addConstraint_V_E(
        const ColElemSide& sideA,
        const ColElemSide& sideB);
    virtual void addConstraint_V_V(
        const ColElemSide& sideA,
        const ColElemSide& sideB);

private:
    double m_restitution = 0.0;  
    double m_friction    = 0.0;  

    bool   m_enableBoundaryCollisions = false;
    double m_compliance = 0.000001;
    bool   m_useCorrectVelocity       = true;
    std::array<double, 2> m_stiffness = { 0.3, 0.3 };
    int m_ccdSubsteps = 25;

    void deleteCollisionConstraints();

    void orderCollisionConstraints();

protected:

    // Enum for binning constraint types to order for solver
    // and defining the order
    enum ConstraintType
    {
        // Rigid on rigid
        BodyBody = 0,
        // Rigid on mesh
        BodyVertex,
        BodyEdge,
        BodyTriangle,
        // mesh on mesh non-ccd constraints
        VertexVertex,
        VertexEdge,
        EdgeEdge,
        VertexTriangle,
        // CCD constraints
        EdgeEdgeCCD,
        NumTypes
    };

    template<class T>
    T* getCachedConstraint(ConstraintType type);

    // Vectors to split out constraint types and allow for ordering
    // A single constraint instance should always either be in a bin or the cache
    // additionally it _may_ be in the m_collisionConstraints structure
    std::vector<PbdConstraint*> m_constraintBins[NumTypes];
    std::vector<PbdConstraint*> m_constraintCache[NumTypes];

    std::vector<PbdConstraint*> m_collisionConstraints; 

    std::unordered_map<PbdCHTableKey, std::function<void(
                                                        const ColElemSide& elemA, const ColElemSide& elemB)>> m_funcTable;
};
} // namespace imstk