#include "Robot/PilotComponent/RobotPilotSO100Component.h"

#include "Actors/MujocoVolumeActor.h"
#include "Kismet/KismetMathLibrary.h"
#include "Kismet/KismetSystemLibrary.h"
#include "Robot/RobotPawn.h"
#include "Serialization/ShaderKeyGenerator.h"

URobotPilotSO100Component::URobotPilotSO100Component()
{
	PrimaryComponentTick.bCanEverTick = true;
	PrimaryComponentTick.bStartWithTickEnabled = true;
}

void URobotPilotSO100Component::BeginPlay()
{
	Super::BeginPlay();
}

void URobotPilotSO100Component::TickComponent(float DeltaTime, enum ELevelTick TickType,
	FActorComponentTickFunction* ThisTickFunction)
{
	// Super::TickComponent(DeltaTime, TickType, ThisTickFunction);
	// const auto Joints = GetCurrentJointsFromPhysicsScene();
	// const auto Controls = GetCurrentControlsFromPhysicScene();
	// UE_LOG(LogTemp, Log, TEXT("Joint: %f - Control: %f - Delta: %f"), Joints.Jaw, Controls.Jaw, Controls.Jaw - Joints.Jaw);
}


void URobotPilotSO100Component::SetTarget(const FTransform& TargetTransformIn)
{
	// Set Base Values
	TargetTransform = TargetTransformIn;
	NextAnimationState();
}



void URobotPilotSO100Component::PrintCurrentActuators() const
{
	PrintActuators(GetCurrentControlsFromPhysicScene());
}

void URobotPilotSO100Component::PrintActuators(FSo100Actuators Actuators)
{
	const auto [Rotation, Pitch, Elbow, WristPitch, WristRoll, Jaw] = Actuators;
	UE_LOG(LogTemp, Display, TEXT("Actuators Rotation %f | Pitch %f | Elbow %f | WristPitch %f | WristRoll %f | Jaw %f"),
		Rotation,
		Pitch,
		Elbow,
		WristPitch,
		WristRoll,
		Jaw
		);
}

void URobotPilotSO100Component::DisableAnim()
{
	AnimationStartTime = 0;
}


FSo100Actuators URobotPilotSO100Component::GetCurrentControlsFromPhysicScene() const
{
	return FSo100Actuators
	{
		RobotOwner->PhysicsSceneProxy->GetActuatorValue(Actuator_Rotation),
		RobotOwner->PhysicsSceneProxy->GetActuatorValue(Actuator_Pitch),
		RobotOwner->PhysicsSceneProxy->GetActuatorValue(Actuator_Elbow),
		RobotOwner->PhysicsSceneProxy->GetActuatorValue(Actuator_WristPitch),
		RobotOwner->PhysicsSceneProxy->GetActuatorValue(Actuator_WristRoll),
		RobotOwner->PhysicsSceneProxy->GetActuatorValue(Actuator_Jaw),
	};
}

FSo100Actuators URobotPilotSO100Component::GetCurrentJointsFromPhysicsScene() const
{
	// Due to Mujoco import plugin not renaming joint differently from controller and Unreal refusing to have 2 variables with the same name
	// Joints get a "1" appended to their name -> refacto so it's "_Joint"
	return FSo100Actuators
	{
		RobotOwner->PhysicsSceneProxy->GetJointValue(FString(Actuator_Rotation).Append("1")),
		RobotOwner->PhysicsSceneProxy->GetJointValue(FString(Actuator_Pitch).Append("1")),
		RobotOwner->PhysicsSceneProxy->GetJointValue(FString(Actuator_Elbow).Append("1")),
		RobotOwner->PhysicsSceneProxy->GetJointValue(FString(Actuator_WristPitch).Append("1")),
		RobotOwner->PhysicsSceneProxy->GetJointValue(FString(Actuator_WristRoll).Append("1")),
		RobotOwner->PhysicsSceneProxy->GetJointValue(FString(Actuator_Jaw).Append("1"))
	};
}

float URobotPilotSO100Component::GetDeltaSumBetweenActuatorValues(const FSo100Actuators& A, const FSo100Actuators& B)
{
	float DeltaSum = 0;

	DeltaSum +=  FMath::Abs(A.Rotation) - FMath::Abs(B.Rotation);
	DeltaSum += FMath::Abs(A.Pitch) - FMath::Abs(B.Pitch);
	DeltaSum += FMath::Abs(A.Elbow) - FMath::Abs(B.Elbow);
	DeltaSum += FMath::Abs(A.WristPitch) - FMath::Abs(B.WristPitch);
	DeltaSum += FMath::Abs(A.WristRoll) - FMath::Abs(B.WristRoll);
	DeltaSum += FMath::Abs(A.Jaw) - FMath::Abs(B.Jaw);

	return DeltaSum;
}

