From 4aa661a52330a38c9f9d0b93d984be4f10c8121c Mon Sep 17 00:00:00 2001
From: Hugo Pezziardi <hugo@pezziardi.net>
Date: Mon, 6 Mar 2023 10:46:19 +0100
Subject: [PATCH] :sparkles: (MotionKit): Add startStabilization method

---
 libs/MotionKit/include/MotionKit.hpp    |  4 +++
 libs/MotionKit/source/MotionKit.cpp     | 21 +++++++++++++
 libs/MotionKit/tests/MotionKit_test.cpp | 40 +++++++++++++++++++++++++
 spikes/lk_motion_kit/main.cpp           |  8 +++++
 4 files changed, 73 insertions(+)

diff --git a/libs/MotionKit/include/MotionKit.hpp b/libs/MotionKit/include/MotionKit.hpp
index 87a7a3aad3..95e115bbd7 100644
--- a/libs/MotionKit/include/MotionKit.hpp
+++ b/libs/MotionKit/include/MotionKit.hpp
@@ -7,6 +7,7 @@
 #include <interface/drivers/Motor.h>
 
 #include "RotationControl.hpp"
+#include "StabilizationControl.hpp"
 #include "interface/drivers/Timeout.h"
 #include "interface/libs/IMUKit.hpp"
 
@@ -23,11 +24,13 @@ class MotionKit
 
 	void startYawRotation(float degrees, Rotation direction,
 						  const std::function<void()> &on_rotation_ended_callback = {});
+	void startStabilization();
 
 	void stop();
 
   private:
 	void processAngleForRotation(const EulerAngles &angles, Rotation direction);
+	void processAngleForStabilization(const EulerAngles &angles);
 
 	void setMotorsSpeedAndDirection(float speed, Rotation direction);
 
@@ -37,6 +40,7 @@ class MotionKit
 	interface::Timeout &_timeout;
 
 	RotationControl _rotation_control;
+	StabilizationControl _stabilization_control;
 	std::function<void()> _on_rotation_ended_callback {};
 
 	bool _target_not_reached	   = false;
diff --git a/libs/MotionKit/source/MotionKit.cpp b/libs/MotionKit/source/MotionKit.cpp
index 9c105e3b02..a05547ea20 100644
--- a/libs/MotionKit/source/MotionKit.cpp
+++ b/libs/MotionKit/source/MotionKit.cpp
@@ -46,6 +46,20 @@ void MotionKit::startYawRotation(float degrees, Rotation direction,
 	_on_rotation_ended_callback = on_rotation_ended_callback;
 }
 
+void MotionKit::startStabilization()
+{
+	stop();
+
+	auto starting_angle = _imukit.getEulerAngles();
+	_stabilization_control.setTarget(starting_angle);
+
+	auto on_euler_angles_rdy_callback = [this](const EulerAngles &euler_angles) {
+		processAngleForStabilization(euler_angles);
+	};
+
+	_imukit.onEulerAnglesReady(on_euler_angles_rdy_callback);
+}
+
 // LCOV_EXCL_START - Dynamic behavior, involving motors and time.
 void MotionKit::processAngleForRotation(const EulerAngles &angles, Rotation direction)
 {
@@ -68,6 +82,13 @@ void MotionKit::processAngleForRotation(const EulerAngles &angles, Rotation dire
 	}
 }
 
+void MotionKit::processAngleForStabilization(const EulerAngles &angles)
+{
+	auto [speed, rotation] = _stabilization_control.processStabilizationAngle(angles);
+
+	setMotorsSpeedAndDirection(speed, rotation);
+}
+
 void MotionKit::setMotorsSpeedAndDirection(float speed, Rotation direction)
 {
 	if (speed == 0.F) {
diff --git a/libs/MotionKit/tests/MotionKit_test.cpp b/libs/MotionKit/tests/MotionKit_test.cpp
index 94a0476f6d..958b486bd6 100644
--- a/libs/MotionKit/tests/MotionKit_test.cpp
+++ b/libs/MotionKit/tests/MotionKit_test.cpp
@@ -134,3 +134,43 @@ TEST_F(MotionKitTest, rotateAndTimeOutOver)
 
 	on_timeout_callback();
 }
+
+TEST_F(MotionKitTest, startStabilization)
+{
+	const EulerAngles angles_quarter_left {0.0F, 0.0F, 90.F};
+
+	EXPECT_CALL(mock_timeout, stop).Times(1);
+	EXPECT_CALL(mock_motor_left, stop).Times(1);
+	EXPECT_CALL(mock_motor_right, stop).Times(1);
+
+	EXPECT_CALL(mock_imukit, getEulerAngles).Times(1);
+
+	EXPECT_CALL(mock_motor_left, spin).Times(AtLeast(1));
+	EXPECT_CALL(mock_motor_right, spin).Times(AtLeast(1));
+
+	motion.startStabilization();
+	mock_imukit.call_angles_ready_callback(angles_quarter_left);
+}
+
+TEST_F(MotionKitTest, startStabilizationAndStop)
+{
+	const EulerAngles angles_quarter_left {0.0F, 0.0F, 90.F};
+
+	EXPECT_CALL(mock_timeout, stop).Times(1);
+	EXPECT_CALL(mock_motor_left, stop).Times(1);
+	EXPECT_CALL(mock_motor_right, stop).Times(1);
+
+	EXPECT_CALL(mock_imukit, getEulerAngles).Times(1);
+
+	EXPECT_CALL(mock_motor_left, spin).Times(AtLeast(1));
+	EXPECT_CALL(mock_motor_right, spin).Times(AtLeast(1));
+
+	motion.startStabilization();
+	mock_imukit.call_angles_ready_callback(angles_quarter_left);
+
+	EXPECT_CALL(mock_timeout, stop).Times(1);
+	EXPECT_CALL(mock_motor_left, stop).Times(1);
+	EXPECT_CALL(mock_motor_right, stop).Times(1);
+
+	motion.stop();
+}
diff --git a/spikes/lk_motion_kit/main.cpp b/spikes/lk_motion_kit/main.cpp
index 193c20b751..5b019e169b 100644
--- a/spikes/lk_motion_kit/main.cpp
+++ b/spikes/lk_motion_kit/main.cpp
@@ -114,6 +114,14 @@ void onMagicCardAvailable(const MagicCard &card)
 		case (MagicCard::number_7.getId()):
 			motionkit.startYawRotation(7, Rotation::counterClockwise);
 			break;
+		case (MagicCard::number_8.getId()):
+			motionkit.startStabilization();
+			rtos::ThisThread::sleep_for(10s);
+			motionkit.stop();
+			break;
+		case (MagicCard::number_9.getId()):
+			motionkit.startStabilization();
+			break;
 		case (MagicCard::number_10.getId()):
 			motionkit.stop();
 			break;