NAK_CVR_Mods/IKFixes/HarmonyPatches.cs

using ABI.CCK.Components;
using ABI_RC.Core.Player;
using ABI_RC.Systems.IK;
using ABI_RC.Systems.IK.SubSystems;
using ABI_RC.Systems.MovementSystem;
using HarmonyLib;
using RootMotion.FinalIK;
using UnityEngine;

namespace NAK.IKFixes.HarmonyPatches;

internal static class BodySystemPatches
{
    static float _ikSimulatedRootAngle = 0f;

    [HarmonyPostfix]
    [HarmonyPatch(typeof(BodySystem), nameof(BodySystem.SetupOffsets))]
    static void Postfix_BodySystem_SetupOffsets(List<TrackingPoint> trackingPoints)
    {
        foreach (TrackingPoint trackingPoint in trackingPoints)
        {
            Transform parent = null;
            float offsetDistance = 0f;

            switch (trackingPoint.assignedRole)
            {
                case TrackingPoint.TrackingRole.LeftKnee:
                    parent = IKSystem.vrik.references.leftCalf;
                    offsetDistance = 0.15f;
                    break;
                case TrackingPoint.TrackingRole.RightKnee:
                    parent = IKSystem.vrik.references.rightCalf;
                    offsetDistance = 0.15f;
                    break;
                case TrackingPoint.TrackingRole.LeftElbow:
                    parent = IKSystem.vrik.references.leftForearm;
                    offsetDistance = -0.15f;
                    break;
                case TrackingPoint.TrackingRole.RightElbow:
                    parent = IKSystem.vrik.references.rightForearm;
                    offsetDistance = -0.15f;
                    break;
                default:
                    break;
            }

            if (parent != null)
            {
                // Set the offset transform's parent and reset its local position and rotation
                trackingPoint.offsetTransform.parent = parent;
                trackingPoint.offsetTransform.localPosition = Vector3.zero;
                trackingPoint.offsetTransform.localRotation = Quaternion.identity;
                trackingPoint.offsetTransform.parent = trackingPoint.referenceTransform;

                // Apply additional offset based on the assigned role
                Vector3 additionalOffset = IKSystem.vrik.references.root.forward * offsetDistance;

                if (IKFixes.EntryAltElbowDirection.Value)
                {
                    switch (trackingPoint.assignedRole)
                    {
                        case TrackingPoint.TrackingRole.LeftElbow:
                            additionalOffset += IKSystem.vrik.references.root.up * -0.15f;
                            break;
                        case TrackingPoint.TrackingRole.RightElbow:
                            additionalOffset += IKSystem.vrik.references.root.up * -0.15f;
                            break;
                        default:
                            break;
                    }
                }

                trackingPoint.offsetTransform.position += additionalOffset;

                // Game originally sets them to about half a meter out, which fucks with slime tracker users and
                // makes the bendGoals less responsive/less accurate. 

                //Funny thing is that IKTweaks specifically made this an option, which should be added to both CVR & Standable for the same reason.

                /// Elbow / knee / chest bend goal offset - controls how far bend goal targets will be away from the actual joint.
                /// Lower values should produce better precision with bent joint, higher values - better stability with straight joint. 
                /// Sensible range of values is between 0 and 1.
            }
        }
    }

