// (CVN) - Candy Valley Network GmbH

#include "Components/LovenseToyAnimationControlComponent.h"

// EXTERNAL INCLUDES
#include <BoneContainer.h>
#include <BoneIndices.h>
#include <BonePose.h>
#include <Animation/AnimationPoseData.h>
#include <Animation/AnimSequence.h>
#include <Animation/Skeleton.h>
#include <Components/SkeletalMeshComponent.h>
#include <Engine/World.h>

// INTERNAL INCLUDES
#include "LovenseIntegration.h"
#include "LovenseManager.h"


ULovenseToyAnimationControlComponent::ULovenseToyAnimationControlComponent() {
	this->PrimaryComponentTick.bCanEverTick = true;
	this->PrimaryComponentTick.bStartWithTickEnabled = false;

	// Realtime evaluation settings
	this->maxAverageVelocity = 50.0f;
	this->maxRecordedVelocityDecaySpeed = 5.0f;
	this->maxRecordedVelocityDecayMinimum = 10.0f;
	this->maxToySpeedInterpolationSpeed = 2.0f;
	this->maxToySpeedTargetDuringTransition = 5.0f;

	// Pattern generation settings
	this->patternToySpeedInputRange = FVector2D(0.0f, 20.0f);

	// Reference data
	this->skeletalMeshComponents = TArray<USkeletalMeshComponent*>();
	this->boneNames = TArray<FName>();

	// Realtime evaluation runtime values.
	this->maxRecordedVelocity = 0.0f;
	this->currentMaxToySpeed = 20.0f;
	this->maxToySpeedTarget = 20.0f;
	this->bRotateClockwise = false;

	// Miscellaneous
	this->bIsEnabled = false;
	this->patternGenerationTimerHandle = FTimerHandle();
	this->getAnimationDataDelegate = FLovenseGetAnimationData();
}

void ULovenseToyAnimationControlComponent::TickComponent(float DeltaTime, ELevelTick TickType, FActorComponentTickFunction* ThisTickFunction) {
	this->RealtimeEvaluationUpdate(DeltaTime);

	Super::TickComponent(DeltaTime, TickType, ThisTickFunction);
}

void ULovenseToyAnimationControlComponent::EndPlay(const EEndPlayReason::Type EndPlayReason) {
	// Disable the component to stock ticking, clear pattern generation timer and clear all commands.
	this->SetEnabled(false);

	Super::EndPlay(EndPlayReason);

	// Cleanup
	this->getAnimationDataDelegate.Clear();
	this->boneNames.Empty();
	this->skeletalMeshComponents.Empty();
}

void ULovenseToyAnimationControlComponent::Initialize(
	const TArray<USkeletalMeshComponent*>& inSkeletalMeshComponents,
	const TArray<FName>& inBoneNames,
	const FLovenseGetAnimationData& inGetAnimationDataDelegate
) {
	// As Initialize() can be called while this component is enabled, we disable it here so the toys don't run with outdated commands.
	this->SetEnabled(false);
	this->skeletalMeshComponents = inSkeletalMeshComponents;
	this->boneNames = inBoneNames;
	this->getAnimationDataDelegate = inGetAnimationDataDelegate;
}

void ULovenseToyAnimationControlComponent::StartBehaviour() {
	// Reset realtime evaluation runtime values.
	this->maxRecordedVelocity = this->maxRecordedVelocityDecayMinimum;
	this->maxToySpeedTarget = 20.0f;
	// Enable the component to start ticking, start pattern generation timer and generate and send a pattern command.
	this->SetEnabled(true);
}

void ULovenseToyAnimationControlComponent::StartTransition() {
	// Set max toy speed to max toy speed during transition and flip rotation direction.
	this->maxToySpeedTarget = this->maxToySpeedTargetDuringTransition;
	this->bRotateClockwise = !this->bRotateClockwise;
}

void ULovenseToyAnimationControlComponent::SetEnabled(bool value) {
	// We purposefully don't check whether we are already enabled/disabled as that allows external code to immediately generate and send a pattern command at any time.
	this->bIsEnabled = value;

	if (!this->bIsEnabled) {
		// When disabled, clear pattern generation timer and clear all currently running commands.
		if (this->GetWorld()) this->GetWorld()->GetTimerManager().ClearTimer(this->patternGenerationTimerHandle);
		this->patternGenerationTimerHandle.Invalidate();
		FLovenseManager::Get()->ClearAllCommands();
	} else {
		// When enabled, (re)start pattern generation timer and generate and send a pattern command.
		this->OnPatternGenerationTimerFinished();
	}

	// Enable/disable component tick, which enables/disables the realtime evaluation.
	this->SetComponentTickEnabled(value);
}

