imstkPbdConstraintFunctor.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 "imstkLineMesh.h"
#include "imstkLogger.h"
#include "imstkParallelUtils.h"
#include "imstkPbdAreaConstraint.h"
#include "imstkPbdBendConstraint.h"
#include "imstkPbdConstantDensityConstraint.h"
#include "imstkPbdConstraint.h"
#include "imstkPbdDihedralConstraint.h"
#include "imstkPbdDistanceConstraint.h"
#include "imstkPbdFemConstraint.h"
#include "imstkPbdFemTetConstraint.h"
#include "imstkPbdVolumeConstraint.h"
#include "imstkPointSet.h"
#include "imstkSurfaceMesh.h"
#include "imstkTetrahedralMesh.h"
#include "imstkPbdConstraintContainer.h"

#include <set>

namespace imstk
{
struct PbdConstraintFunctor
{
    public:
        PbdConstraintFunctor() = default;
        virtual ~PbdConstraintFunctor() = default;

        virtual void operator()(PbdConstraintContainer& constraints) = 0;

        virtual void addConstraints(
            PbdConstraintContainer&                     imstkNotUsed(constraints),
            std::shared_ptr<std::unordered_set<size_t>> imstkNotUsed(vertices)) { }
};

struct PbdConstraintFunctorLambda : public PbdConstraintFunctor
{
    public:
        PbdConstraintFunctorLambda() = default;
        PbdConstraintFunctorLambda(std::function<void(PbdConstraintContainer&)> f) :
            m_func(f) { }

        void operator()(PbdConstraintContainer& constraints) override
        {
            m_func(constraints);
        }

        std::function<void(PbdConstraintContainer&)> m_func = nullptr;
};

struct PbdBodyConstraintFunctor : public PbdConstraintFunctor
{
    public:
        void setGeometry(std::shared_ptr<PointSet> geom)
        {
            m_geom = geom;
        }

        void setBodyIndex(const int bodyIndex) { m_bodyIndex = bodyIndex; }

        int m_bodyIndex = 1;
        std::shared_ptr<PointSet> m_geom = nullptr;
};

struct PbdDistanceConstraintFunctor : public PbdBodyConstraintFunctor
{
    public:
        PbdDistanceConstraintFunctor() = default;
        ~PbdDistanceConstraintFunctor() override = default;

        virtual std::shared_ptr<PbdDistanceConstraint> makeDistConstraint(
            const VecDataArray<double, 3>& vertices,
            int i1, int i2)
        {
            auto constraint = std::make_shared<PbdDistanceConstraint>();
            constraint->initConstraint(vertices[i1], vertices[i2],
                { m_bodyIndex, i1 }, { m_bodyIndex, i2 }, m_stiffness);
            constraint->m_restLength *= m_stretch;
            return constraint;
        }

