Clone Arms Rigging in The Last Clockwinder

AssistantVisuals.cs

// Copyright 2022 John Austin - Pontoco
// License: MIT 
// https://choosealicense.com/licenses/mit/

// This is the rigging code that is used to calculate IK and place arm bones in The Last Clockwinder.
// It began using the Unity Animation Rigging package, but we found that to be too slow and overly separated into too many
// trivial jobs (it was dominated by scheduling time).

// This code is largely the result of manually inlining all of the rigging code we were using from our Animation Rigging setup.
// You'll notice that it's well setup to run as a Job! But we never needed the extra performance. In face, just running this code 25 times
// (once for each clone) was so much faster that this code was no longer a bottleneck. 

using System;
using Assets.Scripts.Utilities;
using Core.Assistants;
using Core.Audio;
using Core.Harness;
using Core.PlayerControl;
using Game.UI.HelpIndicator;
using Global.Utilities.Unity;
using Global.Utilities.Unity.Springs;
using NewtonVR;
using UniRx;
using UniRx.Triggers;
using Unity.Labs.SuperScience;
using UnityEngine;
using UnityEngine.Animations;
using UnityEngine.Animations.Rigging;
using UnityEngine.Playables;
using Utilities;
using Utilities.Unity;

namespace Game
{
    /// <summary>The main class of the assistant rigs. Contains the toplevel for assistant specific code.</summary>
    public class AssistantVisuals : MonoBehaviour
    {
        /// <summary></summary>
        public class Params : ScriptParams<Params>
        {
            public Material BackPlateMaterial1, BackPlateMaterial2, BackPlateMaterial4;

            public int MaxClonesToDisableAnimations = 18;
            public float DisableHandAnimationDistance = 3;
            public float LowFrameRateDistance = 5;
            public int LowFrameRateSkip = 2;
        }

        private class HandAnimation
        {
            private readonly SpringFloat spring = new SpringFloat(10, 3);
            private int timer;
            private bool holdPrev;
            private readonly AnimationClipPlayable clip;

            public HandAnimation(PlayableGraph graph, Animator animator, AnimationClip animClip)
            {
                clip = AnimationClipPlayable.Create(graph, animClip);
                var output = AnimationPlayableOutput.Create(graph, "Animator", animator);
                output.SetSourcePlayable(clip);
            }

            public void Update(bool holdButtonDown)
            {
                if (holdButtonDown)
                {
                    timer = 10;
                }

                spring.TargetValue = timer > 0 || holdButtonDown ? 1 : .3f;
                spring.Step(GameHarness.FixedTime);
                clip.SetTime(spring.CurrentValue);

                holdPrev = holdButtonDown;
                timer = Math.Max(timer - 1, 0);
            }
        }

        [Serializable]
        public class ArmRigging
        {
            private const bool debug = false;
            // Rigging for the skeleton. The skeleton is the intermediate rig that IK is performed on.
            public Transform SkeletonHand;
            public Vector3 SkeletonHandOffset; // Offset of the skeleton hand in local space of the rig hand

            public Transform SkeletonIkHint;
            private Matrix4x4 IkHintChestOffset;

            public Transform SkeletonArmRoot;
            public Transform SkeletonArmMid;
            public Transform SkeletonArmTip;

            private Matrix4x4 SkinArmRootOffset;
            private Matrix4x4 SkinArmMidOffset;

            // Rigging for the skinning. The skinning is the final flat-hierarchy bones that feed the animator.
            public Transform SkinHand;
            public Transform SkinArmRoot;
            public Transform SkinArmMid;
            private Matrix4x4 SkinHandOffset; // offset from the player rig hand positions.

            // Twist correction
            private Quaternion sourceInitialRotationInv;
            private Quaternion bicepInitialRotation;
            private Quaternion forearmInitialRotation;

            public void InitializeRigging(Transform skeletonTorso)
            {
                SkinHandOffset = TransformHelpers.GetTransformBetween(SkeletonHand, SkinHand);
                IkHintChestOffset = TransformHelpers.GetTransformBetween(SkeletonIkHint, skeletonTorso);

                SkinArmRootOffset = TransformHelpers.GetTransformBetween(SkinArmRoot, SkeletonArmRoot);
                SkinArmMidOffset = TransformHelpers.GetTransformBetween(SkinArmMid, SkeletonArmMid);

                // Twist Correction
                sourceInitialRotationInv = Quaternion.Inverse(SkeletonArmTip.localRotation);
                bicepInitialRotation = SkeletonArmRoot.localRotation;
                forearmInitialRotation = SkeletonArmMid.localRotation;
            }