double URobotPilotSO100Component::GetControlJointDeltaForActuator(FString ActuatorName) const
{
	auto const Control = RobotOwner->PhysicsSceneProxy->GetActuatorValue(ActuatorName);
	auto const Joint = RobotOwner->PhysicsSceneProxy->GetJointValue(ActuatorName.Append("1"));
	return Control -Joint;
}

FSo100Actuators URobotPilotSO100Component::LerpActuators(
	const FSo100Actuators& A,
	const FSo100Actuators& B,
	const float Alpha
	)
{
	return FSo100Actuators{
		FMath::Lerp( A.Rotation, B.Rotation, Alpha),
		FMath::Lerp( A.Pitch, B.Pitch, Alpha),
		FMath::Lerp( A.Elbow, B.Elbow, Alpha),
		FMath::Lerp( A.WristPitch, B.WristPitch, Alpha),
		FMath::Lerp( A.WristRoll, B.WristRoll, Alpha),
		FMath::Lerp( A.Jaw, B.Jaw, Alpha),
	};
}

void URobotPilotSO100Component::PostPhysicStepUpdate(const float SimulationTime)
{
	// Do anything that should be done after a mj_step update

	if (AnimationStartTime > 0)
	{
		// Could be moved to if statement, but it becomes confusing - so let's keep a variable
		const bool bHasFinishedAnimation = AnimateActuators(SimulationTime);

		if (bHasFinishedAnimation)
		{
			// AnimationStartTime = 0.; // Only for debug, can be left here but useless in normal operation mode
			NextAnimationState();
		}
	}
}

void URobotPilotSO100Component::NextAnimationState()
{
	const float CurrentPhysicsEngineSceneTime = RobotOwner->PhysicsSceneProxy->GetMujocoData().time;
	AnimationStartTime = CurrentPhysicsEngineSceneTime;
	AnimStartRobotActuators = GetCurrentJointsFromPhysicsScene();

	switch (CurrentAnimationState)
	{
		case 0:
			return RotateToTarget();
		case 1:
			return MoveToTarget();
		case 2:
			return CloseJaw();
		case 3:
			return MoveToDropZone();
		case 4:
			return OpenJaw();
		default:
			return BasePose();
	}
}

void URobotPilotSO100Component::BasePose()
{
	UE_LOG(LogTemp, Log, TEXT("Animate -> BasePose"));
	CurrentAnimationState = 0;
	AnimationDuration = 1.5f;
	AnimTargetRobotActuators = ActuatorsRestPosition;
}

void URobotPilotSO100Component::RotateToTarget()
{
	UE_LOG(LogTemp, Log, TEXT("Animate -> RotateToTarget"));

	if (TargetTransform.GetLocation() == FVector::ZeroVector)
	{
		AnimationStartTime = 0.f;
		CurrentAnimationState = 0;
		return;
	}

	// Compute Pivot Point World Location
	const auto WorldTransform = RobotOwner->RobotActor->GetActorTransform();
	const FVector PivotWorldLocation = WorldTransform.GetLocation() + WorldTransform.GetRotation().RotateVector(PivotOffset);

	// Find Yaw Rotation in degree - cache it for distance to target computation
	RotationToTarget = UKismetMathLibrary::FindLookAtRotation(
		PivotWorldLocation,
		TargetTransform.GetLocation()
		);
	
	// reduce/increase Yaw to not have the fixed jaw colliding with the shape - TODO use middle of the jaw instead of the wall of the jaw
	const auto Dot = FVector::DotProduct(RotationToTarget.Quaternion().GetForwardVector(), WorldTransform.GetRotation().GetForwardVector());
	const auto Mod = .1 * (Dot > 0 ? 1 : -1);

	// Convert to radians
	const auto ActuatorRotation = RotationToTarget.Yaw * (1+Mod) / 180.0f * -PI; // Looks like we are not in the same referential hence the -PI instead of PI ! 

	// Start the animation
	AnimTargetRobotActuators = AnimStartRobotActuators;
	AnimTargetRobotActuators.Rotation = ActuatorRotation;
	CurrentAnimationState = 1;
	AnimationDuration = .7f;
}

