From 9b415a33ae01d422cdfdccfb19f23320ef6765be Mon Sep 17 00:00:00 2001
From: ManifoldFR <wilson.jallet@polytechnique.org>
Date: Tue, 17 Sep 2024 18:59:11 +0200
Subject: [PATCH] great fun exposing boost::unordered_map for CostStack :D

---
 .../src/modelling/expose-cost-stack.cpp       | 77 +++++++++++--------
 tests/python/test_costs.py                    |  5 +-
 2 files changed, 48 insertions(+), 34 deletions(-)

diff --git a/bindings/python/src/modelling/expose-cost-stack.cpp b/bindings/python/src/modelling/expose-cost-stack.cpp
index 9efe06b24..5c4017854 100644
--- a/bindings/python/src/modelling/expose-cost-stack.cpp
+++ b/bindings/python/src/modelling/expose-cost-stack.cpp
@@ -4,6 +4,7 @@
 #include "aligator/modelling/costs/sum-of-costs.hpp"
 
 #include <eigenpy/std-pair.hpp>
+#include <eigenpy/std-map.hpp>
 #include <eigenpy/variant.hpp>
 
 namespace aligator {
@@ -23,40 +24,50 @@ void exposeCostStack() {
   eigenpy::StdPairConverter<CostItem>::registration();
   eigenpy::VariantConverter<CostKey>::registration();
 
-  bp::class_<CostStack, bp::bases<CostAbstract>>(
-      "CostStack", "A weighted sum of other cost functions.", bp::no_init)
-      .def(bp::init<xyz::polymorphic<Manifold>, const int,
-                    const std::vector<PolyCost> &, const std::vector<Scalar> &>(
-          ("self"_a, "space", "nu", "components"_a = bp::list(),
-           "weights"_a = bp::list())))
-      .def(bp::init<const PolyCost &>(("self"_a, "cost")))
-      // .def_readwrite("components", &CostStack::components_,
-      //                "Components of this cost stack.")
-      .def(
-          "addCost",
-          +[](CostStack &self, const PolyCost &cost, const Scalar weight) {
-            // return
-            self.addCost(cost, weight);
-          },
-          ("self"_a, "cost", "weight"_a = 1.),
-          bp::return_internal_reference<>())
-      .def(
-          "addCost",
-          +[](CostStack &self, CostKey key, const PolyCost &cost,
-              const Scalar weight) {
-            // return
-            self.addCost(key, cost, weight);
-          },
-          ("self"_a, "key", "cost", "weight"_a = 1.),
-          bp::return_internal_reference<>())
-      .def("size", &CostStack::size, "Get the number of cost components.")
-      .def(CopyableVisitor<CostStack>())
-      .def(PolymorphicMultiBaseVisitor<CostAbstract>());
+  {
+    bp::scope scope =
+        bp::class_<CostStack, bp::bases<CostAbstract>>(
+            "CostStack", "A weighted sum of other cost functions.", bp::no_init)
+            .def(bp::init<xyz::polymorphic<Manifold>, const int,
+                          const std::vector<PolyCost> &,
+                          const std::vector<Scalar> &>(
+                ("self"_a, "space", "nu", "components"_a = bp::list(),
+                 "weights"_a = bp::list())))
+            .def(bp::init<const PolyCost &>(("self"_a, "cost")))
+            .def_readwrite("components", &CostStack::components_,
+                           "Components of this cost stack.")
+            .def(
+                "addCost",
+                +[](CostStack &self, const PolyCost &cost,
+                    const Scalar weight) {
+                  // return
+                  self.addCost(cost, weight);
+                },
+                ("self"_a, "cost", "weight"_a = 1.),
+                bp::return_internal_reference<>())
+            .def(
+                "addCost",
+                +[](CostStack &self, CostKey key, const PolyCost &cost,
+                    const Scalar weight) {
+                  // return
+                  self.addCost(key, cost, weight);
+                },
+                ("self"_a, "key", "cost", "weight"_a = 1.),
+                bp::return_internal_reference<>())
+            .def("size", &CostStack::size, "Get the number of cost components.")
+            .def(CopyableVisitor<CostStack>())
+            .def(PolymorphicMultiBaseVisitor<CostAbstract>());
+    eigenpy::GenericMapVisitor<CostMap, true>::expose("CostMap");
+  }
 
-  bp::register_ptr_to_python<shared_ptr<CostStackData>>();
-  bp::class_<CostStackData, bp::bases<CostData>>(
-      "CostStackData", "Data struct for CostStack.", bp::no_init)
-      .def_readonly("sub_cost_data", &CostStackData::sub_cost_data);
+  {
+    bp::register_ptr_to_python<shared_ptr<CostStackData>>();
+    bp::scope scope =
+        bp::class_<CostStackData, bp::bases<CostData>>(
+            "CostStackData", "Data struct for CostStack.", bp::no_init)
+            .def_readonly("sub_cost_data", &CostStackData::sub_cost_data);
+    eigenpy::GenericMapVisitor<CostStackData::DataMap, true>::expose("DataMap");
+  }
 }
 
 } // namespace python
diff --git a/tests/python/test_costs.py b/tests/python/test_costs.py
index 9f68d76ea..7ede5e1a0 100644
--- a/tests/python/test_costs.py
+++ b/tests/python/test_costs.py
@@ -145,11 +145,14 @@ def test_stack_error():
     rcost = QuadraticCost(Q, R)
     cost_stack.addCost(rcost)  # optional
 
+    cost_stack.components
+    print(cost_stack.components.todict())
+
     rc2 = QuadraticCost(np.eye(3), np.eye(nu))
     rc3 = QuadraticCost(np.eye(nx), np.eye(nu * 2))
 
     cost_data = cost_stack.createData()
-    print(cost_data.sub_cost_data.tolist())
+    print(cost_data.sub_cost_data.todict())
 
     with pytest.raises(Exception) as e_info:
         cost_stack.addCost(rc2)