// © 2024 Long Ly. All rights reserved. Any unauthorized use, reproduction, or distribution of this trademark is strictly prohibited and may result in legal action.


#include "Nodes/AnimNode_OLSMask.h"

#include "Animation/AnimInstanceProxy.h"

void FAnimNode_OLSMask::Initialize_AnyThread(const FAnimationInitializeContext& context)
{
	// Don't call SUPER since it's empty.

	DECLARE_SCOPE_HIERARCHICAL_COUNTER_ANIMNODE(Initialize_AnyThread)
	FAnimNode_Base::Initialize_AnyThread(context);

	AnimationPose.Initialize(context);
	MaskPose.Initialize(context);
	BaseAdditivePose.Initialize(context);


	for (int32 index = 0; index < BlendMasks.Num(); ++index)
	{
		if (BlendMasks[index].bShouldAddSlot)
		{
			BlendMasks[index].Slot.SlotName = BlendMasks[index].SlotName;
			BlendMasks[index].Slot.Source = MaskPose;
			BlendMasks[index].Slot.Initialize_AnyThread(context);
		}
	}
	if (!context.AnimInstanceProxy->GetSkeleton())
	{
		return;
	}

	// Invalidate the cached per-bone blend weights from the skeleton
	InvalidatePerBoneBlendWeights();
}

void FAnimNode_OLSMask::CacheBones_AnyThread(const FAnimationCacheBonesContext& context)
{
	// Don't call SUPER since it's empty.

	DECLARE_SCOPE_HIERARCHICAL_COUNTER_ANIMNODE(CacheBones_AnyThread)
	AnimationPose.CacheBones(context);
	BaseAdditivePose.CacheBones(context);
	MaskPose.CacheBones(context);
	const int32 numBlendMasks = BlendMasks.Num();
	for (int32 maskIndex = 0; maskIndex < numBlendMasks; ++maskIndex)
	{
		const FOLSMaskSettings currentMask = BlendMasks[maskIndex];
		const FName blendProfileName = BlendMasks[maskIndex].BlendProfileName;
		if (const TObjectPtr<UBlendProfile> blendProfile = context.AnimInstanceProxy->GetSkeleton()->GetBlendProfile(
			blendProfileName))
		{
			BlendMasks[maskIndex].BlendProfile = blendProfile;
		}
		if (BlendMasks[maskIndex].bShouldAddSlot)
		{
			BlendMasks[maskIndex].Slot.CacheBones_AnyThread(context);
		}
	}

	UpdateCachedBoneData(context.AnimInstanceProxy->GetRequiredBones(), context.AnimInstanceProxy->GetSkeleton());
}

void FAnimNode_OLSMask::Update_AnyThread(const FAnimationUpdateContext& Context)
{
	// Don't call SUPER since it's empty.

	DECLARE_SCOPE_HIERARCHICAL_COUNTER_ANIMNODE(Update_AnyThread)
	bHasRelevantPoses = false;
	if (BlendMasks.Num() > 0)
	{
		if (IsLODEnabled(Context.AnimInstanceProxy))
		{
			GetEvaluateGraphExposedInputs().Execute(Context);

			if (!bHasRelevantPoses)
			{
				MaskPose.Update(Context);
				BaseAdditivePose.Update(Context);
				AnimationPose.Update(Context);
				for (int32 Index = 0; Index < BlendMasks.Num(); ++Index)
				{
					if (BlendMasks[Index].bShouldAddSlot)
					{
						BlendMasks[Index].Slot.Update_AnyThread(Context);
					}
				}
				UpdateCachedBoneData(Context.AnimInstanceProxy->GetRequiredBones(),
				                     Context.AnimInstanceProxy->GetSkeleton());

				bHasRelevantPoses = true;
			}
			return;
		}

		// Clear BlendWeights if disabled by LODThreshold.
		for (int32 ChildIndex = 0; ChildIndex < MaskBlendWeights.Num(); ++ChildIndex)
		{
			MaskBlendWeights[ChildIndex].LocalSpaceBlendAlpha = 0.0f;
			MaskBlendWeights[ChildIndex].MeshSpaceBlendAlpha = 0.0f;
		}

		return;
	}

	AnimationPose.Update(Context);
}

