InsertionAlgorithm.h

#pragma once

#include <sofa/collisionAlgorithm/BaseAlgorithm.h>
#include <sofa/collisionAlgorithm/BaseGeometry.h>
#include <sofa/collisionAlgorithm/BaseOperation.h>
#include <sofa/collisionAlgorithm/operations/CreateCenterProximity.h>
#include <sofa/collisionAlgorithm/operations/FindClosestProximity.h>
#include <sofa/collisionAlgorithm/operations/Project.h>
#include <sofa/collisionAlgorithm/proximity/EdgeProximity.h>
#include <sofa/component/constraint/lagrangian/solver/ConstraintSolverImpl.h>
#include <sofa/component/statecontainer/MechanicalObject.h>

namespace sofa::collisionAlgorithm
{

class InsertionAlgorithm : public BaseAlgorithm
{
   public:
    SOFA_CLASS(InsertionAlgorithm, BaseAlgorithm);

    typedef core::behavior::MechanicalState<defaulttype::Vec3Types> MechStateTipType;
    typedef core::objectmodel::SingleLink<InsertionAlgorithm, BaseGeometry,
                                          BaseLink::FLAG_STOREPATH | BaseLink::FLAG_STRONGLINK>
        GeomLink;
    typedef DetectionOutput<BaseProximity, BaseProximity> AlgorithmOutput;
    typedef component::constraint::lagrangian::solver::ConstraintSolverImpl ConstraintSolver;

    GeomLink l_tipGeom, l_surfGeom, l_shaftGeom, l_volGeom;
    Data<AlgorithmOutput> d_collisionOutput, d_insertionOutput;
    Data<bool> d_projective;
    Data<SReal> d_punctureForceThreshold, d_tipDistThreshold;
    ConstraintSolver::SPtr m_constraintSolver;
    std::vector<BaseProximity::SPtr> m_couplingPts;
    Data<bool> d_drawCollision, d_drawPoints;
    Data<SReal> d_drawPointsScale;

    InsertionAlgorithm()
        : l_tipGeom(initLink("tipGeom", "Link to the geometry structure of the needle tip.")),
          l_surfGeom(
              initLink("surfGeom", "Link to the geometry of the surface punctured by the needle.")),
          l_shaftGeom(initLink("shaftGeom", "Link to the geometry structure of the needle shaft.")),
          l_volGeom(initLink("volGeom",
                             "Link to the geometry of volume wherein the needle is inserted.")),
          d_collisionOutput(initData(&d_collisionOutput, "collisionOutput",
                                     "Detected proximities during puncture.")),
          d_insertionOutput(initData(&d_insertionOutput, "insertionOutput",
                                     "Detected proximities during insertion.")),
          d_projective(initData(
              &d_projective, false, "projective",
              "Projection of closest detected proximity back onto the needle tip element.")),
          d_punctureForceThreshold(initData(&d_punctureForceThreshold, -1_sreal,
                                            "punctureForceThreshold",
                                            "Threshold for the force applied to the needle tip. "
                                            "Once exceeded, puncture is initiated.")),
          d_tipDistThreshold(initData(&d_tipDistThreshold, -1_sreal, "tipDistThreshold",
                                      "Threshold for the distance advanced by the needle tip since "
                                      "the last proximity detection. Once exceeded, a new "
                                      "proximity pair is added for the needle-volume coupling.")),
          m_constraintSolver(nullptr),
          m_couplingPts(),
          d_drawCollision(initData(&d_drawCollision, false, "drawcollision", "Draw collision.")),
          d_drawPoints(initData(&d_drawPoints, false, "drawPoints", "Draw detection outputs.")),
          d_drawPointsScale(initData(&d_drawPointsScale, 0.0005, "drawPointsScale",
                                     "Scale the drawing of detection output points."))
    {
    }

