Kraken IK Solver (WIP)

require InlineDrawing;
require Kraken;
require Math;
require Animation;


object TwoBoneStretchyIKSolver : KrakenSolver {
  Xfo initPose[];
};


// Default Constructor
function TwoBoneStretchyIKSolver()
{

}


function TwoBoneStretchyIKSolver(
  Xfo initPose[])
{
  this.initPose = initPose;
}


// Return Arguments for Kraken
function KrakenSolverArg[] TwoBoneStretchyIKSolver.getArguments(){
  KrakenSolverArg args[] = this.parent.getArguments();
  args.push(KrakenSolverArg('rightSide', 'In', 'Boolean'));

  args.push(KrakenSolverArg('ikblend', 'In', 'Scalar'));
  args.push(KrakenSolverArg('softIK', 'In', 'Boolean'));
  args.push(KrakenSolverArg('softRatio', 'In', 'Scalar'));
  args.push(KrakenSolverArg('stretch', 'In', 'Boolean'));
  args.push(KrakenSolverArg('stretchBlend', 'In', 'Scalar'));
  args.push(KrakenSolverArg('slide', 'In', 'Scalar'));
  args.push(KrakenSolverArg('pin', 'In', 'Scalar'));

  args.push(KrakenSolverArg('root', 'In', 'Mat44'));
  args.push(KrakenSolverArg('bone0FK', 'In', 'Mat44'));
  args.push(KrakenSolverArg('bone1FK', 'In', 'Mat44'));
  args.push(KrakenSolverArg('ikHandle', 'In', 'Mat44'));
  args.push(KrakenSolverArg('upV', 'In', 'Mat44'));

  args.push(KrakenSolverArg('bone0Len', 'In', 'Scalar'));
  args.push(KrakenSolverArg('bone1Len', 'In', 'Scalar'));
  args.push(KrakenSolverArg('bone0Out', 'Out', 'Mat44'));
  args.push(KrakenSolverArg('bone1Out', 'Out', 'Mat44'));
  args.push(KrakenSolverArg('bone2Out', 'Out', 'Mat44'));
  return args;
}


