From b531dccad9ef732fdb7f1e341742e3392fc83da7 Mon Sep 17 00:00:00 2001
From: TimQu <tim.quatmann@cs.rwth-aachen.de>
Date: Tue, 23 May 2017 17:55:20 +0200
Subject: [PATCH] .dot output for deterministic models with choice labels

---
 .../models/sparse/DeterministicModel.cpp      | 22 +++++++++++++++++--
 1 file changed, 20 insertions(+), 2 deletions(-)

diff --git a/src/storm/models/sparse/DeterministicModel.cpp b/src/storm/models/sparse/DeterministicModel.cpp
index dc3cdacc4..2cc082fac 100644
--- a/src/storm/models/sparse/DeterministicModel.cpp
+++ b/src/storm/models/sparse/DeterministicModel.cpp
@@ -30,14 +30,32 @@ namespace storm {
             void DeterministicModel<ValueType, RewardModelType>::writeDotToStream(std::ostream& outStream, bool includeLabeling, storm::storage::BitVector const* subsystem, std::vector<ValueType> const* firstValue, std::vector<ValueType> const* secondValue, std::vector<uint_fast64_t> const* stateColoring, std::vector<std::string> const* colors, std::vector<uint_fast64_t>* scheduler, bool finalizeOutput) const {
                 Model<ValueType, RewardModelType>::writeDotToStream(outStream, includeLabeling, subsystem, firstValue, secondValue, stateColoring, colors, scheduler, false);
                 
-                // Simply iterate over all transitions and draw the arrows with probability information attached.
+                // iterate over all transitions and draw the arrows with probability information attached.
                 auto rowIt = this->getTransitionMatrix().begin();
                 for (uint_fast64_t i = 0; i < this->getTransitionMatrix().getRowCount(); ++i, ++rowIt) {
+                    
+                    // Put in an intermediate node if there is a choice labeling
+                    std::string arrowOrigin = std::to_string(i);
+                    if (this->hasChoiceLabeling()) {
+                        arrowOrigin = "\"" + arrowOrigin + "c\"";
+                        outStream << "\t" << arrowOrigin << " [shape = \"point\"]" << std::endl;
+                        outStream << "\t" << i << " -> " << arrowOrigin << " [label= \"{";
+                        bool firstLabel = true;
+                        for (auto const& label : this->getChoiceLabeling().getLabelsOfChoice(i)) {
+                            if (!firstLabel) {
+                                outStream << ", ";
+                            }
+                            firstLabel = false;
+                            outStream << label;
+                        }
+                        outStream << "}\"];" << std::endl;
+                    }
+                    
                     typename storm::storage::SparseMatrix<ValueType>::const_rows row = this->getTransitionMatrix().getRow(i);
                     for (auto const& transition : row) {
                         if (transition.getValue() != storm::utility::zero<ValueType>()) {
                             if (subsystem == nullptr || subsystem->get(transition.getColumn())) {
-                                outStream << "\t" << i << " -> " << transition.getColumn() << " [ label= \"" << transition.getValue() << "\" ];" << std::endl;
+                                outStream << "\t" << arrowOrigin << " -> " << transition.getColumn() << " [ label= \"" << transition.getValue() << "\" ];" << std::endl;
                             }
                         }
                     }