// © 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 "Libraries/OLSLocomotionBPLibrary.h"

#include "SequencePlayerLibrary.h"
#include "Animation/AnimCurveCompressionCodec_UniformIndexable.h"
#include "Animation/AnimNode_SequencePlayer.h"
#include "AnimNodes/AnimNode_SequenceEvaluator.h"

DEFINE_LOG_CATEGORY_STATIC(LogOLSLocomotionLibrary, Verbose, All);

float UOLSLocomotionBPLibrary::GetDistanceCurveValueAtTime(const UAnimSequenceBase* animSequence,
                                                           const float time,
                                                           const FName& curveName)
{
	FAnimCurveBufferAccess bufferCurveAccess(animSequence, curveName);
	if (bufferCurveAccess.IsValid())
	{
		const float clampedTime = FMath::Clamp(time, 0.f, animSequence->GetPlayLength());
		if (animSequence->GetNumberOfSampledKeys() > 2)
		{
			return animSequence->EvaluateCurveData(curveName, clampedTime);
		}
	}

	return 0.f;
}

float UOLSLocomotionBPLibrary::GetTimeAtDistance(const UAnimSequenceBase* animSequence,
                                                 const float& distance, FName curveName)
{
	FAnimCurveBufferAccess bufferCurveAccess(animSequence, curveName);
	if (bufferCurveAccess.IsValid())
	{
		const int32 numKeys = bufferCurveAccess.GetNumSamples();
		if (numKeys < 2)
		{
			return 0.f;
		}

		int32 first = 1;
		int32 last = numKeys - 1;
		int32 count = last - first;

		while (count > 0)
		{
			int32 step = count / 2;
			int32 middle = first + step;

			if (distance > bufferCurveAccess.GetValue(middle))
			{
				first = middle + 1;
				count -= step + 1;
			}
			else
			{
				count = step;
			}
		}

		const float keyAValue = bufferCurveAccess.GetValue(first - 1);
		const float keyBValue = bufferCurveAccess.GetValue(first);
		const float diff = keyBValue - keyAValue;
		const float alpha = !FMath::IsNearlyZero(diff) ? ((distance - keyAValue) / diff) : 0.f;

		const float keyATime = bufferCurveAccess.GetTime(first - 1);
		const float keyBTime = bufferCurveAccess.GetTime(first);
		return FMath::Lerp(keyATime, keyBTime, alpha);
	}

	return 0.f;
}

float UOLSLocomotionBPLibrary::GetDistanceRange(const UAnimSequenceBase* animSequence, const FName& curveName)
{
	FAnimCurveBufferAccess bufferCurveAccess(animSequence, curveName);
	if (bufferCurveAccess.IsValid())
	{
		const int32 numSamples = bufferCurveAccess.GetNumSamples();
		if (numSamples >= 2)
		{
			return (bufferCurveAccess.GetValue(numSamples - 1) - bufferCurveAccess.GetValue(0));
		}
	}
	return 0.f;
}

float UOLSLocomotionBPLibrary::GetTimeAfterDistanceTraveled(const UAnimSequenceBase* animSequence,
															float currentTime,
                                                            float distanceTraveled, FName curveName,
                                                            const bool shouldAllowLooping)
{
	float newTime = currentTime;
	if (!animSequence)
	{
		// Avoid infinite loops if the animation doesn't cover any distance.
		if (!FMath::IsNearlyZero(GetDistanceRange(animSequence, curveName)))
		{
			float accumulatedDistance = 0.f;

			const float sequenceLength = animSequence->GetPlayLength();
			constexpr float stepTime = 1.f / 30.f;

			// Distance Matching expects the distance curve on the animation to increase monotonically. If the curve fails to increase in value
			// after a certain number of iterations, we abandon the algorithm to avoid an infinite loop.

			// Traverse the distance curve, accumulating animated distance until the desired distance is reached.
			while ((accumulatedDistance < distanceTraveled) && (shouldAllowLooping || (newTime + stepTime < sequenceLength)))
			{
				const float currentDistance = GetDistanceCurveValueAtTime(animSequence, newTime, curveName);
				const float distanceAfterStep = GetDistanceCurveValueAtTime(animSequence, newTime + stepTime, curveName);
				const float animationDistanceThisStep = distanceAfterStep - currentDistance;

				if (!FMath::IsNearlyZero(animationDistanceThisStep))
				{
					// Keep advancing if the desired distance hasn't been reached.
					if (accumulatedDistance + animationDistanceThisStep < distanceTraveled)
					{
						FAnimationRuntime::AdvanceTime(shouldAllowLooping, stepTime, newTime, sequenceLength);
						accumulatedDistance += animationDistanceThisStep;
					}
					// Once the desired distance is passed, find the approximate time between samples where the distance will be reached.
					else
					{
						const float DistanceAlpha = (distanceTraveled - accumulatedDistance) / animationDistanceThisStep;
						FAnimationRuntime::AdvanceTime(shouldAllowLooping, DistanceAlpha * stepTime, newTime, sequenceLength);
						break;
					}
						
				}
				else
				{
					FAnimationRuntime::AdvanceTime(shouldAllowLooping, stepTime, newTime, sequenceLength);
					break;
						
				}
			}
		}
		else
		{
			UE_LOG(LogOLSLocomotionLibrary, Warning,
			       TEXT(
				       "Anim sequence (%s) is missing a distance curve or doesn't cover enough distance for GetTimeAfterDistanceTraveled."
			       ), *GetNameSafe(animSequence));
		}
	}
	else
	{
		UE_LOG(LogOLSLocomotionLibrary, Warning, TEXT("Invalid AnimSequence passed to GetTimeAfterDistanceTraveled"));
	}
	
	return newTime;
}