        void operator()(PbdConstraintContainer& constraints) override
        {
            std::shared_ptr<VecDataArray<double, 3>> verticesPtr = m_geom->getVertexPositions();
            const VecDataArray<double, 3>&           vertices    = *verticesPtr;

            auto addDistConstraint = [&](
                std::vector<std::vector<bool>>& E, int i1, int i2)
                                     {
                                         // Make sure i1 is always smaller than i2 for duplicate testing
                                         if (i1 > i2)
                                         {
                                             std::swap(i1, i2);
                                         }
                                         if (E[i1][i2])
                                         {
                                             E[i1][i2] = 0;

                                             auto c = makeDistConstraint(vertices, i1, i2);
                                             constraints.addConstraint(c);
                                         }
                                     };

            if (auto tetMesh = std::dynamic_pointer_cast<TetrahedralMesh>(m_geom))
            {
                std::shared_ptr<VecDataArray<int, 4>> elementsPtr = tetMesh->getCells();
                const VecDataArray<int, 4>&           elements    = *elementsPtr;
                const auto                            nV = tetMesh->getNumVertices();
                std::vector<std::vector<bool>>        E(nV, std::vector<bool>(nV, 1)); // To test for duplicates

                for (int k = 0; k < elements.size(); ++k)
                {
                    auto& tet = elements[k];
                    addDistConstraint(E, tet[0], tet[1]);
                    addDistConstraint(E, tet[0], tet[2]);
                    addDistConstraint(E, tet[0], tet[3]);
                    addDistConstraint(E, tet[1], tet[2]);
                    addDistConstraint(E, tet[1], tet[3]);
                    addDistConstraint(E, tet[2], tet[3]);
                }
            }
            else if (auto triMesh = std::dynamic_pointer_cast<SurfaceMesh>(m_geom))
            {
                std::shared_ptr<VecDataArray<int, 3>> elementsPtr = triMesh->getCells();
                const VecDataArray<int, 3>&           elements    = *elementsPtr;
                const auto                            nV = triMesh->getNumVertices();
                std::vector<std::vector<bool>>        E(nV, std::vector<bool>(nV, 1)); // To test for duplicates

                for (int k = 0; k < elements.size(); ++k)
                {
                    auto& tri = elements[k];
                    addDistConstraint(E, tri[0], tri[1]);
                    addDistConstraint(E, tri[0], tri[2]);
                    addDistConstraint(E, tri[1], tri[2]);
                }
            }
            else if (auto lineMesh = std::dynamic_pointer_cast<LineMesh>(m_geom))
            {
                std::shared_ptr<VecDataArray<int, 2>> elementsPtr = lineMesh->getCells();
                const VecDataArray<int, 2>&           elements    = *elementsPtr;
                const auto&                           nV = lineMesh->getNumVertices();
                std::vector<std::vector<bool>>        E(nV, std::vector<bool>(nV, 1)); // To test for duplicates

                for (int k = 0; k < elements.size(); k++)
                {
                    auto& seg = elements[k];
                    addDistConstraint(E, seg[0], seg[1]);
                }
            }
            else
            {
                LOG(WARNING) << "PbdDistanceConstraint can only be generated with a TetrahedralMesh, SurfaceMesh, or LineMesh";
            }
        }

        void addConstraints(PbdConstraintContainer&                     constraints,
                            std::shared_ptr<std::unordered_set<size_t>> vertices) override
        {
            std::shared_ptr<VecDataArray<double, 3>> initVerticesPtr = m_geom->getInitialVertexPositions();
            const VecDataArray<double, 3>&           initVertices    = *initVerticesPtr;

            std::set<std::pair<int, int>> distanceSet;
            if (auto surfMesh = std::dynamic_pointer_cast<SurfaceMesh>(m_geom))
            {
                std::shared_ptr<VecDataArray<int, 3>> elementsPtr = surfMesh->getCells();
                const VecDataArray<int, 3>&           elements    = *elementsPtr;

                // Build vertex to tri map
                std::vector<std::vector<int>> vertexToCellMap(surfMesh->getNumVertices());
                for (int i = 0; i < elements.size(); i++)
                {
                    const Vec3i& cell = elements[i];
                    vertexToCellMap[cell[0]].push_back(i);
                    vertexToCellMap[cell[1]].push_back(i);
                    vertexToCellMap[cell[2]].push_back(i);
                }
                for (std::vector<int>& faceIds : vertexToCellMap)
                {
                    std::sort(faceIds.begin(), faceIds.end());
                }

                for (const size_t vertIdx : *vertices)
                {
                    const int vid = static_cast<int>(vertIdx);
                    for (const int cellId : vertexToCellMap[vertIdx])
                    {
                        const Vec3i& cell = elements[cellId];
                        int          i1   = 0;
                        int          i2   = 0;
                        for (int i = 0; i < 3; i++)
                        {
                            if (cell[i] == static_cast<int>(vertIdx))
                            {
                                i1 = cell[(i + 1) % 3];
                                i2 = cell[(i + 2) % 3];
                                break;
                            }
                        }
                        auto pair1 = std::make_pair(std::min(vid, i1), std::max(vid, i1));
                        auto pair2 = std::make_pair(std::min(vid, i2), std::max(vid, i2));
                        distanceSet.insert(pair1);
                        distanceSet.insert(pair2);
                    }
                }
            }
            else if (auto lineMesh = std::dynamic_pointer_cast<LineMesh>(m_geom))
            {
                std::shared_ptr<VecDataArray<int, 2>> elementsPtr = lineMesh->getCells();
                const VecDataArray<int, 2>&           elements    = *elementsPtr;

                // Build vertex to cell map
                std::vector<std::vector<int>> vertexToCellMap(lineMesh->getNumVertices());
                for (int k = 0; k < elements.size(); k++)
                {
                    const Vec2i& cell = elements[k];
                    vertexToCellMap[cell[0]].push_back(k);
                    vertexToCellMap[cell[1]].push_back(k);
                }
                for (std::vector<int>& faceIds : vertexToCellMap)
                {
                    std::sort(faceIds.begin(), faceIds.end());
                }

                for (const size_t vertIdx : *vertices)
                {
                    for (const int triIdx : vertexToCellMap[vertIdx])
                    {
                        const Vec2i& cell = elements[triIdx];

                        auto pair1 = std::make_pair(std::min(cell[0], cell[1]), std::max(cell[0], cell[1]));
                        distanceSet.insert(pair1);
                    }
                }
            }
            else
            {
                LOG(FATAL) << "Add element constraints does not support current mesh type";
            }

            constraints.reserve(constraints.getConstraints().size() + distanceSet.size());
            for (auto& c : distanceSet)
            {
                auto constraint = makeDistConstraint(initVertices, c.first, c.second);
                constraints.addConstraint(constraint);
            }
        }

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

