Streamingle_URP/Assets/Scripts/KindRetargeting/TwoBoneIKSolver.cs

using UnityEngine;

namespace KindRetargeting
{
    /// <summary>
    /// 커스텀 Two-Bone IK 솔버.
    /// 소스 무릎 위치를 비율 스케일하여 타겟 무릎을 직접 배치합니다.
    /// cosine law 없이 동작하므로 180° 특이점이 없고 역관절도 자연스럽게 지원합니다.
    /// </summary>
    [System.Serializable]
    public class TwoBoneIKSolver
    {
        [System.Serializable]
        public class LimbIK
        {
            public bool enabled = true;
            public Transform target;
            public Transform bendGoal;

            [Range(0f, 1f)] public float positionWeight = 0f;
            [Range(0f, 1f)] public float rotationWeight = 0f;
            [Range(0f, 1f)] public float bendGoalWeight = 0f;

            [HideInInspector] public Transform upper;
            [HideInInspector] public Transform lower;
            [HideInInspector] public Transform end;
            [HideInInspector] public float upperLength;
            [HideInInspector] public float lowerLength;

            [HideInInspector] public Vector3 localBendNormal;

            // 소스 본 참조 (소스 무릎 위치 기반 IK용)
            [HideInInspector] public Transform sourceUpper;
            [HideInInspector] public Transform sourceLower;
            [HideInInspector] public Transform sourceEnd;
        }

        [HideInInspector] public Animator animator;

        [Header("팔")]
        public LimbIK leftArm = new LimbIK();
        public LimbIK rightArm = new LimbIK();

        [Header("다리")]
        public LimbIK leftLeg = new LimbIK();
        public LimbIK rightLeg = new LimbIK();

        [HideInInspector] public int fabrikIterations = 6;

        private bool isInitialized;

        public void Initialize(Animator targetAnimator)
        {
            animator = targetAnimator;
            if (animator == null || !animator.isHuman) return;

            CacheLimb(leftArm, HumanBodyBones.LeftUpperArm, HumanBodyBones.LeftLowerArm, HumanBodyBones.LeftHand);
            CacheLimb(rightArm, HumanBodyBones.RightUpperArm, HumanBodyBones.RightLowerArm, HumanBodyBones.RightHand);
            CacheLimb(leftLeg, HumanBodyBones.LeftUpperLeg, HumanBodyBones.LeftLowerLeg, HumanBodyBones.LeftFoot);
            CacheLimb(rightLeg, HumanBodyBones.RightUpperLeg, HumanBodyBones.RightLowerLeg, HumanBodyBones.RightFoot);

            isInitialized = true;
        }

        private void CacheLimb(LimbIK limb, HumanBodyBones upperBone, HumanBodyBones lowerBone, HumanBodyBones endBone)
        {
            limb.upper = animator.GetBoneTransform(upperBone);
            limb.lower = animator.GetBoneTransform(lowerBone);
            limb.end = animator.GetBoneTransform(endBone);

            if (limb.upper == null || limb.lower == null || limb.end == null) return;

            limb.upperLength = Vector3.Distance(limb.upper.position, limb.lower.position);
            limb.lowerLength = Vector3.Distance(limb.lower.position, limb.end.position);

            Vector3 ab = limb.lower.position - limb.upper.position;
            Vector3 bc = limb.end.position - limb.lower.position;
            Vector3 bendNormal = Vector3.Cross(ab, bc);

            if (bendNormal.sqrMagnitude < 0.0001f)
            {
                bendNormal = Vector3.Cross(ab.normalized, Vector3.up);
                if (bendNormal.sqrMagnitude < 0.0001f)
                    bendNormal = Vector3.Cross(ab.normalized, Vector3.forward);
            }

            bendNormal.Normalize();
            limb.localBendNormal = Quaternion.Inverse(limb.upper.rotation) * bendNormal;
        }

        /// <summary>
        /// avatarScale 등으로 아바타 크기가 바뀌었을 때 호출하여 본 길이 캐시를 재계산.
        /// </summary>
        public void RefreshLimbLengths()
        {
            if (!isInitialized) return;
            RecacheLength(leftArm);
            RecacheLength(rightArm);
            RecacheLength(leftLeg);
            RecacheLength(rightLeg);
        }

        private void RecacheLength(LimbIK limb)
        {
            if (limb.upper == null || limb.lower == null || limb.end == null) return;
            limb.upperLength = Vector3.Distance(limb.upper.position, limb.lower.position);
            limb.lowerLength = Vector3.Distance(limb.lower.position, limb.end.position);
        }

        public void OnUpdate()
        {
            if (!isInitialized) return;

            SolveLimb(leftArm);
            SolveLimb(rightArm);
            SolveLimb(leftLeg);
            SolveLimb(rightLeg);
        }

