From bfeb0d0b69c026c8703e5c5e95bdbabd264f2c33 Mon Sep 17 00:00:00 2001
From: Dave Parker <dave.x.parker@gmail.com>
Date: Fri, 17 Jul 2015 13:23:05 +0000
Subject: [PATCH] Implement lifting to product for STPG rewards.

git-svn-id: https://www.prismmodelchecker.org/svn/prism/prism/trunk@10356 bbc10eb1-c90d-0410-af57-cb519fbb1720
---
 prism/src/explicit/rewards/STPGRewards.java   | 18 +++-----
 .../explicit/rewards/STPGRewardsSimple.java   | 43 +++++++++++++++++++
 2 files changed, 50 insertions(+), 11 deletions(-)

diff --git a/prism/src/explicit/rewards/STPGRewards.java b/prism/src/explicit/rewards/STPGRewards.java
index a828c282..047c2d86 100644
--- a/prism/src/explicit/rewards/STPGRewards.java
+++ b/prism/src/explicit/rewards/STPGRewards.java
@@ -27,23 +27,16 @@
 
 package explicit.rewards;
 
+import explicit.Model;
+import explicit.Product;
+
 /**
  * Classes that provide (read) access to explicit-state rewards for an STPG.
  * See the {@link explicit.STPG} interface for details of the accompanying model,
  * in particular, for an explanation of nested transitions. 
  */
-public interface STPGRewards extends Rewards
+public interface STPGRewards extends MDPRewards
 {
-	/**
-	 * Get the state reward for state {@code s}.
-	 */
-	public abstract double getStateReward(int s);
-
-	/**
-	 * Get the transition reward for the {@code i}th choice from state {@code s}.
-	 */
-	public abstract double getTransitionReward(int s, int i);
-
 	/**
 	 * Get the transition reward for the {@code i,j}th nested choice from state {@code s}.
 	 */
@@ -53,4 +46,7 @@ public interface STPGRewards extends Rewards
 	 * Build an MDPRewards object containing all the same rewards except for the nested ones.
 	 */
 	public abstract MDPRewards buildMDPRewards();
+
+	@Override
+	public STPGRewards liftFromModel(Product<? extends Model> product);
 }
diff --git a/prism/src/explicit/rewards/STPGRewardsSimple.java b/prism/src/explicit/rewards/STPGRewardsSimple.java
index cc164645..4acc520b 100644
--- a/prism/src/explicit/rewards/STPGRewardsSimple.java
+++ b/prism/src/explicit/rewards/STPGRewardsSimple.java
@@ -30,6 +30,9 @@ package explicit.rewards;
 import java.util.ArrayList;
 import java.util.List;
 
+import explicit.Model;
+import explicit.Product;
+
 public class STPGRewardsSimple extends MDPRewardsSimple implements STPGRewards
 {
 	/** Nested transition rewards */
@@ -46,6 +49,17 @@ public class STPGRewardsSimple extends MDPRewardsSimple implements STPGRewards
 		nestedTransRewards = null;
 	}
 
+	/**
+	 * Constructor: copy MDP rewards, other rewards zero
+	 * @param rews MPD rewards to copy
+	 */
+	public STPGRewardsSimple(MDPRewardsSimple rews)
+	{
+		super(rews);
+		// Initially list is just null (denoting all 0)
+		nestedTransRewards = null;
+	}
+
 	// Mutators
 
 	/**
@@ -123,6 +137,35 @@ public class STPGRewardsSimple extends MDPRewardsSimple implements STPGRewards
 		return list2.get(j);
 	}
 
+	// Converters
+	
+	@Override
+	public STPGRewards liftFromModel(Product<? extends Model> product)
+	{
+		// Lift MDP part
+		MDPRewardsSimple rewardsProdMDP = (MDPRewardsSimple) ((MDPRewardsSimple) this).liftFromModel(product);
+		STPGRewardsSimple rewardsProd = new STPGRewardsSimple(rewardsProdMDP);
+		// Lift nested transition rewards
+		Model modelProd = product.getProductModel();
+		int numStatesProd = modelProd.getNumStates();		
+		if (nestedTransRewards != null) {
+			for (int s = 0; s < numStatesProd; s++) {
+				List<List<Double>> list1 = nestedTransRewards.get(product.getModelState(s));
+				if (list1 != null) {
+					int n1 = list1.size();
+					for (int i = 0; i < n1; i++) {
+						List<Double> list2 = list1.get(i);
+						int n2 = list2.size();
+						for (int j = 0; j < n2; j++) {
+							rewardsProd.setNestedTransitionReward(s, i, j, list2.get(j));
+						}
+					}
+				}
+			}
+		}
+		return rewardsProd;
+	}
+	
 	// Other
 
 	@Override