        double getStretch() const { return m_stretch; }
        void setStretch(double val) { m_stretch = val; }

    protected:
        double m_stiffness = 0.0;
        double m_stretch   = 1.0;
};

struct PbdFemTetConstraintFunctor : public PbdBodyConstraintFunctor
{
    public:
        PbdFemTetConstraintFunctor() = default;
        ~PbdFemTetConstraintFunctor() override = default;

        void operator()(PbdConstraintContainer& constraints) override
        {
            // Check for correct mesh type
            CHECK(std::dynamic_pointer_cast<TetrahedralMesh>(m_geom) != nullptr)
                << "PbdFemTetConstraint can only be generated with a TetrahedralMesh";

            // Create constraints
            auto                                     tetMesh     = std::dynamic_pointer_cast<TetrahedralMesh>(m_geom);
            std::shared_ptr<VecDataArray<double, 3>> verticesPtr = m_geom->getVertexPositions();
            const VecDataArray<double, 3>&           vertices    = *verticesPtr;
            std::shared_ptr<VecDataArray<int, 4>>    elementsPtr = tetMesh->getCells();
            const VecDataArray<int, 4>&              elements    = *elementsPtr;

            std::shared_ptr<VecDataArray<double, 3>> strainParameters;

            if (tetMesh->hasCellAttribute(TetrahedralMesh::StrainParameterName))
            {
                strainParameters = tetMesh->getStrainParameters();
            }

            ParallelUtils::parallelFor(elements.size(),
                [&](const size_t k)
                {
                    const Vec4i& tet = elements[k];

                    PbdFemTetConstraint::MaterialType materialType = m_matType;
                    PbdFemConstraintConfig config = *m_femConfig;
                    if (strainParameters && strainParameters->at(k)[0] >= 0)
                    {
                        materialType = static_cast<PbdFemTetConstraint::MaterialType>((int)strainParameters->at(k)[0]);
                        config.setYoungAndPoisson(strainParameters->at(k)[1], strainParameters->at(k)[2]);
                    }

                    auto c = std::make_shared<PbdFemTetConstraint>(materialType);

                    c->initConstraint(
                        vertices[tet[0]], vertices[tet[1]], vertices[tet[2]], vertices[tet[3]],
                        { m_bodyIndex, tet[0] }, { m_bodyIndex, tet[1] },
                        { m_bodyIndex, tet[2] }, { m_bodyIndex, tet[3] },
                        config);
                    constraints.addConstraint(c);
            }, elements.size() > 100);
        }

        void setMaterialType(const PbdFemTetConstraint::MaterialType materialType) { m_matType = materialType; }
        PbdFemTetConstraint::MaterialType getMaterialType() const { return m_matType; }

        void setFemConfig(std::shared_ptr<PbdFemConstraintConfig> femConfig) { m_femConfig = femConfig; }
        std::shared_ptr<PbdFemConstraintConfig> getFemConfig() const { return m_femConfig; }

    protected:
        PbdFemTetConstraint::MaterialType m_matType = PbdFemTetConstraint::MaterialType::StVK;
        std::shared_ptr<PbdFemConstraintConfig> m_femConfig = nullptr;
};

struct PbdVolumeConstraintFunctor : public PbdBodyConstraintFunctor
{
    public:
        PbdVolumeConstraintFunctor() = default;
        ~PbdVolumeConstraintFunctor() override = default;