FSequenceEvaluatorReference UOLSLocomotionBPLibrary::AdvanceTimeByDistanceMatching(float& outDesiredPlayRate,
																				   const FAnimUpdateContext& updateContext,
																				   const FSequenceEvaluatorReference& sequenceEvaluator,
																				   const float distanceTraveled,
																				   const FName curveName,
																				   const FVector2D playRateClamp /* = FVector2D(0.75f, 1.25f)*/)
{
	sequenceEvaluator.CallAnimNodeFunction<FAnimNode_SequenceEvaluator>(
		TEXT("AdvanceTimeByDistanceMatching"),
		[&outDesiredPlayRate, updateContext, distanceTraveled, curveName, playRateClamp](
		FAnimNode_SequenceEvaluator& inSequenceEvaluator)
		{
			if (const FAnimationUpdateContext* animationUpdateContext = updateContext.GetContext())
			{
				const float deltaTime = animationUpdateContext->GetDeltaTime();

				if (deltaTime > 0 && distanceTraveled > 0)
				{
					if (const UAnimSequenceBase* animSequence = Cast<UAnimSequence>(inSequenceEvaluator.GetSequence()))
					{
						const float currentTime = inSequenceEvaluator.GetExplicitTime();
						const float currentAssetLength = inSequenceEvaluator.GetCurrentAssetLength();
						const bool shouldAllowLooping = inSequenceEvaluator.IsLooping();
						float timeAfterDistanceTraveled = GetTimeAfterDistanceTraveled(
							animSequence,
							currentTime,
							distanceTraveled,
							curveName,
							shouldAllowLooping);

						// Calculate the effective playrate that would result from advancing the animation by the distance traveled.
						// // Account for the animation looping.
						if (timeAfterDistanceTraveled < currentTime)
						{
							timeAfterDistanceTraveled += currentAssetLength;
						}
						float effectivePlayRate = (timeAfterDistanceTraveled - currentTime) / deltaTime;

						outDesiredPlayRate = effectivePlayRate;
						// Clamp the effective play rate.
						if (playRateClamp.X >= 0.f && playRateClamp.X < playRateClamp.Y)
						{
							effectivePlayRate = FMath::Clamp(effectivePlayRate, playRateClamp.X, playRateClamp.Y);
						}

						// Advance animation time by the effective play rate.
						float newTime = currentTime;
						FAnimationRuntime::AdvanceTime(false, effectivePlayRate * deltaTime, newTime,
						                               currentAssetLength);

						if (!inSequenceEvaluator.SetExplicitTime(newTime))
						{
							UE_LOG(LogOLSLocomotionLibrary, Warning,
							       TEXT(
								       "Could not set explicit time on sequence evaluator, value is not dynamic. Set it as Always Dynamic."
							       ));
						}
					}
					else
					{
						UE_LOG(LogOLSLocomotionLibrary, Warning, TEXT("Sequence evaluator does not have an anim sequence to play."));
					}
				}
			}
			else
			{
				UE_LOG(LogOLSLocomotionLibrary, Warning,
				       TEXT("AdvanceTimeByDistanceMatching called with invalid context"));
			}
		});

	return sequenceEvaluator;
}

