imstkPbdBody.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 "imstkMath.h"
#include "imstkVecDataArray.h"

#include <unordered_map>

namespace
{
template<typename T>
void
copyAndAllocate(const std::shared_ptr<T>& src, std::shared_ptr<T>& dest)
{
    if (src != nullptr)
    {
        if (dest == nullptr)
        {
            dest = std::make_shared<T>();
        }
        *dest = *src;
    }
}
} // namespace

namespace imstk
{
class PbdConstraint;
struct PbdBody
{
    public:
        enum class Type
        {
            DEFORMABLE,
            DEFORMABLE_ORIENTED,
            RIGID
        };

        PbdBody() : bodyHandle(-1) { }
        PbdBody(const int handle) : bodyHandle(handle) { }

        void deepCopy(const PbdBody& src)
        {
            fixedNodeInvMass = src.fixedNodeInvMass;
            bodyHandle       = src.bodyHandle;

            copyAndAllocate(src.prevVertices, prevVertices);
            copyAndAllocate(src.vertices, vertices);
            copyAndAllocate(src.velocities, velocities);
            copyAndAllocate(src.masses, masses);
            copyAndAllocate(src.invMasses, invMasses);

            bodyType = src.bodyType;
            if (getOriented())
            {
                copyAndAllocate(src.prevOrientations, prevOrientations);
                copyAndAllocate(src.orientations, orientations);
                copyAndAllocate(src.angularVelocities, angularVelocities);
                copyAndAllocate(src.inertias, inertias);
                copyAndAllocate(src.invInertias, invInertias);
            }

            fixedNodeIds     = src.fixedNodeIds;
            uniformMassValue = src.uniformMassValue;

            externalForce  = src.externalForce;
            externalTorque = src.externalTorque;
        }

        bool getOriented() const
        {
            return (bodyType == Type::DEFORMABLE_ORIENTED || bodyType == Type::RIGID);
        }

        void setRigid(const Vec3d& pos,
                      const double mass = 1.0,
                      const Quatd& orientation = Quatd::Identity(),
                      const Mat3d& inertia     = Mat3d::Identity())
        {
            bodyType = PbdBody::Type::RIGID;

            if (vertices == nullptr)
            {
                vertices = std::make_shared<VecDataArray<double, 3>>();
            }
            * vertices = { pos };

            uniformMassValue = mass;

            if (orientations == nullptr)
            {
                orientations = std::make_shared<StdVectorOfQuatd>();
            }
            * orientations = { orientation };

            if (inertias == nullptr)
            {
                inertias = std::make_shared<StdVectorOfMat3d>();
            }
            * inertias = { inertia };
        }

        void setRigidVelocity(const Vec3d& velocity,
                              const Vec3d& angularVelocity = Vec3d::Zero())
        {
            bodyType = PbdBody::Type::RIGID;

            if (velocities == nullptr)
            {
                velocities = std::make_shared<VecDataArray<double, 3>>();
            }
            * velocities = { velocity };

            if (angularVelocities == nullptr)
            {
                angularVelocities = std::make_shared<VecDataArray<double, 3>>();
            }
            * angularVelocities = { angularVelocity };
        }

        Vec3d getRigidPosition()
        {
            CHECK(bodyType == Type::RIGID) << "Body is not a rigid.";
            return vertices->at(0);
        }

        Quatd getRigidOrientation()
        {
            CHECK(bodyType == Type::RIGID) << "Body is not a rigid.";
            return orientations->at(0);
        }

        void overrideRigidPositionAndOrientation(const Vec3d& pos, const Quatd& orientation)
        {
            CHECK(bodyType == Type::RIGID) << "Body is not a rigid.";
            (*vertices)[0]     = (*prevVertices)[0] = pos;
            (*orientations)[0] = (*prevOrientations)[0] = orientation;
        }

        void overrideLinearAndAngularVelocity(const Vec3d& velocity, const Vec3d& angularVelocity)
        {
            CHECK(bodyType == Type::RIGID) << "Body is not a rigid.";
            (*velocities)[0] = velocity;
            (*angularVelocities)[0] = angularVelocity;
        }

        // Struct Data
        int bodyHandle; 
        Type bodyType = Type::DEFORMABLE;

        std::shared_ptr<VecDataArray<double, 3>> prevVertices;
        std::shared_ptr<VecDataArray<double, 3>> vertices;

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

        std::shared_ptr<DataArray<double>> masses;
        std::shared_ptr<DataArray<double>> invMasses;

        std::shared_ptr<StdVectorOfQuatd> prevOrientations;
        std::shared_ptr<StdVectorOfQuatd> orientations;

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

        std::shared_ptr<StdVectorOfMat3d> inertias;
        std::shared_ptr<StdVectorOfMat3d> invInertias;

        std::vector<int> fixedNodeIds;
        double uniformMassValue = 1.0;

        bool bodyGravity = true;

        Vec3d externalForce  = Vec3d::Zero();
        Vec3d externalTorque = Vec3d::Zero();

        std::unordered_map<int, double> fixedNodeInvMass;

        std::unordered_map<int, std::vector<std::shared_ptr<PbdConstraint>>> cellConstraintMap;
};

struct PbdState
{
    public:
        PbdState() { }

        void deepCopy(const PbdState& src)
        {
            // Copy by value not reference
            m_bodies.resize(src.m_bodies.size());
            for (size_t i = 0; i < m_bodies.size(); i++)
            {
                if (m_bodies[i] == nullptr)
                {
                    m_bodies[i] = std::make_shared<PbdBody>();
                }
                m_bodies[i]->deepCopy(*src.m_bodies[i]);
            }
        }

        inline Vec3d& getPosition(const std::pair<int, int>& bodyParticleId) const { return (*m_bodies[bodyParticleId.first]->vertices)[bodyParticleId.second]; }
        inline Vec3d& getVelocity(const std::pair<int, int>& bodyParticleId) const { return (*m_bodies[bodyParticleId.first]->velocities)[bodyParticleId.second]; }
        inline Quatd& getOrientation(const std::pair<int, int>& bodyParticleId) const { return (*m_bodies[bodyParticleId.first]->orientations)[bodyParticleId.second]; }
        inline Vec3d& getAngularVelocity(const std::pair<int, int>& bodyParticleId) const { return (*m_bodies[bodyParticleId.first]->angularVelocities)[bodyParticleId.second]; }

        inline double getInvMass(const std::pair<int, int>& bodyParticleId) const { return (*m_bodies[bodyParticleId.first]->invMasses)[bodyParticleId.second]; }
        inline Mat3d& getInvInertia(const std::pair<int, int>& bodyParticleId) const { return (*m_bodies[bodyParticleId.first]->invInertias)[bodyParticleId.second]; }

        inline PbdBody::Type getBodyType(const std::pair<int, int>& bodyParticleId) const { return m_bodies[bodyParticleId.first]->bodyType; }

        std::vector<std::shared_ptr<PbdBody>> m_bodies;
};
} // namespace imstk