prism-accumulation/prism/src/explicit/AccumulationProduct.java

package explicit;

import java.util.ArrayList;
import java.util.BitSet;
import java.util.Iterator;
import java.util.Map;

import parser.ast.ExpressionAccumulation;
import parser.ast.AccumulationFactor;
import parser.ast.Expression;
import prism.IntegerBound;
import prism.PrismException;

import common.Dottable;

/**
 * An AccumulationProduct has ProductStates, where the first component is the
 * stateId in the original model, and the second component is the index of an
 * AccumulationTracker.
 * 
 * @author Sascha Wunderlich
 *
 * @param <M>
 */

public abstract class AccumulationProduct<M extends Model,Component> extends ProductWithProductStates<M> implements Dottable
{
	protected final StoragePool<AccumulationTracker<Component>> trackers;
	protected final StoragePool<AccumulationState<Component>> accStates;
	
	protected final ArrayList<BitSet> initStates;
	protected final ArrayList<BitSet> runStates;
	protected final ArrayList<BitSet> goalStates;

	protected final ArrayList<BitSet> labels;
	
	protected int numberOfTracks;
	protected int numberOfWeights;
	
	public AccumulationProduct(M originalModel) {
		super(originalModel);
		trackers = new StoragePool<>();
		accStates = new StoragePool<>();
		
		initStates = new ArrayList<>();
		runStates  = new ArrayList<>();
		goalStates = new ArrayList<>();
		
		labels = new ArrayList<BitSet>();
	}
	
	public final int getNumberOfTracks() {
		return numberOfTracks;
	}

	public final BitSet getInitStates(int track) {
		return initStates.get(track);
	}
	
	public final BitSet getRunStates(int track) {
		return runStates.get(track);
	}
	
	public final BitSet getGoalStates(int track) {
		return goalStates.get(track);
	}
	
	protected void generateTrackInfo(ProductState state, Integer index, ExpressionAccumulation accexp, ProbModelChecker mc) throws PrismException {
		AccumulationState<Component> accState = accStates.getById(state.getSecondState());
		AccumulationTracker<Component> tracker = accState.getTracker(trackers);
		
		for(int trackIdx=0; trackIdx<tracker.getTracks().size(); trackIdx++) {
			AccumulationTrack<Component> track = tracker.getTracks().get(trackIdx);
			
			// Is this an initial
			if (trackIdx==accState.lastRestartNr) {
				initStates.get(trackIdx).set(index);
			}
			
			// Is this a running track?
			if (track != null) {
				runStates.get(trackIdx).set(index);
			}
			
			// Is this a goal track?
			if (isGoalTrack(track, accexp, mc)) {
				goalStates.get(trackIdx).set(index);
			}
		}
	}
	

       /** Checks whether a track is good.
        * 
        * Tracks are good, if they are final and fulfill the constraint in the accumulation expression with context mc.
        * @param track
        * @param accexp
        * @param mc
        * @return
        * @throws PrismException
        */
       protected final boolean isGoalTrack(final AccumulationTrack<Component> track, final ExpressionAccumulation accexp, final ProbModelChecker mc)
                       throws PrismException {
               // Only final tracks can be good
               if (!isFinalTrack(track,accexp,mc)) { return false; }
               boolean isGood = false;
               
               //TODO: these should be double later on
               // Collect the weight linear combination, factor*weight+...     
               int lhs = 0;
               int factorNr = 0;
               for (AccumulationFactor factor : accexp.getConstraint().getFactors()) {
                       lhs += factor.getFactor().evaluateInt(mc.getConstantValues())
                                       * track.getWeight(factorNr);
               }

               // Check the bound
               IntegerBound rhs = IntegerBound.fromTemporalOperatorBound(accexp.getConstraint().getBound(), mc.getConstantValues(), true);

               // For DIA operators, we just check the bound.
               // For BOX operators, we check the INVERTED bound.
               
               // TODO: THIS MAY BE BROKEN
               switch(accexp.getSymbol()) {
                       case ACCBOXMINUS:
                       case ACCBOXPLUS:
                               if (!rhs.isInBounds(lhs)) {
                                       isGood = true;
                               }
                               mc.getLog().println("WARNING: This may be wrong.");
                               break;
                       case ACCDIAMINUS:
                       case ACCDIAPLUS:
                       case ACCUNTIL:
                               if (rhs.isInBounds(lhs)) {
                                       isGood = true;
                               }
                               break;
                       default:
                               throw new RuntimeException("Accumulation symbol cannot be handled...");
                       }
               //if(isGood) {mc.getLog().print("+");} else {mc.getLog().print("-");}
               return isGood;
       }


	
	/** Checks whether a track is final according to the bound in the accumulation expression and the context in mc.
	 * @param track
	 * @param accexp
	 * @param mc
	 * @return
	 * @throws PrismException
	 */
	protected abstract boolean isFinalTrack(final AccumulationTrack<Component> track, final ExpressionAccumulation accexp, final ProbModelChecker mc)
			throws PrismException;

