From 78df5e25a58fbcddb159938304f6f6e71b21dfe3 Mon Sep 17 00:00:00 2001
From: JafarAbdi <jafar.uruc@gmail.com>
Date: Tue, 3 Dec 2024 16:11:50 +0000
Subject: [PATCH] Add support for GenerateRandomPose

---
 core/include/moveit/task_constructor/stages.h |  1 +
 core/python/bindings/src/stages.cpp           | 22 +++++++++++++++++++
 2 files changed, 23 insertions(+)

diff --git a/core/include/moveit/task_constructor/stages.h b/core/include/moveit/task_constructor/stages.h
index 9c3a4429b..0c349221d 100644
--- a/core/include/moveit/task_constructor/stages.h
+++ b/core/include/moveit/task_constructor/stages.h
@@ -51,3 +51,4 @@
 #include "stages/move_relative.h"
 #include "stages/move_to.h"
 #include "stages/predicate_filter.h"
+#include "stages/generate_random_pose.h"
diff --git a/core/python/bindings/src/stages.cpp b/core/python/bindings/src/stages.cpp
index ea8914324..3c9e57fbe 100644
--- a/core/python/bindings/src/stages.cpp
+++ b/core/python/bindings/src/stages.cpp
@@ -57,6 +57,7 @@ PYBIND11_SMART_HOLDER_TYPE_CASTERS(Connect)
 PYBIND11_SMART_HOLDER_TYPE_CASTERS(FixCollisionObjects)
 PYBIND11_SMART_HOLDER_TYPE_CASTERS(GenerateGraspPose)
 PYBIND11_SMART_HOLDER_TYPE_CASTERS(GeneratePlacePose)
+PYBIND11_SMART_HOLDER_TYPE_CASTERS(GenerateRandomPose)
 PYBIND11_SMART_HOLDER_TYPE_CASTERS(GeneratePose)
 PYBIND11_SMART_HOLDER_TYPE_CASTERS(Pick)
 PYBIND11_SMART_HOLDER_TYPE_CASTERS(Place)
@@ -405,6 +406,27 @@ void export_stages(pybind11::module& m) {
 		)")
 	    .def(py::init<const std::string&>(), "name"_a);
 
+	py::enum_<GenerateRandomPose::PoseDimension>(m, "PoseDimension", R"(
+      Define the dimensions of a pose that can be randomized.
+    )")
+      .value("X", GenerateRandomPose::PoseDimension::X, "X dimension")
+      .value("Y", GenerateRandomPose::PoseDimension::Y, "Y dimension")
+      .value("Z", GenerateRandomPose::PoseDimension::Z, "Z dimension")
+      .value("ROLL", GenerateRandomPose::PoseDimension::ROLL, "Roll dimension")
+      .value("PITCH", GenerateRandomPose::PoseDimension::PITCH, "Pitch dimension")
+      .value("YAW", GenerateRandomPose::PoseDimension::YAW, "Yaw dimension");
+
+	properties::class_<GenerateRandomPose, GeneratePose>(m, "GenerateRandomPose", R"(
+			Monitoring generator stage which can be used to generate random poses, based on solutions provided
+			by the monitored stage and the specified pose dimension samplers.
+		)")
+	    .def(py::init<const std::string&>(), "name"_a)
+	    .def("set_max_solutions", &GenerateRandomPose::setMaxSolutions, "max_solutions"_a)
+	    .def("sample_dimension", [](GenerateRandomPose& self, const GenerateRandomPose::PoseDimension pose_dimension,
+                                 const double width) {
+        self.sampleDimension<std::uniform_real_distribution>(pose_dimension, width);
+      }, "pose_dimension"_a, "width"_a);
+
 	properties::class_<Pick, SerialContainer>(m, "Pick", R"(
 			The Pick stage is a specialization of the PickPlaceBase class, which
 			wraps the pipeline to pick or place an object with a given end effector.