From 75ad72a5747de2ca2567285d2dbd6006235198a2 Mon Sep 17 00:00:00 2001
From: Sebastian Junges <sebastian.junges@gmail.com>
Date: Mon, 24 Aug 2020 22:26:00 -0700
Subject: [PATCH] add_memory_labels for memory unfolding added

---
 lib/stormpy/pomdp/__init__.py | 6 +++---
 src/pomdp/transformations.cpp | 6 +++---
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/lib/stormpy/pomdp/__init__.py b/lib/stormpy/pomdp/__init__.py
index 00c7b95..88bc23c 100644
--- a/lib/stormpy/pomdp/__init__.py
+++ b/lib/stormpy/pomdp/__init__.py
@@ -25,7 +25,7 @@ def make_simple(model):
     else:
         return pomdp._make_simple_Double(model)
 
-def unfold_memory(model, memory):
+def unfold_memory(model, memory, add_memory_labels=False):
     """
     Unfold the memory for an FSC into the POMDP
 
@@ -34,9 +34,9 @@ def unfold_memory(model, memory):
     :return: A pomdp that contains states from the product of the original POMDP and the FSC Memory
     """
     if model.supports_parameters:
-        return pomdp._unfold_memory_Rf(model, memory)
+        return pomdp._unfold_memory_Rf(model, memory, add_memory_labels)
     else:
-        return pomdp._unfold_memory_Double(model, memory)
+        return pomdp._unfold_memory_Double(model, memory, add_memory_labels)
 
 def apply_unknown_fsc(model, mode):
     if model.supports_parameters:
diff --git a/src/pomdp/transformations.cpp b/src/pomdp/transformations.cpp
index 9bbafd1..da18cbc 100644
--- a/src/pomdp/transformations.cpp
+++ b/src/pomdp/transformations.cpp
@@ -11,8 +11,8 @@ std::shared_ptr<storm::models::sparse::Pomdp<ValueType>> make_canonic(storm::mod
 }
 
 template<typename ValueType>
-std::shared_ptr<storm::models::sparse::Pomdp<ValueType>> unfold_memory(storm::models::sparse::Pomdp<ValueType> const& pomdp, storm::storage::PomdpMemory const& memory) {
-    storm::transformer::PomdpMemoryUnfolder<ValueType> unfolder(pomdp, memory);
+std::shared_ptr<storm::models::sparse::Pomdp<ValueType>> unfold_memory(storm::models::sparse::Pomdp<ValueType> const& pomdp, storm::storage::PomdpMemory const& memory, bool addMemoryLabels) {
+    storm::transformer::PomdpMemoryUnfolder<ValueType> unfolder(pomdp, memory, addMemoryLabels);
     return unfolder.transform();
 }
 
@@ -42,7 +42,7 @@ void define_transformations_nt(py::module &m) {
 template<typename ValueType>
 void define_transformations(py::module& m, std::string const& vtSuffix) {
     m.def(("_make_canonic_" + vtSuffix).c_str(), &make_canonic<ValueType>, "Return a canonicly-ordered POMDP", py::arg("pomdp"));
-    m.def(("_unfold_memory_" + vtSuffix).c_str(), &unfold_memory<ValueType>, "Unfold memory into a POMDP", py::arg("pomdp"), py::arg("memorystructure"));
+    m.def(("_unfold_memory_" + vtSuffix).c_str(), &unfold_memory<ValueType>, "Unfold memory into a POMDP", py::arg("pomdp"), py::arg("memorystructure"), py::arg("memorylabels") = false);
     m.def(("_make_simple_"+ vtSuffix).c_str(), &make_simple<ValueType>, "Make POMDP simple", py::arg("pomdp"));
     m.def(("_apply_unknown_fsc_" + vtSuffix).c_str(), &apply_unknown_fsc<ValueType>, "Apply unknown FSC",py::arg("pomdp"), py::arg("application_mode")=storm::transformer::PomdpFscApplicationMode::SIMPLE_LINEAR);
 }