            public void UpdateSkeleton(PlayerHand hand, Vector3 skeletonTorsoPosition, Quaternion skeletonTorsoRotation,
                                       float twistCorrectionAll, float twistCorrectionBicep,
                                       float twistCorrectionForearm)
            {
                TwistCorrection(SkeletonArmTip, SkeletonArmRoot, SkeletonArmMid, twistCorrectionAll,
                    twistCorrectionBicep, twistCorrectionForearm);
                
                // todo: we only need positions of these
                var armHint = Matrix4x4.TRS(skeletonTorsoPosition, skeletonTorsoRotation, Vector3.one) *
                              IkHintChestOffset;

                if (debug)
                {
                    GizmoModule.DrawSphere(armHint.ExtractTranslation(), .01f, Color.blue);
                }

                // Skeleton hand positions
                var skeletonHandPosition = hand.transform.localToWorldMatrix.MultiplyPoint3x4(SkeletonHandOffset);
                var skeletonHandRotation = hand.transform.rotation;

                // Skin hand positions
                // todo: can be optimized if not inlined/simplified.
                SkinHand.position = skeletonHandPosition + SkinHandOffset.ExtractTranslation();
                SkinHand.rotation = skeletonHandRotation * SkinHandOffset.ExtractRotation();

                // Skeleton Arms (IK)
                SolveTwoBoneIk(SkeletonArmRoot, SkeletonArmMid, SkeletonArmTip,
                    armHint.ExtractTranslation(), SkinHand.position);

                // Skin arms 
                SkinArmRoot.position =
                    (SkeletonArmRoot.localToWorldMatrix * SkinArmRootOffset).ExtractTranslation();
                SkinArmRoot.rotation =
                    (SkeletonArmRoot.localToWorldMatrix * SkinArmRootOffset).ExtractRotation();

                SkinArmMid.position =
                    (SkeletonArmMid.localToWorldMatrix * SkinArmMidOffset).ExtractTranslation();
                SkinArmMid.rotation =
                    (SkeletonArmMid.localToWorldMatrix * SkinArmMidOffset).ExtractRotation();
            }

            private void SolveTwoBoneIk( /*Vector3 rootPosition, Vector3 midPosition, Vector3 tipPosition, */
                Transform rootBone, Transform midBone, Transform tipBone,
                Vector3 hintPosition,
                Vector3 targetPosition)
            {
                const float k_SqrEpsilon = 1e-8f;
                Vector3 aPosition = rootBone.position;
                Vector3 bPosition = midBone.position;
                Vector3 cPosition = tipBone.position;

                Vector3 ab = bPosition - aPosition;
                Vector3 bc = cPosition - bPosition;
                Vector3 ac = cPosition - aPosition;
                Vector3 at = targetPosition - aPosition;

                float abLen = ab.magnitude;
                float bcLen = bc.magnitude;
                float acLen = ac.magnitude;
                float atLen = at.magnitude;

                float oldAbcAngle = TriangleAngle(acLen, abLen, bcLen);
                float newAbcAngle = TriangleAngle(atLen, abLen, bcLen);

                // Rotate the forearm to be the correct angle:

                // Bend normal strategy is to take whatever has been provided in the animation
                // stream to minimize configuration changes, however if this is collinear
                // try computing a bend normal given the desired target position.
                // If this also fails, try resolving axis using hint if provided.
                Vector3 axis = Vector3.Cross(ab, bc);
                if (axis.sqrMagnitude < k_SqrEpsilon)
                {
                    axis = Vector3.Cross(hintPosition - aPosition, bc);

                    if (axis.sqrMagnitude < k_SqrEpsilon)
                        axis = Vector3.Cross(at, bc);

                    if (axis.sqrMagnitude < k_SqrEpsilon)
                        axis = Vector3.up;
                }

                axis = Vector3.Normalize(axis);

                // todo replace with quaternion-axis-angle constructor
                float a = 0.5f * (oldAbcAngle - newAbcAngle);
                float sin = Mathf.Sin(a);
                float cos = Mathf.Cos(a);
                Quaternion deltaR = new Quaternion(axis.x * sin, axis.y * sin, axis.z * sin, cos);
                midBone.rotation = deltaR * midBone.rotation;

                // Rotate the root to align the forearm to the target:

                cPosition = tipBone.position; // This updates the positive based on the rotated mid bone.
                ac = cPosition - aPosition;
                rootBone.rotation = QuaternionExt.FromToRotation(ac, at) * rootBone.rotation;

                float acSqrMag = ac.sqrMagnitude;
                if (acSqrMag > 0f)
                {
                    bPosition = midBone.position;
                    cPosition = tipBone.position;
                    ab = bPosition - aPosition;
                    ac = cPosition - aPosition;

                    Vector3 acNorm = ac / Mathf.Sqrt(acSqrMag);
                    Vector3 ah = hintPosition - aPosition;
                    Vector3 abProj = ab - acNorm * Vector3.Dot(ab, acNorm);
                    Vector3 ahProj = ah - acNorm * Vector3.Dot(ah, acNorm);

                    float maxReach = abLen + bcLen;
                    if (abProj.sqrMagnitude > (maxReach * maxReach * 0.001f) && ahProj.sqrMagnitude > 0f)
                    {
                        Quaternion hintR = QuaternionExt.FromToRotation(abProj, ahProj);
                        rootBone.rotation = hintR * rootBone.rotation;
                    }
                }
            }