        private void SolveLimb(LimbIK limb)
        {
            if (!limb.enabled || limb.target == null) return;
            if (limb.upper == null || limb.lower == null || limb.end == null) return;

            float weight = limb.positionWeight;
            if (weight < 0.0001f) return; // 완전히 0일 때만 스킵

            // FK 회전 저장 (IK 결과와 블렌딩용)
            Quaternion fkUpperRot = limb.upper.rotation;
            Quaternion fkLowerRot = limb.lower.rotation;
            Quaternion fkEndRot = limb.end.rotation;

            // 본 길이는 Initialize/RefreshLimbLengths에서 캐싱됨 (avatarScale 변경 시 갱신)
            float upperLen = limb.upperLength;
            float lowerLen = limb.lowerLength;

            Vector3 targetPos = limb.target.position;

            // --- 무릎 위치 결정 ---
            Vector3 kneePos;

            if (limb.sourceUpper != null && limb.sourceLower != null && limb.sourceEnd != null)
            {
                kneePos = ComputeKneePosFromSource(limb, upperLen, lowerLen, targetPos);
            }
            else
            {
                kneePos = ComputeKneePosFromBendGoal(limb, upperLen, lowerLen, targetPos);
            }

            // --- IK 회전 계산 ---
            Vector3 currentUpperDir = (limb.lower.position - limb.upper.position).normalized;
            Vector3 desiredUpperDir = (kneePos - limb.upper.position).normalized;
            Quaternion ikUpperRot = Quaternion.FromToRotation(currentUpperDir, desiredUpperDir) * limb.upper.rotation;

            // upper 적용 후 lower 계산 (자식 위치가 바뀌므로 임시 적용)
            limb.upper.rotation = ikUpperRot;
            Vector3 currentLowerDir = (limb.end.position - limb.lower.position).normalized;
            Vector3 desiredLowerDir = (targetPos - limb.lower.position).normalized;
            Quaternion ikLowerRot = Quaternion.FromToRotation(currentLowerDir, desiredLowerDir) * limb.lower.rotation;

            // --- FK/IK 블렌딩: weight로 부드럽게 전환 ---
            limb.upper.rotation = Quaternion.Slerp(fkUpperRot, ikUpperRot, weight);
            limb.lower.rotation = Quaternion.Slerp(fkLowerRot, ikLowerRot, weight);

            if (limb.rotationWeight > 0.0001f)
            {
                Quaternion ikEndRot = limb.target.rotation;
                limb.end.rotation = Quaternion.Slerp(fkEndRot, ikEndRot, limb.rotationWeight);
            }
        }

        /// <summary>
        /// 소스 무릎 위치를 타겟 프레임으로 옮긴 raw 위치를 초기값으로 두고,
        /// FABRIK 2회 반복으로 hip↔knee=upperLen, knee↔target=lowerLen 을 모두 만족시킵니다.
        ///
        /// 정규화(normalize)에 의한 부호 증폭이 없어 입력 노이즈가 출력에 그대로 비례 전달됩니다.
        /// → 다리 거의 펴진 상태에서 모캡 노이즈로 굽힘 부호가 흔들려도 점프 없이 연속적.
        /// → 의도된 역관절도 자연스럽게 따라감 (소스 무릎이 어느 쪽이든 raw로 받음).
        /// </summary>
        private Vector3 ComputeKneePosFromSource(LimbIK limb, float upperLen, float lowerLen, Vector3 targetPos)
        {
            float chainLength = upperLen + lowerLen;

            // 소스 hip→foot 방향
            Vector3 sourceHipToFoot = limb.sourceEnd.position - limb.sourceUpper.position;
            float sourceHipToFootMag = sourceHipToFoot.magnitude;
            if (sourceHipToFootMag < 0.001f) return limb.lower.position;
            Vector3 sourceHipToFootDir = sourceHipToFoot / sourceHipToFootMag;

            // 소스 hip→knee (굽힘 정보 raw, 정규화 X → 노이즈 증폭 없음)
            Vector3 sourceHipToKnee = limb.sourceLower.position - limb.sourceUpper.position;
            float sourceUpperLen = sourceHipToKnee.magnitude;
            float sourceLowerLen = Vector3.Distance(limb.sourceLower.position, limb.sourceEnd.position);
            float sourceChain = sourceUpperLen + sourceLowerLen;
            if (sourceChain < 0.001f) return limb.lower.position;

            // 타겟 hip→target 방향 / 거리
            Vector3 toTarget = targetPos - limb.upper.position;
            float targetDist = toTarget.magnitude;
            if (targetDist < 0.001f) return limb.lower.position;
            Vector3 toTargetDir = toTarget / targetDist;

            // 도달 불가 거리는 클램프 (FABRIK 진동 방지)
            // |upper-lower| 보다 가깝거나 chainLength 보다 멀면 본 길이 제약을 만족하는 무릎 위치 없음
            float effectiveDist = Mathf.Clamp(targetDist, Mathf.Abs(upperLen - lowerLen) + 0.001f, chainLength - 0.001f);
            Vector3 effectiveTarget = limb.upper.position + toTargetDir * effectiveDist;

            // 소스 → 타겟 프레임 회전 + 다리 길이 비율 스케일 → 초기 무릎 위치
            Quaternion frameRotation = Quaternion.FromToRotation(sourceHipToFootDir, toTargetDir);
            float scale = (upperLen + lowerLen) / sourceChain;
            Vector3 kneePos = limb.upper.position + frameRotation * (sourceHipToKnee * scale);

            // FABRIK: knee 를 lowerLen 구면(target 중심) ↔ upperLen 구면(hip 중심) 사이에서 번갈아 투영
            for (int i = 0; i < fabrikIterations; i++)
            {
                // 1) target 중심 lowerLen 구면에 투영 → knee↔target = lowerLen 보장
                Vector3 fromTarget = kneePos - effectiveTarget;
                float distFromTarget = fromTarget.magnitude;
                if (distFromTarget > 0.001f)
                    kneePos = effectiveTarget + (fromTarget / distFromTarget) * lowerLen;

                // 2) hip 중심 upperLen 구면에 투영 → hip↔knee = upperLen 보장
                Vector3 fromHip = kneePos - limb.upper.position;
                float distFromHip = fromHip.magnitude;
                if (distFromHip > 0.001f)
                    kneePos = limb.upper.position + (fromHip / distFromHip) * upperLen;
            }

            return kneePos;
        }