void ULovenseToyAnimationControlComponent::OnGamePaused() {
	// Disable the component when the game has been paused to stop the toys.
	// You will need to manually call OnGamePaused() as there are no global events for that in unreal.
	this->SetEnabled(false);
}

void ULovenseToyAnimationControlComponent::OnGameUnpaused() {
	// Enable the component when the game has been unpaused.
	// You will need to manually call OnGameUnpaused() as there are no global events for that in unreal.
	this->SetEnabled(true);
}

void ULovenseToyAnimationControlComponent::OnPatternGenerationTimerFinished() {
	if (!this->GetWorld()) return;
	// Restart the timer.
	// We don't use the loop functionality of the timer manager so this can be called before the timer finishes.
	this->GetWorld()->GetTimerManager().SetTimer(this->patternGenerationTimerHandle, this, &ULovenseToyAnimationControlComponent::OnPatternGenerationTimerFinished, 4.0f);

	// If delegate is not bound, there's nothing to do.
	if (!this->getAnimationDataDelegate.IsBound()) return;

	// Execute the delegate to fetch required animation data.
	TArray<UAnimSequence*> animations = TArray<UAnimSequence*>();
	TArray<float> times = TArray<float>();
	this->getAnimationDataDelegate.Execute(animations, times);

	// Generate a pattern based on the returned animation data.
	FLovensePattern pattern = FLovensePattern();
	if (!this->GeneratePattern(animations, times, pattern)) return;

	// If generation was successful, send pattern command.
	FLovenseManager::Get()->SendCommand_Pattern(nullptr, pattern,nullptr);
}

void ULovenseToyAnimationControlComponent::RealtimeEvaluationUpdate(float deltaTime) {
	// Don't update if the integration is not running or we are disabled or we don't have any skeletal meshes or bones.
	if (!FLovenseManager::Get()->IsLovenseRunning() || !this->bIsEnabled || this->skeletalMeshComponents.Num() <= 0 || this->boneNames.Num() <= 0) return;

	// Calculate velocity based on the average physics linear velocity of the bones.
	float velocitySum = 0.0f;
	for (int32 i = 0; i < this->skeletalMeshComponents.Num(); i++) {
		// Calculate sum of bone physics linear velocities for each mesh.
		float meshBoneVelocitySum = 0.0f;
		for (FName boneName : this->boneNames) {
			// Get bone location and extract physics linear velocity at that location.
			FVector boneLocation = this->skeletalMeshComponents[i]->GetSocketLocation(boneName);
			meshBoneVelocitySum += this->skeletalMeshComponents[i]->GetPhysicsLinearVelocityAtPoint(boneLocation, boneName).Size();
		}

		// Calculate average velocity of the bones and add to our velocity sum.
		velocitySum += meshBoneVelocitySum / static_cast<float>(this->boneNames.Num());
	}

	// Calculate average velocity between all meshes.
	float averageVelocity = velocitySum / static_cast<float>(this->skeletalMeshComponents.Num());
	// Cap averageVelocity or we might end up in a situation where it takes a long time for maxRecordedVelocity to decay.
	averageVelocity = FMath::Min(averageVelocity, this->maxAverageVelocity);
	// Slowly decay max recorded velocity unless new velocity is greater.
	// This is so we don't get stuck with a very high max velocity when the animation slows down significantly,
	// causing a much lower speed resolution as everything would be around 1 or 2 toy speed.
	if (averageVelocity <= this->maxRecordedVelocity) {
		this->maxRecordedVelocity = FMath::Max(this->maxRecordedVelocity - this->maxRecordedVelocityDecaySpeed * deltaTime, this->maxRecordedVelocityDecayMinimum);
	} else {
		this->maxRecordedVelocity = averageVelocity;
	}

	// The max toy speed gets reduced to maxToySpeedTargetDuringTransition during a transition (see StartTransition()).
	// To avoid abrupt toy speed changes, we interpolate the max toy speed.
	this->currentMaxToySpeed = FMath::FInterpTo(this->currentMaxToySpeed, this->maxToySpeedTarget, deltaTime, this->maxToySpeedInterpolationSpeed);
	// Convert the average velocity into a toy speed the Lovense commands will accept.
	float fToySpeed = FMath::GetMappedRangeValueClamped(
		FVector2D(0.0f, this->maxRecordedVelocity),
		FVector2D(1.0f, this->currentMaxToySpeed),
		averageVelocity
	);
	// It's better to apply the toy strength multiplier here, so we don't lose the decimal precision for the multiplication,
	// as the speed from here on forwards will be an integer.
	fToySpeed = FMath::Max(fToySpeed * FLovenseManager::Get()->GetToyStrengthMultiplier(), 1.0f);
	// Remap the scaled toy command speed to the toy pump speed range.
	float fPumpSpeed = FMath::GetMappedRangeValueClamped(
		FVector2D(1.0f, this->currentMaxToySpeed),
		FVector2D(1.0f, 3.0f),
		fToySpeed
	);
	int32 toySpeed = static_cast<int32>(FMath::RoundToInt(fToySpeed));
	int32 pumpSpeed = static_cast<int32>(FMath::RoundToInt(fPumpSpeed));

	// Send commands with the calculated speeds.
	//FLovenseManager::Get()->SendCommand_Vibrate(nullptr, toySpeed);
	//this->bRotateClockwise
	//	? FLovenseManager::Get()->SendCommand_RotateClockwise(nullptr, toySpeed)
	//	: FLovenseManager::Get()->SendCommand_RotateAntiClockwise(nullptr, toySpeed);
	//FLovenseManager::Get()->SendCommand_AirAuto(nullptr, pumpSpeed);
}