	/** Get the initial Component of the accumulation monitor.
	 * @return
	 */
	protected abstract Component getInitialComponent();
	
	/** Update the Component of the accumulation monitor when leaving modelFromStateId.
	 * @param modelFromStateId
	 * @param track
	 * @param accexp
	 * @param mc
	 * @return
	 */
	protected abstract Component updateComponent(Integer modelFromStateId, final AccumulationTrack<Component> track,
			final ExpressionAccumulation accexp, final StateModelChecker mc);
	
	/** Updates a single AccumulationTrack. Used in updateAccumulationState.
	 * 
	 * Uses updateComponent to get the next Component, accumulates the current weights
	 * and makes a new track.
	 * @param modelFromStateId
	 * @param track
	 * @param accexp
	 * @param weights
	 * @param mc
	 * @return
	 */
	protected final AccumulationTrack<Component> updateTrack(Integer modelFromStateId, final AccumulationTrack<Component> track,
			final ExpressionAccumulation accexp, final double[] weights, final StateModelChecker mc) {
		Component nextComponent = updateComponent(modelFromStateId, track, accexp, mc);

		// If we are done, return null-Track
		if (nextComponent == null) { return null; }
		
		// Otherwise, we update the weights and increase the step.
		double[] newweights = new double[weights.length];
		for (int i = 0; i < weights.length; i++) {
			newweights[i] = weights[i] + track.getWeights()[i];
		}
		
		return new AccumulationTrack<Component>(newweights, nextComponent);
	}
	
	/** Generates a new accumulation state from an old one.
	 * 
	 * To do so, it reads the necessary information from the model and its rewards
	 * and updates all tracks and their goodness accordingly.
	 * @param modelFromStateId
	 * @param accstate
	 * @param accexp
	 * @param weights
	 * @param mc
	 * @return
	 * @throws PrismException
	 */
	protected final AccumulationState<Component> updateAccumulationState(final int modelFromStateId,
			final AccumulationState<Component> accstate, final ExpressionAccumulation accexp,
			final double[] weights, final ProbModelChecker mc) throws PrismException {
		
		// Check if we even need to fire here.
		if(accexp.hasFireOn()) {
			boolean stutter = true;
			for(Expression f : accexp.getFireOn()) {
				if(mc.checkExpression(originalModel, f, null).getBitSet().get(modelFromStateId)) {
					//mc.getLog().println("f:" + f);
					stutter = false;
					break;
				}
			}
			// If this is a stutter action, we can return the same accstate. Copies are made on modification.
			if(stutter) { return accstate; }
		}
		// ...otherwise proceed.
		
		// Get the old tracker and tracks.
		AccumulationTracker<Component> oldTracker = accstate.getTracker(trackers);
		ArrayList<AccumulationTrack<Component>> oldTracks = oldTracker.getTracks();
		
		// This restart will be...
		int newLastRestartNr = accstate.getNextRestartNr();
		//mc.getLog().print(newLastRestartNr);
		
		// Build the new tracks and determine their goodness;
		ArrayList<AccumulationTrack<Component>> newTracks = new ArrayList<>();
		
		int trackNr = 0;
		for(AccumulationTrack<Component> oldTrack : oldTracks) {
			
			AccumulationTrack<Component> newTrack;
			
			// restart or advance
			if(trackNr == newLastRestartNr) {
				// If we have a restart, we produce an initial track and let it advance.
				AccumulationTrack<Component> freshTrack = new AccumulationTrack<Component>(numberOfWeights, getInitialComponent());
				//newTrack = updateTrack(modelFromStateId, freshTrack, accexp, weights, mc);
				newTrack = freshTrack;
			} else if (oldTrack == null) {
				// If the old track is undefined, the new track is as well.
				newTrack = null;
			} else {
				// Otherwise, the track is defined and advances.
				assert oldTrack != null;
				newTrack = updateTrack(modelFromStateId, oldTrack, accexp, weights, mc);
			}
			
			newTracks.add(newTrack);
			trackNr++;
		}
		
		//Create a new tracker with the right tracks and add it to storage.
		AccumulationTracker<Component> newTracker = new AccumulationTracker<>(newTracks);
		int newTrackerId = trackers.findOrAdd(newTracker);
		
		return new AccumulationState<>(newTrackerId, newLastRestartNr, numberOfTracks);
	}
	
