From 557dd4f57e1bcfae59b4b5dd34ef6f9c161e7f6c Mon Sep 17 00:00:00 2001
From: Thomas Knoll <thomas.knolł@student.tugraz.at>
Date: Tue, 8 Aug 2023 15:40:50 +0200
Subject: [PATCH] extended smg / stg / simulator python classes

---
 lib/stormpy/simulator.py |  3 +++
 src/core/result.cpp      |  2 +-
 src/storage/model.cpp    | 16 ++++++++++++++--
 3 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/lib/stormpy/simulator.py b/lib/stormpy/simulator.py
index 8c8c6ef..6e0444d 100644
--- a/lib/stormpy/simulator.py
+++ b/lib/stormpy/simulator.py
@@ -204,6 +204,9 @@ class SparseSimulator(Simulator):
                 raise RuntimeError("Program level observations require model with state valuations")
         self._state_valuations = self._model.state_valuations
 
+    def get_current_state(self):
+        return self._engine.get_current_state()
+
 
 
 def create_simulator(model, seed = None):
diff --git a/src/core/result.cpp b/src/core/result.cpp
index 961626d..cd3dbc9 100644
--- a/src/core/result.cpp
+++ b/src/core/result.cpp
@@ -40,7 +40,7 @@ void define_result(py::module& m) {
         .def_property_readonly("_pareto_curve", &storm::modelchecker::CheckResult::isParetoCurveCheckResult, "Flag if result is a pareto curve")
         .def_property_readonly("result_for_all_states", &storm::modelchecker::CheckResult::isResultForAllStates, "Flag if result is for all states")
         .def_property_readonly("has_scheduler", &storm::modelchecker::CheckResult::hasScheduler, "Flag if a scheduler is present")
-        .def_property_readonly("has_schield", &storm::modelchecker::CheckResult::hasShield, "Flag if a schield is present")
+        .def_property_readonly("has_shield", &storm::modelchecker::CheckResult::hasShield, "Flag if a schield is present")
 
         .def("as_explicit_qualitative", [](storm::modelchecker::CheckResult const& result) {
                 return result.asExplicitQualitativeCheckResult();
diff --git a/src/storage/model.cpp b/src/storage/model.cpp
index 73ab53a..ef482a1 100644
--- a/src/storage/model.cpp
+++ b/src/storage/model.cpp
@@ -6,6 +6,7 @@
 #include "storm/models/sparse/Dtmc.h"
 #include "storm/models/sparse/Mdp.h"
 #include "storm/models/sparse/Smg.h"
+#include "storm/models/sparse/StochasticTwoPlayerGame.h"
 #include "storm/models/sparse/Pomdp.h"
 #include "storm/models/sparse/Ctmc.h"
 #include "storm/models/sparse/MarkovAutomaton.h"
@@ -35,6 +36,7 @@ template<typename ValueType> using SparseModel = storm::models::sparse::Model<Va
 template<typename ValueType> using SparseDtmc = storm::models::sparse::Dtmc<ValueType>;
 template<typename ValueType> using SparseMdp = storm::models::sparse::Mdp<ValueType>;
 template<typename ValueType> using Smg = storm::models::sparse::Smg<ValueType>;
+template<typename ValueType> using StochasticTwoPlayerGame = storm::models::sparse::StochasticTwoPlayerGame<ValueType>;
 template<typename ValueType> using SparsePomdp = storm::models::sparse::Pomdp<ValueType>;
 template<typename ValueType> using SparseCtmc = storm::models::sparse::Ctmc<ValueType>;
 template<typename ValueType> using SparseMarkovAutomaton = storm::models::sparse::MarkovAutomaton<ValueType>;
@@ -226,10 +228,20 @@ void define_sparse_model(py::module& m, std::string const& vtSuffix) {
     py::class_<Smg<ValueType>, std::shared_ptr<Smg<ValueType>>> smg(m, ("Sparse" + vtSuffix + "Smg").c_str(), "SMG in sparse representation", model);
     smg
         .def(py::init<Smg<ValueType>>(), py::arg("other_model"))
-        
+        .def(py::init<ModelComponents<ValueType> const&>(), py::arg("components"))
+        .def("get_player_of_state", &Smg<ValueType>::getPlayerOfState, py::arg("state_index"))
+        .def("get_player_index", &Smg<ValueType>::getPlayerIndex, py::arg("player_name"))
+    ;
 
+    py::class_<StochasticTwoPlayerGame<ValueType>, std::shared_ptr<StochasticTwoPlayerGame<ValueType>>> stg(m, ("Sparse" + vtSuffix + "Stg").c_str(), "STG in sparse representation", model);
+    stg
+        .def(py::init<StochasticTwoPlayerGame<ValueType>>(), py::arg("other_model"))
+        .def(py::init<ModelComponents<ValueType> const&>(), py::arg("components"))
+        .def_property_readonly("player_1_matrix", &StochasticTwoPlayerGame<ValueType>::getPlayer1Matrix)
+        .def_property_readonly("player_2_matrix", &StochasticTwoPlayerGame<ValueType>::getPlayer2Matrix)
+        .def_property_readonly("has_player_2_choice_labeling", &StochasticTwoPlayerGame<ValueType>::hasPlayer2ChoiceLabeling)
+        .def_property_readonly("player_2_choice_labeling", &StochasticTwoPlayerGame<ValueType>::getPlayer2ChoiceLabeling)
     ;
-        
 
     py::class_<SparsePomdp<ValueType>, std::shared_ptr<SparsePomdp<ValueType>>>(m, ("Sparse" + vtSuffix + "Pomdp").c_str(), "POMDP in sparse representation", mdp)
         .def(py::init<SparsePomdp<ValueType>>(), py::arg("other_model"))