    void init() override
    {
        BaseAlgorithm::init();
        this->getContext()->get<ConstraintSolver>(m_constraintSolver);
        msg_warning_when(!m_constraintSolver)
            << "No constraint solver found in context. Insertion algorithm is disabled.";

        if (d_punctureForceThreshold.getValue() < 0)
        {
            msg_warning() << d_punctureForceThreshold.getName() +
                                 " parameter not defined or set to negative value." msgendl
                          << "Puncture will not function properly; provide a positive value";
            d_punctureForceThreshold.setValue(std::numeric_limits<double>::max());
        }

        if (d_tipDistThreshold.getValue() < 0)
        {
            msg_warning() << d_tipDistThreshold.getName() +
                                 " parameter not defined or set to negative value." msgendl
                          << "Needle-volume coupling is disabled; provide a positive value";
            d_tipDistThreshold.setValue(std::numeric_limits<double>::max());
        }
    }

    void draw(const core::visual::VisualParams* vparams)
    {
        if (!vparams->displayFlags().getShowCollisionModels() && !d_drawCollision.getValue())
            return;
        vparams->drawTool()->disableLighting();

        const AlgorithmOutput& collisionOutput = d_collisionOutput.getValue();
        for (const auto& it : collisionOutput)
        {
            vparams->drawTool()->drawLine(it.first->getPosition(), it.second->getPosition(),
                                          type::RGBAColor(0, 1, 0, 1));
        }

        const AlgorithmOutput& insertionOutput = d_insertionOutput.getValue();
        for (const auto& it : insertionOutput)
        {
            vparams->drawTool()->drawSphere(it.first->getPosition(), d_drawPointsScale.getValue(),
                                            type::RGBAColor(1, 0, 1, 0.9));
            vparams->drawTool()->drawSphere(it.second->getPosition(), d_drawPointsScale.getValue(),
                                            type::RGBAColor(0, 0, 1, 0.9));
            vparams->drawTool()->drawLine(it.first->getPosition(), it.second->getPosition(),
                                          type::RGBAColor(1, 1, 0, 1));
        }
    }