            private void TwistCorrection(Transform source, Transform bicep, Transform forearm,
                                         float twistCorrectionAll, float twistCorrectionBicep,
                                         float twistCorrectionForearm)
            {
                float w = twistCorrectionAll;

                if (w > 0f)
                {
                    Quaternion twistRot =
                        TwistRotation(Vector3.forward, sourceInitialRotationInv * source.localRotation);

                    // bicep

                    Quaternion rot = Quaternion.Lerp(Quaternion.identity, twistRot, twistCorrectionBicep);
                    bicep.localRotation = Quaternion.Lerp(bicepInitialRotation, rot, w);

                    rot = Quaternion.Lerp(Quaternion.identity, twistRot, twistCorrectionForearm);
                    forearm.localRotation = Quaternion.Lerp(forearmInitialRotation, rot, w);
                }
            }

            static Quaternion TwistRotation(Vector3 axis, Quaternion rot)
            {
                return new Quaternion(axis.x * rot.x, axis.y * rot.y, axis.z * rot.z, rot.w);
            }
        }

        /// <summary>The renderers to modify when the Assistant is selected.</summary>
        public Renderer BodyRenderer;
        public Renderer BackPlateRenderer;
        public Animator HandAnimator;
        public PlayerRig rig;
        public AnimationClip LHandClose, RHandClose;

        public Transform SkeletonTorso;
        public Vector3 TorsoHeadOffset; // Offset of the torso relative to the rig head

        public ArmRigging LeftArmRig, RightArmRig;

        // Rigging for the skinning. The skinning is the final flat-hierarchy bones that feed the animator.
        public Transform SkinTorso;
        public Transform SkinHead;
        private Matrix4x4 SkinTorsoOffset; // offset from the player rig hand positions.
        private Vector3 SkinHeadPositionOffset;
        private Quaternion SkinHeadRotationOffset;

        [Range(0, 1)]
        public float TwistCorrectionStrengthAll;
        [Range(0, 1)]
        public float TwistCorrectionStrengthBicep;
        [Range(0, 1)]
        public float TwistCorrectionStrengthForearm;

        private RigInputPlayback playback;
        private MaterialPropertyBlock materialPropertyBlock;
        private Assistant assistant;

        private static readonly int RingProgressProperty = Shader.PropertyToID("_RingProgress");

        private void Awake()
        {
            rig = GetComponent<PlayerRig>();
            playback = rig.Input.Playback.ValueOr(null);

            assistant = GetComponent<Assistant>();
            BackPlateRenderer.material = assistant.Data.MinLoopMeasures switch
            {
                1 => Params.Instance.BackPlateMaterial1,
                2 => Params.Instance.BackPlateMaterial2,
                4 => Params.Instance.BackPlateMaterial4,
                _ => Params.Instance.BackPlateMaterial2
            };

            materialPropertyBlock = new MaterialPropertyBlock();

            HandAnimator.applyRootMotion = false;
            graph = PlayableGraph.Create("HandAnimationGraph");
            graph.SetTimeUpdateMode(DirectorUpdateMode.GameTime);

            lHand = new HandAnimation(graph, HandAnimator, LHandClose);
            rHand = new HandAnimation(graph, HandAnimator, RHandClose);

            graph.Play();

            InitializeRigging();
        }

        public void UpdateEffects()
        {
            // Run purely visual logic in Update().
            if (playback == null)
            {
                return;
            }

            BackPlateRenderer.GetPropertyBlock(materialPropertyBlock);
            materialPropertyBlock.SetFloat(RingProgressProperty, playback.CurrentLoopProgress);
            BackPlateRenderer.SetPropertyBlock(materialPropertyBlock);
        }

        public NVRHand LeftHand, RightHand;

        private PlayableGraph graph;
        private HandAnimation lHand, rHand;