	/**
	 * Creates the initial state for the accumulation part of this of this accumulation product. Returns its ID.
	 * 
	 * The initial state has a tracker filled with null-tracks, none of which is good
	 * where the first restart is the first (0th) track.
	 * @return
	 */
	protected final int createInitialStateId() {
		Component initialComponent = getInitialComponent();
		
		// The initial active track is the first one, all tracks are non-good by default
		int initialActiveTrack = 0;
		
		// Generate the initial track and product state
		AccumulationTracker<Component> initialTracker = new AccumulationTracker<Component>(numberOfTracks, numberOfWeights, initialComponent);
		int initialTrackerId  = trackers.findOrAdd(initialTracker);
		AccumulationState<Component> initialAccState = new AccumulationState<Component>(initialTrackerId, initialActiveTrack, numberOfTracks);
		int initialAccStateId = accStates.findOrAdd(initialAccState);
		
		return initialAccStateId;
	}
	
	@Override
	public String toDot() {
		StringBuffer result = new StringBuffer();
		
		result.append("digraph " + originalModel.getModelType() + " {\n");
		
		for(int i = 0; i < prod_states.size(); i++) {
			ProductState fromState = prod_states.get(i);
			AccumulationState<Component> accState = accStates.getById(fromState.getSecondState());
			AccumulationTracker<Component> tracker = accState.getTracker(trackers);
			result.append(""
					+ i
					+ "[shape=box, color=black"
					+ " label= < <TABLE BORDER=\"0\">"
					+ "<TR>"
					+ "<TD>" + i + "=" + fromState + "</TD>"
					+ "</TR><TR>"
					+ "<TD> " + accState	+ "</TD>"
					+ "</TR><TR>"
					+ "<TD>\"" + Dottable.quoteForDot(tracker.toString()) + "\"</TD>"
					+ "</TR>"
					+ " </TABLE> >]\n");
			
			switch(productModel.getModelType()) {
			case DTMC:
				DTMCExplicit castDTMC = (DTMCExplicit)productModel;
				Iterator<Map.Entry<Integer, Double>> dtmcIter = castDTMC.getTransitionsIterator(i);
				while (dtmcIter.hasNext()) {
					Map.Entry<Integer, Double> e = dtmcIter.next();
					result.append(i + " -> " + e.getKey() + " [ label=\"");
					result.append(e.getValue() + "\" ];\n");
				}
				break;
			case MDP:
				MDPExplicit castMDP = (MDPExplicit)productModel;
				for(int c = 0; c < castMDP.getNumChoices(i); c++) {
					Iterator<Map.Entry<Integer, Double>> mdpIter = castMDP.getTransitionsIterator(i, c);
					while (mdpIter.hasNext()) {
						Map.Entry<Integer, Double> e = mdpIter.next();
						result.append(i + " -> " + e.getKey() + " [ label=\"");
						result.append(c + "," + e.getValue() + "\" ];\n");
					}
				}
				break;
			default:
				break;
			}
		}
		
		result.append("}");
		
		return result.toString();
	}
}