From 7cee208518e777d2daa74b21a1cdb8b5587fb667 Mon Sep 17 00:00:00 2001
From: Sebastian Junges <sebastian.junges@gmail.com>
Date: Mon, 24 Aug 2020 23:41:57 -0700
Subject: [PATCH] keep state valuations

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

diff --git a/lib/stormpy/pomdp/__init__.py b/lib/stormpy/pomdp/__init__.py
index 88bc23c..8f93f3e 100644
--- a/lib/stormpy/pomdp/__init__.py
+++ b/lib/stormpy/pomdp/__init__.py
@@ -13,7 +13,7 @@ def make_canonic(model):
     else:
         return pomdp._make_canonic_Double(model)
 
-def make_simple(model):
+def make_simple(model, keep_state_valuations=False):
     """
     Make the POMDP simple (aka alternating), i.e., each state has at most two actions, and if there is nondeterminism, then there is no probabilistic branching,
 
@@ -21,11 +21,11 @@ def make_simple(model):
     :return:
     """
     if model.supports_parameters:
-        return pomdp._make_simple_Rf(model)
+        return pomdp._make_simple_Rf(model, keep_state_valuations)
     else:
-        return pomdp._make_simple_Double(model)
+        return pomdp._make_simple_Double(model, keep_state_valuations)
 
-def unfold_memory(model, memory, add_memory_labels=False):
+def unfold_memory(model, memory, add_memory_labels=False, keep_state_valuations=False):
     """
     Unfold the memory for an FSC into the POMDP
 
@@ -34,9 +34,9 @@ def unfold_memory(model, memory, add_memory_labels=False):
     :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, add_memory_labels)
+        return pomdp._unfold_memory_Rf(model, memory, add_memory_labels, keep_state_valuations)
     else:
-        return pomdp._unfold_memory_Double(model, memory, add_memory_labels)
+        return pomdp._unfold_memory_Double(model, memory, add_memory_labels, keep_state_valuations)
 
 def apply_unknown_fsc(model, mode):
     if model.supports_parameters:
diff --git a/src/pomdp/transformations.cpp b/src/pomdp/transformations.cpp
index da18cbc..7fd2287 100644
--- a/src/pomdp/transformations.cpp
+++ b/src/pomdp/transformations.cpp
@@ -11,15 +11,15 @@ 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, bool addMemoryLabels) {
-    storm::transformer::PomdpMemoryUnfolder<ValueType> unfolder(pomdp, memory, addMemoryLabels);
+std::shared_ptr<storm::models::sparse::Pomdp<ValueType>> unfold_memory(storm::models::sparse::Pomdp<ValueType> const& pomdp, storm::storage::PomdpMemory const& memory, bool addMemoryLabels, bool keepStateValuations) {
+    storm::transformer::PomdpMemoryUnfolder<ValueType> unfolder(pomdp, memory, addMemoryLabels, keepStateValuations);
     return unfolder.transform();
 }
 
 template<typename ValueType>
-std::shared_ptr<storm::models::sparse::Pomdp<ValueType>> make_simple(storm::models::sparse::Pomdp<ValueType> const& pomdp) {
+std::shared_ptr<storm::models::sparse::Pomdp<ValueType>> make_simple(storm::models::sparse::Pomdp<ValueType> const& pomdp, bool keepStateValuations) {
     storm::transformer::BinaryPomdpTransformer<ValueType> transformer;
-    return transformer.transform(pomdp,true);
+    return transformer.transform(pomdp,true, keepStateValuations);
 }
 
 template<typename ValueType>
@@ -42,8 +42,8 @@ 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"), py::arg("memorylabels") = false);
-    m.def(("_make_simple_"+ vtSuffix).c_str(), &make_simple<ValueType>, "Make POMDP simple", py::arg("pomdp"));
+    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, py::arg("keep_state_valuations")=false);
+    m.def(("_make_simple_"+ vtSuffix).c_str(), &make_simple<ValueType>, "Make POMDP simple", py::arg("pomdp"), py::arg("keep_state_valuations")=false);
     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);
 }