/*=========================================================================

   Library: iMSTK

   Copyright (c) Kitware, Inc. & Center for Modeling, Simulation,
   & Imaging in Medicine, Rensselaer Polytechnic Institute.

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

      http://www.apache.org/licenses/LICENSE-2.0.txt

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.

=========================================================================*/

#include "imstkLevelSetCH.h"
#include "imstkCollisionData.h"
#include "imstkImageData.h"
#include "imstkLevelSetDeformableObject.h"
#include "imstkLevelSetModel.h"
#include "imstkRbdConstraint.h"
#include "imstkRigidObject2.h"

namespace imstk
{
LevelSetCH::LevelSetCH()
{
    setKernel(m_kernelSize, m_kernelSigma);
}

LevelSetCH::~LevelSetCH()
{
    if (m_kernelWeights != nullptr)
    {
        delete[] m_kernelWeights;
    }
}

void
LevelSetCH::setInputLvlSetObj(std::shared_ptr<LevelSetDeformableObject> lvlSetObj)
{
    setInputObjectA(lvlSetObj);
}

void
LevelSetCH::setInputRigidObj(std::shared_ptr<RigidObject2> rbdObj)
{
    setInputObjectB(rbdObj);
    maskAllPoints();
}

std::shared_ptr<LevelSetDeformableObject>
LevelSetCH::getLvlSetObj()
{
    return std::dynamic_pointer_cast<LevelSetDeformableObject>(getInputObjectA());
}

std::shared_ptr<RigidObject2>
LevelSetCH::getRigidObj()
{
    return std::dynamic_pointer_cast<RigidObject2>(getInputObjectB());
}

void
LevelSetCH::setKernel(const int size, const double sigma)
{
    m_kernelSize  = size;
    m_kernelSigma = sigma;
    if (size % 2 == 0)
    {
        LOG(WARNING) << "LevelSetCH kernel size must be odd, increasing by 1";
        m_kernelSize++;
    }
    if (m_kernelWeights != nullptr)
    {
        delete[] m_kernelWeights;
    }
    m_kernelWeights = new double[size * size * size];

    const double invDiv   = 1.0 / (2.0 * sigma * sigma);
    const int    halfSize = static_cast<int>(size * 0.5);
    int          i = 0;
    for (int z = -halfSize; z < halfSize + 1; z++)
    {
        for (int y = -halfSize; y < halfSize + 1; y++)
        {
            for (int x = -halfSize; x < halfSize + 1; x++)
            {
                const double dist = Vec3i(x, y, z).cast<double>().norm();
                m_kernelWeights[i++] = std::exp(-dist * invDiv);
            }
        }
    }
}

void
LevelSetCH::handle(
    const CDElementVector<CollisionElement>& elementsA,
    const CDElementVector<CollisionElement>& elementsB)
{
    std::shared_ptr<LevelSetDeformableObject> lvlSetObj = getLvlSetObj();
    std::shared_ptr<RigidObject2>             rbdObj    = getRigidObj();

    if (lvlSetObj == nullptr || rbdObj == nullptr)
    {
        return;
    }

    std::shared_ptr<LevelSetModel> lvlSetModel = lvlSetObj->getLevelSetModel();
    std::shared_ptr<ImageData>     grid = std::dynamic_pointer_cast<SignedDistanceField>(lvlSetModel->getModelGeometry())->getImage();

    if (grid == nullptr)
    {
        LOG(FATAL) << "Error: level set model geometry is not ImageData";
        return;
    }

    //const Vec3i& dim = grid->getDimensions();
    const Vec3d& invSpacing = grid->getInvSpacing();
    const Vec3d& origin     = grid->getOrigin();

    // LevelSetCH requires both sides
    if (elementsA.getSize() != elementsB.getSize())
    {
        return;
    }

    if (m_useProportionalForce)
    {
        // Apply impulses at points of contacts
        for (size_t i = 0; i < elementsA.getSize(); i++)
        {
            const CollisionElement& lsmContactElement = elementsA[i];
            const CollisionElement& rbdContactElement = elementsB[i];

            if (lsmContactElement.m_type != CollisionElementType::PointDirection
                || rbdContactElement.m_type != CollisionElementType::PointIndexDirection)
            {
                continue;
            }

            // If the point is in the mask, let it apply impulses
            if (m_ptIdMask.count(rbdContactElement.m_element.m_PointIndexDirectionElement.ptIndex) != 0)
            {
                const Vec3d& pos    = lsmContactElement.m_element.m_PointDirectionElement.pt;
                const Vec3d& normal = lsmContactElement.m_element.m_PointDirectionElement.dir;
                const Vec3i  coord  = (pos - origin).cwiseProduct(invSpacing).cast<int>();

                // Scale the applied impulse by the normal force
                const double fN = normal.normalized().dot(rbdObj->getRigidBody()->getForce()) / rbdObj->getRigidBody()->getForce().norm();
                const double S  = std::max(fN, 0.0) * m_velocityScaling;

                const int halfSize = static_cast<int>(m_kernelSize * 0.5);
                int       j = 0;
                for (int z = -halfSize; z < halfSize + 1; z++)
                {
                    for (int y = -halfSize; y < halfSize + 1; y++)
                    {
                        for (int x = -halfSize; x < halfSize + 1; x++)
                        {
                            const Vec3i fCoord = coord + Vec3i(x, y, z);
                            lvlSetModel->addImpulse(fCoord, S * m_kernelWeights[j++]);
                        }
                    }
                }
            }
        }
    }
    else
    {
        // Apply impulses at points of contacts
        for (size_t i = 0; i < elementsA.getSize(); i++)
        {
            const CollisionElement& lsmContactElement = elementsA[i];
            const CollisionElement& rbdContactElement = elementsB[i];

            if (lsmContactElement.m_type != CollisionElementType::PointDirection
                || rbdContactElement.m_type != CollisionElementType::PointIndexDirection)
            {
                continue;
            }

            // If the point is in the mask, let it apply impulses
            if (m_ptIdMask.count(rbdContactElement.m_element.m_PointIndexDirectionElement.ptIndex) != 0)
            {
                const Vec3d& pos = lsmContactElement.m_element.m_PointDirectionElement.pt;
                //const Vec3d& normal = pdColData[i].dirAtoB;
                const Vec3i  coord = (pos - origin).cwiseProduct(invSpacing).cast<int>();
                const double S     = m_velocityScaling;

                const int halfSize = static_cast<int>(m_kernelSize * 0.5);
                int       j = 0;
                for (int z = -halfSize; z < halfSize + 1; z++)
                {
                    for (int y = -halfSize; y < halfSize + 1; y++)
                    {
                        for (int x = -halfSize; x < halfSize + 1; x++)
                        {
                            const Vec3i fCoord = coord + Vec3i(x, y, z);
                            lvlSetModel->addImpulse(fCoord, S * m_kernelWeights[j++]);
                        }
                    }
                }
            }
        }
    }
}

void
LevelSetCH::maskAllPoints()
{
    std::shared_ptr<PointSet> pointSet = std::dynamic_pointer_cast<PointSet>(getRigidObj()->getCollidingGeometry());
    for (int i = 0; i < pointSet->getNumVertices(); i++)
    {
        m_ptIdMask.insert(i);
    }
}
}