#include "Robot/PilotComponent/RobotPilotSO100Component.h"

#include "LuckyDataTransferSubsystem.h"
#include "Actors/MujocoVolumeActor.h"
#include "Kismet/KismetMathLibrary.h"
#include "Robot/RobotPawn.h"

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

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

	// Start the RestPost animation 0.1s after loading
	AnimationStartTime = .1f; 
	bBreakAfterAnimation = true;
	RestPose();
}

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);
}

FTransform URobotPilotSO100Component::GetReachableTransform()
{
	const auto RobotTransform = RobotOwner->RobotActor->GetActorTransform();

	// Robot actor is built Y axis forward, like everything in Unreal - Rotate to get Arm facing X world axis aka forward vector
	const auto ArmWorldRotation = FRotator{0,90,0}.Quaternion() * RobotTransform.GetRotation();

	// Find Arm Pivot Location
	const auto ArmPivotLocation = RobotTransform.GetLocation() + RobotTransform.GetRotation().RotateVector(PivotOffset) + FVector{0,0,5};

	// Compute a random Yaw
	const float RandomYaw = MaxYaw * FMath::RandRange(0., 1.) * (FMath::RandBool() ? 1 : -1);
	const FQuat RandomRotation = FRotator{0,RandomYaw,0}.Quaternion() * ArmWorldRotation;



	// Compute Random Range within reach of the arm and add this to pivot location
	// Add a bit more than the Jaw Offset - TODO Offsets must be better computed
	const float RandomRange = JawOffset.X + MaxRange * FMath::RandRange(0.1f, 1.f);
	const FVector RandomLocation = ArmPivotLocation + RandomRotation.GetForwardVector() * RandomRange;

	// Find Look at rotation from location to Pivot
	auto RewardAxis = RandomLocation-ArmPivotLocation;
	RewardAxis.Z = 0; // Nullify Z to keep a 2D vector -> ensure the geometry roll/pitch are 0
	const FRotator TowardPivotRotation = UKismetMathLibrary::MakeRotFromXZ(RewardAxis, FVector::UpVector);

	// Return the Object Transform
	return FTransform(TowardPivotRotation, RandomLocation);
}

bool URobotPilotSO100Component::GetIsReadyForTraining()
{
	const auto CurrentJoints = GetCurrentJointsFromPhysicsScene();
	return AreActuatorsAlmostEqual(CurrentJoints, ActuatorsRestPosition); 
}

bool URobotPilotSO100Component::GetIsInRestState()
{
	return CurrentAnimationState == 0 && GetIsReadyForTraining();
}


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

void URobotPilotSO100Component::SetRobotCurrentRewardZone(const FTransform& RewardTransformIn)
{
	const auto DirectionToTarget = (RewardTransformIn.GetLocation() - RobotOwner->RobotActor->GetActorLocation()).GetSafeNormal();

	// TODO This is wrong way to do because GetRightVector is the arm forward vector due to robot actor being rotated -90 yaw at spawn
	bDropZoneIsRight = FVector::DotProduct(RobotOwner->RobotActor->GetActorRotation().Quaternion().GetRightVector(), DirectionToTarget) > 0.f;
}

void URobotPilotSO100Component::ReceiveRemoteCommand(const FRemoteControlPayload& RemoteRobotPayload)
{
	for (const auto& [ActuatorName, ActuatorValue] : RemoteRobotPayload.Commands)
	{
		// Will print an error message if it doesn't exists
		RobotOwner->PhysicsSceneProxy->SetActuatorValue(ActuatorName, ActuatorValue);
	}
}

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;
}