        public HelpBubble ShowDeleteHint()
        {
            Assistant a = GetComponent<Assistant>();

            Transform rigHead = GetComponentInParent<PlayerRig>().Head.transform;
            var bubble = HelpBubble.Spawn(rigHead, Vector3.up * .2f,
                HelpBubble.Params.Instance.HelpBubblePrefab_DeleteClone,
                // Lower the target position slightly to match the head visual rather than rig head position.
                targetOffsetWorld: new Vector3(0, -0.15f, 0));

            // Destroy if the clone is destroyed.
            a.OnDestroyAsObservable().First().Subscribe(() =>
            {
                if (bubble != null) // can happen during scene destruction. the bubble might get destroyed first.
                {
                    bubble.Hide();
                }
            });

            return bubble;
        }

        public void InitializeRigging()
        {
            SkinTorsoOffset = TransformHelpers.GetTransformBetween(SkinTorso, SkeletonTorso);
            SkinHeadPositionOffset = SkinHead.position - rig.Head.transform.position;
            SkinHeadRotationOffset = Quaternion.Inverse(rig.Head.transform.rotation) * SkinHead.rotation;

            LeftArmRig.InitializeRigging(SkeletonTorso);
            RightArmRig.InitializeRigging(SkeletonTorso);
        }

        public void UpdateRigging()
        {
            
            Params param = Params.Instance;
            bool highCloneNumber = Assistant.AllAssistants.Count > param.MaxClonesToDisableAnimations;

            // After ~18 clones we start to stagger/disable animation updates:

            // Disable the hand animations past a certain distance:
            bool disabledHandAnimations = highCloneNumber &&
                                          Vector3.Distance(rig.Head.transform.position,
                                              RealPlayer.Instance.Rig.Head.transform.position) >
                                          param.DisableHandAnimationDistance;
            if (!disabledHandAnimations)
            {
                HandAnimator.enabled = true;
                lHand.Update(LeftHand.HoldButtonPressed);
                rHand.Update(RightHand.HoldButtonPressed);
            }
            else
            {
                HandAnimator.enabled = false;
            }

            // Update the animation poses for clones every N frames, staggered to spread out the workload.
            // Only applies to clones far enough away from the player.
            bool lowFramerateAnimation = highCloneNumber &&
                                         Vector3.Distance(rig.Head.transform.position,
                                             RealPlayer.Instance.Rig.Head.transform.position) >
                                         param.LowFrameRateDistance;

            bool isMyFrame = Time.frameCount % param.LowFrameRateSkip ==
                             assistant.Identity % param.LowFrameRateSkip;
            if (!lowFramerateAnimation || isMyFrame)
            {
                UpdateSkeleton();
            }
        }

        public void UpdateSkeleton()
        {
            // Torso 
            var skeletonTorsoPosition = rig.Head.transform.position + TorsoHeadOffset;
            var skeletonTorsoRotation = Quaternion.Euler(0, rig.Head.transform.rotation.eulerAngles.y, 0);

            SkeletonTorso.position = skeletonTorsoPosition;
            SkeletonTorso.rotation = skeletonTorsoRotation;

            Matrix4x4 skeletonTorsoTransform = Matrix4x4.TRS(skeletonTorsoPosition, skeletonTorsoRotation, Vector3.one);

            SkinTorso.position = (skeletonTorsoTransform * SkinTorsoOffset).ExtractTranslation();
            SkinTorso.rotation = (skeletonTorsoTransform * SkinTorsoOffset).ExtractRotation();

            // Head
            SkinHead.position = rig.Head.transform.position + SkinHeadPositionOffset;
            SkinHead.rotation = rig.Head.transform.rotation * SkinHeadRotationOffset;

            LeftArmRig.UpdateSkeleton(rig.LeftHand, skeletonTorsoPosition, skeletonTorsoRotation, 
                TwistCorrectionStrengthAll, TwistCorrectionStrengthBicep, TwistCorrectionStrengthForearm);
            RightArmRig.UpdateSkeleton(rig.RightHand, skeletonTorsoPosition, skeletonTorsoRotation,
            TwistCorrectionStrengthAll, TwistCorrectionStrengthBicep, TwistCorrectionStrengthForearm);
        }

        private static float TriangleAngle(float aLen, float aLen1, float aLen2)
        {
            float c = Mathf.Clamp((aLen1 * aLen1 + aLen2 * aLen2 - aLen * aLen) / (aLen1 * aLen2) / 2.0f, -1.0f,
                1.0f);
            return Mathf.Acos(c);
        }

        private void OnDestroy()
        {
            graph.Destroy();
        }
    }
}