FSequenceEvaluatorReference UOLSLocomotionBPLibrary::DistanceMatchToTarget(const FAnimUpdateContext& updateContext,
                                                                           const FSequenceEvaluatorReference& sequenceEvaluator,
                                                                           float distanceToTarget,
                                                                           bool shouldDistanceMatchStop,
                                                                           float stopDistanceThreshHold,
                                                                           float animEndTime,
                                                                           FName curveName)
{
	sequenceEvaluator.CallAnimNodeFunction<FAnimNode_SequenceEvaluator>(
		TEXT("DistanceMatchToTarget"),
		[updateContext,sequenceEvaluator,distanceToTarget, shouldDistanceMatchStop,stopDistanceThreshHold,animEndTime,
			curveName](
		FAnimNode_SequenceEvaluator& inSequenceEvaluator)
		{
			if (const UAnimSequenceBase* animSequence = inSequenceEvaluator.GetSequence())
			{
				if (GetDistanceCurveValueAtTime(animSequence,
				                                USequenceEvaluatorLibrary::GetAccumulatedTime(sequenceEvaluator),
				                                curveName) > stopDistanceThreshHold && !shouldDistanceMatchStop)
				{
					const float newTime = GetTimeAtDistance(animSequence, -distanceToTarget, curveName);
					if (!inSequenceEvaluator.SetExplicitTime(newTime))
					{
						UE_LOG(LogOLSLocomotionLibrary, Warning,
						       TEXT(
							       "Could not set explicit time on sequence evaluator, value is not dynamic. Set it as Always Dynamic."
						       ));
					}
				}
				else
				{
					USequenceEvaluatorLibrary::AdvanceTime(updateContext, sequenceEvaluator, 1.0f);
					if (animEndTime > 0)
					{
						const float desiredTime = FMath::Clamp(
							USequenceEvaluatorLibrary::GetAccumulatedTime(sequenceEvaluator), 0, animEndTime);
						USequenceEvaluatorLibrary::SetExplicitTime(sequenceEvaluator, desiredTime);
					}
				}
			}
			else
			{
				UE_LOG(LogOLSLocomotionLibrary, Warning,
				       TEXT("Sequence evaluator does not have an anim sequence to play."));
			}
		});

	return sequenceEvaluator;
}

FSequencePlayerReference UOLSLocomotionBPLibrary::SetPlayRateToMatchSpeed(const FSequencePlayerReference& sequencePlayer,
																		  float speedToMatch,
																		  FVector2D playRateClamp)
{
	sequencePlayer.CallAnimNodeFunction<FAnimNode_SequencePlayer>(
		TEXT("SetPlayrateToMatchSpeed"),
		[speedToMatch, playRateClamp](FAnimNode_SequencePlayer& sequencePlayer)
		{
			if (const UAnimSequence* animSequence = Cast<UAnimSequence>(sequencePlayer.GetSequence()))
			{
				const float animLength = animSequence->GetPlayLength();
				if (!FMath::IsNearlyZero(animLength))
				{
					// Calculate the speed as: (distance traveled by the animation) / (length of the animation)
					const FVector rootMotionTranslation = animSequence->ExtractRootMotionFromRange(0.0f, animLength).
					                                                    GetTranslation();
					const float rootMotionDistance = rootMotionTranslation.Size2D();
					if (!FMath::IsNearlyZero(rootMotionDistance))
					{
						const float animationSpeed = rootMotionDistance / animLength;
						float desiredPlayRate = speedToMatch / animationSpeed;
						if (playRateClamp.X >= 0.0f && playRateClamp.X < playRateClamp.Y)
						{
							desiredPlayRate = FMath::Clamp(desiredPlayRate, playRateClamp.X, playRateClamp.Y);
						}

						if (!sequencePlayer.SetPlayRate(desiredPlayRate))
						{
							UE_LOG(LogOLSLocomotionLibrary, Warning,
							       TEXT(
								       "Could not set play rate on sequence player, value is not dynamic. Set it as Always Dynamic."
							       ));
						}
					}
					else
					{
						UE_LOG(LogOLSLocomotionLibrary, Warning,
						       TEXT("Unable to adjust playrate for animation with no root motion delta (%s)."),
						       *GetNameSafe(animSequence));
					}
				}
				else
				{
					UE_LOG(LogOLSLocomotionLibrary, Warning,
					       TEXT("Unable to adjust playrate for zero length animation (%s)."),
					       *GetNameSafe(animSequence));
				}
			}
			else
			{
				UE_LOG(LogOLSLocomotionLibrary, Warning,
				       TEXT("Sequence player does not have an anim sequence to play."));
			}
		});

	return sequencePlayer;
}