bool URobotPilotSO100Component::AreActuatorsAlmostEqual(const FSo100Actuators& A, const FSo100Actuators& B)
{
	if (FMath::Abs(A.Rotation) - FMath::Abs(B.Rotation) > 0.001) return false;
	if (FMath::Abs(A.Pitch) - FMath::Abs(B.Pitch) > 0.001) return false;
	if (FMath::Abs(A.Elbow) - FMath::Abs(B.Elbow) > 0.001) return false;
	if (FMath::Abs(A.WristPitch) - FMath::Abs(B.WristPitch) > 0.001) return false;
	if (FMath::Abs(A.WristRoll) - FMath::Abs(B.WristRoll) > 0.001) return false;
	if (FMath::Abs(A.Jaw) - FMath::Abs(B.Jaw) > 0.001) return false;

	return true;
}

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)
		{
			if (bBreakAfterAnimation)
			{
				bBreakAfterAnimation = false;
				AnimationStartTime = 0;
				return;
			}
			
			// AnimationStartTime = 0.; // Only for debug, can be left here but useless in normal operation mode
			return NextAnimationState();
		}
	}
}

bool URobotPilotSO100Component::IsJawOverCurrent() const
{
	const auto DeltaActJntJaw = GetControlJointDeltaForActuator(Actuator_Jaw); 
	return DeltaActJntJaw > OverCurrentThreshold;
}

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 RestPose();
	}
}

void URobotPilotSO100Component::RestPose()
{
	UE_LOG(LogTemp, Log, TEXT("Animate -> BasePose"));
	CurrentAnimationState = 0;
	AnimationDuration = .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());

	// TODO Better computation of Jaw Center would avoid that ugliness
	// TODO This is not working for values Dot<0.15 - we need a better solution
	const auto ModBaseAlpha = FMath::Abs(Dot) / 0.3;
	const auto ModBase = 0.1 * FMath::Lerp(5.f, 1.f, FMath::Min(ModBaseAlpha, 1)); // TODO Dot Product below 0.5 needs more yaw modification 
	
	const auto Mod = ModBase * (Dot > 0 ? 1 : -1); // TODO Hardcoded value - compute better Jaw Offset and robot geometry awareness 
	UE_LOG(LogTemp, Log, TEXT("Dot : %f | ModBaseAlpha: %f | ModBase: %f | Mod: %f"), Dot, ModBaseAlpha, ModBase, Mod);

	// 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 = .33f;
}

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);

	// TODO Better computations
	// 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());

	// TODO Compute correct ranges and avoid to add a hardcoded value to adjust the arm extension
	const auto AlphaExtend = FMath::Clamp(Distance / (MaxRange + 4), 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 = .66f;
}

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 = .5f;
}

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

	AnimTargetRobotActuators = ActuatorsDropZone;
	AnimTargetRobotActuators.Rotation = ActuatorsDropZone.Rotation * (bDropZoneIsRight ? 1. : -1.); 

	// Here the Jaw should keep being target to closed
	AnimTargetRobotActuators.Jaw = ClosedJaw;
	
	CurrentAnimationState = 4;
	AnimationDuration = 1.5f;
}

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

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

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 bool bIsArrivedAtActuatorsConfig = AreActuatorsAlmostEqual(AnimTargetRobotActuators, GetCurrentJointsFromPhysicsScene());

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

	// UE_LOG(LogTemp, Log, TEXT("Animate -> AnimateActuators %f - %f"), AnimAlpha, GetControlJointDeltaForActuator(Actuator_Jaw));
	
	// Stop the animation if we reached the target
	if (AnimAlpha >= 1. && bIsArrivedAtActuatorsConfig) 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
	// If the target Jaw is closed position, aka we are grabbing, but we are over-current, then we don't hold more to not squeeze the object
	// But if the over-current stops, aka the object rotated a bit, then we hold tighter
	if (AnimTargetRobotActuators.Jaw != ClosedJaw || !IsJawOverCurrent())
	{
		RobotOwner->PhysicsSceneProxy->SetActuatorValue(Actuator_Jaw, FMath::Lerp(
			AnimStartRobotActuators.Jaw,
			AnimTargetRobotActuators.Jaw,
			AnimAlpha
		));
	}

	return false;
}