openmw/apps/openmw/mwrender/easyik.cpp

// Authors: Joar Engberg, GPT4o, Maksim Eremenko :)
// Ported with a significant GPT4o (Copilot) assistance from a C# (Unity script) Inverse Kinematic solver implementation
// by Joar Engberg (EasyIK) https://github.com/joaen/EasyIK/blob/master/EasyIK/Assets/Scripts/EasyIK.cs

#include "easyik.hpp"
#include <algorithm>
#include <cmath>
#include <optional>
#include <osg/MatrixTransform>
#include <osg/Quat>
#include <osg/Vec3>
#include <vector>

// Note: armatureBones are in local space, poleDirection and ikTarget are in world space.
EasyIK::EasyIK(const std::vector<osg::ref_ptr<osg::MatrixTransform>>& armatureBones, const osg::Vec3f& ikTarget,
    int iterations, float tolerance, std::optional<osg::Vec3f> poleDirection, Debug::DebugDrawer* debugDrawer)
    : mNumberOfJoints(armatureBones.size()) // Number of joints is one more than the number of bones
    , mIkTarget(ikTarget)
    , mIterations(iterations)
    , mTolerance(tolerance)
    , mPoleDirection(poleDirection)
    , mDebugDrawer(debugDrawer)
    , mArmatureBones(armatureBones)
{
    initialize();
}

void EasyIK::initialize()
{
    mJointChainLength = 0;
    mJointPositions.resize(mNumberOfJoints);
    mLocalJointPositions.resize(mNumberOfJoints);
    mBoneLengths.resize(mNumberOfJoints - 1);

    for (int i = 0; i < mNumberOfJoints - 1; ++i)
    {
        osg::ref_ptr<osg::MatrixTransform> current = mArmatureBones[i];
        osg::ref_ptr<osg::MatrixTransform> child = mArmatureBones[i + 1];
        mBoneLengths[i] = child->getMatrix().getTrans().length();
        mJointChainLength += mBoneLengths[i];
    }
}

void EasyIK::poleConstraint()
{
    if (mNumberOfJoints != 3 || !mPoleDirection.has_value())
        return;

    osg::Vec3 limbAxis = mJointPositions[2] - mJointPositions[0];
    limbAxis.normalize();
    osg::Vec3 poleDirection = mPoleDirection.value();
    poleDirection.normalize();
    osg::Vec3 boneDirection = mJointPositions[1] - mJointPositions[0];
    boneDirection.normalize();

    // Orthonormalize the vectors
    poleDirection -= limbAxis * (poleDirection * limbAxis);
    poleDirection.normalize();
    boneDirection -= limbAxis * (boneDirection * limbAxis);
    boneDirection.normalize();

    osg::Quat angle;
    angle.makeRotate(boneDirection, poleDirection);

    mJointPositions[1] = angle * (mJointPositions[1] - mJointPositions[0]) + mJointPositions[0];
}

void EasyIK::backward()
{
    for (int i = mNumberOfJoints - 1; i >= 0; --i)
    {
        if (i == mNumberOfJoints - 1)
        {
            mJointPositions[i] = mIkTarget;
        }
        else
        {
            osg::Vec3 pos = mJointPositions[i] - mJointPositions[i + 1];
            pos.normalize();
            mJointPositions[i] = mJointPositions[i + 1] + pos * mBoneLengths[i];
        }
    }
}

void EasyIK::forward()
{
    for (int i = 0; i < mNumberOfJoints; ++i)
    {
        if (i == 0)
        {
            mJointPositions[i] = osg::computeLocalToWorld(mArmatureBones[0]->getParentalNodePaths()[0]).getTrans();
        }
        else
        {
            osg::Vec3 pos = mJointPositions[i] - mJointPositions[i - 1];
            pos.normalize();
            mJointPositions[i] = mJointPositions[i - 1] + pos * mBoneLengths[i - 1];
        }
    }
}

