From 0742736fc46f03bfc922f2304b9691142b516e76 Mon Sep 17 00:00:00 2001
From: Jens Vanhooydonck <jensvanhooydonck@gmail.com>
Date: Mon, 27 Nov 2023 14:17:17 +0100
Subject: [PATCH] Allow editing allowed collision matrix in python + fix
 get_entry function (#2551)

Co-authored-by: Sebastian Jahr <sebastian.jahr@picknik.ai>
---
 .../moveit/planning_scene/planning_scene.h    |  4 +++
 .../planning_scene/src/planning_scene.cpp     |  5 ++++
 .../collision_detection/collision_matrix.cpp  | 28 +++++++++++++++----
 .../planning_scene/planning_scene.cpp         |  5 ++--
 4 files changed, 35 insertions(+), 7 deletions(-)

diff --git a/moveit_core/planning_scene/include/moveit/planning_scene/planning_scene.h b/moveit_core/planning_scene/include/moveit/planning_scene/planning_scene.h
index 472c63d68a..70717c5d63 100644
--- a/moveit_core/planning_scene/include/moveit/planning_scene/planning_scene.h
+++ b/moveit_core/planning_scene/include/moveit/planning_scene/planning_scene.h
@@ -291,9 +291,13 @@ class MOVEIT_PLANNING_SCENE_EXPORT PlanningScene : public std::enable_shared_fro
   {
     return acm_ ? *acm_ : parent_->getAllowedCollisionMatrix();
   }
+
   /** \brief Get the allowed collision matrix */
   collision_detection::AllowedCollisionMatrix& getAllowedCollisionMatrixNonConst();
 
+  /** \brief Set the allowed collision matrix */
+  void setAllowedCollisionMatrix(const collision_detection::AllowedCollisionMatrix& acm);
+
   /**@}*/
 
   /**
diff --git a/moveit_core/planning_scene/src/planning_scene.cpp b/moveit_core/planning_scene/src/planning_scene.cpp
index 20d8364ff9..73d9e6801f 100644
--- a/moveit_core/planning_scene/src/planning_scene.cpp
+++ b/moveit_core/planning_scene/src/planning_scene.cpp
@@ -558,6 +558,11 @@ collision_detection::AllowedCollisionMatrix& PlanningScene::getAllowedCollisionM
   return *acm_;
 }
 
+void PlanningScene::setAllowedCollisionMatrix(const collision_detection::AllowedCollisionMatrix& acm)
+{
+  getAllowedCollisionMatrixNonConst() = acm;
+}
+
 const moveit::core::Transforms& PlanningScene::getTransforms()
 {
   // Trigger an update of the robot transforms
diff --git a/moveit_py/src/moveit/moveit_core/collision_detection/collision_matrix.cpp b/moveit_py/src/moveit/moveit_core/collision_detection/collision_matrix.cpp
index fb5434b8b1..b0ea225050 100644
--- a/moveit_py/src/moveit/moveit_core/collision_detection/collision_matrix.cpp
+++ b/moveit_py/src/moveit/moveit_core/collision_detection/collision_matrix.cpp
@@ -40,16 +40,34 @@ namespace moveit_py
 {
 namespace bind_collision_detection
 {
-std::pair<bool, collision_detection::AllowedCollision::Type>
-getEntry(const std::shared_ptr<collision_detection::AllowedCollisionMatrix>& acm, const std::string& name1,
-         const std::string& name2)
+// TODO: Create a custom typecaster/revise the current implementation to return std::pair<bool,
+// collision_detection::AllowedCollision::Type>
+std::pair<bool, std::string> getEntry(const collision_detection::AllowedCollisionMatrix& acm, const std::string& name1,
+                                      const std::string& name2)
 {
   // check acm for collision
   collision_detection::AllowedCollision::Type type;
-  bool collision_allowed = acm->getEntry(name1, name2, type);
+  bool collision_allowed = acm.getEntry(name1, name2, type);
+  std::string type_str;
+  if (type == collision_detection::AllowedCollision::Type::NEVER)
+  {
+    type_str = "NEVER";
+  }
+  else if (type == collision_detection::AllowedCollision::Type::ALWAYS)
+  {
+    type_str = "ALWAYS";
+  }
+  else if (type == collision_detection::AllowedCollision::Type::CONDITIONAL)
+  {
+    type_str = "CONDITIONAL";
+  }
+  else
+  {
+    type_str = "UNKNOWN";
+  }
 
   // should return a tuple true/false and the allowed collision type
-  std::pair<bool, collision_detection::AllowedCollision::Type> result = std::make_pair(collision_allowed, type);
+  std::pair<bool, std::string> result = std::make_pair(collision_allowed, type_str);
   return result;
 }
 
diff --git a/moveit_py/src/moveit/moveit_core/planning_scene/planning_scene.cpp b/moveit_py/src/moveit/moveit_core/planning_scene/planning_scene.cpp
index 22cef42453..5df1c18f30 100644
--- a/moveit_py/src/moveit/moveit_core/planning_scene/planning_scene.cpp
+++ b/moveit_py/src/moveit/moveit_core/planning_scene/planning_scene.cpp
@@ -111,8 +111,9 @@ void initPlanningScene(py::module& m)
                     py::return_value_policy::move)
 
       .def_property("transforms", py::overload_cast<>(&planning_scene::PlanningScene::getTransforms), nullptr)
-      .def_property("allowed_collision_matrix", &planning_scene::PlanningScene::getAllowedCollisionMatrix, nullptr,
-                    py::return_value_policy::move)
+      .def_property("allowed_collision_matrix", &planning_scene::PlanningScene::getAllowedCollisionMatrix,
+                    &planning_scene::PlanningScene::setAllowedCollisionMatrix,
+                    py::return_value_policy::reference_internal)
       // methods
       .def("__copy__",
            [](const planning_scene::PlanningScene* self) {