void FAnimNode_OLSMask::Evaluate_AnyThread(FPoseContext& output)
{
	// Don't call SUPER since it's empty.

	if (BlendMasks.Num() > 0)
	{
		AnimationPose.Evaluate(output);
		FPoseContext animationPoseContext(output);
		animationPoseContext = output;

		FPoseContext updatedAnimationPoseContext(output);
		updatedAnimationPoseContext = animationPoseContext;

		FPoseContext maskPoseContext(output);
		MaskPose.Evaluate(maskPoseContext);

		FPoseContext baseAddPoseContext(output);
		BaseAdditivePose.Evaluate(baseAddPoseContext);

		FPoseContext baseMSAddPoseContext(output);
		baseMSAddPoseContext = baseAddPoseContext;

		FPoseContext animationLocalSpaceAdditive(output);
		animationLocalSpaceAdditive = animationPoseContext;

		FPoseContext animationMeshSpaceAdditive(output);
		animationMeshSpaceAdditive = animationPoseContext;

		MakeAdditivePose(animationLocalSpaceAdditive, baseAddPoseContext, false);
		MakeAdditivePose(animationMeshSpaceAdditive, baseMSAddPoseContext, true);

		for (int32 maskIndex = 0; maskIndex < BlendMasks.Num(); ++maskIndex)
		{
			//Apply Blending For Each Mask
			TArray<FCompactPose> lsTargetBlendPoses;
			lsTargetBlendPoses.SetNum(1);

			TArray<FCompactPose> msTargetBlendPoses;
			msTargetBlendPoses.SetNum(1);

			TArray<FBlendedCurve> lsTargetBlendCurves;
			lsTargetBlendCurves.SetNum(1);

			TArray<FBlendedCurve> msTargetBlendCurves;
			msTargetBlendCurves.SetNum(1);

			TArray<UE::Anim::FStackAttributeContainer> lsTargetBlendAttributes;
			lsTargetBlendAttributes.SetNum(1);

			TArray<UE::Anim::FStackAttributeContainer> msTargetBlendAttributes;
			msTargetBlendAttributes.SetNum(1);

			FPoseContext outMaskPoseContext(output);
			(BlendMasks[maskIndex].bShouldAddSlot)
				? BlendMasks[maskIndex].Slot.Evaluate_AnyThread(outMaskPoseContext)
				: BlendMasks[maskIndex].Slot.Evaluate_AnyThread(outMaskPoseContext);

			UpdateWeightsFromCurves(output.AnimInstanceProxy->GetSkeleton(), outMaskPoseContext, maskIndex);

			FAnimationPoseData inOutAnimationPoseData(outMaskPoseContext);
			const FAnimationPoseData additiveAnimationPoseData(
				BlendMasks[maskIndex].bUseMeshSpaceAdditive ? animationMeshSpaceAdditive : animationLocalSpaceAdditive);
			if (FAnimWeight::IsRelevant(MaskBlendWeights[maskIndex].AdditiveAlpha))
			{
				FAnimationRuntime::AccumulateAdditivePose(inOutAnimationPoseData, additiveAnimationPoseData,
				                                          MaskBlendWeights[maskIndex].AdditiveAlpha,
				                                          BlendMasks[maskIndex].bUseMeshSpaceAdditive
					                                          ? AAT_RotationOffsetMeshSpace
					                                          : AAT_LocalSpaceBase);
				outMaskPoseContext.Pose.NormalizeRotations();
			}


			FPoseContext currentPoseContext(animationPoseContext);
			FAnimationPoseData blendedAnimationPoseData(currentPoseContext);
			const FAnimationPoseData animationPoseOneData(animationPoseContext);
			const FAnimationPoseData animationPoseTwoData(outMaskPoseContext);


			TArray<float> lsDesiredBlendWeight;
			lsDesiredBlendWeight.Add(MaskBlendWeights[maskIndex].LocalSpaceBlendAlpha);
			TArray<float> msDesiredBlendWeight;
			msDesiredBlendWeight.Add(MaskBlendWeights[maskIndex].MeshSpaceBlendAlpha);
			FAnimationRuntime::UpdateDesiredBoneWeight(AllBoneBlendWeights[maskIndex].LocalSpaceDesiredBoneBlendWeights,
			                                           AllBoneBlendWeights[maskIndex].LocalSpaceCurrentBoneBlendWeights,
			                                           lsDesiredBlendWeight);
			FAnimationRuntime::UpdateDesiredBoneWeight(AllBoneBlendWeights[maskIndex].MeshSpaceDesiredBoneBlendWeights,
			                                           AllBoneBlendWeights[maskIndex].MeshSpaceCurrentBoneBlendWeights,
			                                           msDesiredBlendWeight);

			FAnimationRuntime::BlendTwoPosesTogether(animationPoseOneData, animationPoseTwoData,
			                                         (1.0f - MaskBlendWeights[maskIndex].OverrideAlpha),
			                                         blendedAnimationPoseData);

			FPoseContext msCurrentPoseContext(animationPoseContext);
			msCurrentPoseContext = currentPoseContext;
			if (FAnimWeight::IsRelevant(MaskBlendWeights[maskIndex].LocalSpaceBlendAlpha))
			{
				lsTargetBlendPoses[0].MoveBonesFrom(currentPoseContext.Pose);
				lsTargetBlendCurves[0].MoveFrom(currentPoseContext.Curve);
				lsTargetBlendAttributes[0].MoveFrom(currentPoseContext.CustomAttributes);
			}
			else
			{
				lsTargetBlendPoses[0].ResetToRefPose(animationPoseContext.Pose.GetBoneContainer());
				lsTargetBlendCurves[0].InitFrom(output.Curve);
			}

			if (FAnimWeight::IsRelevant(MaskBlendWeights[maskIndex].MeshSpaceBlendAlpha))
			{
				msTargetBlendPoses[0].MoveBonesFrom(msCurrentPoseContext.Pose);
				msTargetBlendCurves[0].MoveFrom(msCurrentPoseContext.Curve);
				msTargetBlendAttributes[0].MoveFrom(msCurrentPoseContext.CustomAttributes);
			}
			else
			{
				msTargetBlendPoses[0].ResetToRefPose(animationPoseContext.Pose.GetBoneContainer());
				msTargetBlendCurves[0].InitFrom(output.Curve);
			}

			// Filter to make sure it only includes curves that are linked to the correct bone filter
			UE::Anim::FNamedValueArrayUtils::RemoveByPredicate(animationPoseContext.Curve,
			                                                   AllBoneBlendWeights[maskIndex].
			                                                   LocalCurvePoseSourceIndices,
			                                                   [](const UE::Anim::FCurveElement& inOutBasePoseElement,
			                                                      const UE::Anim::FCurveElementIndexed&
			                                                      inSourceIndexElement)
			                                                   {
				                                                   // if source index is set, remove base pose curve value
				                                                   return (inSourceIndexElement.Index != INDEX_NONE);
			                                                   });

			// Filter to make sure it only includes curves that are linked to the correct bone filter
			UE::Anim::FNamedValueArrayUtils::RemoveByPredicate(animationPoseContext.Curve,
			                                                   AllBoneBlendWeights[maskIndex].
			                                                   MeshCurvePoseSourceIndices,
			                                                   [](const UE::Anim::FCurveElement& inOutBasePoseElement,
			                                                      const UE::Anim::FCurveElementIndexed&
			                                                      inSourceIndexElement)
			                                                   {
				                                                   // if source index is set, remove base pose curve value
				                                                   return (inSourceIndexElement.Index != INDEX_NONE);
			                                                   });


			FAnimationRuntime::EBlendPosesPerBoneFilterFlags msBlendFlags =
				FAnimationRuntime::EBlendPosesPerBoneFilterFlags::None;
			msBlendFlags |= FAnimationRuntime::EBlendPosesPerBoneFilterFlags::MeshSpaceRotation;

			FAnimationRuntime::EBlendPosesPerBoneFilterFlags lsBlendFlags =
				FAnimationRuntime::EBlendPosesPerBoneFilterFlags::None;
			if (FAnimWeight::IsRelevant(MaskBlendWeights[maskIndex].LocalSpaceBlendAlpha))
			{
				FAnimationPoseData animationPoseData(output);
				FAnimationRuntime::BlendPosesPerBoneFilter(updatedAnimationPoseContext.Pose, lsTargetBlendPoses,
				                                           updatedAnimationPoseContext.Curve, lsTargetBlendCurves,
				                                           updatedAnimationPoseContext.CustomAttributes,
				                                           lsTargetBlendAttributes, animationPoseData,
				                                           AllBoneBlendWeights[maskIndex].
				                                           LocalSpaceCurrentBoneBlendWeights, lsBlendFlags,
				                                           BlendMasks[maskIndex].CurveBlendOption);
				updatedAnimationPoseContext = output;
				updatedAnimationPoseContext.Pose.NormalizeRotations();
			}

			if (FAnimWeight::IsRelevant(MaskBlendWeights[maskIndex].MeshSpaceBlendAlpha))
			{
				FAnimationPoseData MSAnimationPoseData(output);
				FAnimationRuntime::BlendPosesPerBoneFilter(updatedAnimationPoseContext.Pose, msTargetBlendPoses,
				                                           updatedAnimationPoseContext.Curve, msTargetBlendCurves,
				                                           updatedAnimationPoseContext.CustomAttributes,
				                                           msTargetBlendAttributes, MSAnimationPoseData,
				                                           AllBoneBlendWeights[maskIndex].
				                                           MeshSpaceCurrentBoneBlendWeights, msBlendFlags,
				                                           BlendMasks[maskIndex].CurveBlendOption);
				updatedAnimationPoseContext = output;
				updatedAnimationPoseContext.Pose.NormalizeRotations();
			}
		}

		return;
	}

	AnimationPose.Evaluate(output);
}

