// Copyright 2022, Algoryx Simulation AB.

#include "YumiArmController.h"

// AGX Dynamics for Unreal includes.
#include "Constraints/AGX_Constraint1DofComponent.h"
#include "AGX_Simulation.h"

UYumiArmController::UYumiArmController()
{
	PrimaryComponentTick.bCanEverTick = false;
}

void UYumiArmController::BeginPlay()
{
	Super::BeginPlay();
	UAGX_Simulation* Simulation = UAGX_Simulation::GetFrom(this);
	if (Simulation == nullptr)
	{
		UE_LOG(
			LogTemp, Error,
			TEXT("YumiArmController could not get the AGX Dynamics for Unreal Simulation. "
				 "The controller will not be able to step."));
		return;
	}

	PreviousSimTime = Simulation->GetTimeStamp();
	Simulation->PostStepForward.AddDynamic(this, &UYumiArmController::TickController);
}

void UYumiArmController::Init(
	TArray<UAGX_Constraint1DofComponent*> InJointsArm,
	TArray<UAGX_Constraint1DofComponent*> InJointsGripper)
{
	JointsArm = InJointsArm;
	JointsGripper = InJointsGripper;

	for (UAGX_Constraint1DofComponent* Constraint : JointsArm)
	{
		if (Constraint)
		{
			Constraint->LockController.SetEnable(true);
		}
	}

	for (UAGX_Constraint1DofComponent* Constraint : JointsGripper)
	{
		if (Constraint)
		{
			Constraint->LockController.SetEnable(true);
		}
	}
}

void UYumiArmController::StartRunSequence(const FRobotSequence& InSequence)
{
	CurrentSequence.Sequence = InSequence;
	CurrentSequence.TargetElapsedTime = 0.0;
	CurrentSequence.TargetIndex = 0;
	CurrentSequence.TargetStartArmJoints = GetCurrentJointPos(JointsArm);
	CurrentSequence.TargetStartGripperJoints = GetCurrentJointPos(JointsGripper);
	CurrentSequence.IsExecuting = true;
}

void UYumiArmController::TickController(double SimTime)
{
	if (CurrentSequence.IsExecuting)
	{
		const double DeltaTime = SimTime - PreviousSimTime;
		RunCurrentSequence(DeltaTime);
	}

	PreviousSimTime = SimTime;
}

void UYumiArmController::RunCurrentSequence(double DeltaTime)
{
	CurrentSequence.TargetElapsedTime += DeltaTime;
	const TArray<FRobotTarget>& Targets = CurrentSequence.Sequence.GetTargets();
	if (CurrentSequence.TargetIndex >= Targets.Num())
	{
		SetSequenceFinished();
		return;
	}

	const FRobotTarget Target = Targets[CurrentSequence.TargetIndex];
	if (Target.ArmJointPos.Num() != JointsArm.Num())
	{
		UE_LOG(
			LogTemp, Error,
			TEXT("YumiArmController cannot execute sequence because the number of arm joints found "
				 "in the current Target (%d) does not match the number of arm joints of the robot "
				 "(%d)."),
			Target.ArmJointPos.Num(), JointsArm.Num());
		SetSequenceFinished();
		return;
	}

	if (Target.GripperJointPos.Num() != JointsGripper.Num())
	{
		UE_LOG(
			LogTemp, Error,
			TEXT("YumiArmController cannot execute sequence because the number of gripper joints "
				 "found in the current Target (%d) does not match the number of gripper joints of "
				 "the robot (%d)."),
			Target.GripperJointPos.Num(), JointsGripper.Num());
		SetSequenceFinished();
		return;
	}

	// Target completion ratio.
	const double Tcr = CurrentSequence.TargetElapsedTime / Target.Duration;

	auto UpdateJoints = [](TArray<UAGX_Constraint1DofComponent*>& Joints,
						   const TArray<double>& StartJoints, const TArray<double>& TargetJoints,
						   double Tcr) -> bool
	{
		for (int i = 0; i < Joints.Num(); i++)
		{
			UAGX_Constraint1DofComponent* Constraint = Joints[i];
			if (Constraint == nullptr)
			{
				UE_LOG(
					LogTemp, Error,
					TEXT("YumiArmController cannot execute sequence because "
						 "one of the joints was nullptr."));
				return false;
			}

			if (Tcr > 1.0)
			{
				// We have reached the target. Set the joint position exactly.
				Constraint->LockController.SetPosition(TargetJoints[i]);
			}
			else
			{
				// We have not yet reached the target, interpolate.
				const double JointTargetDelta = TargetJoints[i] - StartJoints[i];
				const double JointPosThisTick = StartJoints[i] + JointTargetDelta * Tcr;
				Constraint->LockController.SetPosition(JointPosThisTick);
			}
		}
		return true;
	};

	// Update arm joints.
	if (!UpdateJoints(JointsArm, CurrentSequence.TargetStartArmJoints, Target.ArmJointPos, Tcr))
	{
		// Logging done in UpdateJoints.
		SetSequenceFinished();
		return;
	}

	// Update gripper joints.
	if (!UpdateJoints(
			JointsGripper, CurrentSequence.TargetStartGripperJoints, Target.GripperJointPos, Tcr))
	{
		// Logging done in UpdateJoints.
		SetSequenceFinished();
		return;
	}

	if (Tcr > 1.0)
	{
		// Target reached.
		SetTargetReached();
	}
}

void UYumiArmController::SetSequenceFinished()
{
	CurrentSequence = SequenceExecutionData();
}

void UYumiArmController::SetTargetReached()
{
	const int Index = CurrentSequence.TargetIndex;
	if (Index < CurrentSequence.Sequence.GetTargets().Num())
	{
		CurrentSequence.TargetStartArmJoints =
			CurrentSequence.Sequence.GetTargets()[Index].ArmJointPos;
		CurrentSequence.TargetStartGripperJoints =
			CurrentSequence.Sequence.GetTargets()[Index].GripperJointPos;
	}

	CurrentSequence.TargetElapsedTime = 0.0;
	CurrentSequence.TargetIndex++;
}

TArray<double> UYumiArmController::GetCurrentJointPos(
	const TArray<UAGX_Constraint1DofComponent*>& Joints)
{
	TArray<double> JointPos;

	for (UAGX_Constraint1DofComponent* Constraint : Joints)
	{
		if (Constraint == nullptr)
		{
			continue;
		}

		JointPos.Add(Constraint->GetAngle());
	}

	return JointPos;
}