        /// <summary>
        /// bendGoal 기반 fallback (소스 참조가 없는 사지용).
        /// 기존 cosine law + bendNormal 방식.
        /// </summary>
        private Vector3 ComputeKneePosFromBendGoal(LimbIK limb, float upperLen, float lowerLen, Vector3 targetPos)
        {
            float chainLength = upperLen + lowerLen;
            Vector3 toTarget = targetPos - limb.upper.position;
            float targetDist = toTarget.magnitude;
            if (targetDist < 0.001f) return limb.lower.position;
            Vector3 toTargetDir = toTarget / targetDist;

            // bendGoal의 수직 성분(rejection)으로 팔꿈치/무릎 방향 직접 결정
            Vector3 bendDir;
            if (limb.bendGoal != null && limb.bendGoalWeight > 0.001f)
            {
                Vector3 toBendGoal = limb.bendGoal.position - limb.upper.position;
                Vector3 rejection = toBendGoal - Vector3.Dot(toBendGoal, toTargetDir) * toTargetDir;
                if (rejection.sqrMagnitude > 0.0001f)
                {
                    // T-포즈 기반 기본 방향
                    Vector3 defaultBendDir = GetDefaultBendDir(limb, toTargetDir);
                    bendDir = Vector3.Lerp(defaultBendDir, rejection.normalized, limb.bendGoalWeight);
                }
                else
                {
                    bendDir = GetDefaultBendDir(limb, toTargetDir);
                }
            }
            else
            {
                bendDir = GetDefaultBendDir(limb, toTargetDir);
            }
            bendDir.Normalize();

            // 코사인 법칙으로 upper 각도 계산
            float clampedDist = Mathf.Clamp(targetDist, Mathf.Abs(upperLen - lowerLen) + 0.001f, chainLength - 0.001f);
            float cosAngle = (clampedDist * clampedDist + upperLen * upperLen - lowerLen * lowerLen)
                           / (2f * clampedDist * upperLen);
            cosAngle = Mathf.Clamp(cosAngle, -1f, 1f);
            float angle = Mathf.Acos(cosAngle);

            // 무릎/팔꿈치 위치: toTargetDir + bendDir 방향으로 angle만큼 오프셋
            // sin(angle) = 수직 성분, cos(angle) = 직선 성분
            Vector3 kneePos = limb.upper.position
                            + toTargetDir * (upperLen * Mathf.Cos(angle))
                            + bendDir * (upperLen * Mathf.Sin(angle));

            return kneePos;
        }

        /// <summary>
        /// T-포즈 기반 기본 굽힘 방향 (bendGoal이 없거나 불안정할 때)
        /// </summary>
        private Vector3 GetDefaultBendDir(LimbIK limb, Vector3 toTargetDir)
        {
            Vector3 bendNormal = limb.upper.rotation * limb.localBendNormal;
            Vector3 bendDir = Vector3.Cross(bendNormal, toTargetDir);
            if (bendDir.sqrMagnitude < 0.0001f)
                bendDir = Vector3.Cross(Vector3.up, toTargetDir);
            return bendDir.normalized;
        }
    }
}