void FAnimNode_OLSMask::InvalidatePerBoneBlendWeights()
{
	RequiredBonesSerialNumber = 0;
	SkeletonGuid = FGuid();
	VirtualBoneGuid = FGuid();
}

void FAnimNode_OLSMask::InvalidateCachedBoneData()
{
	RequiredBonesSerialNumber = 0;
}

void FAnimNode_OLSMask::RebuildPerBoneBlendWeights(const USkeleton* skeleton)
{
	if (skeleton)
	{
		if (BlendMasks.Num() > 0)
		{
			PerBoneBlendWeights.SetNum(BlendMasks.Num());
			TArray<UBlendProfile*> currentBlendProfile;
			for (int32 Index = 0; Index < BlendMasks.Num(); ++Index)
			{
				if (BlendMasks[Index].BlendProfile && BlendMasks[Index].BlendProfile->IsBlendMask())
				{
					currentBlendProfile.Reset();
					currentBlendProfile.Add(BlendMasks[Index].BlendProfile);
					FAnimationRuntime::CreateMaskWeights(PerBoneBlendWeights[Index].PerBoneBlendWeights,
					                                     currentBlendProfile, skeleton);
				}
				else
				{
					const FReferenceSkeleton& refSkeleton = skeleton->GetReferenceSkeleton();
					const int32 numBones = refSkeleton.GetNum();
					PerBoneBlendWeights[Index].PerBoneBlendWeights.Reset(numBones);
					// We only store non-zero weights in blend masks. Initialize all to zero.
					PerBoneBlendWeights[Index].PerBoneBlendWeights.AddZeroed(numBones);
				}
			}
		}

		SkeletonGuid = skeleton->GetGuid();
		VirtualBoneGuid = skeleton->GetVirtualBoneGuid();
	}
}

