From 1021d242bd6a5bc01cd02239c658732fb6252924 Mon Sep 17 00:00:00 2001
From: Roman Andriushchenko <iandri@fit.vut.cz>
Date: Fri, 8 Dec 2023 15:13:53 +0100
Subject: [PATCH] fix not using optional coloring

---
 paynt/quotient/coloring.py | 12 +++++++++++-
 paynt/quotient/quotient.py | 18 ++++--------------
 2 files changed, 15 insertions(+), 15 deletions(-)

diff --git a/paynt/quotient/coloring.py b/paynt/quotient/coloring.py
index 721d3d88c..93e77ebce 100644
--- a/paynt/quotient/coloring.py
+++ b/paynt/quotient/coloring.py
@@ -56,7 +56,7 @@ def __init__(self, mdp, holes, action_to_hole_options):
     @property
     def num_choices(self):
         return len(self.action_to_hole_options)
-        
+
     def select_actions(self, family):
         '''
         Select non-default actions relevant in the provided design space.
@@ -112,3 +112,13 @@ def select_actions(self, family):
         selected_actions_bv = stormpy.synthesis.construct_selection(self.default_actions, selected_actions)
         
         return hole_selected_actions,selected_actions,selected_actions_bv
+
+    
+    def choices_to_hole_selection(self, choices):
+        hole_selection = [set() for hole_index in self.holes.hole_indices]
+        for choice in choices:
+            choice_options = self.action_to_hole_options[choice]
+            for hole_index,option in choice_options.items():
+                hole_selection[hole_index].add(option)
+        hole_selection = [list(options) for options in hole_selection]
+        return hole_selection
diff --git a/paynt/quotient/quotient.py b/paynt/quotient/quotient.py
index 073993115..c5e5321ac 100644
--- a/paynt/quotient/quotient.py
+++ b/paynt/quotient/quotient.py
@@ -185,25 +185,15 @@ def state_to_choice_to_choices(self, state_to_choice):
                 choices.set(choice,True)
         return choices
 
-
-    def choices_to_hole_selection(self, choices, coloring=None):
-        if coloring is None:
-            coloring = self.coloring
-        hole_selection = [set() for hole_index in self.design_space.hole_indices]
-        for choice in choices:
-            choice_options = self.coloring.action_to_hole_options[choice]
-            for hole_index,option in choice_options.items():
-                hole_selection[hole_index].add(option)
-        hole_selection = [list(options) for options in hole_selection]
-        return hole_selection
-
     
     def scheduler_selection(self, mdp, scheduler, coloring=None):
         ''' Get hole options involved in the scheduler selection. '''
         assert scheduler.memoryless and scheduler.deterministic
-        state_to_choice = self.scheduler_to_state_to_choice(mdp, scheduler, keep_reachable=True)
+        state_to_choice = self.scheduler_to_state_to_choice(mdp, scheduler)
         choices = self.state_to_choice_to_choices(state_to_choice)
-        hole_selection = self.choices_to_hole_selection(choices,coloring)
+        if coloring is None:
+            coloring = self.coloring
+        hole_selection = coloring.choices_to_hole_selection(choices)
         return hole_selection