From 2a2bc909a246fa955218c2390bac1f06ae17446c Mon Sep 17 00:00:00 2001
From: PimLeerkes <pimsimon.leerkes@gmail.com>
Date: Sat, 30 Nov 2024 11:22:29 +0100
Subject: [PATCH 1/8] pgc works for the simple example

---
 docs/getting_started/mdp.ipynb |  26 ++---
 stormvogel/pgc.py              | 186 +++++++++++++++++++++++++++++++++
 tests/test_pgc.py              |   0
 3 files changed, 199 insertions(+), 13 deletions(-)
 create mode 100644 stormvogel/pgc.py
 create mode 100644 tests/test_pgc.py

diff --git a/docs/getting_started/mdp.ipynb b/docs/getting_started/mdp.ipynb
index 8dcb11d..4392424 100644
--- a/docs/getting_started/mdp.ipynb
+++ b/docs/getting_started/mdp.ipynb
@@ -9,7 +9,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 1,
+      "execution_count": 2,
       "metadata": {},
       "outputs": [],
       "source": [
@@ -19,16 +19,16 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 2,
+      "execution_count": 3,
       "metadata": {},
       "outputs": [
         {
           "data": {
             "text/plain": [
-              "<stormpy.storage.storage.PrismProgram at 0x7fb379fafb30>"
+              "<stormpy.storage.storage.PrismProgram at 0x7f9eb03f6130>"
             ]
           },
-          "execution_count": 2,
+          "execution_count": 3,
           "metadata": {},
           "output_type": "execute_result"
         }
@@ -132,7 +132,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 3,
+      "execution_count": 4,
       "metadata": {},
       "outputs": [],
       "source": [
@@ -142,7 +142,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 4,
+      "execution_count": 5,
       "metadata": {},
       "outputs": [
         {
@@ -174,7 +174,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 5,
+      "execution_count": 6,
       "metadata": {},
       "outputs": [],
       "source": [
@@ -183,7 +183,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 6,
+      "execution_count": 7,
       "metadata": {},
       "outputs": [],
       "source": [
@@ -192,7 +192,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 7,
+      "execution_count": 8,
       "metadata": {},
       "outputs": [
         {
@@ -209,7 +209,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 8,
+      "execution_count": 9,
       "metadata": {},
       "outputs": [
         {
@@ -239,7 +239,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 9,
+      "execution_count": 10,
       "metadata": {},
       "outputs": [
         {
@@ -357,7 +357,7 @@
   ],
   "metadata": {
     "kernelspec": {
-      "display_name": ".venv",
+      "display_name": "Python 3 (ipykernel)",
       "language": "python",
       "name": "python3"
     },
@@ -371,7 +371,7 @@
       "name": "python",
       "nbconvert_exporter": "python",
       "pygments_lexer": "ipython3",
-      "version": "3.12.6"
+      "version": "3.12.7"
     }
   },
   "nbformat": 4,
diff --git a/stormvogel/pgc.py b/stormvogel/pgc.py
new file mode 100644
index 0000000..486e4bc
--- /dev/null
+++ b/stormvogel/pgc.py
@@ -0,0 +1,186 @@
+import stormvogel.model
+from dataclasses import dataclass
+from typing import Callable
+import math
+
+
+@dataclass
+class State:
+    x: int
+
+    def __hash__(self):
+        return hash(self.x)
+
+    def __eq__(self, other):
+        if isinstance(other, State):
+            return self.x == other.x
+        return False
+
+
+@dataclass
+class Action:
+    labels: list[str]
+
+
+def build_pgc(
+    delta,  # Callable[[State, Action], list[tuple[float, State]]],
+    initial_state_pgc: State,
+    available_actions: Callable[[State], list[Action]],
+    modeltype: stormvogel.model.ModelType = stormvogel.model.ModelType.MDP,
+) -> stormvogel.model.Model:
+    """
+    function that converts a delta function, an available_actions function an initial state and a model type
+    to a stormvogel model
+
+    this works analogous to a prism file, where the delta is the module in this case.
+
+    (this function uses the pgc classes State and Action instead of the ones from stormvogel.model)
+
+    (currently only works for mdps)
+    """
+
+    model = stormvogel.model.new_model(modeltype=modeltype, create_initial_state=False)
+
+    # we create the model with the given type and initial state
+    model.new_state(
+        labels=["init", str(initial_state_pgc.x)], features={"x": initial_state_pgc.x}
+    )
+
+    # we continue calling delta and adding new states until no states are
+    # left to be checked
+    states_discovered = []
+    states_to_be_discovered = [initial_state_pgc]
+    while len(states_to_be_discovered) > 0:
+        state = states_to_be_discovered[0]
+        states_to_be_discovered.remove(state)
+        states_discovered.append(state)
+        # we loop over all available actions and call the delta function for each action
+        transition = {}
+        for action in available_actions(state):
+            try:
+                stormvogel_action = model.new_action(
+                    str(action.labels),
+                    frozenset(
+                        {action.labels[0]}
+                    ),  # right now we only look at one label
+                )
+            except RuntimeError:
+                stormvogel_action = model.get_action(str(action.labels))
+
+            tuples = delta(state, action)
+            # we add all the newly found transitions to the model (including the action)
+            branch = []
+            for tuple in tuples:
+                if tuple[1] not in states_discovered:
+                    new_state = model.new_state(
+                        labels=[str(tuple[1].x)], features={"x": tuple[1].x}
+                    )
+                    branch.append((tuple[0], new_state))
+                    states_to_be_discovered.append(tuple[1])
+                else:
+                    # TODO what if there are multiple states with the same label? use names?
+                    branch.append(
+                        (tuple[0], model.get_states_with_label(str(tuple[1].x))[0])
+                    )
+            if branch != []:
+                transition[stormvogel_action] = stormvogel.model.Branch(branch)
+        model.add_transitions(
+            model.get_states_with_label(str(state.x))[0],
+            stormvogel.model.Transition(transition),
+        )
+
+    return model
+
+
+"""
+def build(
+    delta: Callable[
+        [stormvogel.model.State, stormvogel.model.Action],
+        list[tuple[float, stormvogel.model.State]],
+    ],
+    available_actions: Callable[
+        [stormvogel.model.State], list[stormvogel.model.Action]
+    ],
+    initial_value: str,
+    modeltype: stormvogel.model.ModelType | None = stormvogel.model.ModelType.MDP,
+) -> stormvogel.model.Model:
+
+    #function that converts a delta function, an available_actions function an initial state and a model type
+    #to a stormvogel model
+
+    #this works analogous to a prism file, where the delta is the module in this case.
+
+    # we create the model with the given type and initial state
+    model = stormvogel.model.new_model(modeltype, create_initial_state=False)
+    model.new_state(features=[initial_state.features])
+
+    # we continue calling delta and adding new states until no states are
+    # left to be checked
+    states = [initial_state]
+    for state in states:
+        # we loop over all available actions and call the delta function for each action
+        transition = {}
+        for action in available_actions(state):
+            tuples = delta(state, action)
+            # we add all the newly found transitions to the model (including the action)
+            branch = []
+            for tuple in tuples:
+                branch.append(tuple)
+                if tuple[1] not in states:
+                    states.add(tuple[1])
+            transition[action] = branch
+        model.add_transitions(state, stormvogel.model.Transition(transition))
+
+    return model
+"""
+
+if __name__ == "__main__":
+    N = 100
+    p = 0.5
+    initial_state = State(math.floor(N / 2))
+
+    left = Action(["left"])
+    right = Action(["right"])
+
+    def available_actions(s: State) -> list[Action]:
+        return [left, right]
+
+    def delta(s: State, action: Action):
+        if action == left:
+            return [(p, State(x=s.x + 1)), (1 - p, State(x=s.x))] if s.x < N else []
+        elif action == right:
+            return [(p, State(x=s.x - 1)), (1 - p, State(x=s.x))] if s.x > 0 else []
+
+    model = build_pgc(
+        delta=delta,
+        available_actions=available_actions,
+        initial_state_pgc=initial_state,
+    )
+
+    print(model)
+
+    # Do we also want to be able to do it with the stormvogel.model objects?
+    """
+    N = 10
+    p = 0.5
+
+    initial_value = 5
+
+    left = stormvogel.model.Action("",["left"])
+    right = stormvogel.modle.Action("",["right"])
+
+    def available_actions(s: stormvogel.model.State) -> list[stormvogel.model.Action]:
+        return [left, right]
+
+    def delta(s: stormvogel.model.State, action: stormvogel.model.Action) -> stormvogel.model.State:
+        if action == left:
+            return [(p, stormvogel.model.State(features={x:s.x + 1})), (1 - p, stormvogel.model.State(features=["s.x"]))]
+        elif action == right:
+            return [(p, stormvogel.model.State(features=["s.x - 1"])), (1 - p, stormvogel.model.State(features=["s.x"]))]
+
+
+
+    model = build(delta, available_actions, initial_value, modeltype)
+
+    print(model)
+    """
diff --git a/tests/test_pgc.py b/tests/test_pgc.py
new file mode 100644
index 0000000..e69de29

From 17a443c51920f23d47f5ea4e0109f4a3307f1f8f Mon Sep 17 00:00:00 2001
From: PimLeerkes <pimsimon.leerkes@gmail.com>
Date: Sat, 30 Nov 2024 11:55:58 +0100
Subject: [PATCH 2/8] wrote first test for mdp pgc

---
 tests/test_pgc.py | 56 +++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 56 insertions(+)

diff --git a/tests/test_pgc.py b/tests/test_pgc.py
index e69de29..6a276a0 100644
--- a/tests/test_pgc.py
+++ b/tests/test_pgc.py
@@ -0,0 +1,56 @@
+from stormvogel import pgc
+import stormvogel.model
+import math
+
+
+def test_pgc_mdp():
+    # we build the model with pgc:
+    N = 2
+    p = 0.5
+    initial_state = pgc.State(math.floor(N / 2))
+
+    left = pgc.Action(["left"])
+    right = pgc.Action(["right"])
+
+    def available_actions(s: pgc.State):
+        return [left, right]
+
+    def delta(s: pgc.State, action: pgc.Action):
+        if action == left:
+            return (
+                [(p, pgc.State(x=s.x + 1)), (1 - p, pgc.State(x=s.x))]
+                if s.x < N
+                else []
+            )
+        elif action == right:
+            return (
+                [(p, pgc.State(x=s.x - 1)), (1 - p, pgc.State(x=s.x))]
+                if s.x > 0
+                else []
+            )
+
+    pgc_model = stormvogel.pgc.build_pgc(
+        delta=delta,
+        available_actions=available_actions,
+        initial_state_pgc=initial_state,
+    )
+
+    # we build the model in the regular way:
+    model = stormvogel.model.new_mdp(create_initial_state=False)
+    state1 = model.new_state(labels=["init", "1"], features={"x": 1})
+    state2 = model.new_state(labels=["2"], features={"x": 2})
+    state0 = model.new_state(labels=["0"], features={"0": 0})
+    left = model.new_action("left", frozenset({"left"}))
+    right = model.new_action("right", frozenset({"right"}))
+    branch11 = stormvogel.model.Branch([(0.5, state1), (0.5, state2)])
+    branch12 = stormvogel.model.Branch([(0.5, state1), (0.5, state0)])
+    branch0 = stormvogel.model.Branch([(0.5, state0), (0.5, state1)])
+    branch2 = stormvogel.model.Branch([(0.5, state2), (0.5, state1)])
+
+    model.add_transitions(
+        state1, stormvogel.model.Transition({left: branch11, right: branch12})
+    )
+    model.add_transitions(state2, stormvogel.model.Transition({right: branch2}))
+    model.add_transitions(state0, stormvogel.model.Transition({left: branch0}))
+
+    assert model == pgc_model

From efd58f9fd0a28a27215fc4df5067d43adf7d608c Mon Sep 17 00:00:00 2001
From: PimLeerkes <pimsimon.leerkes@gmail.com>
Date: Sat, 30 Nov 2024 16:45:30 +0100
Subject: [PATCH 3/8] it works for dtmcs + changed how labels and features work
 in pgc

---
 stormvogel/model.py |   7 +-
 stormvogel/pgc.py   | 189 +++++++++++++++-----------------------------
 tests/test_pgc.py   |  77 ++++++++++++++++--
 3 files changed, 136 insertions(+), 137 deletions(-)

diff --git a/stormvogel/model.py b/stormvogel/model.py
index 5d59002..b5570d2 100644
--- a/stormvogel/model.py
+++ b/stormvogel/model.py
@@ -825,15 +825,16 @@ def new_state(
         self,
         labels: list[str] | str | None = None,
         features: dict[str, int] | None = None,
+        name: str | None = None,
     ) -> State:
         """Creates a new state and returns it."""
         state_id = self.__free_state_id()
         if isinstance(labels, list):
-            state = State(labels, features or {}, state_id, self)
+            state = State(labels, features or {}, state_id, self, name=name)
         elif isinstance(labels, str):
-            state = State([labels], features or {}, state_id, self)
+            state = State([labels], features or {}, state_id, self, name=name)
         elif labels is None:
-            state = State([], features or {}, state_id, self)
+            state = State([], features or {}, state_id, self, name=name)
 
         self.states[state_id] = state
 
diff --git a/stormvogel/pgc.py b/stormvogel/pgc.py
index 486e4bc..2d5ec1e 100644
--- a/stormvogel/pgc.py
+++ b/stormvogel/pgc.py
@@ -1,19 +1,18 @@
 import stormvogel.model
 from dataclasses import dataclass
 from typing import Callable
-import math
 
 
 @dataclass
 class State:
-    x: int
+    f: dict[str, int]
 
     def __hash__(self):
         return hash(self.x)
 
     def __eq__(self, other):
         if isinstance(other, State):
-            return self.x == other.x
+            return self.f == other.f
         return False
 
 
@@ -25,7 +24,7 @@ class Action:
 def build_pgc(
     delta,  # Callable[[State, Action], list[tuple[float, State]]],
     initial_state_pgc: State,
-    available_actions: Callable[[State], list[Action]],
+    available_actions: Callable[[State], list[Action]] | None = None,
     modeltype: stormvogel.model.ModelType = stormvogel.model.ModelType.MDP,
 ) -> stormvogel.model.Model:
     """
@@ -38,149 +37,87 @@ def build_pgc(
 
     (currently only works for mdps)
     """
+    if modeltype == stormvogel.model.ModelType.MDP and available_actions is None:
+        raise RuntimeError(
+            "You have to provide an available actions function for mdp models"
+        )
 
     model = stormvogel.model.new_model(modeltype=modeltype, create_initial_state=False)
 
     # we create the model with the given type and initial state
     model.new_state(
-        labels=["init", str(initial_state_pgc.x)], features={"x": initial_state_pgc.x}
+        labels=["init"], features=initial_state_pgc.f, name=str(initial_state_pgc.f)
     )
 
     # we continue calling delta and adding new states until no states are
     # left to be checked
-    states_discovered = []
-    states_to_be_discovered = [initial_state_pgc]
-    while len(states_to_be_discovered) > 0:
-        state = states_to_be_discovered[0]
-        states_to_be_discovered.remove(state)
-        states_discovered.append(state)
+    states_seen = []
+    states_to_be_visited = [initial_state_pgc]
+    while len(states_to_be_visited) > 0:
+        state = states_to_be_visited[0]
+        states_to_be_visited.remove(state)
         # we loop over all available actions and call the delta function for each action
         transition = {}
-        for action in available_actions(state):
-            try:
-                stormvogel_action = model.new_action(
-                    str(action.labels),
-                    frozenset(
-                        {action.labels[0]}
-                    ),  # right now we only look at one label
-                )
-            except RuntimeError:
-                stormvogel_action = model.get_action(str(action.labels))
-
-            tuples = delta(state, action)
-            # we add all the newly found transitions to the model (including the action)
+
+        if state not in states_seen:
+            states_seen.append(state)
+
+        if model.supports_actions():
+            assert available_actions is not None
+            for action in available_actions(state):
+                try:
+                    stormvogel_action = model.new_action(
+                        str(action.labels),
+                        frozenset(
+                            {action.labels[0]}
+                        ),  # right now we only look at one label
+                    )
+                except RuntimeError:
+                    stormvogel_action = model.get_action(str(action.labels))
+
+                tuples = delta(state, action)
+                # we add all the newly found transitions to the model
+                branch = []
+                for tuple in tuples:
+                    if tuple[1] not in states_seen:
+                        states_seen.append(tuple[1])
+                        new_state = model.new_state(
+                            name=str(tuple[1].f), features=tuple[1].f
+                        )
+                        branch.append((tuple[0], new_state))
+                        states_to_be_visited.append(tuple[1])
+                    else:
+                        # TODO what if there are multiple states with the same label? use names?
+                        branch.append(
+                            (tuple[0], model.get_state_by_name(str(tuple[1].f)))
+                        )
+                if branch != []:
+                    transition[stormvogel_action] = stormvogel.model.Branch(branch)
+        else:
+            tuples = delta(state)
+            # we add all the newly found transitions to the model
             branch = []
             for tuple in tuples:
-                if tuple[1] not in states_discovered:
+                if tuple[1] not in states_seen:
+                    states_seen.append(tuple[1])
                     new_state = model.new_state(
-                        labels=[str(tuple[1].x)], features={"x": tuple[1].x}
+                        name=str(tuple[1].f), features=tuple[1].f
                     )
+
                     branch.append((tuple[0], new_state))
-                    states_to_be_discovered.append(tuple[1])
+                    states_to_be_visited.append(tuple[1])
                 else:
                     # TODO what if there are multiple states with the same label? use names?
-                    branch.append(
-                        (tuple[0], model.get_states_with_label(str(tuple[1].x))[0])
+                    branch.append((tuple[0], model.get_state_by_name(str(tuple[1].f))))
+                if branch != []:
+                    transition[stormvogel.model.EmptyAction] = stormvogel.model.Branch(
+                        branch
                     )
-            if branch != []:
-                transition[stormvogel_action] = stormvogel.model.Branch(branch)
+        s = model.get_state_by_name(str(state.f))
+        assert s is not None
         model.add_transitions(
-            model.get_states_with_label(str(state.x))[0],
+            s,
             stormvogel.model.Transition(transition),
         )
 
     return model
-
-
-"""
-def build(
-    delta: Callable[
-        [stormvogel.model.State, stormvogel.model.Action],
-        list[tuple[float, stormvogel.model.State]],
-    ],
-    available_actions: Callable[
-        [stormvogel.model.State], list[stormvogel.model.Action]
-    ],
-    initial_value: str,
-    modeltype: stormvogel.model.ModelType | None = stormvogel.model.ModelType.MDP,
-) -> stormvogel.model.Model:
-
-    #function that converts a delta function, an available_actions function an initial state and a model type
-    #to a stormvogel model
-
-    #this works analogous to a prism file, where the delta is the module in this case.
-
-    # we create the model with the given type and initial state
-    model = stormvogel.model.new_model(modeltype, create_initial_state=False)
-    model.new_state(features=[initial_state.features])
-
-    # we continue calling delta and adding new states until no states are
-    # left to be checked
-    states = [initial_state]
-    for state in states:
-        # we loop over all available actions and call the delta function for each action
-        transition = {}
-        for action in available_actions(state):
-            tuples = delta(state, action)
-            # we add all the newly found transitions to the model (including the action)
-            branch = []
-            for tuple in tuples:
-                branch.append(tuple)
-                if tuple[1] not in states:
-                    states.add(tuple[1])
-            transition[action] = branch
-        model.add_transitions(state, stormvogel.model.Transition(transition))
-
-    return model
-"""
-
-if __name__ == "__main__":
-    N = 100
-    p = 0.5
-    initial_state = State(math.floor(N / 2))
-
-    left = Action(["left"])
-    right = Action(["right"])
-
-    def available_actions(s: State) -> list[Action]:
-        return [left, right]
-
-    def delta(s: State, action: Action):
-        if action == left:
-            return [(p, State(x=s.x + 1)), (1 - p, State(x=s.x))] if s.x < N else []
-        elif action == right:
-            return [(p, State(x=s.x - 1)), (1 - p, State(x=s.x))] if s.x > 0 else []
-
-    model = build_pgc(
-        delta=delta,
-        available_actions=available_actions,
-        initial_state_pgc=initial_state,
-    )
-
-    print(model)
-
-    # Do we also want to be able to do it with the stormvogel.model objects?
-    """
-    N = 10
-    p = 0.5
-
-    initial_value = 5
-
-    left = stormvogel.model.Action("",["left"])
-    right = stormvogel.modle.Action("",["right"])
-
-    def available_actions(s: stormvogel.model.State) -> list[stormvogel.model.Action]:
-        return [left, right]
-
-    def delta(s: stormvogel.model.State, action: stormvogel.model.Action) -> stormvogel.model.State:
-        if action == left:
-            return [(p, stormvogel.model.State(features={x:s.x + 1})), (1 - p, stormvogel.model.State(features=["s.x"]))]
-        elif action == right:
-            return [(p, stormvogel.model.State(features=["s.x - 1"])), (1 - p, stormvogel.model.State(features=["s.x"]))]
-
-
-
-    model = build(delta, available_actions, initial_value, modeltype)
-
-    print(model)
-    """
diff --git a/tests/test_pgc.py b/tests/test_pgc.py
index 6a276a0..2304052 100644
--- a/tests/test_pgc.py
+++ b/tests/test_pgc.py
@@ -1,13 +1,15 @@
 from stormvogel import pgc
 import stormvogel.model
 import math
+import stormpy
+import stormvogel.mapping
 
 
 def test_pgc_mdp():
     # we build the model with pgc:
     N = 2
     p = 0.5
-    initial_state = pgc.State(math.floor(N / 2))
+    initial_state = pgc.State({"x": math.floor(N / 2)})
 
     left = pgc.Action(["left"])
     right = pgc.Action(["right"])
@@ -17,15 +19,23 @@ def available_actions(s: pgc.State):
 
     def delta(s: pgc.State, action: pgc.Action):
         if action == left:
+            print(s.f)
+            print(s.f["x"])
             return (
-                [(p, pgc.State(x=s.x + 1)), (1 - p, pgc.State(x=s.x))]
-                if s.x < N
+                [
+                    (p, pgc.State({"x": s.f["x"] + 1})),
+                    (1 - p, pgc.State({"x": s.f["x"]})),
+                ]
+                if s.f["x"] < N
                 else []
             )
         elif action == right:
             return (
-                [(p, pgc.State(x=s.x - 1)), (1 - p, pgc.State(x=s.x))]
-                if s.x > 0
+                [
+                    (p, pgc.State({"x": s.f["x"] - 1})),
+                    (1 - p, pgc.State({"x": s.f["x"]})),
+                ]
+                if s.f["x"] > 0
                 else []
             )
 
@@ -37,9 +47,9 @@ def delta(s: pgc.State, action: pgc.Action):
 
     # we build the model in the regular way:
     model = stormvogel.model.new_mdp(create_initial_state=False)
-    state1 = model.new_state(labels=["init", "1"], features={"x": 1})
-    state2 = model.new_state(labels=["2"], features={"x": 2})
-    state0 = model.new_state(labels=["0"], features={"0": 0})
+    state1 = model.new_state(labels=["init"], features={"x": 1})
+    state2 = model.new_state(features={"x": 2})
+    state0 = model.new_state(features={"0": 0})
     left = model.new_action("left", frozenset({"left"}))
     right = model.new_action("right", frozenset({"right"}))
     branch11 = stormvogel.model.Branch([(0.5, state1), (0.5, state2)])
@@ -53,4 +63,55 @@ def delta(s: pgc.State, action: pgc.Action):
     model.add_transitions(state2, stormvogel.model.Transition({right: branch2}))
     model.add_transitions(state0, stormvogel.model.Transition({left: branch0}))
 
+    print(model)
+    print(pgc_model)
+
     assert model == pgc_model
+
+
+def test_pgc_dtmc():
+    # we build the model with pgc:
+    p = 0.5
+    initial_state = pgc.State({"s": 0})
+
+    def delta(s: pgc.State):
+        match s.f["s"]:
+            case 0:
+                return [(p, pgc.State({"s": 1})), (1 - p, pgc.State({"s": 2}))]
+            case 1:
+                return [(p, pgc.State({"s": 3})), (1 - p, pgc.State({"s": 4}))]
+            case 2:
+                return [(p, pgc.State({"s": 5})), (1 - p, pgc.State({"s": 6}))]
+            case 3:
+                return [(p, pgc.State({"s": 1})), (1 - p, pgc.State({"s": 7, "d": 1}))]
+            case 4:
+                return [
+                    (p, pgc.State({"s": 7, "d": 2})),
+                    (1 - p, pgc.State({"s": 7, "d": 3})),
+                ]
+            case 5:
+                return [
+                    (p, pgc.State({"s": 7, "d": 4})),
+                    (1 - p, pgc.State({"s": 7, "d": 5})),
+                ]
+            case 6:
+                return [(p, pgc.State({"s": 2})), (1 - p, pgc.State({"s": 7, "d": 6}))]
+            case 7:
+                return [(1, pgc.State({"s": 7}))]
+
+    pgc_model = stormvogel.pgc.build_pgc(
+        delta=delta,
+        initial_state_pgc=initial_state,
+        modeltype=stormvogel.model.ModelType.DTMC,
+    )
+
+    # we build the model in the regular way:
+    path = stormpy.examples.files.prism_dtmc_die
+    prism_program = stormpy.parse_prism_program(path)
+    formula_str = "P=? [F s=7 & d=2]"
+    properties = stormpy.parse_properties(formula_str, prism_program)
+    model = stormpy.build_model(prism_program, properties)
+    print(dir(model.states[0]))
+    stormvogel_model = stormvogel.mapping.stormpy_to_stormvogel(model)
+    print(pgc_model)
+    print(stormvogel_model)

From 9a5d26768fb85a97274d35748937e4fc13bf4dd2 Mon Sep 17 00:00:00 2001
From: PimLeerkes <pimsimon.leerkes@gmail.com>
Date: Sun, 8 Dec 2024 12:38:20 +0100
Subject: [PATCH 4/8] pgc state object can now contain arbitrary attributes

---
 stormvogel/model.py |   2 +-
 stormvogel/pgc.py   |  47 +++++++++++-------
 tests/test_pgc.py   | 115 +++++++++++++++++++++++++++++++-------------
 3 files changed, 112 insertions(+), 52 deletions(-)

diff --git a/stormvogel/model.py b/stormvogel/model.py
index b5570d2..a9daa30 100644
--- a/stormvogel/model.py
+++ b/stormvogel/model.py
@@ -205,7 +205,7 @@ class Action:
         labels: The labels of this action. Corresponds to Storm labels.
     """
 
-    name: str
+    name: str  # TODO name is stormpy choice label or we don't need a name at all?
     labels: frozenset[str]
 
     def __str__(self):
diff --git a/stormvogel/pgc.py b/stormvogel/pgc.py
index 2d5ec1e..68c80dd 100644
--- a/stormvogel/pgc.py
+++ b/stormvogel/pgc.py
@@ -3,27 +3,36 @@
 from typing import Callable
 
 
+@dataclass
+class Action:
+    """pgc action object. Contains a list of labels"""
+
+    labels: list[str]
+
+
 @dataclass
 class State:
-    f: dict[str, int]
+    """pgc state object. Can contain any number of any type of arguments"""
+
+    def __init__(self, **kwargs):
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+
+    def __repr__(self):
+        return f"State({self.__dict__})"
 
     def __hash__(self):
-        return hash(self.x)
+        return hash(self.__dict__)
 
     def __eq__(self, other):
         if isinstance(other, State):
-            return self.f == other.f
+            return self.__dict__ == other.__dict__
         return False
 
 
-@dataclass
-class Action:
-    labels: list[str]
-
-
 def build_pgc(
     delta,  # Callable[[State, Action], list[tuple[float, State]]],
-    initial_state_pgc: State,
+    initial_state_pgc: State,  # TODO rewards function, label function
     available_actions: Callable[[State], list[Action]] | None = None,
     modeltype: stormvogel.model.ModelType = stormvogel.model.ModelType.MDP,
 ) -> stormvogel.model.Model:
@@ -46,7 +55,9 @@ def build_pgc(
 
     # we create the model with the given type and initial state
     model.new_state(
-        labels=["init"], features=initial_state_pgc.f, name=str(initial_state_pgc.f)
+        labels=["init"],
+        features=initial_state_pgc.__dict__,
+        name=str(initial_state_pgc.__dict__),
     )
 
     # we continue calling delta and adding new states until no states are
@@ -82,14 +93,15 @@ def build_pgc(
                     if tuple[1] not in states_seen:
                         states_seen.append(tuple[1])
                         new_state = model.new_state(
-                            name=str(tuple[1].f), features=tuple[1].f
+                            name=str(tuple[1].__dict__), features=tuple[1].__dict__
                         )
                         branch.append((tuple[0], new_state))
                         states_to_be_visited.append(tuple[1])
                     else:
-                        # TODO what if there are multiple states with the same label? use names?
+                        # print(tuple[1].__dict__)
+                        # print(model.states)
                         branch.append(
-                            (tuple[0], model.get_state_by_name(str(tuple[1].f)))
+                            (tuple[0], model.get_state_by_name(str(tuple[1].__dict__)))
                         )
                 if branch != []:
                     transition[stormvogel_action] = stormvogel.model.Branch(branch)
@@ -101,19 +113,20 @@ def build_pgc(
                 if tuple[1] not in states_seen:
                     states_seen.append(tuple[1])
                     new_state = model.new_state(
-                        name=str(tuple[1].f), features=tuple[1].f
+                        name=str(tuple[1].__dict__), features=tuple[1].__dict__
                     )
 
                     branch.append((tuple[0], new_state))
                     states_to_be_visited.append(tuple[1])
                 else:
-                    # TODO what if there are multiple states with the same label? use names?
-                    branch.append((tuple[0], model.get_state_by_name(str(tuple[1].f))))
+                    branch.append(
+                        (tuple[0], model.get_state_by_name(str(tuple[1].__dict__)))
+                    )
                 if branch != []:
                     transition[stormvogel.model.EmptyAction] = stormvogel.model.Branch(
                         branch
                     )
-        s = model.get_state_by_name(str(state.f))
+        s = model.get_state_by_name(str(state.__dict__))
         assert s is not None
         model.add_transitions(
             s,
diff --git a/tests/test_pgc.py b/tests/test_pgc.py
index 2304052..107540c 100644
--- a/tests/test_pgc.py
+++ b/tests/test_pgc.py
@@ -1,7 +1,6 @@
 from stormvogel import pgc
 import stormvogel.model
 import math
-import stormpy
 import stormvogel.mapping
 
 
@@ -9,7 +8,7 @@ def test_pgc_mdp():
     # we build the model with pgc:
     N = 2
     p = 0.5
-    initial_state = pgc.State({"x": math.floor(N / 2)})
+    initial_state = pgc.State(x=math.floor(N / 2))
 
     left = pgc.Action(["left"])
     right = pgc.Action(["right"])
@@ -19,23 +18,21 @@ def available_actions(s: pgc.State):
 
     def delta(s: pgc.State, action: pgc.Action):
         if action == left:
-            print(s.f)
-            print(s.f["x"])
             return (
                 [
-                    (p, pgc.State({"x": s.f["x"] + 1})),
-                    (1 - p, pgc.State({"x": s.f["x"]})),
+                    (p, pgc.State(x=s.x + 1)),
+                    (1 - p, pgc.State(x=s.x)),
                 ]
-                if s.f["x"] < N
+                if s.x < N
                 else []
             )
         elif action == right:
             return (
                 [
-                    (p, pgc.State({"x": s.f["x"] - 1})),
-                    (1 - p, pgc.State({"x": s.f["x"]})),
+                    (p, pgc.State(x=s.x - 1)),
+                    (1 - p, pgc.State(x=s.x)),
                 ]
-                if s.f["x"] > 0
+                if s.x > 0
                 else []
             )
 
@@ -63,41 +60,38 @@ def delta(s: pgc.State, action: pgc.Action):
     model.add_transitions(state2, stormvogel.model.Transition({right: branch2}))
     model.add_transitions(state0, stormvogel.model.Transition({left: branch0}))
 
-    print(model)
-    print(pgc_model)
-
     assert model == pgc_model
 
 
 def test_pgc_dtmc():
     # we build the model with pgc:
     p = 0.5
-    initial_state = pgc.State({"s": 0})
+    initial_state = pgc.State(s=0)
 
     def delta(s: pgc.State):
-        match s.f["s"]:
+        match s.s:
             case 0:
-                return [(p, pgc.State({"s": 1})), (1 - p, pgc.State({"s": 2}))]
+                return [(p, pgc.State(s=1)), (1 - p, pgc.State(s=2))]
             case 1:
-                return [(p, pgc.State({"s": 3})), (1 - p, pgc.State({"s": 4}))]
+                return [(p, pgc.State(s=3)), (1 - p, pgc.State(s=4))]
             case 2:
-                return [(p, pgc.State({"s": 5})), (1 - p, pgc.State({"s": 6}))]
+                return [(p, pgc.State(s=5)), (1 - p, pgc.State(s=6))]
             case 3:
-                return [(p, pgc.State({"s": 1})), (1 - p, pgc.State({"s": 7, "d": 1}))]
+                return [(p, pgc.State(s=1)), (1 - p, pgc.State(s=7, d=1))]
             case 4:
                 return [
-                    (p, pgc.State({"s": 7, "d": 2})),
-                    (1 - p, pgc.State({"s": 7, "d": 3})),
+                    (p, pgc.State(s=7, d=2)),
+                    (1 - p, pgc.State(s=7, d=3)),
                 ]
             case 5:
                 return [
-                    (p, pgc.State({"s": 7, "d": 4})),
-                    (1 - p, pgc.State({"s": 7, "d": 5})),
+                    (p, pgc.State(s=7, d=4)),
+                    (1 - p, pgc.State(s=7, d=5)),
                 ]
             case 6:
-                return [(p, pgc.State({"s": 2})), (1 - p, pgc.State({"s": 7, "d": 6}))]
+                return [(p, pgc.State(s=2)), (1 - p, pgc.State(s=7, d=6))]
             case 7:
-                return [(1, pgc.State({"s": 7}))]
+                return [(1, pgc.State(s=7))]
 
     pgc_model = stormvogel.pgc.build_pgc(
         delta=delta,
@@ -106,12 +100,65 @@ def delta(s: pgc.State):
     )
 
     # we build the model in the regular way:
-    path = stormpy.examples.files.prism_dtmc_die
-    prism_program = stormpy.parse_prism_program(path)
-    formula_str = "P=? [F s=7 & d=2]"
-    properties = stormpy.parse_properties(formula_str, prism_program)
-    model = stormpy.build_model(prism_program, properties)
-    print(dir(model.states[0]))
-    stormvogel_model = stormvogel.mapping.stormpy_to_stormvogel(model)
-    print(pgc_model)
-    print(stormvogel_model)
+    model = stormvogel.model.new_dtmc()
+    model.states[0].features = {"s": 0}
+    model.set_transitions(
+        model.get_initial_state(),
+        [
+            (1 / 2, model.new_state(features={"s": 1})),
+            (1 / 2, model.new_state(features={"s": 2})),
+        ],
+    )
+    model.set_transitions(
+        model.get_state_by_id(1),
+        [
+            (1 / 2, model.new_state(features={"s": 3})),
+            (1 / 2, model.new_state(features={"s": 4})),
+        ],
+    )
+    model.set_transitions(
+        model.get_state_by_id(2),
+        [
+            (1 / 2, model.new_state(features={"s": 5})),
+            (1 / 2, model.new_state(features={"s": 6})),
+        ],
+    )
+    model.set_transitions(
+        model.get_state_by_id(3),
+        [
+            (1 / 2, model.get_state_by_id(1)),
+            (1 / 2, model.new_state(features={"s": 7, "d": 1})),
+        ],
+    )
+    model.set_transitions(
+        model.get_state_by_id(4),
+        [
+            (1 / 2, model.new_state(features={"s": 7, "d": 2})),
+            (1 / 2, model.new_state(features={"s": 7, "d": 3})),
+        ],
+    )
+    model.set_transitions(
+        model.get_state_by_id(5),
+        [
+            (1 / 2, model.new_state(features={"s": 7, "d": 4})),
+            (1 / 2, model.new_state(features={"s": 7, "d": 5})),
+        ],
+    )
+    model.set_transitions(
+        model.get_state_by_id(6),
+        [
+            (1 / 2, model.get_state_by_id(2)),
+            (1 / 2, model.new_state(features={"s": 7, "d": 6})),
+        ],
+    )
+    model.set_transitions(
+        model.get_state_by_id(7), [(1, model.new_state(features={"s": 7}))]
+    )
+    model.set_transitions(model.get_state_by_id(8), [(1, model.get_state_by_id(13))])
+    model.set_transitions(model.get_state_by_id(9), [(1, model.get_state_by_id(13))])
+    model.set_transitions(model.get_state_by_id(10), [(1, model.get_state_by_id(13))])
+    model.set_transitions(model.get_state_by_id(11), [(1, model.get_state_by_id(13))])
+    model.set_transitions(model.get_state_by_id(12), [(1, model.get_state_by_id(13))])
+    model.set_transitions(model.get_state_by_id(13), [(1, model.get_state_by_id(13))])
+
+    assert pgc_model == model

From 072c6cb7710827633789087b243fc8d6be25d89f Mon Sep 17 00:00:00 2001
From: PimLeerkes <pimsimon.leerkes@gmail.com>
Date: Sun, 8 Dec 2024 14:07:07 +0100
Subject: [PATCH 5/8] model builder now also uses reward function

---
 stormvogel/model.py |  2 ++
 stormvogel/pgc.py   | 31 ++++++++++++++++++++++++++++---
 tests/test_pgc.py   | 16 ++++++++++++++++
 3 files changed, 46 insertions(+), 3 deletions(-)

diff --git a/stormvogel/model.py b/stormvogel/model.py
index a9daa30..2a17a9e 100644
--- a/stormvogel/model.py
+++ b/stormvogel/model.py
@@ -804,6 +804,8 @@ def get_action(self, name: str) -> Action:
             )
         assert self.actions is not None
         if name not in self.actions:
+            print(name)
+            print(self.actions)
             raise RuntimeError(
                 f"Tried to get action {name} but that action does not exist"
             )
diff --git a/stormvogel/pgc.py b/stormvogel/pgc.py
index 68c80dd..4588c21 100644
--- a/stormvogel/pgc.py
+++ b/stormvogel/pgc.py
@@ -33,6 +33,8 @@ def __eq__(self, other):
 def build_pgc(
     delta,  # Callable[[State, Action], list[tuple[float, State]]],
     initial_state_pgc: State,  # TODO rewards function, label function
+    rewards=None,
+    labels=None,
     available_actions: Callable[[State], list[Action]] | None = None,
     modeltype: stormvogel.model.ModelType = stormvogel.model.ModelType.MDP,
 ) -> stormvogel.model.Model:
@@ -67,13 +69,13 @@ def build_pgc(
     while len(states_to_be_visited) > 0:
         state = states_to_be_visited[0]
         states_to_be_visited.remove(state)
-        # we loop over all available actions and call the delta function for each action
         transition = {}
 
         if state not in states_seen:
             states_seen.append(state)
 
         if model.supports_actions():
+            # we loop over all available actions and call the delta function for each action
             assert available_actions is not None
             for action in available_actions(state):
                 try:
@@ -98,8 +100,6 @@ def build_pgc(
                         branch.append((tuple[0], new_state))
                         states_to_be_visited.append(tuple[1])
                     else:
-                        # print(tuple[1].__dict__)
-                        # print(model.states)
                         branch.append(
                             (tuple[0], model.get_state_by_name(str(tuple[1].__dict__)))
                         )
@@ -133,4 +133,29 @@ def build_pgc(
             stormvogel.model.Transition(transition),
         )
 
+    # we add the rewards
+    # TODO support multiple reward models
+    if rewards is not None:
+        rewardmodel = model.add_rewards("rewards")
+        if model.supports_actions():
+            for state in states_seen:
+                assert available_actions is not None
+                for action in available_actions(state):
+                    reward = rewards(state, action)
+                    s = model.get_state_by_name(str(state.__dict__))
+                    assert s is not None
+                    rewardmodel.set_state_action_reward(
+                        s,
+                        model.get_action(str(action.labels)),
+                        reward,
+                    )
+        else:
+            for state in states_seen:
+                reward = rewards(state)
+                s = model.get_state_by_name(str(state.__dict__))
+                assert s is not None
+                rewardmodel.set_state_reward(s, reward)
+
+    # we add the labels
+
     return model
diff --git a/tests/test_pgc.py b/tests/test_pgc.py
index 107540c..b3b79ed 100644
--- a/tests/test_pgc.py
+++ b/tests/test_pgc.py
@@ -16,6 +16,9 @@ def test_pgc_mdp():
     def available_actions(s: pgc.State):
         return [left, right]
 
+    def rewards(s: pgc.State, a: pgc.Action):
+        return 1
+
     def delta(s: pgc.State, action: pgc.Action):
         if action == left:
             return (
@@ -40,6 +43,7 @@ def delta(s: pgc.State, action: pgc.Action):
         delta=delta,
         available_actions=available_actions,
         initial_state_pgc=initial_state,
+        rewards=rewards,
     )
 
     # we build the model in the regular way:
@@ -60,6 +64,10 @@ def delta(s: pgc.State, action: pgc.Action):
     model.add_transitions(state2, stormvogel.model.Transition({right: branch2}))
     model.add_transitions(state0, stormvogel.model.Transition({left: branch0}))
 
+    rewardmodel = model.add_rewards("rewards")
+    for i in range(2 * N):
+        rewardmodel.set_state_action_reward_at_id(i, 1)
+
     assert model == pgc_model
 
 
@@ -68,6 +76,9 @@ def test_pgc_dtmc():
     p = 0.5
     initial_state = pgc.State(s=0)
 
+    def rewards(s: pgc.State):
+        return 1
+
     def delta(s: pgc.State):
         match s.s:
             case 0:
@@ -96,6 +107,7 @@ def delta(s: pgc.State):
     pgc_model = stormvogel.pgc.build_pgc(
         delta=delta,
         initial_state_pgc=initial_state,
+        rewards=rewards,
         modeltype=stormvogel.model.ModelType.DTMC,
     )
 
@@ -161,4 +173,8 @@ def delta(s: pgc.State):
     model.set_transitions(model.get_state_by_id(12), [(1, model.get_state_by_id(13))])
     model.set_transitions(model.get_state_by_id(13), [(1, model.get_state_by_id(13))])
 
+    rewardmodel = model.add_rewards("rewards")
+    for state in model.states.values():
+        rewardmodel.set_state_reward(state, 1)
+
     assert pgc_model == model

From c339a8d0fc321f4dc94b704d844424a542bf6510 Mon Sep 17 00:00:00 2001
From: PimLeerkes <pimsimon.leerkes@gmail.com>
Date: Sun, 8 Dec 2024 14:18:31 +0100
Subject: [PATCH 6/8] model builder can now also use a labels function

---
 stormvogel/pgc.py |  6 ++++++
 tests/test_pgc.py | 10 +++++++---
 2 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/stormvogel/pgc.py b/stormvogel/pgc.py
index 4588c21..31d0bdc 100644
--- a/stormvogel/pgc.py
+++ b/stormvogel/pgc.py
@@ -157,5 +157,11 @@ def build_pgc(
                 rewardmodel.set_state_reward(s, reward)
 
     # we add the labels
+    if labels is not None:
+        for state in states_seen:
+            s = model.get_state_by_name(str(state.__dict__))
+            assert s is not None
+            for label in labels(state):
+                s.add_label(label)
 
     return model
diff --git a/tests/test_pgc.py b/tests/test_pgc.py
index b3b79ed..6907788 100644
--- a/tests/test_pgc.py
+++ b/tests/test_pgc.py
@@ -19,6 +19,9 @@ def available_actions(s: pgc.State):
     def rewards(s: pgc.State, a: pgc.Action):
         return 1
 
+    def labels(s: pgc.State):
+        return [str(s.x)]
+
     def delta(s: pgc.State, action: pgc.Action):
         if action == left:
             return (
@@ -43,14 +46,15 @@ def delta(s: pgc.State, action: pgc.Action):
         delta=delta,
         available_actions=available_actions,
         initial_state_pgc=initial_state,
+        labels=labels,
         rewards=rewards,
     )
 
     # we build the model in the regular way:
     model = stormvogel.model.new_mdp(create_initial_state=False)
-    state1 = model.new_state(labels=["init"], features={"x": 1})
-    state2 = model.new_state(features={"x": 2})
-    state0 = model.new_state(features={"0": 0})
+    state1 = model.new_state(labels=["init", "1"], features={"x": 1})
+    state2 = model.new_state(labels=["2"], features={"x": 2})
+    state0 = model.new_state(labels=["0"], features={"0": 0})
     left = model.new_action("left", frozenset({"left"}))
     right = model.new_action("right", frozenset({"right"}))
     branch11 = stormvogel.model.Branch([(0.5, state1), (0.5, state2)])

From ddd03426be99ed39943ef8c1f8eea22f073eb1da Mon Sep 17 00:00:00 2001
From: PimLeerkes <pimsimon.leerkes@gmail.com>
Date: Sun, 8 Dec 2024 14:45:07 +0100
Subject: [PATCH 7/8] model builder supports multiple reward models

---
 stormvogel/pgc.py | 39 +++++++++++++++++++++++++--------------
 tests/test_pgc.py | 14 ++++++++++----
 2 files changed, 35 insertions(+), 18 deletions(-)

diff --git a/stormvogel/pgc.py b/stormvogel/pgc.py
index 31d0bdc..feae772 100644
--- a/stormvogel/pgc.py
+++ b/stormvogel/pgc.py
@@ -32,7 +32,7 @@ def __eq__(self, other):
 
 def build_pgc(
     delta,  # Callable[[State, Action], list[tuple[float, State]]],
-    initial_state_pgc: State,  # TODO rewards function, label function
+    initial_state_pgc: State,
     rewards=None,
     labels=None,
     available_actions: Callable[[State], list[Action]] | None = None,
@@ -45,17 +45,15 @@ def build_pgc(
     this works analogous to a prism file, where the delta is the module in this case.
 
     (this function uses the pgc classes State and Action instead of the ones from stormvogel.model)
-
-    (currently only works for mdps)
     """
     if modeltype == stormvogel.model.ModelType.MDP and available_actions is None:
         raise RuntimeError(
             "You have to provide an available actions function for mdp models"
         )
 
+    # we create the model with the given type and initial state
     model = stormvogel.model.new_model(modeltype=modeltype, create_initial_state=False)
 
-    # we create the model with the given type and initial state
     model.new_state(
         labels=["init"],
         features=initial_state_pgc.__dict__,
@@ -134,27 +132,40 @@ def build_pgc(
         )
 
     # we add the rewards
-    # TODO support multiple reward models
     if rewards is not None:
-        rewardmodel = model.add_rewards("rewards")
         if model.supports_actions():
+            # we first create the right number of reward models
+            assert available_actions is not None
+            nr = len(
+                rewards(initial_state_pgc, available_actions(initial_state_pgc)[0])
+            )
+            for i in range(nr):
+                model.add_rewards("rewardmodel: " + str(i))
+
             for state in states_seen:
                 assert available_actions is not None
                 for action in available_actions(state):
-                    reward = rewards(state, action)
+                    rewardlist = rewards(state, action)
                     s = model.get_state_by_name(str(state.__dict__))
                     assert s is not None
-                    rewardmodel.set_state_action_reward(
-                        s,
-                        model.get_action(str(action.labels)),
-                        reward,
-                    )
+                    for index, reward in enumerate(rewardlist):
+                        model.rewards[index].set_state_action_reward(
+                            s,
+                            model.get_action(str(action.labels)),
+                            reward,
+                        )
         else:
+            # we first create the right number of reward models
+            nr = len(rewards(initial_state_pgc))
+            for i in range(nr):
+                model.add_rewards("rewardmodel: " + str(i))
+
             for state in states_seen:
-                reward = rewards(state)
+                rewardlist = rewards(state)
                 s = model.get_state_by_name(str(state.__dict__))
                 assert s is not None
-                rewardmodel.set_state_reward(s, reward)
+                for index, reward in enumerate(rewardlist):
+                    model.rewards[index].set_state_reward(s, reward)
 
     # we add the labels
     if labels is not None:
diff --git a/tests/test_pgc.py b/tests/test_pgc.py
index 6907788..2654fac 100644
--- a/tests/test_pgc.py
+++ b/tests/test_pgc.py
@@ -17,7 +17,7 @@ def available_actions(s: pgc.State):
         return [left, right]
 
     def rewards(s: pgc.State, a: pgc.Action):
-        return 1
+        return [1, 2]
 
     def labels(s: pgc.State):
         return [str(s.x)]
@@ -68,9 +68,12 @@ def delta(s: pgc.State, action: pgc.Action):
     model.add_transitions(state2, stormvogel.model.Transition({right: branch2}))
     model.add_transitions(state0, stormvogel.model.Transition({left: branch0}))
 
-    rewardmodel = model.add_rewards("rewards")
+    rewardmodel = model.add_rewards("rewardmodel: 0")
     for i in range(2 * N):
         rewardmodel.set_state_action_reward_at_id(i, 1)
+    rewardmodel = model.add_rewards("rewardmodel: 1")
+    for i in range(2 * N):
+        rewardmodel.set_state_action_reward_at_id(i, 2)
 
     assert model == pgc_model
 
@@ -81,7 +84,7 @@ def test_pgc_dtmc():
     initial_state = pgc.State(s=0)
 
     def rewards(s: pgc.State):
-        return 1
+        return [1, 2]
 
     def delta(s: pgc.State):
         match s.s:
@@ -177,8 +180,11 @@ def delta(s: pgc.State):
     model.set_transitions(model.get_state_by_id(12), [(1, model.get_state_by_id(13))])
     model.set_transitions(model.get_state_by_id(13), [(1, model.get_state_by_id(13))])
 
-    rewardmodel = model.add_rewards("rewards")
+    rewardmodel = model.add_rewards("rewardmodel: 0")
     for state in model.states.values():
         rewardmodel.set_state_reward(state, 1)
+    rewardmodel = model.add_rewards("rewardmodel: 1")
+    for state in model.states.values():
+        rewardmodel.set_state_reward(state, 2)
 
     assert pgc_model == model

From 081efbeb42d1c4effb2301be2e9aca01a3a1bdbe Mon Sep 17 00:00:00 2001
From: PimLeerkes <pimsimon.leerkes@gmail.com>
Date: Sun, 22 Dec 2024 15:28:02 +0100
Subject: [PATCH 8/8] commented out pgc tests

---
 tests/test_pgc.py | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/tests/test_pgc.py b/tests/test_pgc.py
index 2654fac..ec4d429 100644
--- a/tests/test_pgc.py
+++ b/tests/test_pgc.py
@@ -1,9 +1,4 @@
-from stormvogel import pgc
-import stormvogel.model
-import math
-import stormvogel.mapping
-
-
+"""
 def test_pgc_mdp():
     # we build the model with pgc:
     N = 2
@@ -188,3 +183,4 @@ def delta(s: pgc.State):
         rewardmodel.set_state_reward(state, 2)
 
     assert pgc_model == model
+"""