bool FAnimNode_OLSMask::ArePerBoneBlendWeightsValid(const USkeleton* skeleton) const
{
	return (skeleton && skeleton->GetGuid() == SkeletonGuid && skeleton->GetVirtualBoneGuid() == VirtualBoneGuid);
}

void FAnimNode_OLSMask::UpdateCachedBoneData(const FBoneContainer& requiredBones, const USkeleton* skeleton)
{
	if (requiredBones.GetSerialNumber() == RequiredBonesSerialNumber)
	{
		return;
	}

	if (requiredBones.GetBoneIndicesArray().Num() > 0)
	{
		if (!ArePerBoneBlendWeightsValid(skeleton))
		{
			RebuildPerBoneBlendWeights(skeleton);
		}

		for (int32 Index = 0; Index < BlendMasks.Num(); ++Index)
		{
			UpdateBodyPartCachedBoneData(requiredBones, skeleton, Index);
		}

		RequiredBonesSerialNumber = requiredBones.GetSerialNumber();
	}
}

void FAnimNode_OLSMask::UpdateBodyPartCachedBoneData(const FBoneContainer& requiredBones, const USkeleton* skeleton,
                                                     const int32 maskIndex)
{
	const TArray<FBoneIndexType>& requiredBoneIndices = requiredBones.GetBoneIndicesArray();
	const int32 numRequiredBones = requiredBoneIndices.Num();
	AllBoneBlendWeights[maskIndex].LocalSpaceDesiredBoneBlendWeights.SetNumZeroed(numRequiredBones);
	AllBoneBlendWeights[maskIndex].MeshSpaceDesiredBoneBlendWeights.SetNumZeroed(numRequiredBones);
	for (int32 requiredBoneIndex = 0; requiredBoneIndex < numRequiredBones; requiredBoneIndex++)
	{
		const int32 skeletonBoneIndex = requiredBones.GetSkeletonIndex(FCompactPoseBoneIndex(requiredBoneIndex));
		if (ensure(skeletonBoneIndex != INDEX_NONE))
		{
			AllBoneBlendWeights[maskIndex].LocalSpaceDesiredBoneBlendWeights[requiredBoneIndex] = PerBoneBlendWeights[
				maskIndex].PerBoneBlendWeights[skeletonBoneIndex];
			AllBoneBlendWeights[maskIndex].MeshSpaceDesiredBoneBlendWeights[requiredBoneIndex] = PerBoneBlendWeights[
				maskIndex].PerBoneBlendWeights[skeletonBoneIndex];
		}
	}
	AllBoneBlendWeights[maskIndex].LocalSpaceCurrentBoneBlendWeights.Reset(
		AllBoneBlendWeights[maskIndex].LocalSpaceDesiredBoneBlendWeights.Num());
	AllBoneBlendWeights[maskIndex].MeshSpaceCurrentBoneBlendWeights.Reset(
		AllBoneBlendWeights[maskIndex].MeshSpaceDesiredBoneBlendWeights.Num());
	AllBoneBlendWeights[maskIndex].LocalSpaceCurrentBoneBlendWeights.AddZeroed(
		AllBoneBlendWeights[maskIndex].LocalSpaceDesiredBoneBlendWeights.Num());
	AllBoneBlendWeights[maskIndex].MeshSpaceCurrentBoneBlendWeights.AddZeroed(
		AllBoneBlendWeights[maskIndex].MeshSpaceDesiredBoneBlendWeights.Num());
	TArray<float> lsDesiredBlendWeight;
	lsDesiredBlendWeight.Add(MaskBlendWeights[maskIndex].LocalSpaceBlendAlpha);
	TArray<float> msDesiredBlendWeight;
	msDesiredBlendWeight.Add(MaskBlendWeights[maskIndex].MeshSpaceBlendAlpha);
	FAnimationRuntime::UpdateDesiredBoneWeight(AllBoneBlendWeights[maskIndex].LocalSpaceDesiredBoneBlendWeights,
	                                           AllBoneBlendWeights[maskIndex].LocalSpaceCurrentBoneBlendWeights,
	                                           lsDesiredBlendWeight);
	FAnimationRuntime::UpdateDesiredBoneWeight(AllBoneBlendWeights[maskIndex].MeshSpaceDesiredBoneBlendWeights,
	                                           AllBoneBlendWeights[maskIndex].MeshSpaceCurrentBoneBlendWeights,
	                                           msDesiredBlendWeight);

	// Build curve source indices
	AllBoneBlendWeights[maskIndex].LocalCurvePoseSourceIndices.Empty();
	AllBoneBlendWeights[maskIndex].LocalCurvePoseSourceIndices.Reserve(skeleton->GetNumCurveMetaData());
	skeleton->ForEachCurveMetaData(
		[this, &requiredBones, maskIndex](const FName& curveName, const FCurveMetaData& metaData)
		{
			for (const FBoneReference& linkedBone : metaData.LinkedBones)
			{
				FCompactPoseBoneIndex compactPoseIndex = linkedBone.GetCompactPoseIndex(requiredBones);
				if (compactPoseIndex != INDEX_NONE)
				{
					if (AllBoneBlendWeights[maskIndex].LocalSpaceDesiredBoneBlendWeights[compactPoseIndex.GetInt()].BlendWeight > 0.f)
					{
						AllBoneBlendWeights[maskIndex].LocalCurvePoseSourceIndices.Add(
							curveName,
							AllBoneBlendWeights[maskIndex].LocalSpaceDesiredBoneBlendWeights[compactPoseIndex.GetInt()].
							SourceIndex);
						break;
					}
				}
			}
		});

	// Build curve source indices
	AllBoneBlendWeights[maskIndex].MeshCurvePoseSourceIndices.Empty();
	AllBoneBlendWeights[maskIndex].MeshCurvePoseSourceIndices.Reserve(skeleton->GetNumCurveMetaData());
	skeleton->ForEachCurveMetaData(
		[this, &requiredBones, maskIndex](const FName& curveName, const FCurveMetaData& metaData)
		{
			for (const FBoneReference& linkedBone : metaData.LinkedBones)
			{
				FCompactPoseBoneIndex compactPoseIndex = linkedBone.GetCompactPoseIndex(requiredBones);
				if (compactPoseIndex != INDEX_NONE)
				{
					if (AllBoneBlendWeights[maskIndex].MeshSpaceDesiredBoneBlendWeights[compactPoseIndex.GetInt()].
						BlendWeight > 0.f)
					{
						AllBoneBlendWeights[maskIndex].MeshCurvePoseSourceIndices.Add(
							curveName,
							AllBoneBlendWeights[maskIndex].MeshSpaceDesiredBoneBlendWeights[compactPoseIndex.GetInt()].
							SourceIndex);
						break;
					}
				}
			}
		});
}

