From 06da755a74fdd2106e3b9175095c85cbab533fc8 Mon Sep 17 00:00:00 2001
From: jpdefrutos <javier.perezdefrutos@sintef.no>
Date: Thu, 7 Mar 2024 15:58:07 +0100
Subject: [PATCH] Added pickling support to the Robot class

---
 source/pyromocc/source/pyromocc.cpp           | 29 ++++++++++++++++++-
 source/romocc/Robot.cpp                       |  5 ++++
 source/romocc/Robot.h                         |  2 ++
 .../communication/CommunicationInterface.cpp  |  5 ++++
 .../communication/CommunicationInterface.h    |  6 ++++
 source/romocc/robotics/RobotState.cpp         |  5 ++++
 source/romocc/robotics/RobotState.h           |  2 ++
 7 files changed, 53 insertions(+), 1 deletion(-)

diff --git a/source/pyromocc/source/pyromocc.cpp b/source/pyromocc/source/pyromocc.cpp
index 50e6a11..9c38d52 100644
--- a/source/pyromocc/source/pyromocc.cpp
+++ b/source/pyromocc/source/pyromocc.cpp
@@ -23,7 +23,34 @@ PYBIND11_MODULE(pyromocc, m) {
         .def("configure", &Robot::configure)
         .def("get_state", &Robot::getCurrentState)
         .def("stop_move", &Robot::stopMove)
-        .def("_connect", &Robot::connect, "Connects to the robot.");
+        .def("_connect", &Robot::connect, "Connects to the robot.")
+        .def("get_comms_config", &Robot::getCommuncationConfiguration)
+        .def(py::pickle(
+            /* Get object state */
+            [](Robot &r) {
+                Manipulator manipulator = r.getCurrentState()->getManipulator();
+                connectionConfiguration commsConfig = r.getCommuncationConfiguration();
+                if (r.isConnected())
+                    r.disconnect();
+                return py::make_tuple(manipulator.manipulator, manipulator.sw_version, commsConfig.host, commsConfig.port);
+            },
+
+            /* Create a new object and restore the state */
+            [](py::tuple pickled_state) {
+                if (pickled_state.size() != 4)
+                    throw std::runtime_error("Invalid state");
+
+                Manipulator manipulator = { pickled_state[0].cast<ManipulatorType>(), pickled_state[1].cast<std::string>() };
+                /* Create a new instance of Robot */
+                std::unique_ptr<Robot> r = std::unique_ptr<Robot>(new Robot());
+                r->configure(manipulator, pickled_state[2].cast<std::string>(), pickled_state[3].cast<int>());
+                if (!r->connect()) {
+                    throw std::runtime_error("Failed to unpickle, cannot connect to robot");
+                }
+
+                return r;
+            }
+        ));
 
     robot.def("_movej", [](Robot& self, Eigen::Ref<const Eigen::RowVectorXd> target, double acc, double vel,
             double time = 0, double blendRad = 0, bool wait=false){
diff --git a/source/romocc/Robot.cpp b/source/romocc/Robot.cpp
index 6b84c6e..337787d 100644
--- a/source/romocc/Robot.cpp
+++ b/source/romocc/Robot.cpp
@@ -55,6 +55,11 @@ RobotState::pointer Robot::getCurrentState() const
     return mCurrentState;
 }
 
+connectionConfiguration Robot::getCommuncationConfiguration() const
+{
+    return mCommunicationInterface->getConnectionConfig();
+}
+
 void Robot::runMotionQueue(MotionQueue queue)
 {
     mMotionQueue = std::move(queue);
diff --git a/source/romocc/Robot.h b/source/romocc/Robot.h
index f17c92e..3f25c87 100644
--- a/source/romocc/Robot.h
+++ b/source/romocc/Robot.h
@@ -28,6 +28,8 @@ class ROMOCC_EXPORT Robot : public Object
         void shutdown();
 
         RobotState::pointer getCurrentState() const;
+        connectionConfiguration getCommuncationConfiguration() const;
+
         void addUpdateSubscription(std::function<void()> updateFunction);
 
         template <class Target>
diff --git a/source/romocc/communication/CommunicationInterface.cpp b/source/romocc/communication/CommunicationInterface.cpp
index 5f9d903..ae003d1 100644
--- a/source/romocc/communication/CommunicationInterface.cpp
+++ b/source/romocc/communication/CommunicationInterface.cpp
@@ -100,6 +100,11 @@ void CommunicationInterface::setManipulator(romocc::Manipulator manipulator)
     mCurrentState->setManipulator(manipulator);
 }
 
+connectionConfiguration CommunicationInterface::getConnectionConfig()
+{
+    return connectionConfiguration{ mHost, mPort };
+}
+
 RobotState::pointer CommunicationInterface::getRobotState()
 {
     return mCurrentState;
diff --git a/source/romocc/communication/CommunicationInterface.h b/source/romocc/communication/CommunicationInterface.h
index 17d740e..ee15363 100644
--- a/source/romocc/communication/CommunicationInterface.h
+++ b/source/romocc/communication/CommunicationInterface.h
@@ -13,6 +13,11 @@
 
 namespace romocc
 {
+struct connectionConfiguration
+{
+    std::string host;
+    int port;
+};
 
 class ROMOCC_EXPORT CommunicationInterface : public Object
 {
@@ -24,6 +29,7 @@ class ROMOCC_EXPORT CommunicationInterface : public Object
 
         void configConnection(std::string host, int port);
         void setManipulator(Manipulator manipulator);
+        connectionConfiguration getConnectionConfig();
 
         RobotState::pointer getRobotState();
 
diff --git a/source/romocc/robotics/RobotState.cpp b/source/romocc/robotics/RobotState.cpp
index ed1d696..438d1c1 100644
--- a/source/romocc/robotics/RobotState.cpp
+++ b/source/romocc/robotics/RobotState.cpp
@@ -246,4 +246,9 @@ Transform3d RobotState::jointConfigToOperationalConfig(romocc::Vector6d jointCon
     return transform_to_joint(jointConfig);
 }
 
+Manipulator RobotState::getManipulator()
+{
+    return mManipulator;
+}
+
 }
\ No newline at end of file
diff --git a/source/romocc/robotics/RobotState.h b/source/romocc/robotics/RobotState.h
index bfa3297..b84935f 100644
--- a/source/romocc/robotics/RobotState.h
+++ b/source/romocc/robotics/RobotState.h
@@ -52,6 +52,8 @@ class ROMOCC_EXPORT RobotState : public Object
         Vector6d operationalConfigToJointConfig(Transform3d transform);
         Transform3d jointConfigToOperationalConfig(Vector6d jointConfig);
 
+        Manipulator getManipulator();
+
         std::shared_ptr<FKSolver> getFKSolver();
         std::shared_ptr<IKSolver> getIKSolver();