void EasyIK::solveIK()
{
    // Initialize joint positions in world space
    for (int i = 0; i < mNumberOfJoints; ++i)
    {
        mJointPositions[i] = osg::computeLocalToWorld(mArmatureBones[i]->getParentalNodePaths()[0]).getTrans();
    }

    // Calculate distance to target
    mDistanceToTarget = (mJointPositions[0] - mIkTarget).length();

    // If the target is out of reach, stretch the chain towards the target
    if (mDistanceToTarget > mJointChainLength)
    {
        osg::Vec3 direction = mIkTarget - mJointPositions[0];
        direction.normalize();
        for (int i = 1; i < mNumberOfJoints; ++i)
        {
            mJointPositions[i] = mJointPositions[i - 1] + direction * mBoneLengths[i - 1];
        }
    }
    else
    {
        // Iteratively solve the IK problem
        float distToTarget = (mJointPositions[mNumberOfJoints - 1] - mIkTarget).length();
        int counter = 0;
        while (distToTarget > mTolerance)
        {
            backward();
            forward();
            distToTarget = (mJointPositions[mNumberOfJoints - 1] - mIkTarget).length();
            counter++;
            if (counter > mIterations)
            {
                break;
            }
        }
    }

    // Apply pole constraint
    poleConstraint();

    // Apply the calculated positions to the bones in local space
    for (int i = 0; i < mNumberOfJoints - 1; ++i)
    {
        auto boneName = mArmatureBones[i]->getName();

        // World/local conversion data
        osg::Matrix worldToLocal = osg::computeWorldToLocal(mArmatureBones[i]->getParentalNodePaths()[0]);
        osg::Quat worldToLocalRot = worldToLocal.getRotate();

        // Directions should be in local space, so after deriving them from joint positions - we convert them to the
        // local space of the bone. Start direction IS a direction in which the bone is
        // currently pointing (in its own local space!), an assumption here is made that the bone is always pointing
        // along 1.0 0 0 in its own local space, so far - this assumption seem to stand correct.
        osg::Vec3f startDirection = osg::Vec3f(1.0f, 0.0f, 0.0f);
        osg::Vec3f solvedDirection = worldToLocalRot * (mJointPositions[i + 1] - mJointPositions[i]);
        solvedDirection.normalize();

        // Find a rotation from startDirection to solvedDirection
        osg::Quat targetRotation;
        targetRotation.makeRotate(startDirection, solvedDirection);

        // Rotate the bone
        osg::Matrix currentMatrix = mArmatureBones[i]->getMatrix();
        currentMatrix.preMultRotate(targetRotation);
        mArmatureBones[i]->setMatrix(currentMatrix);
    }
}

void EasyIK::alignFoot(const osg::Vec3f& normal)
{
    // This alignes a foot to the ground based on provided normal. Probably angular limits should be introduced.
    if (mNumberOfJoints < 1)
        return;

    osg::Vec3f invertedNormal = -normal;
    osg::Matrix worldToLocal = osg::computeWorldToLocal(mArmatureBones.back()->getParentalNodePaths()[0]);
    osg::Vec3f localNormal = worldToLocal.getRotate() * invertedNormal;

    osg::Vec3f startDirection = osg::Vec3f(1.0f, 0.0f, 0.0f);
    osg::Quat rotation;
    rotation.makeRotate(startDirection, localNormal);

    osg::Matrix currentMatrix = mArmatureBones.back()->getMatrix();
    currentMatrix.preMultRotate(rotation);
    mArmatureBones.back()->setMatrix(currentMatrix);
}

void EasyIK::debugDraw()
{
    for (const auto& jointPosition : mJointPositions)
    {
        mDebugDrawer->drawCube(jointPosition, osg::Vec3f(10.0f, 10.0f, 10.0f), Debug::colorCyan);
    }

    if (mPoleDirection.has_value())
    {
        osg::Vec3f polePosition
            = mJointPositions[0] + mPoleDirection.value() * 10.0f; // Arbitrary distance for visualization
        mDebugDrawer->drawCube(polePosition, osg::Vec3f(5.0f, 5.0f, 5.0f), Debug::colorYellow);
    }

    mDebugDrawer->drawCube(mIkTarget, osg::Vec3f(5.0f, 5.0f, 5.0f), Debug::colorGreen);
}