        void operator()(PbdConstraintContainer& constraints) override
        {
            // Check for correct mesh type
            CHECK(std::dynamic_pointer_cast<TetrahedralMesh>(m_geom) != nullptr)
                << "PbdVolumeConstraint can only be generated with a TetrahedralMesh";

            // Create constraints
            auto                                     tetMesh     = std::dynamic_pointer_cast<TetrahedralMesh>(m_geom);
            std::shared_ptr<VecDataArray<double, 3>> verticesPtr = m_geom->getVertexPositions();
            const VecDataArray<double, 3>&           vertices    = *verticesPtr;
            std::shared_ptr<VecDataArray<int, 4>>    elementsPtr = tetMesh->getCells();
            const VecDataArray<int, 4>&              elements    = *elementsPtr;

            ParallelUtils::parallelFor(elements.size(),
                [&](const size_t k)
                {
                    auto& tet = elements[k];
                    auto c    = std::make_shared<PbdVolumeConstraint>();
                    c->initConstraint(
                        vertices[tet[0]], vertices[tet[1]], vertices[tet[2]], vertices[tet[3]],
                        { m_bodyIndex, tet[0] }, { m_bodyIndex, tet[1] },
                        { m_bodyIndex, tet[2] }, { m_bodyIndex, tet[3] },
                        m_stiffness);
                    constraints.addConstraint(c);
            });
        }

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

    protected:
        double m_stiffness = 0.0;
};

struct PbdAreaConstraintFunctor : public PbdBodyConstraintFunctor
{
    public:
        PbdAreaConstraintFunctor() = default;
        ~PbdAreaConstraintFunctor() override = default;

        void operator()(PbdConstraintContainer& constraints) override
        {
            // Check for correct mesh type
            CHECK(std::dynamic_pointer_cast<SurfaceMesh>(m_geom) != nullptr)
                << "PbdAreaConstraint can only be generated with a SurfaceMesh";

            // ok, now create constraints
            auto                                     triMesh     = std::dynamic_pointer_cast<SurfaceMesh>(m_geom);
            std::shared_ptr<VecDataArray<double, 3>> verticesPtr = m_geom->getVertexPositions();
            const VecDataArray<double, 3>&           vertices    = *verticesPtr;
            std::shared_ptr<VecDataArray<int, 3>>    elemenstPtr = triMesh->getCells();
            const VecDataArray<int, 3>&              elements    = *elemenstPtr;

            ParallelUtils::parallelFor(elements.size(),
                [&](const size_t k)
                {
                    auto& tri = elements[k];
                    auto c    = std::make_shared<PbdAreaConstraint>();
                    c->initConstraint(vertices[tri[0]], vertices[tri[1]], vertices[tri[2]],
                        { m_bodyIndex, tri[0] }, { m_bodyIndex, tri[1] }, { m_bodyIndex, tri[2] },
                        m_stiffness);
                    constraints.addConstraint(c);
            });
        }

        void addConstraints(PbdConstraintContainer&                     constraints,
                            std::shared_ptr<std::unordered_set<size_t>> vertices) override
        {
            // Check for correct mesh type
            CHECK(std::dynamic_pointer_cast<SurfaceMesh>(m_geom) != nullptr)
                << "Add element constraints does not support current mesh type.";

            auto                                     triMesh = std::dynamic_pointer_cast<SurfaceMesh>(m_geom);
            std::shared_ptr<VecDataArray<double, 3>> initVerticesPtr = m_geom->getInitialVertexPositions();
            std::shared_ptr<VecDataArray<int, 3>>    elementsPtr     = triMesh->getCells();

            const VecDataArray<double, 3>& initVertices = *initVerticesPtr;
            const VecDataArray<int, 3>&    elements     = *elementsPtr;

            // Build vertex to tri map
            std::vector<std::vector<size_t>> vertexToTriMap(triMesh->getNumVertices());
            for (int k = 0; k < elements.size(); k++)
            {
                const Vec3i& tri = elements[k];
                vertexToTriMap[tri[0]].push_back(k);
                vertexToTriMap[tri[1]].push_back(k);
                vertexToTriMap[tri[2]].push_back(k);
            }
            for (std::vector<size_t>& faceIds : vertexToTriMap)
            {
                std::sort(faceIds.begin(), faceIds.end());
            }

            std::unordered_set<size_t> areaSet;
            for (const size_t& vertIdx : *vertices)
            {
                for (const size_t& triIdx : vertexToTriMap[vertIdx])
                {
                    areaSet.insert(triIdx);
                }
            }

            auto addAreaConstraint =
                [&](size_t k, double stiffness)
                {
                    const Vec3i& tri = elements[k];
                    auto         c   = std::make_shared<PbdAreaConstraint>();
                    c->initConstraint(initVertices[tri[0]], initVertices[tri[1]], initVertices[tri[2]],
                        { m_bodyIndex, tri[0] }, { m_bodyIndex, tri[1] }, { m_bodyIndex, tri[2] },
                        stiffness);
                    constraints.addConstraint(c);
                };

            constraints.reserve(constraints.getConstraints().size() + areaSet.size());
            for (auto& c : areaSet)
            {
                addAreaConstraint(c, m_stiffness);
            }
        }

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

