// Copyright (c) 2012-2024 Wojciech Figat. All rights reserved.

#include "InverseKinematics.h"

void InverseKinematics::SolveAimIK(const Transform& node, const Vector3& target, Quaternion& outNodeCorrection)
{
    Vector3 toTarget = target - node.Translation;
    toTarget.Normalize();
    const Vector3 fromNode = Vector3::Forward;
    Quaternion::FindBetween(fromNode, toTarget, outNodeCorrection);
}


void InverseKinematics::SolveTwoBoneIK(Transform& rootTransform, Transform& midJointTransform, Transform& endEffectorTransform, const Vector3& targetPosition, const Vector3& poleVector, bool allowStretching, float maxStretchScale)
{
    // Calculate limb segment lengths
    Real lowerLimbLength = (endEffectorTransform.Translation - midJointTransform.Translation).Length();
    Real upperLimbLength = (midJointTransform.Translation - rootTransform.Translation).Length();
    Vector3 midJointPos = midJointTransform.Translation;

    // Calculate the direction and length towards the target
    Vector3 toTargetVector = targetPosition - rootTransform.Translation;
    Real toTargetLength = toTargetVector.Length();
    Real totalLimbLength = lowerLimbLength + upperLimbLength;

    // Normalize the direction vector or set a default direction if too small
    Vector3 toTargetDir;
    if (toTargetLength < ZeroTolerance)
    {
        toTargetLength = ZeroTolerance;
        toTargetDir = Vector3(1, 0, 0);
    }
    else
    {
        toTargetDir = toTargetVector.GetNormalized();
    }

    // Calculate the pole vector direction
    Vector3 poleVectorDelta = poleVector - rootTransform.Translation;
    const Real poleVectorLengthSqr = poleVectorDelta.LengthSquared();

    Vector3 jointPlaneNormal, bendDirection;
    if (poleVectorLengthSqr < ZeroTolerance * ZeroTolerance)
    {
        bendDirection = Vector3::Forward;
        jointPlaneNormal = Vector3::Up;
    }
    else
    {
        jointPlaneNormal = toTargetDir ^ poleVectorDelta;
        if (jointPlaneNormal.LengthSquared() < ZeroTolerance * ZeroTolerance)
        {
            toTargetDir.FindBestAxisVectors(jointPlaneNormal, bendDirection);
        }
        else
        {
            jointPlaneNormal.Normalize();
            bendDirection = poleVectorDelta - (poleVectorDelta | toTargetDir) * toTargetDir;
            bendDirection.Normalize();
        }
    }

    // Handle limb stretching if allowed
    if (allowStretching)
    {
        const Real initialStretchRatio = 1.0f;
        const Real stretchRange = maxStretchScale - initialStretchRatio;
        if (stretchRange > ZeroTolerance && totalLimbLength > ZeroTolerance)
        {
            const Real reachRatio = toTargetLength / totalLimbLength;
            const Real scalingFactor = (maxStretchScale - 1.0f) * Math::Saturate((reachRatio - initialStretchRatio) / stretchRange);
            if (scalingFactor > ZeroTolerance)
            {
                lowerLimbLength *= 1.0f + scalingFactor;
                upperLimbLength *= 1.0f + scalingFactor;
                totalLimbLength *= 1.0f + scalingFactor;
            }
        }
    }

    // Calculate new positions for joint and end effector
    Vector3 newEndEffectorPos = targetPosition;
    Vector3 newMidJointPos = midJointPos;

    if (toTargetLength >= totalLimbLength) {
        // Target is beyond the reach of the limb
        Vector3 rootToEnd = (targetPosition - rootTransform.Translation).GetNormalized();

        // Calculate the slight offset towards the pole vector
        Vector3 rootToPole = (poleVector - rootTransform.Translation).GetNormalized();
        Vector3 slightBendDirection = Vector3::Cross(rootToEnd, rootToPole);
        if (slightBendDirection.LengthSquared() < ZeroTolerance * ZeroTolerance) {
            slightBendDirection = Vector3::Up;
        }
        else {
            slightBendDirection.Normalize();
        }

        // Calculate the direction from root to mid joint with a slight offset towards the pole vector
        Vector3 midJointDirection = Vector3::Cross(slightBendDirection, rootToEnd).GetNormalized();
        Real slightOffset = upperLimbLength * 0.01f; // Small percentage of the limb length for slight offset
        newMidJointPos = rootTransform.Translation + rootToEnd * (upperLimbLength - slightOffset) + midJointDirection * slightOffset;
    }
    else
    {
        // Target is within reach, calculate joint position
        const Real twoAb = 2.0f * upperLimbLength * toTargetLength;
        const Real cosAngle = twoAb > ZeroTolerance ? (upperLimbLength * upperLimbLength + toTargetLength * toTargetLength - lowerLimbLength * lowerLimbLength) / twoAb : 0.0f;
        const bool reverseUpperBone = cosAngle < 0.0f;
        const Real angle = Math::Acos(cosAngle);
        const Real jointLineDist = upperLimbLength * Math::Sin(angle);
        const Real projJointDistSqr = upperLimbLength * upperLimbLength - jointLineDist * jointLineDist;
        Real projJointDist = projJointDistSqr > 0.0f ? Math::Sqrt(projJointDistSqr) : 0.0f;
        if (reverseUpperBone)
            projJointDist *= -1.0f;
        newMidJointPos = rootTransform.Translation + projJointDist * toTargetDir + jointLineDist * bendDirection;
    }

    // Update root joint orientation
    {
        // Vector from root joint to mid joint (local Y-axis direction)
        Vector3 localY = (newMidJointPos - rootTransform.Translation).GetNormalized();

        // Vector from mid joint to end effector (used to calculate plane normal)
        Vector3 midToEnd = (newEndEffectorPos - newMidJointPos).GetNormalized();

        // Calculate the plane normal (local Z-axis direction)
        Vector3 localZ = Vector3::Cross(localY, midToEnd).GetNormalized();

        // Calculate the local X-axis direction, should be perpendicular to the Y and Z axes
        Vector3 localX = Vector3::Cross(localY, localZ).GetNormalized();

        // Correct the local Z-axis direction based on the cross product of X and Y to ensure orthogonality
        localZ = Vector3::Cross(localX, localY).GetNormalized();

        // Construct a rotation from the orthogonal basis vectors
        Quaternion newRootJointOrientation = Quaternion::LookRotation(localZ, localY);

        // Apply the new rotation to the root joint
        rootTransform.Orientation = newRootJointOrientation;
    }


    // Update mid joint orientation to point Y-axis towards the end effector and Z-axis perpendicular to the IK plane
    {
        // Vector from mid joint to end effector (local Y-axis direction after rotation)
        Vector3 midToEnd = (newEndEffectorPos - newMidJointPos).GetNormalized();

        // Calculate the plane normal using the root, mid joint, and end effector positions (will be the local Z-axis direction)
        Vector3 rootToMid = (newMidJointPos - rootTransform.Translation).GetNormalized();
        Vector3 planeNormal = Vector3::Cross(rootToMid, midToEnd).GetNormalized();


        // Vector from mid joint to end effector (local Y-axis direction)
        Vector3 localY = (newEndEffectorPos - newMidJointPos).GetNormalized();

        // Calculate the plane normal using the root, mid joint, and end effector positions (local Z-axis direction)
        Vector3 localZ = Vector3::Cross(rootToMid, localY).GetNormalized();

        //// Calculate the local X-axis direction, should be perpendicular to the Y and Z axes
        Vector3 localX = Vector3::Cross(localY, localZ).GetNormalized();

        // Correct the local Z-axis direction based on the cross product of X and Y to ensure orthogonality
        localZ = Vector3::Cross(localX, localY).GetNormalized();


        // Construct a rotation from the orthogonal basis vectors
        // The axes are used differently here than a standard LookRotation to align Z towards the end and Y perpendicular
        Quaternion newMidJointOrientation = Quaternion::LookRotation(localZ, localY); // Assuming FromLookRotation creates a rotation with the first vector as forward and the second as up

        // Apply the new rotation to the mid joint
        midJointTransform.Orientation = newMidJointOrientation;
        midJointTransform.Translation = newMidJointPos;
    }

    // Update end effector transform
    endEffectorTransform.Translation = newEndEffectorPos;
}