#if ENGINE_MAJOR_VERSION > 5 || ENGINE_MAJOR_VERSION == 5 && ENGINE_MINOR_VERSION >= 2
/**
 * @brief Extracts the world transform of a bone in an animation at a specific time.
 * \n Copied from MovieSceneSkeletalAnimationTrack.cpp.
 */
static FTransform GetWorldTransformForBone(UAnimSequence* animSequence, USkeletalMeshComponent* meshComponent, const FName& boneName, float seconds) {
	FTransform worldTransform = FTransform::Identity;
	//AnimSequence->GetBoneTransform doesn't seem to be as accurate as GetAnimationPose
	FMemMark mark(FMemStack::Get());
	FCompactPose outPose;
	TArray<FBoneIndexType> requiredBoneIndexArray;
#if ENGINE_MAJOR_VERSION > 5 || ENGINE_MAJOR_VERSION == 5 && ENGINE_MINOR_VERSION >= 3
	const UE::Anim::FCurveFilterSettings curveEvalOption;
#else
	const FCurveEvaluationOption curveEvalOption;
#endif
	requiredBoneIndexArray.AddUninitialized(animSequence->GetSkeleton()->GetReferenceSkeleton().GetNum());
	for (int32 boneIndex = 0; boneIndex < requiredBoneIndexArray.Num(); ++boneIndex) {
		requiredBoneIndexArray[boneIndex] = boneIndex;
	}

	FBoneContainer boneContainer(requiredBoneIndexArray, curveEvalOption, *animSequence->GetSkeleton());
	outPose.ResetToRefPose(boneContainer);
	FBlendedCurve outCurve;
	outCurve.InitFrom(boneContainer);
	UE::Anim::FStackAttributeContainer tempAttributes;
	FAnimationPoseData animationPoseData(outPose, outCurve, tempAttributes);
	FAnimExtractContext extractionContext(seconds, false);
	animSequence->GetAnimationPose(animationPoseData, extractionContext);
	int32 meshIndex = animationPoseData.GetPose().GetBoneContainer().GetPoseBoneIndexForBoneName(boneName);
	if (meshIndex != INDEX_NONE) {
		FCompactPoseBoneIndex compactPoseIndex = animationPoseData.GetPose().GetBoneContainer().MakeCompactPoseIndex(FMeshPoseBoneIndex(meshIndex));
		if (compactPoseIndex != INDEX_NONE) {
			FTransform boneTransform = FTransform::Identity;
			worldTransform *= boneTransform;
			do {
				boneTransform = animationPoseData.GetPose()[compactPoseIndex];
				worldTransform *= boneTransform;
				compactPoseIndex = animationPoseData.GetPose().GetBoneContainer().GetParentBoneIndex(compactPoseIndex);
			} while (compactPoseIndex.IsValid());
		}
	}

	return worldTransform;
}
#else
/**
 * @brief Extracts the world transform of a bone in an animation at a specific time.
 * \n Copied from MovieSceneSkeletalAnimationTrack.cpp.
 */
