From a7191e24ba38fc316cd1f92438fcce9e7668f27c Mon Sep 17 00:00:00 2001
From: Sebastian Junges <sebastian.junges@rwth-aachen.de>
Date: Wed, 28 Mar 2018 18:15:20 +0200
Subject: [PATCH] support for queries on schedulers

---
 examples/schedulers/01-schedulers.py |  5 +++++
 src/mod_storage.cpp                  |  4 +++-
 src/storage/distribution.cpp         | 17 +++++++++++++++++
 src/storage/distribution.h           |  6 ++++++
 src/storage/scheduler.cpp            | 27 +++++++++++++++++++++++----
 src/storage/scheduler.h              |  3 ++-
 6 files changed, 56 insertions(+), 6 deletions(-)
 create mode 100644 src/storage/distribution.cpp
 create mode 100644 src/storage/distribution.h

diff --git a/examples/schedulers/01-schedulers.py b/examples/schedulers/01-schedulers.py
index 75a8cc7..6a0b2bd 100644
--- a/examples/schedulers/01-schedulers.py
+++ b/examples/schedulers/01-schedulers.py
@@ -18,6 +18,11 @@ def example_schedulers_01():
     result = stormpy.model_checking(model, formulas[0], extract_scheduler = True)
     assert result.has_scheduler
     print(result.scheduler)
+    assert result.scheduler.memoryless
+    assert result.scheduler.deterministic
+
+    for i in range(0,model.nr_states):
+        print("In state {} choose action {}".format(i,result.scheduler.get_choice(i).get_deterministic_choice()))
 
 
 
diff --git a/src/mod_storage.cpp b/src/mod_storage.cpp
index 2e0c3e1..4fd3cbd 100644
--- a/src/mod_storage.cpp
+++ b/src/mod_storage.cpp
@@ -3,6 +3,7 @@
 #include "storage/bitvector.h"
 #include "storage/model.h"
 #include "storage/matrix.h"
+#include "storage/distribution.h"
 #include "storage/scheduler.h"
 #include "storage/prism.h"
 #include "storage/state.h"
@@ -24,5 +25,6 @@ PYBIND11_MODULE(storage, m) {
     define_prism(m);
     define_labeling(m);
     define_expressions(m);
-    define_scheduler(m);
+    define_scheduler<double>(m, "Double");
+    define_distribution<double>(m, "Double");
 }
diff --git a/src/storage/distribution.cpp b/src/storage/distribution.cpp
new file mode 100644
index 0000000..79a03a2
--- /dev/null
+++ b/src/storage/distribution.cpp
@@ -0,0 +1,17 @@
+#include "distribution.h"
+#include "src/helpers.h"
+
+#include "storm/storage/Distribution.h"
+
+template<typename ValueType>
+void define_distribution(py::module& m, std::string vt_suffix) {
+    using Distrib = storm::storage::Distribution<ValueType, uint32_t>;
+
+    std::string distributionClassName = std::string("Distribution") + vt_suffix;
+    py::class_<Distrib> distribution(m, distributionClassName.c_str(), "Finite Support Distribution");
+    distribution
+            .def("__str__", &streamToString<Distrib>);
+}
+
+
+template void define_distribution<double>(py::module&, std::string vt_suffix);
\ No newline at end of file
diff --git a/src/storage/distribution.h b/src/storage/distribution.h
new file mode 100644
index 0000000..3396827
--- /dev/null
+++ b/src/storage/distribution.h
@@ -0,0 +1,6 @@
+#pragma once
+
+#include "common.h"
+
+template<typename ValueType>
+void define_distribution(py::module& m, std::string vt_suffix);
\ No newline at end of file
diff --git a/src/storage/scheduler.cpp b/src/storage/scheduler.cpp
index 065e888..5d79f09 100644
--- a/src/storage/scheduler.cpp
+++ b/src/storage/scheduler.cpp
@@ -3,16 +3,35 @@
 
 #include "storm/storage/Scheduler.h"
 
-void define_scheduler(py::module& m) {
-    using Scheduler = storm::storage::Scheduler<double>;
+template<typename ValueType>
+void define_scheduler(py::module& m, std::string vt_suffix) {
+    using Scheduler = storm::storage::Scheduler<ValueType>;
+    using SchedulerChoice = storm::storage::SchedulerChoice<ValueType>;
 
-    py::class_<Scheduler, std::shared_ptr<storm::storage::Scheduler<double>>> scheduler(m, "Scheduler", "A Finite Memory Scheduler");
+    std::string schedulerClassName = std::string("Scheduler") + vt_suffix;
+    py::class_<Scheduler, std::shared_ptr<storm::storage::Scheduler<ValueType>>> scheduler(m, schedulerClassName.c_str(), "A Finite Memory Scheduler");
     scheduler
             .def("__str__", [](Scheduler const& s) {
                 std::stringstream str;
                 s.printToStream(str);
                 return str.str();
             })
+            .def_property_readonly("memoryless", &Scheduler::isMemorylessScheduler, "Is the scheduler memoryless?")
+            .def_property_readonly("memory_size", &Scheduler::getNumberOfMemoryStates, "How much memory does the scheduler take?")
+            .def_property_readonly("deterministic", &Scheduler::isDeterministicScheduler, "Is the scheduler deterministic?")
+            .def("get_choice", &Scheduler::getChoice, py::arg("state_index"), py::arg("memory_index") = 0)
     ;
 
-}
\ No newline at end of file
+    std::string schedulerChoiceClassName = std::string("SchedulerChoice") + vt_suffix;
+    py::class_<SchedulerChoice> schedulerChoice(m, schedulerChoiceClassName.c_str(), "A choice of a finite memory scheduler");
+    schedulerChoice
+        .def_property_readonly("defined", &SchedulerChoice::isDefined, "Is the choice defined by the scheduler?")
+        .def_property_readonly("deterministic", &SchedulerChoice::isDeterministic, "Is the choice deterministic (given by a Dirac distribution)?")
+        .def("get_deterministic_choice", &SchedulerChoice::getDeterministicChoice, "Get the deterministic choice")
+        .def("get_choice", &SchedulerChoice::getChoiceAsDistribution, "Get the distribution over the actions")
+        .def("__str__", &streamToString<SchedulerChoice>);
+
+}
+
+
+template void define_scheduler<double>(py::module& m, std::string vt_suffix);
\ No newline at end of file
diff --git a/src/storage/scheduler.h b/src/storage/scheduler.h
index b66000d..24cd4cf 100644
--- a/src/storage/scheduler.h
+++ b/src/storage/scheduler.h
@@ -3,4 +3,5 @@
 #include "common.h"
 
 
-void define_scheduler(py::module& m);
+template<typename ValueType>
+void define_scheduler(py::module& m, std::string vt_suffix);