    [HarmonyPrefix]
    [HarmonyPatch(typeof(BodySystem), nameof(BodySystem.Update))]
    static bool Prefix_BodySystem_Update(ref BodySystem __instance)
    {
        void SetArmWeight(IKSolverVR.Arm arm, float weight)
        {
            arm.positionWeight = weight;
            arm.rotationWeight = weight;
            arm.shoulderRotationWeight = weight;
            arm.shoulderTwistWeight = weight;
            // assumed fix of bend goal weight if arms disabled with elbows (havent tested)
            // why is there no "usingElbowTracker" flag like knees? where is the consistancy???
            arm.bendGoalWeight = arm.bendGoal != null ? weight : 0f;
        }
        void SetLegWeight(IKSolverVR.Leg leg, float weight)
        {
            leg.positionWeight = weight;
            leg.rotationWeight = weight;
            // fixes knees bending to tracker if feet disabled (running anim)
            leg.bendGoalWeight = leg.usingKneeTracker ? weight : 0f;
        }
        void SetPelvisWeight(IKSolverVR.Spine spine, float weight)
        {
            // looks better when hips are disabled while running
            spine.pelvisPositionWeight = weight;
            spine.pelvisRotationWeight = weight;
        }

        if (IKSystem.vrik != null)
        {
            IKSolverVR solver = IKSystem.vrik.solver;

            if (BodySystem.TrackingEnabled)
            {
                IKSystem.vrik.enabled = true;
                solver.IKPositionWeight = BodySystem.TrackingPositionWeight;
                solver.locomotion.weight = BodySystem.TrackingLocomotionEnabled ? 1f : 0f;

                // fixes arm weights not being set if leftArm & rightArm targets are null
                // game handles TrackingLegs in PlayerSetup, but not for knee goals
                SetArmWeight(solver.leftArm, BodySystem.TrackingLeftArmEnabled && solver.leftArm.target != null ? 1f : 0f);
                SetArmWeight(solver.rightArm, BodySystem.TrackingRightArmEnabled && solver.rightArm.target != null ? 1f : 0f);
                SetLegWeight(solver.leftLeg, BodySystem.TrackingLeftLegEnabled && solver.leftLeg.target != null ? 1f : 0f);
                SetLegWeight(solver.rightLeg, BodySystem.TrackingRightLegEnabled && solver.leftLeg.target != null ? 1f : 0f);
                SetPelvisWeight(solver.spine, solver.spine.pelvisTarget != null ? 1f : 0f);

                // makes running animation look better
                if (BodySystem.isCalibratedAsFullBody)
                {
                    bool isRunning = BodySystem.PlayRunningAnimationInFullBody && MovementSystem.Instance.movementVector.magnitude > 0f;
                    if (isRunning) SetPelvisWeight(solver.spine, 0f);
                }
            }
            else
            {
                IKSystem.vrik.enabled = false;
                solver.IKPositionWeight = 0f;
                solver.locomotion.weight = 0f;

                SetArmWeight(solver.leftArm, 0f);
                SetArmWeight(solver.rightArm, 0f);
                SetLegWeight(solver.leftLeg, 0f);
                SetLegWeight(solver.rightLeg, 0f);
                SetPelvisWeight(solver.spine, 0f);
            }

            float maxRootAngle = BodySystem.isCalibratedAsFullBody || IKFixes.EntryUseFakeRootAngle.Value ? (PlayerSetup.Instance._emotePlaying ? 180f : 0f) : (PlayerSetup.Instance._emotePlaying ? 180f : 25f);
            solver.spine.maxRootAngle = maxRootAngle;

            if (IKFixes.EntryUseFakeRootAngle.Value)
            {
                // Emulate maxRootAngle because CVR doesn't have the player controller set up ideally for VRIK.
                // I believe they'd need to change which object vrik.references.root is, as using avatar object is bad!
                // This is a small small fix, but makes it so the feet dont point in the direction of the head
                // when turning. It also means turning with joystick & turning IRL make feet behave the same and follow behind.
                float weightedAngleLimit = IKFixes.EntryFakeRootAngleLimit.Value * solver.locomotion.weight;
                float pivotAngle = MovementSystem.Instance.rotationPivot.eulerAngles.y;
                float deltaAngleRoot = Mathf.DeltaAngle(pivotAngle, _ikSimulatedRootAngle);
                float absDeltaAngleRoot = Mathf.Abs(deltaAngleRoot);

                if (absDeltaAngleRoot > weightedAngleLimit)
                {
                    deltaAngleRoot = Mathf.Sign(deltaAngleRoot) * weightedAngleLimit;
                    _ikSimulatedRootAngle = Mathf.MoveTowardsAngle(_ikSimulatedRootAngle, pivotAngle, absDeltaAngleRoot - weightedAngleLimit);
                }
                solver.spine.rootHeadingOffset = deltaAngleRoot;
            }

            // custom IK settings
            solver.spine.neckStiffness = IKFixes.EntryNeckStiffness.Value;
            solver.spine.bodyRotStiffness = IKFixes.EntryBodyRotStiffness.Value;
            solver.spine.rotateChestByHands = IKFixes.EntryRotateChestByHands.Value;
        }

        int count = IKSystem.Instance.AllTrackingPoints.FindAll((TrackingPoint m) => m.isActive && m.isValid && m.suggestedRole > TrackingPoint.TrackingRole.Invalid).Count;

        // fixes having all tracking points disabled forcing calibration
        if (count == 0)
        {
            __instance._fbtAvailable = false;
            return false;
        }

        // solid body count block
        int num = 0;
        if (BodySystem.enableLeftFootTracking) num++;
        if (BodySystem.enableRightFootTracking) num++;
        if (BodySystem.enableHipTracking) num++;
        if (BodySystem.enableLeftKneeTracking) num++;
        if (BodySystem.enableRightKneeTracking) num++;
        if (BodySystem.enableChestTracking) num++;
        if (BodySystem.enableLeftElbowTracking) num++;
        if (BodySystem.enableRightElbowTracking) num++;

        __instance._fbtAvailable = (count >= num);

        return false;
    }

