From b0d87d997b2db1032757aeaeb9550d26d404048c Mon Sep 17 00:00:00 2001
From: Christoph Froehlich <christoph.froehlich@ait.ac.at>
Date: Wed, 11 Oct 2023 11:36:02 +0000
Subject: [PATCH] Move wraparound logic into separate function

---
 .../trajectory.hpp                            | 19 ++++++++--
 .../src/trajectory.cpp                        | 36 +++++++++++--------
 2 files changed, 38 insertions(+), 17 deletions(-)

diff --git a/joint_trajectory_controller/include/joint_trajectory_controller/trajectory.hpp b/joint_trajectory_controller/include/joint_trajectory_controller/trajectory.hpp
index 28cfbd6f44..a994b7fe5e 100644
--- a/joint_trajectory_controller/include/joint_trajectory_controller/trajectory.hpp
+++ b/joint_trajectory_controller/include/joint_trajectory_controller/trajectory.hpp
@@ -44,18 +44,31 @@ class Trajectory
     const trajectory_msgs::msg::JointTrajectoryPoint & current_point,
     std::shared_ptr<trajectory_msgs::msg::JointTrajectory> joint_trajectory);
 
-  /* Set the point before the trajectory message is replaced/appended
+  /**
+   *  Set the point before the trajectory message is replaced/appended
    * Example: if we receive a new trajectory message and it's first point is 0.5 seconds
    * from the current one, we call this function to log the current state, then
    * append/replace the current trajectory
-   * \param normalize_joint_error Vector of boolean where true value corresponds to a joint that
+   * \param joints_angle_wraparound Vector of boolean where true value corresponds to a joint that
    * wrap around (ie. is continuous).
    */
   JOINT_TRAJECTORY_CONTROLLER_PUBLIC
   void set_point_before_trajectory_msg(
     const rclcpp::Time & current_time,
     const trajectory_msgs::msg::JointTrajectoryPoint & current_point,
-    const std::vector<bool> & normalize_joint_error = std::vector<bool>());
+    const std::vector<bool> & joints_angle_wraparound = std::vector<bool>());
+
+  /**
+   * \param current_position The current position given from the controller, which will be adapted.
+   * \param next_position Next position from which to compute the wraparound offset, i.e.,
+   *      the first trajectory point
+   * \param joints_angle_wraparound Vector of boolean where true value corresponds to a joint that
+   * wrap around (ie. is continuous).
+   */
+  JOINT_TRAJECTORY_CONTROLLER_PUBLIC
+  void wraparound_joint(
+    std::vector<double> & current_position, const std::vector<double> next_position,
+    const std::vector<bool> & joints_angle_wraparound);
 
   JOINT_TRAJECTORY_CONTROLLER_PUBLIC
   void update(std::shared_ptr<trajectory_msgs::msg::JointTrajectory> joint_trajectory);
diff --git a/joint_trajectory_controller/src/trajectory.cpp b/joint_trajectory_controller/src/trajectory.cpp
index f7e8118091..143d547a55 100644
--- a/joint_trajectory_controller/src/trajectory.cpp
+++ b/joint_trajectory_controller/src/trajectory.cpp
@@ -46,28 +46,36 @@ Trajectory::Trajectory(
 void Trajectory::set_point_before_trajectory_msg(
   const rclcpp::Time & current_time,
   const trajectory_msgs::msg::JointTrajectoryPoint & current_point,
-  const std::vector<bool> & normalize_joint_error)
+  const std::vector<bool> & joints_angle_wraparound)
 {
   time_before_traj_msg_ = current_time;
   state_before_traj_msg_ = current_point;
 
-  if (!normalize_joint_error.empty())
+  // Compute offsets due to wrapping joints
+  wraparound_joint(
+    state_before_traj_msg_.positions, trajectory_msg_->points[0].positions,
+    joints_angle_wraparound);
+}
+
+void Trajectory::wraparound_joint(
+  std::vector<double> & current_position, const std::vector<double> next_position,
+  const std::vector<bool> & joints_angle_wraparound)
+{
+  double dist;
+  // joints_angle_wraparound is even empty, or has the same size as the number of joints
+  for (size_t i = 0; i < joints_angle_wraparound.size(); i++)
   {
-    for (size_t i = 0; i < current_point.positions.size(); i++)
+    if (joints_angle_wraparound[i])
     {
-      if (normalize_joint_error[i])
+      dist = angles::shortest_angular_distance(current_position[i], next_position[i]);
+
+      // Deal with singularity at M_PI shortest distance
+      if (std::abs(std::abs(dist) - M_PI) < 1e-9)
       {
-        double prev_position = state_before_traj_msg_.positions[i];
-        double next_position = trajectory_msg_->points[0].positions[i];
-        double dist = angles::shortest_angular_distance(prev_position, next_position);
-
-        // Deal with singularity at M_PI shortest distance
-        if (std::abs(std::abs(dist) - M_PI) < 1e-9)
-        {
-          dist = next_position > prev_position ? std::abs(dist) : -std::abs(dist);
-        }
-        state_before_traj_msg_.positions[i] = next_position - dist;
+        dist = next_position[i] > current_position[i] ? std::abs(dist) : -std::abs(dist);
       }
+
+      current_position[i] = next_position[i] - dist;
     }
   }
 }