From b53f8496d1d9be9e4e5ebda130111388a7d241bc Mon Sep 17 00:00:00 2001
From: Claudio Chies <61051109+Claudio-Chies@users.noreply.github.com>
Date: Tue, 19 Nov 2024 10:45:23 +0100
Subject: [PATCH] Included turning radius calculations for vertical changes and
 removed the requirement for the same altitude in the 2D turning radius logic.

---
 src/lib/motion_planning/PositionSmoothing.cpp | 25 +++++++++++++++++--
 .../motion_planning/TrajectoryConstraints.hpp |  6 +----
 2 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/src/lib/motion_planning/PositionSmoothing.cpp b/src/lib/motion_planning/PositionSmoothing.cpp
index 28c70c690d60..ea93708bedc9 100644
--- a/src/lib/motion_planning/PositionSmoothing.cpp
+++ b/src/lib/motion_planning/PositionSmoothing.cpp
@@ -106,15 +106,36 @@ float PositionSmoothing::_getMaxXYSpeed(const Vector3f(&waypoints)[3]) const
 
 float PositionSmoothing::_getMaxZSpeed(const Vector3f(&waypoints)[3]) const
 {
-
+	const auto &start_position = waypoints[0];
 	const auto &target = waypoints[1];
+	const auto &next_target = waypoints[2];
 
 	Vector3f pos_traj(_trajectory[0].getCurrentPosition(),
 			  _trajectory[1].getCurrentPosition(),
 			  _trajectory[2].getCurrentPosition());
 
 	const float distance_start_target = fabs(target(2) - pos_traj(2));
-	const float arrival_z_speed = 0.f;
+	float arrival_z_speed = 0.0f;
+
+	const float distance_target_next = fabsf(target(2) - next_target(2));
+
+	const bool target_next_different = distance_target_next  > 0.001f;
+
+	Vector2f target_xy_z = {target.xy().norm(), target(2)};
+	Vector2f next_target_xy_z = {next_target.xy().norm(), next_target(2)};
+	Vector2f start_position_xy_z = {start_position.xy().norm(), start_position(2)};
+
+	if (target_next_different
+	   ) {
+		const float alpha = acosf(Vector2f((target_xy_z - start_position_xy_z)).unit_or_zero().dot(
+						  Vector2f((target_xy_z - next_target_xy_z)).unit_or_zero()));
+
+		const float safe_alpha = math::constrain(alpha, 0.f, M_PI_F - FLT_EPSILON);
+		float accel_tmp = _trajectory[2].getMaxAccel();
+		float max_speed_in_turn = math::trajectory::computeMaxSpeedInWaypoint(safe_alpha, accel_tmp,
+					  _vertical_acceptance_radius);
+		arrival_z_speed = math::min(max_speed_in_turn, _trajectory[2].getMaxVel());
+	}
 
 	float max_speed = math::min(_trajectory[2].getMaxVel(), math::trajectory::computeMaxSpeedFromDistance(
 					    _trajectory[2].getMaxJerk(), _trajectory[2].getMaxAccel(),
diff --git a/src/lib/motion_planning/TrajectoryConstraints.hpp b/src/lib/motion_planning/TrajectoryConstraints.hpp
index 6d52e1861e6e..ffe68840cb12 100644
--- a/src/lib/motion_planning/TrajectoryConstraints.hpp
+++ b/src/lib/motion_planning/TrajectoryConstraints.hpp
@@ -74,15 +74,11 @@ inline float computeStartXYSpeedFromWaypoints(const Vector3f &start_position, co
 
 	const bool target_next_different = distance_target_next  > 0.001f;
 	const bool waypoint_overlap = distance_target_next < config.xy_accept_rad;
-	const bool has_reached_altitude = fabsf(target(2) - start_position(2)) < config.z_accept_rad;
-	const bool altitude_stays_same = fabsf(next_target(2) - target(2)) < config.z_accept_rad;
 
 	float speed_at_target = 0.0f;
 
 	if (target_next_different &&
-	    !waypoint_overlap &&
-	    has_reached_altitude &&
-	    altitude_stays_same
+	    !waypoint_overlap
 	   ) {
 		const float alpha = acosf(Vector2f((target - start_position).xy()).unit_or_zero().dot(
 						  Vector2f((target - next_target).xy()).unit_or_zero()));