    [HarmonyPrefix]
    [HarmonyPatch(typeof(BodySystem), nameof(BodySystem.AssignRemainingTrackers))]
    static bool Prefix_BodySystem_AssignRemainingTrackers()
    {
        return IKFixes.EntryAssignRemainingTrackers.Value;
    }

    [HarmonyPrefix]
    [HarmonyPatch(typeof(BodySystem), nameof(BodySystem.MuscleUpdate))]
    static void Prefix_BodySystem_MuscleUpdate()
    {
        if (BodySystem.isCalibrating)
        {
            IKSystem.Instance.humanPose.bodyRotation = Quaternion.identity;
            IKSystem.vrik.solver.spine.maxRootAngle = 0f; // idk, testing
        }

        if (BodySystem.isCalibratedAsFullBody && BodySystem.TrackingPositionWeight > 0f)
        {
            bool isRunning = BodySystem.PlayRunningAnimationInFullBody && MovementSystem.Instance.movementVector.magnitude > 0f;
            if (!isRunning)
            {
                // Resetting bodyRotation made running animations look funky
                IKSystem.Instance.applyOriginalHipPosition = true;
                IKSystem.Instance.applyOriginalHipRotation = false;
                IKSystem.Instance.humanPose.bodyRotation = Quaternion.identity;
            }
            else
            {
                // This looks much better when running
                IKSystem.Instance.applyOriginalHipPosition = true;
                IKSystem.Instance.applyOriginalHipRotation = true;
            }
        }

        // TODO: Rewrite to exclude setting T-pose to limbs that are not tracked
    }

    [HarmonyPostfix]
    [HarmonyPatch(typeof(BodySystem), nameof(BodySystem.Calibrate))]
    static void Postfix_BodySystem_Calibrate()
    {
        IKSystem.Instance.applyOriginalHipPosition = false;
        IKSystem.Instance.applyOriginalHipRotation = false;
        IKSystem.vrik.solver.leftLeg.bendToTargetWeight = 0.25f;
        IKSystem.vrik.solver.rightLeg.bendToTargetWeight = 0.25f;
    }
}

internal static class IKSystemPatches
{
    [HarmonyPrefix]
    [HarmonyPatch(typeof(IKSystem), nameof(IKSystem.InitializeAvatar))]
    static void Prefix_IKSystem_InitializeAvatar(ref IKSystem __instance)
    {
        __instance.applyOriginalHipPosition = true;
        __instance.applyOriginalHipRotation = true;
    }
}