FVector UOLSLocomotionBPLibrary::PredictGroundMovementStopLocation(const FVector& velocity,
                                                                   bool shouldUseSeparateBrakingFriction,
                                                                   float brakingFriction, float groundFriction,
                                                                   float brakingFrictionFactor,
                                                                   float brakingDecelerationWalking)
{
	FVector predictedStopLocation = FVector::ZeroVector;

	// Determine the actual braking friction
	float actualBrakingFriction = shouldUseSeparateBrakingFriction ? brakingFriction : groundFriction;
	actualBrakingFriction = FMath::Max(0.f, actualBrakingFriction * FMath::Max(0.f, brakingFrictionFactor));

	// Calculate 2D velocity and speed
	const FVector velocity2D = FVector(velocity.X, velocity.Y, 0.f);
	const float speed2D = velocity2D.Size();

	// Check if there's movement to stop
	if (speed2D > 0.f)
	{
		// Calculate braking divisor
		const float divisor = actualBrakingFriction * speed2D + FMath::Max(0.f, brakingDecelerationWalking);

		// Check if stopping is possible
		if (divisor > 0.f)
		{
			// Calculate time to stop
			const float timeToStop = speed2D / divisor;

			// Calculate predicted stop location
			predictedStopLocation = velocity2D * timeToStop + 0.5f * ((-actualBrakingFriction) * velocity2D -
				brakingDecelerationWalking * velocity2D.GetSafeNormal()) * FMath::Square(timeToStop);
		}
	}

	return predictedStopLocation;
}

FVector UOLSLocomotionBPLibrary::PredictGroundMovementPivotLocation(const FVector& acceleration,
                                                                    const FVector& velocity,
                                                                    float groundFriction)
{
	FVector predictedPivotLocation = FVector::ZeroVector;

	const FVector acceleration2D = acceleration * FVector(1.f, 1.f, 0.f);

	const float accelerationSize2D = acceleration2D.Size();

	// Check if velocity is along the opposite direction of acceleration
	if ((velocity | acceleration2D) < 0.0f)
	{
		const float speedAlongAcceleration = -(velocity | acceleration2D);
		const float divisor = accelerationSize2D + 2.f * speedAlongAcceleration * groundFriction;

		// Check if stopping is possible
		if (divisor > 0.f)
		{
			const float timeToDirectionChange = speedAlongAcceleration / divisor;

			// Calculate the acceleration force
			const FVector accelerationForce = acceleration - (velocity - acceleration2D * velocity.Size2D()) * groundFriction;

			// Calculate the predicted pivot location
			predictedPivotLocation = velocity * timeToDirectionChange + 0.5f * accelerationForce * timeToDirectionChange * timeToDirectionChange;
		}
	}

	return predictedPivotLocation;
}

EOLSCardinalDirection UOLSLocomotionBPLibrary::SelectCardinalDirectionFromAngle(float angle,
                                                                                float deadZone,
                                                                                EOLSCardinalDirection currentDirection,
                                                                                bool useCurrentDirection /* = false */)
{

	const float absAngle = FMath::Abs(angle);
	float fwdDeadZone = deadZone;
	float bwdDeadZone = deadZone;

	if (useCurrentDirection)
	{
		if (currentDirection == EOLSCardinalDirection::EForward)
		{
			fwdDeadZone *= 2.f;
		}
		else if (currentDirection == EOLSCardinalDirection::EBackward)
		{
			bwdDeadZone *= 2.f;
		}
	}

	if(absAngle <= (45 + fwdDeadZone))
	{
		return EOLSCardinalDirection::EForward;
	}
	else if (absAngle >= (135 - bwdDeadZone))
	{
		return EOLSCardinalDirection::EBackward;
	}
	else if (angle > 0)
	{
		return EOLSCardinalDirection::ERight;
	}

	return EOLSCardinalDirection::ELeft;
}

EOLSCardinalDirection UOLSLocomotionBPLibrary::GetOppositeCardinalDirectional(EOLSCardinalDirection currentDirection)
{
	switch (currentDirection)
	{
		case EOLSCardinalDirection::EForward: {return EOLSCardinalDirection::EBackward;}
		case EOLSCardinalDirection::EBackward: {return EOLSCardinalDirection::EForward;}
		case EOLSCardinalDirection::ELeft: {return EOLSCardinalDirection::ERight;}
		case EOLSCardinalDirection::ERight: {return EOLSCardinalDirection::ELeft;}
	}

	return EOLSCardinalDirection::EForward;
}

EOLSHipDirection UOLSLocomotionBPLibrary::GetOppositeHipDirection(EOLSHipDirection currentHipDirection)
{
	return (currentHipDirection == EOLSHipDirection::EForward ? EOLSHipDirection::EBackward : EOLSHipDirection::EForward);
}

void UOLSLocomotionBPLibrary::TryLinkAnimLayer(USkeletalMeshComponent* mesh,
                                               TSubclassOf<UAnimInstance> animClass,
                                               FName groupName,
                                               bool shouldUnlinkGroupIfInvalid)
{
	if (!animClass->IsValidLowLevelFast())
	{
		if (shouldUnlinkGroupIfInvalid)
		{
			if (const TObjectPtr<UAnimInstance> linkedAnimInstance = mesh->GetLinkedAnimLayerInstanceByGroup(groupName))
			{
				mesh->UnlinkAnimClassLayers(linkedAnimInstance.GetClass());
			}
		}
		
		return;
	}

	mesh->LinkAnimClassLayers(animClass);
}