    protected:
        double m_stiffness = 0.0;
};

struct PbdBendConstraintFunctor : public PbdBodyConstraintFunctor
{
    public:
        PbdBendConstraintFunctor() = default;
        ~PbdBendConstraintFunctor() override = default;

        void operator()(PbdConstraintContainer& constraints) override
        {
            // Check for correct mesh type
            CHECK(std::dynamic_pointer_cast<LineMesh>(m_geom) != nullptr)
                << "PbdBendConstraint can only be generated with a LineMesh";

            auto                                     lineMesh    = std::dynamic_pointer_cast<LineMesh>(m_geom);
            std::shared_ptr<VecDataArray<double, 3>> verticesPtr = m_geom->getVertexPositions();
            const VecDataArray<double, 3>&           vertices    = *verticesPtr;
            /*std::shared_ptr< VecDataArray<int, 2>> indicesPtr = lineMesh->getCells();
            const VecDataArray<int, 2>& indices = *indicesPtr*/

            auto addBendConstraint =
                [&](const double k, int i1, int i2, int i3)
                {
                    // i1 should always come first
                    if (i2 < i1)
                    {
                        std::swap(i1, i2);
                    }
                    // i3 should always come last
                    if (i2 > i3)
                    {
                        std::swap(i2, i3);
                    }

                    auto c = std::make_shared<PbdBendConstraint>();
                    if (m_restLengthOverride >= 0.0)
                    {
                        c->initConstraint({ m_bodyIndex, i1 }, { m_bodyIndex, i2 }, { m_bodyIndex, i3 },
                            m_restLengthOverride, k);
                    }
                    else
                    {
                        c->initConstraint(vertices[i1], vertices[i2], vertices[i3],
                            { m_bodyIndex, i1 }, { m_bodyIndex, i2 }, { m_bodyIndex, i3 }, k);
                    }
                    constraints.addConstraint(c);
                };

            // Iterate sets of stride # of segments
            for (int k = 0; k < vertices.size() - m_stride * 2; k++)
            {
                addBendConstraint(m_stiffness, k, k + m_stride, k + 2 * m_stride);
            }
        }

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

        void setStride(const int stride)
        {
            m_stride = stride;
            CHECK(m_stride >= 1) << "Stride should be at least 1.";
        }

        int getStride() const { return m_stride; }

        double getRestLength() const { return m_restLengthOverride; }
        void setRestLength(const double restLength) { m_restLengthOverride = restLength; }

    protected:
        double m_stiffness = 0.0;
        int m_stride       = 3;
        double m_restLengthOverride = -1.0;
};

struct PbdDihedralConstraintFunctor : public PbdBodyConstraintFunctor
{
    public:
        PbdDihedralConstraintFunctor() = default;
        ~PbdDihedralConstraintFunctor() override = default;