    void doDetection()
    {
        if (!l_tipGeom || !l_surfGeom || !l_shaftGeom || !l_volGeom) return;

        auto& collisionOutput = *d_collisionOutput.beginEdit();
        auto& insertionOutput = *d_insertionOutput.beginEdit();

        insertionOutput.clear();
        collisionOutput.clear();

        if (m_couplingPts.empty())
        {
            // 1. Puncture algorithm
            auto createTipProximity =
                Operations::CreateCenterProximity::Operation::get(l_tipGeom->getTypeInfo());
            auto findClosestProxOnSurf =
                Operations::FindClosestProximity::Operation::get(l_surfGeom);
            auto projectOnSurf = Operations::Project::Operation::get(l_surfGeom);
            auto projectOnTip = Operations::Project::Operation::get(l_tipGeom);

            const bool isProjective = d_projective.getValue();
            const SReal punctureForceThreshold = d_punctureForceThreshold.getValue();
            for (auto itTip = l_tipGeom->begin(); itTip != l_tipGeom->end(); itTip++)
            {
                BaseProximity::SPtr tipProx = createTipProximity(itTip->element());
                if (!tipProx) continue;
                const BaseProximity::SPtr surfProx = findClosestProxOnSurf(
                    tipProx, l_surfGeom.get(), projectOnSurf, getFilterFunc());
                if (surfProx)
                {
                    surfProx->normalize();

                    // 1.1 Check whether puncture is happening - if so, create coupling point
                    if (m_constraintSolver)
                    {
                        const MechStateTipType::SPtr mstate =
                            l_tipGeom->getContext()->get<MechStateTipType>();
                        const auto& lambda =
                            m_constraintSolver->getLambda()[mstate.get()].read()->getValue();
                        SReal norm{0_sreal};
                        for (const auto& l : lambda) {
                            norm += l.norm();
                        }
                        if (norm > punctureForceThreshold)
                        {
                            m_couplingPts.push_back(surfProx);
                            continue;
                        }
                    }

                    // 1.2 If not, create a proximity pair for the tip-surface collision
                    if (isProjective)
                    {
                        tipProx = projectOnTip(surfProx->getPosition(), itTip->element()).prox;
                        if (!tipProx) continue;
                        tipProx->normalize();
                    }
                    collisionOutput.add(tipProx, surfProx);
                }
            }

            // 1.3 Collision with the shaft geometry
            if (collisionOutput.size())
            {
            auto createShaftProximity =
                Operations::CreateCenterProximity::Operation::get(l_shaftGeom->getTypeInfo());
            auto projectOnShaft = Operations::Project::Operation::get(l_shaftGeom);
                for (auto itShaft = l_shaftGeom->begin(); itShaft != l_shaftGeom->end(); itShaft++)
                {
                    BaseProximity::SPtr shaftProx = createShaftProximity(itShaft->element());
                    if (!shaftProx) continue;
                    const BaseProximity::SPtr surfProx = findClosestProxOnSurf(
                        shaftProx, l_surfGeom.get(), projectOnSurf, getFilterFunc());
                    if (surfProx)
                    {
                        surfProx->normalize();
    
                        // 1.2 If not, create a proximity pair for the tip-surface collision
                        if (d_projective.getValue())
                        {
                            shaftProx = projectOnShaft(surfProx->getPosition(), itShaft->element()).prox;
                            if (!shaftProx) continue;
                            shaftProx->normalize();
                        }
                        collisionOutput.add(shaftProx, surfProx);
                    }
                }
            }
        }
        else
        {
            // 2. Needle insertion algorithm
            ElementIterator::SPtr itTip = l_tipGeom->begin();
            auto createTipProximity =
                Operations::CreateCenterProximity::Operation::get(itTip->getTypeInfo());
            const BaseProximity::SPtr tipProx = createTipProximity(itTip->element());

            // 2.1 Check whether coupling point should be added
            const type::Vec3 tip2Pt = m_couplingPts.back()->getPosition() - tipProx->getPosition();
            if (tip2Pt.norm() > d_tipDistThreshold.getValue())
            {
                auto findClosestProxOnVol =
                    Operations::FindClosestProximity::Operation::get(l_volGeom);
                auto projectOnVol = Operations::Project::Operation::get(l_volGeom);
                const BaseProximity::SPtr volProx =
                    findClosestProxOnVol(tipProx, l_volGeom.get(), projectOnVol, getFilterFunc());
                if (volProx)
                {
                    volProx->normalize();
                    m_couplingPts.push_back(volProx);
                }
            }
            else // Don't bother with removing the point that was just added
            {
                // 2.2. Check whether coupling point should be removed 
                ElementIterator::SPtr itShaft = l_shaftGeom->begin(l_shaftGeom->getSize() - 2);
                auto createShaftProximity =
                    Operations::CreateCenterProximity::Operation::get(itShaft->getTypeInfo());
                const BaseProximity::SPtr shaftProx = createShaftProximity(itShaft->element());
                const EdgeProximity::SPtr edgeProx = dynamic_pointer_cast<EdgeProximity>(shaftProx);
                const type::Vec3 normal = (edgeProx->element()->getP1()->getPosition() -
                                           edgeProx->element()->getP0()->getPosition())
                                              .normalized();
                if (dot(tip2Pt, normal) > 0_sreal) {
                    m_couplingPts.pop_back();
                }
            }
        }

        if (!m_couplingPts.empty())
        {
            // 3. Re-project proximities on the shaft geometry
            auto findClosestProxOnShaft =
                Operations::FindClosestProximity::Operation::get(l_shaftGeom);
            auto projectOnShaft = Operations::Project::Operation::get(l_shaftGeom);
            for (int i = 0; i < m_couplingPts.size(); i++)
            {
                const BaseProximity::SPtr shaftProx = findClosestProxOnShaft(
                    m_couplingPts[i], l_shaftGeom.get(), projectOnShaft, getFilterFunc());
                insertionOutput.add(shaftProx, m_couplingPts[i]);
            }
        }

        d_collisionOutput.endEdit();
        d_insertionOutput.endEdit();
    }
};

}  // namespace sofa::collisionAlgorithm