static FTransform GetWorldTransformForBone(UAnimSequence* animSequence, USkeletalMeshComponent* meshComponent, const FName& boneName, float seconds) {
	FName currentBoneName = boneName;
	FTransform worldTransform = FTransform::Identity;

	do {
		int32 boneIndex = meshComponent->GetBoneIndex(currentBoneName);
		FTransform boneTransform = FTransform::Identity;
		int32 trackIndex = 0;

#if ENGINE_MAJOR_VERSION >= 5
#if WITH_EDITOR
		const UAnimDataModel* model = animSequence->GetDataModel();
		if (const FBoneAnimationTrack* trackData = model->FindBoneTrackByName(boneName)) {
			trackIndex = trackData->BoneTreeIndex;
		}
#else
		trackIndex = animSequence->GetCompressedTrackToSkeletonMapTable().IndexOfByPredicate([boneIndex](const FTrackToSkeletonMap& mapping) {
			return mapping.BoneTreeIndex == boneIndex;
		});
#endif

		if (trackIndex == INDEX_NONE) {
			UE_LOG(LogLovenseIntegration, Warning, TEXT("Couldn't find raw animation track for bone \"%s\" with index %d"), *currentBoneName.ToString(), boneIndex);
			break;
		}
#else
		// Check if there is an animation track for the bone.
		for (; trackIndex < animSequence->GetRawTrackToSkeletonMapTable().Num(); ++trackIndex) {
			if (animSequence->GetRawTrackToSkeletonMapTable()[trackIndex].BoneTreeIndex == boneIndex) {
				break;
			}
		}

		if (trackIndex == animSequence->GetRawTrackToSkeletonMapTable().Num()) {
			UE_LOG(LogLovenseIntegration, Warning, TEXT("Couldn't find raw animation track for bone \"%s\" with index %d"), *currentBoneName.ToString(), boneIndex);
			break;
		}
#endif

		// Extract the bone transform from the animation at the specified time.
		// Raw data is only available in the editor, so bUseRawData should only be true if WITH_EDITOR.
		animSequence->GetBoneTransform(boneTransform, trackIndex, seconds, WITH_EDITOR);
		// Compose bone transform onto the world transform.
		worldTransform *= boneTransform;

		// Walk up the hierarchy, so we can calculate the world transform.
		currentBoneName = meshComponent->GetParentBone(currentBoneName);
	} while (currentBoneName != NAME_None);

	return worldTransform;
}
#endif

