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;
using UnityEngine.Events;

namespace NAK.IKFixes.HarmonyPatches;

internal static class BodySystemPatches
{
    private static float _ikSimulatedRootAngle;

    [HarmonyPostfix]
    [HarmonyPatch(typeof(BodySystem), nameof(BodySystem.SetupOffsets))]
    private 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;
            }

            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;
                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.
            }
        }
    }

    private static 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;
    }

    private static 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;
    }

    private static void SetPelvisWeight(IKSolverVR.Spine spine, float weight)
    {
        // looks better when hips are disabled while running
        spine.pelvisPositionWeight = weight;
        spine.pelvisRotationWeight = weight;
    }

    [HarmonyPrefix]
    [HarmonyPatch(typeof(BodySystem), nameof(BodySystem.Update))]
    private static bool Prefix_BodySystem_Update(ref BodySystem __instance)
    {
        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);
            }
            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 = 25f;
            float rootHeadingOffset = 0f;

            if (BodySystem.isCalibratedAsFullBody || IKFixes.EntryUseFakeRootAngle.Value)
                maxRootAngle = 0f;

            if (PlayerSetup.Instance._emotePlaying)
                maxRootAngle = 180f;

            if (IKFixes.EntryUseFakeRootAngle.Value && !BodySystem.isCalibratedAsFullBody)
            {
                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);

                deltaAngleRoot = Mathf.Clamp(deltaAngleRoot, -weightedAngleLimit, weightedAngleLimit);
                _ikSimulatedRootAngle = Mathf.MoveTowardsAngle(_ikSimulatedRootAngle, pivotAngle, absDeltaAngleRoot - weightedAngleLimit);

                rootHeadingOffset = deltaAngleRoot;
            }

            solver.spine.maxRootAngle = maxRootAngle;
            solver.spine.rootHeadingOffset = rootHeadingOffset;

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

            if (!IKSystem.vrik.solver.leftLeg.usingKneeTracker)
                IKSystem.vrik.solver.leftLeg.bendToTargetWeight = IKFixes.EntryBendToTargetWeight.Value;
            if (!IKSystem.vrik.solver.rightLeg.usingKneeTracker)
                IKSystem.vrik.solver.rightLeg.bendToTargetWeight = IKFixes.EntryBendToTargetWeight.Value;
        }
        
        int count = IKSystem.Instance.AllTrackingPoints.FindAll(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))]
    private static bool Prefix_BodySystem_AssignRemainingTrackers()
    {
        return IKFixes.EntryAssignRemainingTrackers.Value;
    }

    [HarmonyPrefix]
    [HarmonyPatch(typeof(BodySystem), nameof(BodySystem.MuscleUpdate))]
    private 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 = MovementSystem.Instance.movementVector.magnitude > 0f;
            bool isGrounded = MovementSystem.Instance._isGrounded;
            bool isFlying = MovementSystem.Instance.flying;
            bool playRunningAnimation = BodySystem.PlayRunningAnimationInFullBody;

            if ((playRunningAnimation && (isRunning || !isGrounded && !isFlying)))
            {
                SetPelvisWeight(IKSystem.vrik.solver.spine, 0f);
                IKSystem.Instance.applyOriginalHipPosition = true;
                IKSystem.Instance.applyOriginalHipRotation = true;
            }
            else
            {
                IKSystem.Instance.applyOriginalHipPosition = true;
                IKSystem.Instance.applyOriginalHipRotation = false;
                IKSystem.Instance.humanPose.bodyRotation = Quaternion.identity;
            }
        }

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

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

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

    [HarmonyPostfix]
    [HarmonyPatch(typeof(IKSystem), nameof(IKSystem.InitializeHalfBodyIK))]
    private static void Postfix_IKSystem_InitializeHalfBodyIK(ref IKSystem __instance)
    {
        if (!IKFixes.EntryUseIKPose.Value) return;

        __instance._poseHandler.GetHumanPose(ref __instance.humanPose);

        for (int i = 0; i < IKPoseMuscles.Length; i++)
        {
            __instance.ApplyMuscleValue((MuscleIndex)i, IKPoseMuscles[i], ref __instance.humanPose.muscles);
        }
        __instance.humanPose.bodyPosition = Vector3.up;
        __instance.humanPose.bodyRotation = Quaternion.identity;
        __instance._poseHandler.SetHumanPose(ref __instance.humanPose);

        // recentering avatar so it doesnt need to step from random place on switch
        IKSystem.vrik.transform.localPosition = Vector3.zero;
        IKSystem.vrik.transform.localRotation = Quaternion.identity;
        // janky fix, initializing early with correct pose
        IKSystem.vrik.solver.Initiate(IKSystem.vrik.transform);
    }

    private static readonly float[] IKPoseMuscles = new float[]
    {
            0.00133321f,
            8.195831E-06f,
            8.537738E-07f,
            -0.002669832f,
            -7.651234E-06f,
            -0.001659694f,
            0f,
            0f,
            0f,
            0.04213953f,
            0.0003007996f,
            -0.008032114f,
            -0.03059979f,
            -0.0003182998f,
            0.009640567f,
            0f,
            0f,
            0f,
            0f,
            0f,
            0f,
            0.5768794f,
            0.01061097f,
            -0.1127839f,
            0.9705755f,
            0.07972051f,
            -0.0268422f,
            0.007237188f,
            0f,
            0.5768792f,
            0.01056608f,
            -0.1127519f,
            0.9705756f,
            0.07971933f,
            -0.02682396f,
            0.007229362f,
            0f,
            -5.651802E-06f,
            -3.034899E-07f,
            0.4100508f,
            0.3610304f,
            -0.0838329f,
            0.9262537f,
            0.1353517f,
            -0.03578902f,
            0.06005657f,
            -4.95989E-06f,
            -1.43007E-06f,
            0.4096187f,
            0.363263f,
            -0.08205152f,
            0.9250782f,
            0.1345718f,
            -0.03572125f,
            0.06055461f,
            -1.079177f,
            0.2095419f,
            0.6140652f,
            0.6365265f,
            0.6683931f,
            -0.4764312f,
            0.8099416f,
            0.8099371f,
            0.6658203f,
            -0.7327053f,
            0.8113618f,
            0.8114051f,
            0.6643661f,
            -0.40341f,
            0.8111364f,
            0.8111367f,
            0.6170399f,
            -0.2524227f,
            0.8138723f,
            0.8110135f,
            -1.079171f,
            0.2095456f,
            0.6140658f,
            0.6365255f,
            0.6683878f,
            -0.4764301f,
            0.8099402f,
            0.8099376f,
            0.6658241f,
            -0.7327023f,
            0.8113653f,
            0.8113793f,
            0.664364f,
            -0.4034042f,
            0.811136f,
            0.8111364f,
            0.6170469f,
            -0.2524345f,
            0.8138595f,
            0.8110138f
     };
}

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

    [HarmonyPrefix]
    [HarmonyPatch(typeof(PlayerSetup), nameof(PlayerSetup.ResetIk))]
    private 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;
    }
}