        void operator()(PbdConstraintContainer& constraints) override
        {
            // Check for correct mesh type
            CHECK(std::dynamic_pointer_cast<SurfaceMesh>(m_geom) != nullptr)
                << "PbdDihedralConstraint can only be generated with a SurfaceMesh";

            // Create constraints
            auto                                     triMesh     = std::dynamic_pointer_cast<SurfaceMesh>(m_geom);
            std::shared_ptr<VecDataArray<double, 3>> verticesPtr = triMesh->getVertexPositions();
            const VecDataArray<double, 3>&           vertices    = *verticesPtr;
            std::shared_ptr<VecDataArray<int, 3>>    elementsPtr = triMesh->getCells();
            const VecDataArray<int, 3>&              elements    = *elementsPtr;
            const int                                nV = triMesh->getNumVertices();
            std::vector<std::vector<int>>            vertIdsToTriangleIds(nV);

            for (int k = 0; k < elements.size(); ++k)
            {
                const Vec3i& tri = elements[k];
                vertIdsToTriangleIds[tri[0]].push_back(k);
                vertIdsToTriangleIds[tri[1]].push_back(k);
                vertIdsToTriangleIds[tri[2]].push_back(k);
            }

            // Used to resolve duplicates
            std::vector<std::vector<bool>> E(nV, std::vector<bool>(nV, 1));

            auto addDihedralConstraint =
                [&](const std::vector<int>& r1, const std::vector<int>& r2,
                    const int k, int i1, int i2)
                {
                    if (i1 > i2) // Make sure i1 is always smaller than i2
                    {
                        std::swap(i1, i2);
                    }
                    if (E[i1][i2])
                    {
                        E[i1][i2] = 0;

                        // Find the shared edge
                        std::vector<size_t> rs(2);
                        auto                it = std::set_intersection(r1.begin(), r1.end(), r2.begin(), r2.end(), rs.begin());
                        rs.resize(static_cast<size_t>(it - rs.begin()));
                        if (rs.size() > 1)
                        {
                            int          idx  = (rs[0] == k) ? 1 : 0;
                            const Vec3i& tri0 = elements[k];
                            const Vec3i& tri1 = elements[rs[idx]];
                            int          idx0 = 0;
                            int          idx1 = 0;
                            for (int i = 0; i < 3; i++)
                            {
                                if (tri0[i] != i1 && tri0[i] != i2)
                                {
                                    idx0 = tri0[i];
                                }
                                if (tri1[i] != i1 && tri1[i] != i2)
                                {
                                    idx1 = tri1[i];
                                }
                            }
                            auto c = std::make_shared<PbdDihedralConstraint>();
                            c->initConstraint(vertices[idx0], vertices[idx1], vertices[i1], vertices[i2],
                                { m_bodyIndex, idx0 }, { m_bodyIndex, idx1 },
                                { m_bodyIndex, i1 }, { m_bodyIndex, i2 }, m_stiffness);
                            constraints.addConstraint(c);
                        }
                    }
                };

            for (size_t i = 0; i < vertIdsToTriangleIds.size(); i++)
            {
                std::sort(vertIdsToTriangleIds[i].begin(), vertIdsToTriangleIds[i].end());
            }

            // For every triangle
            for (int k = 0; k < elements.size(); ++k)
            {
                const Vec3i& tri = elements[k];

                // Get all the neighbor triangles (to the vertices)
                std::vector<int>& neighborTriangles0 = vertIdsToTriangleIds[tri[0]];
                std::vector<int>& neighborTriangles1 = vertIdsToTriangleIds[tri[1]];
                std::vector<int>& neighborTriangles2 = vertIdsToTriangleIds[tri[2]];

                // Add constraints between all the triangles
                addDihedralConstraint(neighborTriangles0, neighborTriangles1, k, tri[0], tri[1]);
                addDihedralConstraint(neighborTriangles0, neighborTriangles2, k, tri[0], tri[2]);
                addDihedralConstraint(neighborTriangles1, neighborTriangles2, k, tri[1], tri[2]);
            }
        }

