imstkPbdObjectGrasping.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 "imstkCollisionInteraction.h"
#include "imstkMacros.h"
#include "imstkMath.h"
#include "imstkPbdConstraint.h"

#include <unordered_map>

namespace imstk
{
class AnalyticalGeometry;
class PbdObject;
class PickingAlgorithm;
class PointwiseMap;

class PbdObjectGrasping : public SceneObject
{
public:
    enum class GraspMode
    {
        Vertex,     // Grab a vertex (most performant)
        Cell,       // Grab an entire cell/s
        RayPoint,   // Grab a point on the nearest cell along the ray
        RayCell     // Grab a cell along the ray
    };

    PbdObjectGrasping(std::shared_ptr<PbdObject> graspedObject,
                      std::shared_ptr<PbdObject> grasperObject = nullptr);
    ~PbdObjectGrasping() override = default;

    IMSTK_TYPE_NAME(PbdObjectGrasping)

    
    void setStiffness(const double stiffness) { m_stiffness = stiffness; }
    double getStiffness() const { return m_stiffness; }

    void setCompliance(const double compliance) { m_compliance = compliance; }
    double getCompliance() const { return m_compliance; }

    void beginVertexGrasp(std::shared_ptr<AnalyticalGeometry> geometry);

    void beginCellGrasp(std::shared_ptr<AnalyticalGeometry> geometry, std::string cdType = "");

    void beginRayPointGrasp(std::shared_ptr<AnalyticalGeometry> geometry,
                            const Vec3d& rayStart, const Vec3d& rayDir, const double maxDist = -1.0);

    void beginRayCellGrasp(std::shared_ptr<AnalyticalGeometry> geometry,
                           const Vec3d& rayStart, const Vec3d& rayDir, const double maxDist = -1.0);

    void regrasp()
    {
        m_regrasp = m_isGrasping;
    }

    void endGrasp();

    void addPickConstraints();

    void removePickConstraints();

    virtual void addPointToPointConstraint(
        const std::vector<PbdParticleId>& ptsA,
        const std::vector<double>& weightsA,
        const std::vector<PbdParticleId>& ptsB,
        const std::vector<double>& weightsB,
        const double stiffnessA, const double stiffnessB);

    virtual void addBodyToBodyConstraint(
        const PbdParticleId& graspedBodyId,
        const PbdParticleId& grasperBodyId,
        const Vec3d&         pointOnBodies,
        const double         compliance);

    virtual void addPointToBodyConstraint(
        const PbdParticleId& graspedParticleId,
        const PbdParticleId& grasperBodyId,
        const Vec3d&         pointOnBody,
        const double         compliance);

    void setPickingAlgorithm(std::shared_ptr<PickingAlgorithm> pickMethod) { m_pickMethod = pickMethod; }
    std::shared_ptr<PickingAlgorithm> getPickingAlgorithm() const { return m_pickMethod; }

    void setGeometryToPick(std::shared_ptr<Geometry> geomToPick, std::shared_ptr<PointwiseMap> map)
    {
        m_geomToPick = geomToPick;
        m_geometryToPickMap = map;
    }

    std::shared_ptr<TaskNode> getPickingNode() const { return m_pickingNode; }

    void initGraphEdges(std::shared_ptr<TaskNode> source, std::shared_ptr<TaskNode> sink) override;

    bool hasConstraints() const;

    bool getGraspState() const { return m_isGrasping; }

protected:
    virtual void updatePicking();

    void updateConstraints();

    std::shared_ptr<TaskNode> m_pickingNode = nullptr;

    std::shared_ptr<Geometry>     m_geomToPick = nullptr;
    std::shared_ptr<PointwiseMap> m_geometryToPickMap = nullptr;

    std::shared_ptr<PbdObject> m_objectToGrasp      = nullptr;
    std::shared_ptr<PbdObject> m_grasperObject      = nullptr; // Optional
    std::shared_ptr<AnalyticalGeometry> m_graspGeom = nullptr;

    std::shared_ptr<PickingAlgorithm> m_pickMethod = nullptr;
    GraspMode m_graspMode = GraspMode::Cell;

    bool m_isGrasping     = false;
    bool m_regrasp        = false; 
    bool m_isPrevGrasping = false;

    double m_stiffness  = 0.4;    // For deformables
    double m_compliance = 0.0001; // For rigid bodies (inverse of stiffness)

    std::vector<std::tuple<PbdParticleId, Vec3d>> m_constraintPts;
    std::vector<std::shared_ptr<PbdConstraint>>   m_constraints; 

private:
    std::vector<PbdConstraint*> m_collisionConstraints;
};
} // namespace imstk