LuckyWorld/Source/LuckyWorldV2/Public/Robot/PilotComponent/RobotPilotSO100Component.h

#pragma once
#include "CoreMinimal.h"
#include "Robot/PilotComponent/RobotPilotComponent.h"
#include "RobotPilotSO100Component.generated.h"

struct FRemoteControlPayload;

USTRUCT(BlueprintType)
struct FSo100Actuators
{
	GENERATED_BODY()
	UPROPERTY(EditAnywhere, BlueprintReadWrite)
	double Rotation = 0.;
	
	UPROPERTY(EditAnywhere, BlueprintReadWrite)
	double Pitch = 0.;

	UPROPERTY(EditAnywhere, BlueprintReadWrite)
	double Elbow = 0.;

	UPROPERTY(EditAnywhere, BlueprintReadWrite)
	double WristPitch = 0.;

	UPROPERTY(EditAnywhere, BlueprintReadWrite)
	double WristRoll = 0.;
	
	UPROPERTY(EditAnywhere, BlueprintReadWrite)
	double Jaw = 0.;
};

UCLASS(Blueprintable)
class LUCKYWORLDV2_API URobotPilotSO100Component : public URobotPilotComponent
{
	GENERATED_BODY()

public:
	URobotPilotSO100Component();
	
	virtual void BeginPlay() override;
	virtual void TickComponent(float DeltaTime, enum ELevelTick TickType, FActorComponentTickFunction* ThisTickFunction) override;
	virtual FTransform GetReachableTransform() override;
	virtual bool GetIsReadyForTraining() override;
	virtual bool GetIsInRestState() override;

	virtual void SetRobotTarget(const FTransform& TargetTransformIn) override;
	virtual void SetRobotCurrentRewardZone(const FTransform& RewardTransformIn) override;
	virtual void ReceiveRemoteCommand(const FRemoteControlPayload& RemoteRobotPayload) override;

	// ------------------
	// ------ DATA ------
	// ------------------
	virtual FTrainingEpisodeData GetTrainingEpisodeData() override;
	virtual FJsonObject GetBufferedControlsData() override;
	virtual FJsonObject GetBufferedJointsData() override;
	TArray<FSo100Actuators> ControlsDataBuffer;
	TArray<FSo100Actuators> JointsDataBuffer;
	static FJsonObject GetStats(const TArray<FSo100Actuators>& ActuatorStates);
	
private:
	FTransform TargetTransform;
	bool bDropZoneIsRight = false;

	//---------------------
	//------- DEBUG -------
	//---------------------
public:
	UFUNCTION(BlueprintCallable)
	void PrintCurrentActuators() const;
	static void PrintActuators(FSo100Actuators Actuators);

	UFUNCTION(BlueprintCallable)
	void DisableAnim();

private:
	// SO100 Controls by name
	FString Actuator_Rotation = FString("Rotation");
	FString Actuator_Pitch = FString("Pitch");
	FString Actuator_Elbow = FString("Elbow");
	FString Actuator_WristPitch = FString("Wrist_Pitch");
	FString Actuator_WristRoll = FString("Wrist_Roll");
	FString Actuator_Jaw = FString("Jaw");

	// SO100 Static Variables
	// TODO Those values must be more precise, and probably that the way we compute the rotation is not good enough
	// TODO Let's discuss how to improve that algorithm
	FVector PivotOffset = FVector{-0.000030, 4.520021, 1.650041}; // From the Robot Location
	FVector JawOffset =  FVector{23, 2, 9}; // From the Pivot
	float MaxRange = 20.757929; // fixed_jaw_pad_3 ForwardVectorLength Delta between Rest and MaxExtend
	float MaxYaw = 80.f;
	
	// Actuators Joints and Controls are expressed in doubles
	double ClosedJaw = -0.01;
	double OpenedJaw = -2.0;
	int32 JawState = 0; // 0 - Opened || 1 - Grabbing
	
	/**
	 * Query the physic proxy on the RobotOwner to get the SO100 actuators values
	 * @return 
	 */
public:
	FSo100Actuators GetCurrentControlsFromPhysicScene() const;
private:
	FSo100Actuators GetCurrentJointsFromPhysicsScene() const;
	double GetControlJointDeltaForActuator(FString ActuatorName) const;
	static float GetDeltaSumBetweenActuatorValues(const FSo100Actuators& A, const FSo100Actuators& B);
	static bool AreActuatorsAlmostEqual(const FSo100Actuators& A, const FSo100Actuators& B);
	static FSo100Actuators LerpActuators(const FSo100Actuators& A, const FSo100Actuators& B, const float Alpha); 
	
	// Called after every physic step - this is a substep tick
	virtual void PostPhysicStepUpdate(const float SimulationTime) override;
	bool AnimateActuators(float SimulationTime); // Bound to the PhysicProxy post-update delegate
	FSo100Actuators CurrentRobotActuators; // This will be updated by the post-physic delegate 
	FSo100Actuators AnimStartRobotActuators;
	FSo100Actuators AnimTargetRobotActuators;
	FRotator RotationToTarget = FRotator::ZeroRotator;

	// ------------------------
	// ----- OVER-CURRENT -----
	// ------------------------
	bool bDetectOverCurrent = false;
	const float OverCurrentThreshold = 0.1;
	bool IsJawOverCurrent() const;

	// Quick and dirty sequence of moves
	// -1 -> Start Game, extended
	// 0 -> Retract - base pose
	// 1 -> Rotate towards target
	// 2 -> Move to target
	// 3 -> Close the jaw
	// 4 -> Go to drop zone
	// 5 -> open jaw
	int32 CurrentAnimationState = -1;
	void NextAnimationState();
	bool bBreakAfterAnimation = false;
	
	void RestPose();
	void RotateToTarget();
	void MoveToTarget();
	void CloseJaw();
	void MoveToDropZone();
	void OpenJaw();
	
	
	
	// Here let's write the code trying to match with Constantin class
	// After both classes have been designed around specific needs, see what can be migrated in the parent class and update both children

	FSo100Actuators ActuatorsRestPosition {
		0.,
		-1.54,
		3.105,
		-1.5,
		1.47,
		-1.39
	};

	FSo100Actuators ActuatorsMaxExtendPosition {
		0.0000001,
		0.081027,
		-0.01707,
		-0.075,
		1.469020,
		-1.389073
	};
	
	FSo100Actuators ActuatorsDropZone {
		PI / 2 + 0.25,
		-2.17,
		0.805,
		1.345,
		1.61,
		0
	};
};