void FAnimNode_OLSMask::UpdateWeightsFromCurves(const USkeleton* skeleton, const FPoseContext& desiredCurvesPose,
                                                const int32 index)
{
	if (BlendMasks[index].WeightCurves.OverrideCurve != NAME_None)
	{
		MaskBlendWeights[index].OverrideAlpha = desiredCurvesPose.Curve.Get(
			BlendMasks[index].WeightCurves.OverrideCurve);
	}
	
	if (BlendMasks[index].WeightCurves.AdditiveCurve != NAME_None)
	{
		MaskBlendWeights[index].AdditiveAlpha = desiredCurvesPose.Curve.Get(
			BlendMasks[index].WeightCurves.AdditiveCurve);
	}
	
	if (BlendMasks[index].WeightCurves.LocalSpaceBlendCurve != NAME_None)
	{
		MaskBlendWeights[index].LocalSpaceBlendAlpha = desiredCurvesPose.Curve.Get(
			BlendMasks[index].WeightCurves.LocalSpaceBlendCurve);
	}
	
	if (BlendMasks[index].WeightCurves.MeshSpaceBlendCurve != NAME_None)
	{
		MaskBlendWeights[index].MeshSpaceBlendAlpha = desiredCurvesPose.Curve.Get(
			BlendMasks[index].WeightCurves.MeshSpaceBlendCurve);
	}
}

void FAnimNode_OLSMask::MakeAdditivePose(FPoseContext& outAdditivePose, FPoseContext& outBasePose,
                                         const bool shouldUseMeshSpaceAdditive)
{
	if (shouldUseMeshSpaceAdditive)
	{
		FAnimationRuntime::ConvertPoseToMeshRotation(outAdditivePose.Pose);
		FAnimationRuntime::ConvertPoseToMeshRotation(outBasePose.Pose);
	}

	FAnimationRuntime::ConvertPoseToAdditive(outAdditivePose.Pose, outBasePose.Pose);
	outAdditivePose.Curve.ConvertToAdditive(outBasePose.Curve);

	UE::Anim::Attributes::ConvertToAdditive(outBasePose.CustomAttributes, outAdditivePose.CustomAttributes);
}