void URobotPilotSO100Component::MoveToTarget()
{
	UE_LOG(LogTemp, Log, TEXT("Animate -> MoveToTarget"));

	// Get Pivot World
	const auto WorldTransform = RobotOwner->RobotActor->GetActorTransform();
	const FVector PivotWorldLocation = WorldTransform.GetLocation() + WorldTransform.GetRotation().RotateVector(PivotOffset);
	
	// Rotate Jaw offset towards target
	// Get pure 2d rotation
	RotationToTarget.Pitch = 0;
	RotationToTarget.Roll = 0;
	const auto JawOffsetToPivot = JawOffset - PivotOffset;
	const auto JawPositionWorld = RotationToTarget.RotateVector(JawOffsetToPivot) + PivotWorldLocation;
	const auto Distance = FVector::Distance(JawPositionWorld, TargetTransform.GetLocation());

	const auto AlphaExtend = FMath::Clamp(Distance / MaxReach, 0., 1.);
	
	// Set the target actuators values 
	AnimTargetRobotActuators = GetCurrentJointsFromPhysicsScene();
	AnimTargetRobotActuators = LerpActuators(ActuatorsRestPosition, ActuatorsMaxExtendPosition, AlphaExtend);

	
	AnimTargetRobotActuators.Rotation = AnimStartRobotActuators.Rotation;
	PrintActuators(AnimTargetRobotActuators);
	
	// Start the animation 	
	CurrentAnimationState = 2;
	AnimationDuration = 2.f;
	
	DrawDebugLine(
		this->GetWorld(),
		JawPositionWorld,
		JawPositionWorld + (TargetTransform.GetLocation() - JawPositionWorld).GetSafeNormal() * Distance,
		FColor::Green,
		true
	);
}

void URobotPilotSO100Component::CloseJaw()
{
	UE_LOG(LogTemp, Log, TEXT("Animate -> CloseJaw"));

	// Here we need overcurrent detection - not here actually
	AnimTargetRobotActuators.Jaw = ClosedJaw;

	// Reset TargetTransform
	TargetTransform = FTransform();
	
	// Start the animation
	bDetectOverCurrent = true;
	CurrentAnimationState = 3;
	AnimationDuration = 2.f;
}

void URobotPilotSO100Component::MoveToDropZone()
{
	UE_LOG(LogTemp, Log, TEXT("Animate -> MoveToDropZone"));

	AnimTargetRobotActuators = ActuatorsDropZone;
	AnimTargetRobotActuators.Rotation = ActuatorsDropZone.Rotation * (FMath::RandBool() ? 1. : -1.); 
	AnimTargetRobotActuators.Jaw = GetCurrentJointsFromPhysicsScene().Jaw;
	CurrentAnimationState = 4;
	AnimationDuration = 3.f;
}

void URobotPilotSO100Component::OpenJaw()
{
	UE_LOG(LogTemp, Log, TEXT("Animate -> OpenJaw"));

	AnimTargetRobotActuators.Jaw = OpenedJaw;
	CurrentAnimationState = 5;
	AnimationDuration = 0.6f;
}

bool URobotPilotSO100Component::AnimateActuators(const float SimulationTime)
{
	const double AnimAlpha = FMath::Clamp((SimulationTime - AnimationStartTime) / AnimationDuration, 0., 1.);

	// Need to wait for the joints to be in the right position before switching to next animation
	const auto DeltaSum = GetDeltaSumBetweenActuatorValues(AnimTargetRobotActuators, GetCurrentJointsFromPhysicsScene());

	// Alternative Animation completion event - checking for over-current
	if (bDetectOverCurrent && GetControlJointDeltaForActuator(Actuator_Jaw) > OverCurrentThreshold)
	{
		bDetectOverCurrent = false;
		return true;
	}

	
	// UE_LOG(LogTemp, Log, TEXT("AnimationAlpha: %f - Delta: %f"), AnimAlpha, DeltaSum);

	// Stop the animation if we reached the target
	if (AnimAlpha >= 1. && DeltaSum <= .001) return true;

	// Rotation
	RobotOwner->PhysicsSceneProxy->SetActuatorValue(Actuator_Rotation, FMath::Lerp(
		AnimStartRobotActuators.Rotation,
		AnimTargetRobotActuators.Rotation,
		AnimAlpha
	));

	// Pitch
	RobotOwner->PhysicsSceneProxy->SetActuatorValue(Actuator_Pitch, FMath::Lerp(
		AnimStartRobotActuators.Pitch,
		AnimTargetRobotActuators.Pitch,
		AnimAlpha
	));

	// Elbow
	RobotOwner->PhysicsSceneProxy->SetActuatorValue(Actuator_Elbow, FMath::Lerp(
		AnimStartRobotActuators.Elbow,
		AnimTargetRobotActuators.Elbow,
		AnimAlpha
	));
		
	// WristPitch
	RobotOwner->PhysicsSceneProxy->SetActuatorValue(Actuator_WristPitch, FMath::Lerp(
		AnimStartRobotActuators.WristPitch,
		AnimTargetRobotActuators.WristPitch,
		AnimAlpha
	));

	// WristRoll
	RobotOwner->PhysicsSceneProxy->SetActuatorValue(Actuator_WristRoll, FMath::Lerp(
		AnimStartRobotActuators.WristRoll,
		AnimTargetRobotActuators.WristRoll,
		AnimAlpha
	));

	// Jaw
	RobotOwner->PhysicsSceneProxy->SetActuatorValue(Actuator_Jaw, FMath::Lerp(
		AnimStartRobotActuators.Jaw,
		AnimTargetRobotActuators.Jaw,
		AnimAlpha
	));

	return false;
}