internal static class VRIKPatches
{
    /**
        Leg solver uses virtual bone calf and foot, plus world tracked knee position for normal maths.
        This breaks as you playspace up, because calf and foot position aren't offset yet in solve order.
    **/

    // Add ApplyBendGoal(); to the second line of RootMotionNew.FinalIK.IKSolverVR.Leg.Solve(bool)
    // https://github.com/knah/VRCMods/tree/b6c4198fb8e06174ea511fe1f8a3257dfef2fdd2 IKTweaks

    [HarmonyPostfix]
    [HarmonyPatch(typeof(IKSolverVR.Leg), nameof(IKSolverVR.Leg.Stretching))]
    static void Postfix_IKSolverVR_Leg_Stretching(ref IKSolverVR.Leg __instance)
    {
        // I am patching here because Stretching() is always called
        // and i am not doing a transpiler to place it on the second line
        __instance.ApplyBendGoal();
    }

    // IKTweaks original method does not do this?
    // idk i prefer it, upper leg rotates weird otherwise

    [HarmonyPrefix]
    [HarmonyPatch(typeof(IKSolverVR.Leg), nameof(IKSolverVR.Leg.ApplyOffsets))]
    static bool Prefix_IKSolverVR_Leg_ApplyOffsets(ref IKSolverVR.Leg __instance)
    {
        //This is the second part of the above fix, preventing the solver from calculating a bad bendNormal
        //when it doesn't need to. The knee tracker should dictate the bendNormal completely.

        //TODO: investigate lower leg not bending towards knee direction

        if (__instance.usingKneeTracker)
        {
            __instance.ApplyPositionOffset(__instance.footPositionOffset, 1f);
            __instance.ApplyRotationOffset(__instance.footRotationOffset, 1f);
            Quaternion quaternion = Quaternion.FromToRotation(__instance.footPosition - __instance.position, __instance.footPosition + __instance.heelPositionOffset - __instance.position);
            __instance.footPosition = __instance.position + quaternion * (__instance.footPosition - __instance.position);
            __instance.footRotation = quaternion * __instance.footRotation;
            return false;
        }

        // run full method like normal otherwise
        float num = __instance.bendGoalWeight;
        __instance.bendGoalWeight = 0f;
        __instance.ApplyOffsetsOld();
        __instance.bendGoalWeight = num;
        return false;
    }
}

internal static class PlayerSetupPatches
{
    // Last Movement Parent Info
    static CVRMovementParent lastMovementParent;
    static Vector3 lastMovementPosition;
    static Quaternion lastMovementRotation;

    [HarmonyPrefix]
    [HarmonyPatch(typeof(PlayerSetup), nameof(PlayerSetup.ResetIk))]
    static bool Prefix_PlayerSetup_ResetIk()
    {
        if (IKSystem.vrik == null) return false;

        CVRMovementParent currentParent = MovementSystem.Instance._currentParent;
        if (currentParent != null && currentParent._referencePoint != null)
        {
            // Get current position, VR pivots around VR camera
            Vector3 currentPosition = MovementSystem.Instance.rotationPivot.transform.position;
            currentPosition.y = IKSystem.vrik.transform.position.y; // set pivot to floor
            Quaternion currentRotation = Quaternion.Euler(0f, currentParent.transform.rotation.eulerAngles.y, 0f);

            // Convert to delta position (how much changed since last frame)
            Vector3 deltaPosition = currentPosition - lastMovementPosition;
            Quaternion deltaRotation = Quaternion.Inverse(lastMovementRotation) * currentRotation;

            // Prevent targeting other parent position
            if (lastMovementParent == currentParent || lastMovementParent == null)
            {
                // Add platform motion to IK solver
                IKSystem.vrik.solver.AddPlatformMotion(deltaPosition, deltaRotation, currentPosition);
            }

            // Store for next frame
            lastMovementParent = currentParent;
            lastMovementPosition = currentPosition;
            lastMovementRotation = currentRotation;
            return false;
        }

        return true;
    }
}