From 2e5a4bfd56dd7992865bb7b0da02f9488d584a4b Mon Sep 17 00:00:00 2001
From: Dave Parker <dave.x.parker@gmail.com>
Date: Thu, 25 Feb 2021 13:32:50 +0000
Subject: [PATCH] Implement (extended) .tra file export for POMDPs.

---
 prism/src/explicit/POMDP.java | 33 +++++++++++++++++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/prism/src/explicit/POMDP.java b/prism/src/explicit/POMDP.java
index 475dad33..cb5eda1b 100644
--- a/prism/src/explicit/POMDP.java
+++ b/prism/src/explicit/POMDP.java
@@ -27,8 +27,13 @@
 
 package explicit;
 
+import java.util.Iterator;
+import java.util.Map;
+import java.util.TreeMap;
+
 import explicit.rewards.MDPRewards;
 import prism.ModelType;
+import prism.PrismLog;
 import prism.PrismUtils;
 
 /**
@@ -44,6 +49,34 @@ public interface POMDP extends MDP, PartiallyObservableModel
 		return ModelType.POMDP;
 	}
 
+	@Override
+	default void exportToPrismExplicitTra(PrismLog out)
+	{
+		// Output transitions to .tra file
+		int numStates = getNumStates();
+		out.print(numStates + " " + getNumChoices() + " " + getNumTransitions() + " " + getNumObservations() + "\n");
+		TreeMap<Integer, Double> sorted = new TreeMap<Integer, Double>();
+		for (int i = 0; i < numStates; i++) {
+			int numChoices = getNumChoices(i);
+			for (int j = 0; j < numChoices; j++) {
+				// Extract transitions and sort by destination state index (to match PRISM-exported files)
+				Iterator<Map.Entry<Integer, Double>> iter = getTransitionsIterator(i, j);
+				while (iter.hasNext()) {
+					Map.Entry<Integer, Double> e = iter.next();
+					sorted.put(e.getKey(), e.getValue());
+				}
+				// Print out (sorted) transitions
+				for (Map.Entry<Integer, Double> e : sorted.entrySet()) {
+					// Note use of PrismUtils.formatDouble to match PRISM-exported files
+					out.print(i + " " + j + " " + e.getKey() + " " + PrismUtils.formatDouble(e.getValue()) + " " + getObservation(e.getKey()));
+					Object action = getAction(i, j);
+					out.print(action == null ? "\n" : (" " + action + "\n"));
+				}
+				sorted.clear();
+			}
+		}
+	}
+
 	@Override
 	default String infoString()
 	{