// Solve
function TwoBoneStretchyIKSolver.solve!
(
  Boolean drawDebug,
  Scalar rigScale,
  Boolean rightSide,

  Scalar ikblend,
  Boolean softIK,
  Scalar softRatio,
  Boolean stretch,
  Scalar stretchBlend,
  Scalar slide,
  Scalar pin,

  Mat44 root,
  Mat44 bone0FK,
  Mat44 bone1FK,
  Mat44 ikHandle,
  Mat44 upV,

  Scalar bone0Len,
  Scalar bone1Len,
  io Mat44 bone0Out,
  io Mat44 bone1Out,
  io Mat44 bone2Out
){

  Xfo bone0FkXfo = Xfo(bone0FK);
  Xfo bone1FkXfo = Xfo(bone1FK);

  Xfo bone0Xfo;
  Xfo bone1Xfo;
  Xfo bone2Xfo;

  // TODO: Fix FK bone0 translate (could lock this in the component?)
  // TODO: Fix FK/IK blend "translation" blending; instead blend lengths

  Scalar outBone0Len = bone0Len;
  Scalar outBone1Len = bone1Len;

  // Scale to global rig scale
  Scalar scaledBone0Len = bone0Len * rigScale;
  Scalar scaledBone1Len = bone1Len * rigScale;

  // Solve IK
  if(ikblend > 0.0)
  {

    Vec3 ikHandlePos = ikHandle.translation();
    Vec3 rootPos = bone0Xfo.tr;
    Scalar distanceToIK = (rootPos - ikHandlePos).length();

    // Lock mid to the upVector (pole vector)
    // e.g. could be used to lock an elbow on a table
    if (pin > 0.0)
    {
        Vec3 pinPt = upV.translation();
        scaledBone0Len = Math_linearInterpolate(scaledBone0Len, (pinPt - rootPos).length(), pin);
        scaledBone1Len = Math_linearInterpolate(scaledBone1Len, (pinPt - ikHandlePos).length(), pin);

        outBone0Len = scaledBone0Len / rigScale;
        outBone1Len = scaledBone1Len / rigScale;
    }

    if (pin != 1.0)
    {

      // TODO: Allow scale of zero and the evaluation to finish
      Scalar chainLen = scaledBone0Len + scaledBone1Len;
      if (chainLen == 0) return;

      // Slide mid to end (+1) or to start (-1)
      if (slide != 0.0)
      {
        Scalar shift;
        if (slide > 0.0)
          shift = slide * scaledBone1Len;
        else
          shift = slide * scaledBone0Len;

        // Update the bone lengths
        scaledBone0Len += shift;
        scaledBone1Len -= shift;

        outBone0Len = scaledBone0Len / rigScale;
        outBone1Len = scaledBone1Len / rigScale;
        chainLen = scaledBone0Len + scaledBone1Len;
      }

      // Soft IK scale
      Scalar finalRatio = 1.0;
      if (softIK && softRatio > 0.0)
      {
        // For this we drag the IK handle behind after the softDist
        // See: http://www.softimageblog.com/archives/108
        Scalar softLen = softRatio * chainLen;
        Scalar hardLen = chainLen - softLen;

        // If we left the hardDistArea we entered the soft area
        if (distanceToIK > hardLen)
        {
          Scalar exponent = -(distanceToIK - hardLen) / softLen;
          Scalar softTargetDistance = softLen * (1 - exp(exponent)) + hardLen;

          finalRatio = softTargetDistance / distanceToIK;

          if (!stretch || stretchBlend < 1.0)
          {
            // Compute soft IK location for non-stretchy
            Scalar nonStretchyFinalRatio = Math_linearInterpolate(finalRatio, 1.0, stretchBlend);
            Vec3 direction = ikHandlePos - rootPos;
            direction = direction.multiplyScalar(nonStretchyFinalRatio);
            ikHandlePos = rootPos + direction;
            distanceToIK *= nonStretchyFinalRatio;

          }

          if (stretch && stretchBlend > 0.0)
          {
            // Scale bones by ratio between IK and Soft IK so that we preserve the soft IK
            // motion while hitting the IK handle. See: http://www.softimageblog.com/archives/109
            Scalar stretchyFinalRatio = 1.0 / Math_linearInterpolate(1.0, finalRatio, stretchBlend);

            chainLen *= stretchyFinalRatio;
            scaledBone0Len *= stretchyFinalRatio;
            scaledBone1Len *= stretchyFinalRatio;
            outBone0Len *= stretchyFinalRatio;
            outBone1Len *= stretchyFinalRatio;
          }

        }
      }

      // Stretchy
      if (stretch && stretchBlend > 0.0)
      {
          if (chainLen < distanceToIK)
          {
              Scalar diff = distanceToIK / chainLen;
              diff = Math_linearInterpolate(1.0, diff, stretchBlend);

              scaledBone0Len *= diff;
              scaledBone1Len *= diff;
              outBone0Len *= diff;
              outBone1Len *= diff;
          }

      }
    }

    solve2BoneIK(
      scaledBone0Len,
      scaledBone1Len,
      root.translation(),
      upV.translation(),
      ikHandlePos,
      bone0Xfo,
      bone1Xfo
    );

    bone2Xfo = bone1Xfo;
    bone2Xfo.tr = bone1Xfo.transformVector(Vec3(outBone1Len, 0.0, 0.0));

    // Set IK  scaling
    bone0Xfo.sc = Vec3(scaledBone0Len / (rigScale * bone0Len), rigScale, rigScale);
    bone1Xfo.sc = Vec3(scaledBone1Len / (rigScale * bone1Len), rigScale, rigScale);
    bone2Xfo.sc = Vec3(rigScale, rigScale, rigScale);

  }

  // Solve FK
  if (ikblend < 1.0)
  {
    // Project bone2 to the end of bone 1
    Xfo bone2FkXfo = bone1FkXfo;
    bone2FkXfo.tr = bone1FkXfo.transformVector(Vec3(outBone1Len, 0.0, 0.0));

    // Set FK scale
    bone0FkXfo.sc = bone0FkXfo.sc.multiplyScalar(rigScale);
    bone1FkXfo.sc = bone1FkXfo.sc.multiplyScalar(rigScale);
    bone2FkXfo.sc = bone2FkXfo.sc.multiplyScalar(rigScale);

    // Only FK
    if (ikblend == 0.0)
    {
      bone0Xfo = bone0FkXfo;
      bone1Xfo = bone1FkXfo;
      bone2Xfo = bone2FkXfo;
    }
    // Solve IK/FK blend
    else
    {
      bone0Xfo.ori = bone0FkXfo.ori.sphericalLinearInterpolate(bone0Xfo.ori, ikblend);
      bone0Xfo.sc =  bone0FkXfo.sc.linearInterpolate(bone0Xfo.sc, ikblend);

      bone1Xfo.tr =  bone1FkXfo.tr.linearInterpolate(bone1Xfo.tr, ikblend);
      bone1Xfo.ori = bone1FkXfo.ori.sphericalLinearInterpolate(bone1Xfo.ori, ikblend);
      bone1Xfo.sc =  bone1FkXfo.sc.linearInterpolate(bone1Xfo.sc, ikblend);

      bone2Xfo.tr =  bone2FkXfo.tr.linearInterpolate(bone2Xfo.tr, ikblend);
      bone2Xfo.ori = bone2FkXfo.ori.sphericalLinearInterpolate(bone2Xfo.ori, ikblend);
      bone2Xfo.sc =  bone2FkXfo.sc.linearInterpolate(bone2Xfo.sc, ikblend);
    }
  }

  bone0Out = bone0Xfo.toMat44();
  bone1Out = bone1Xfo.toMat44();
  bone2Out = bone2Xfo.toMat44();

  // Set debugging visibility.
  this.setDebug(drawDebug);
  if(this.drawDebug){

    Color boneColor(1.0, 1.0, 0);
    etDrawBone(this.handle.rootTransform, 'bone0', bone0Xfo, scaledBone0Len, scaledBone0Len * 0.15, boneColor);
    etDrawBone(this.handle.rootTransform, 'bone1', bone1Xfo, scaledBone1Len, scaledBone1Len * 0.15, boneColor);
  }
}