bool ULovenseToyAnimationControlComponent::GeneratePattern(const TArray<UAnimSequence*>& animations, const TArray<float>& times, FLovensePattern& outPattern) {
	outPattern = FLovensePattern();
	bool bIsInputDataValid =
		this->skeletalMeshComponents.Num() > 0
		&& this->skeletalMeshComponents[0]
		&& this->boneNames.Num() > 0
		&& animations.Num() > 0
		&& animations[0]
		// There needs to be an animation for each skeletal mesh component.
		&& this->skeletalMeshComponents.Num() == animations.Num()
		// Pattern command will fail if there are more than 50 pattern values. First time value will be swallowed, so we need max 51 values in total.
		&& times.Num() <= 51;

	if (!bIsInputDataValid) return false;

	// 100ms is the minimum interval the pattern command will accept. We use it for maximum pattern resolution.
	CONSTEXPR int32 patternInterval = 100;
	// The pattern command only accepts up to 50 pattern elements, therefore 5s is the maximum time we can have with an interval of 100ms.
	// We use the longest time possible to minimize the amount of commands being sent, reducing hiccups in command execution.
	CONSTEXPR int32 patternTime = 5;
	CONSTEXPR int32 numPatternElements = patternTime * 1000 / patternInterval;
	CONSTEXPR float intervalSeconds = static_cast<float>(patternInterval) / 1000.0f;

	outPattern.bVibrate = true;
	outPattern.bRotate = true;
	outPattern.bPump = true;
	outPattern.bThrust = true;
	outPattern.bSuck = true;
	outPattern.bFinger = true;
	outPattern.interval = patternInterval;
	outPattern.time = patternTime;
	outPattern.pattern.Reserve(numPatternElements);

	TArray<TArray<FTransform>> previousTransforms = TArray<TArray<FTransform>>();
	previousTransforms.Init(TArray<FTransform>(), this->skeletalMeshComponents.Num());
	for (TArray<FTransform>& previousTransform : previousTransforms) {
		previousTransform.Init(FTransform::Identity, this->boneNames.Num());
	}

	auto FillPattern = [this, &outPattern, &animations, &previousTransforms, intervalSeconds](float time) ->void {
		// The animations in the animations array all need to have the same sequence length, therefore we can just use the animation in the first index here to adjust the input time.
		// Take the animation play rate into account.
		time *= animations[0]->RateScale;
		// The time can be multiples of the animation sequence length, so we modulo here to create a seamless loop in the result pattern.
#if ENGINE_MAJOR_VERSION >= 5
		time = FMath::Fmod(time, animations[0]->GetPlayLength());
#else
		time = FMath::Fmod(time, animations[0]->SequenceLength);
#endif

		// Calculate velocity based on the delta bone locations extracted from the animations.
		float velocitySum = 0.0f;
		for (int32 meshIndex = 0; meshIndex < this->skeletalMeshComponents.Num(); ++meshIndex) {
			// Calculate sum of delta bone locations.
			float distanceSum = 0.0f;
			for (int32 boneIndex = 0; boneIndex < this->boneNames.Num(); ++boneIndex) {
				// Extract bone transform from the animation.
				FTransform currentTransform = GetWorldTransformForBone(animations[meshIndex], this->skeletalMeshComponents[meshIndex], boneNames[boneIndex], time);

				// If previous transform is identity, it means this is the first iteration and we cannot get a delta location.
				if (!previousTransforms[meshIndex][boneIndex].Equals(FTransform::Identity)) {
					distanceSum += FVector::Dist(currentTransform.GetTranslation(), previousTransforms[meshIndex][boneIndex].GetTranslation());
				}

				// Cache current bone transform for next iteration.
				previousTransforms[meshIndex][boneIndex] = currentTransform;
			}
			// Calculate average velocity and add to our velocity sum.
			velocitySum += distanceSum / static_cast<float>(this->boneNames.Num()) / intervalSeconds;
		}

		// If this is the first iteration, velocitySum will be 0.0f -> skip.
		// velocitySum could also be 0.0f if the same time value was passed in right after the other or if the animations have not moved during the delta time.
		// In that case there'll be a missing pattern element.
		if (velocitySum <= 0.0f) return;

		// Calculate average velocity between all meshes.
		// If there's only one skeletal mesh component, the velocity sum will generally be greater than if there were multiple,
		// as one mesh might not move very much if there are multiple, so we max to 2 here to account for that.
		float averageVelocity = velocitySum / FMath::Max(static_cast<float>(this->skeletalMeshComponents.Num()), 2.0f);
		// Convert the speed sum into a toy speed the Lovense commands will accept.
		float fToySpeed = FMath::GetMappedRangeValueClamped(this->patternToySpeedInputRange, FVector2D(1.0f, 20.0f), averageVelocity);
		// It's better to apply the toy strength multiplier here, so we don't lose the decimal precision for the multiplication,
		// as the speed from here on forwards will be an integer.
		int32 toySpeed = FMath::Max(static_cast<int32>(FMath::RoundToInt(fToySpeed * FLovenseManager::Get()->GetToyStrengthMultiplier())), 1);
		// Add our final toy speed to the pattern.
		outPattern.pattern.Add(toySpeed);
	};

	if (times.Num() > 0) {
		// If a custom times array was provided by calculateAnimationTimeDelegate, use it to generate the pattern.
		// Users will want to provide a custom times array to take e.g. dynamic animation play rate variations into account.
		for (float time : times) {
			FillPattern(time);
		}
	} else {
		// If no times array was provided, use a constant interval of 0.1 seconds instead.
		// The toy delay accounts for the time it takes the toy to start executing the command, which is heavily dependent on the bluetooth connection quality,
		// so we let the user adjust the delay.
		int32 toyDelay = FLovenseManager::Get()->GetToyDelay();

		// Iterate until we have 50 time values as we can only have 50 pattern values.
		for (int32 milliseconds = toyDelay; outPattern.pattern.Num() < 50; milliseconds += outPattern.interval) {
			float time = static_cast<float>(milliseconds) / 1000.0f;
			FillPattern(time);
		}
	}

	return true;
}