        void addConstraints(PbdConstraintContainer&                     constraints,
                            std::shared_ptr<std::unordered_set<size_t>> vertices) override
        {
            // Check for correct mesh type
            CHECK(std::dynamic_pointer_cast<SurfaceMesh>(m_geom) != nullptr)
                << "Add element constraints does not support current mesh type.";

            auto                                     triMesh = std::dynamic_pointer_cast<SurfaceMesh>(m_geom);
            std::shared_ptr<VecDataArray<double, 3>> initVerticesPtr = m_geom->getInitialVertexPositions();
            std::shared_ptr<VecDataArray<int, 3>>    elementsPtr     = triMesh->getCells();

            const VecDataArray<double, 3>& initVertices = *initVerticesPtr;
            const VecDataArray<int, 3>&    elements     = *elementsPtr;

            // Build vertex to tri map
            std::vector<std::vector<size_t>> vertexToTriMap(triMesh->getNumVertices());
            for (int k = 0; k < elements.size(); k++)
            {
                const Vec3i& tri = elements[k];
                vertexToTriMap[tri[0]].push_back(k);
                vertexToTriMap[tri[1]].push_back(k);
                vertexToTriMap[tri[2]].push_back(k);
            }
            for (std::vector<size_t>& faceIds : vertexToTriMap)
            {
                std::sort(faceIds.begin(), faceIds.end());
            }

            std::map<std::pair<size_t, size_t>, std::pair<size_t, size_t>> dihedralSet;
            for (const size_t& vertIdx : *vertices)
            {
                for (const size_t& triIdx : vertexToTriMap[vertIdx])
                {
                    const Vec3i& tri = elements[triIdx];
                    for (size_t i = 0; i < 3; i++)
                    {
                        size_t j  = (i + 1) % 3;
                        size_t i0 = tri[i];
                        size_t i1 = tri[j];
                        if (i0 > i1)
                        {
                            std::swap(i0, i1);
                        }
                        const std::vector<size_t>& r0 = vertexToTriMap[i0];
                        const std::vector<size_t>& r1 = vertexToTriMap[i1];
                        std::vector<size_t>        rs(2);
                        auto                       it = std::set_intersection(r0.begin(), r0.end(), r1.begin(), r1.end(), rs.begin());
                        rs.resize(static_cast<size_t>(it - rs.begin()));
                        if (rs.size() > 1)
                        {
                            dihedralSet[std::make_pair(i0, i1)] = std::make_pair(rs[0], rs[1]);
                        }
                    }
                }
            }

            auto addDihedralConstraint =
                [&](int t0, int t1, int i1, int i2, double stiffness)
                {
                    const Vec3i& tri0 = elements[t0];
                    const Vec3i& tri1 = elements[t1];
                    int          idx0 = 0;
                    int          idx1 = 0;
                    for (int i = 0; i < 3; i++)
                    {
                        if (tri0[i] != i1 && tri0[i] != i2)
                        {
                            idx0 = tri0[i];
                        }
                        if (tri1[i] != i1 && tri1[i] != i2)
                        {
                            idx1 = tri1[i];
                        }
                    }
                    auto c = std::make_shared<PbdDihedralConstraint>();
                    c->initConstraint(initVertices[idx0], initVertices[idx1], initVertices[i1], initVertices[i2],
                        { m_bodyIndex, idx0 }, { m_bodyIndex, idx1 },
                        { m_bodyIndex, i1 }, { m_bodyIndex, i2 },
                        stiffness);
                    constraints.addConstraint(c);
                };

            constraints.reserve(constraints.getConstraints().size() + dihedralSet.size());
            for (auto& c : dihedralSet)
            {
                addDihedralConstraint(
                    static_cast<int>(c.second.first), static_cast<int>(c.second.second),
                    static_cast<int>(c.first.first), static_cast<int>(c.first.second),
                    m_stiffness);
            }
        }

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

    protected:
        double m_stiffness = 0.0;
};

struct PbdConstantDensityConstraintFunctor : public PbdBodyConstraintFunctor
{
    public:
        PbdConstantDensityConstraintFunctor() = default;
        ~PbdConstantDensityConstraintFunctor() override = default;

        void operator()(PbdConstraintContainer& constraints) override
        {
            // Check for correct mesh type
            CHECK(std::dynamic_pointer_cast<PointSet>(m_geom) != nullptr)
                << "PbdConstantDensityConstraint can only be generated with a PointSet";

            auto c = std::make_shared<PbdConstantDensityConstraint>();
            c->initConstraint(m_geom->getNumVertices(), m_bodyIndex, m_particleRadius, m_restDensity);
            constraints.addConstraint(c);
        }

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

        void setParticleRadius(const double particleRadius) { m_particleRadius = particleRadius; }
        double getParticleRadius() const { return m_particleRadius; }

        void setRestDensity(const double restDensity) { m_restDensity = restDensity; }
        double getRestDensity() const { return m_restDensity; }

    protected:
        double m_stiffness      = 0.0;
        double m_particleRadius = 0.2;
        double m_restDensity    = 6378.0;
};
} // namespace imstk