From edbd7e827bf10d7851ce34ac662b402842be9a24 Mon Sep 17 00:00:00 2001
From: sjunges <sebastian.junges@rwth-aachen.de>
Date: Fri, 25 Sep 2015 14:16:35 +0200
Subject: [PATCH] MDPs, restrict choices, get choiceindex

Former-commit-id: 6d1f2ff37d12fd3bfb7d3d0db42d46ba7bab30f8
---
 src/models/sparse/Mdp.cpp                   | 11 ++++++++---
 src/models/sparse/Mdp.h                     |  8 +++++++-
 src/models/sparse/NondeterministicModel.cpp |  4 ++--
 src/models/sparse/NondeterministicModel.h   |  2 +-
 4 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/src/models/sparse/Mdp.cpp b/src/models/sparse/Mdp.cpp
index f6e4b4800..6396af48d 100644
--- a/src/models/sparse/Mdp.cpp
+++ b/src/models/sparse/Mdp.cpp
@@ -75,15 +75,20 @@ namespace storm {
             }
 
             template <typename ValueType, typename RewardModelType>
-            Mdp<ValueType, RewardModelType> Mdp<ValueType, RewardModelType>::restrictActions(storm::storage::BitVector const& enabledActions) const {
-                storm::storage::SparseMatrix<ValueType> restrictedTransitions = this->getTransitionMatrix().restrictRows(enabledActions);
+            Mdp<ValueType, RewardModelType> Mdp<ValueType, RewardModelType>::restrictChoices(storm::storage::BitVector const& enabledChoices) const {
+                storm::storage::SparseMatrix<ValueType> restrictedTransitions = this->getTransitionMatrix().restrictRows(enabledChoices);
                 std::unordered_map<std::string, RewardModelType> newRewardModels;
                 for (auto const& rewardModel : this->getRewardModels()) {
-                    newRewardModels.emplace(rewardModel.first, rewardModel.second.restrictActions(enabledActions));
+                    newRewardModels.emplace(rewardModel.first, rewardModel.second.restrictActions(enabledChoices));
                 }
                 return Mdp<ValueType, RewardModelType>(restrictedTransitions, this->getStateLabeling(), newRewardModels, this->getOptionalChoiceLabeling());
             }
 
+            template<typename ValueType, typename RewardModelType>
+            uint_least64_t Mdp<ValueType, RewardModelType>::getChoiceIndex(storm::storage::StateActionPair const& stateactPair) const {
+                return this->getNondeterministicChoiceIndices()[stateactPair.getState()]+stateactPair.getAction();
+            }
+
             template class Mdp<double>;
             template class Mdp<float>;
 
diff --git a/src/models/sparse/Mdp.h b/src/models/sparse/Mdp.h
index 9bc29bbc9..a28a84929 100644
--- a/src/models/sparse/Mdp.h
+++ b/src/models/sparse/Mdp.h
@@ -1,6 +1,7 @@
 #ifndef STORM_MODELS_SPARSE_MDP_H_
 #define STORM_MODELS_SPARSE_MDP_H_
 
+#include <src/storage/StateActionPair.h>
 #include "src/models/sparse/NondeterministicModel.h"
 #include "src/utility/OsDetection.h"
 
@@ -65,7 +66,12 @@ namespace storm {
                  * @param enabledActions A BitVector of lenght numberOfChoices(), which is one iff the action should be kept.
                  * @return A subMDP.
                  */
-                Mdp<ValueType, RewardModelType> restrictActions(storm::storage::BitVector const& enabledActions) const;
+                Mdp<ValueType, RewardModelType> restrictChoices(storm::storage::BitVector const& enabledActions) const;
+
+                /*!
+                 *  For a state/action pair, get the choice index referring to the state-action pair.
+                 */
+                uint_fast64_t getChoiceIndex(storm::storage::StateActionPair const& stateactPair) const;
             };
             
         } // namespace sparse
diff --git a/src/models/sparse/NondeterministicModel.cpp b/src/models/sparse/NondeterministicModel.cpp
index 9cfbe75f9..ecc350d8b 100644
--- a/src/models/sparse/NondeterministicModel.cpp
+++ b/src/models/sparse/NondeterministicModel.cpp
@@ -47,7 +47,7 @@ namespace storm {
                 return indices[state+1] - indices[state];
             }
             
-            template<typename ValueType, typename RewardModelType>
+            /*template<typename ValueType, typename RewardModelType>
             void NondeterministicModel<ValueType, RewardModelType>::modifyStateActionRewards(RewardModelType& rewardModel, std::map<std::pair<uint_fast64_t, LabelSet>, typename RewardModelType::ValueType> const& modifications) const {
                 STORM_LOG_THROW(rewardModel.hasStateActionRewards(), storm::exceptions::InvalidOperationException, "Cannot modify state-action rewards, because the reward model does not have state-action rewards.");
                 STORM_LOG_THROW(this->hasChoiceLabeling(), storm::exceptions::InvalidOperationException, "Cannot modify state-action rewards, because the model does not have an action labeling.");
@@ -61,7 +61,7 @@ namespace storm {
                         }
                     }
                 }
-            }
+            }*/
             
             template<typename ValueType, typename RewardModelType>
             void NondeterministicModel<ValueType, RewardModelType>::reduceToStateBasedRewards() {
diff --git a/src/models/sparse/NondeterministicModel.h b/src/models/sparse/NondeterministicModel.h
index 5bacf6dd7..ad09143c7 100644
--- a/src/models/sparse/NondeterministicModel.h
+++ b/src/models/sparse/NondeterministicModel.h
@@ -80,7 +80,7 @@ namespace storm {
                  * @param rewardModel The reward model whose state-action rewards to modify.
                  * @param modifications A mapping from state-action pairs to the their new reward values.
                  */
-                void modifyStateActionRewards(RewardModelType& rewardModel, std::map<std::pair<uint_fast64_t, LabelSet>, typename RewardModelType::ValueType> const& modifications) const;
+                //void modifyStateActionRewards(RewardModelType& rewardModel, std::map<std::pair<uint_fast64_t, LabelSet>, typename RewardModelType::ValueType> const& modifications) const;
                 
                 virtual void reduceToStateBasedRewards() override;