code: //==============================================================================
  //	
  //	Copyright (c) 2016-
  //	Authors:
  //	* Steffen Maercker <maercker@tcs.inf.tu-dresden.de> (TU Dresden)
  //	* Joachim Klein <klein@tcs.inf.tu-dresden.de> (TU Dresden)
  //	
  //------------------------------------------------------------------------------
  //	
  //	This file is part of PRISM.
  //	
  //	PRISM is free software; you can redistribute it and/or modify
  //	it under the terms of the GNU General Public License as published by
  //	the Free Software Foundation; either version 2 of the License, or
  //	(at your option) any later version.
  //	
  //	PRISM is distributed in the hope that it will be useful,
  //	but WITHOUT ANY WARRANTY; without even the implied warranty of
  //	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  //	GNU General Public License for more details.
  //	
  //	You should have received a copy of the GNU General Public License
  //	along with PRISM; if not, write to the Free Software Foundation,
  //	Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
  //	
  //==============================================================================
  
  package explicit.modelviews;
  
  import java.util.AbstractMap;
  import java.util.BitSet;
  import java.util.Collections;
  import java.util.Iterator;
  import java.util.List;
  import java.util.Map.Entry;
  import java.util.Set;
  import java.util.function.IntFunction;
  import java.util.function.IntPredicate;
  
  import common.iterable.SingletonIterator;
  import explicit.MDP;
  import parser.State;
  import parser.Values;
  import parser.VarList;
  
  /**
   * An MDPView that takes an existing MDP and
   * adds additional choices to certain states.
   */
  public class MDPAdditionalChoices extends MDPView
  {
  	private MDP model;
  	private IntFunction<List<Iterator<Entry<Integer, Double>>>> choices;
  	private IntFunction<List<Object>> actions;
  
  	/**
  	 * If {@code choices} returns {@code null} for a state and a choice, no additional choice is added.
  	 * If {@code actions} is {@code null} or returns {@code null} for a state, no additional action is attached.
  	 *
  	 * @param model
  	 * @param choices
  	 * @param actions
  	 */
  	public MDPAdditionalChoices(final MDP model, final IntFunction<List<Iterator<Entry<Integer, Double>>>> choices,
  			IntFunction<List<Object>> actions)
  	{
  		this.model = model;
  		this.choices = choices;
  		this.actions = actions;
  	}
  
  	public MDPAdditionalChoices(final MDPAdditionalChoices additional)
  	{
  		super(additional);
  		model = additional.model;
  		choices = additional.choices;
  		actions = additional.actions;
  	}
  
  
  
  	//--- Cloneable ---
  
  	@Override
  	public MDPView clone()
  	{
  		return new MDPAdditionalChoices(this);
  	}
  
  
  
  	//--- Model ---
  
  	@Override
  	public int getNumStates()
  	{
  		return model.getNumStates();
  	}
  
  	@Override
  	public int getNumInitialStates()
  	{
  		return model.getNumInitialStates();
  	}
  
  	@Override
  	public Iterable<Integer> getInitialStates()
  	{
  		return model.getInitialStates();
  	}
  
  	@Override
  	public int getFirstInitialState()
  	{
  		return model.getFirstInitialState();
  	}
  
  	@Override
  	public boolean isInitialState(final int state)
  	{
  		return model.isInitialState(state);
  	}
  
  	@Override
  	public List<State> getStatesList()
  	{
  		return model.getStatesList();
  	}
  
  	@Override
  	public VarList getVarList()
  	{
  		return model.getVarList();
  	}
  
  	@Override
  	public Values getConstantValues()
  	{
  		return model.getConstantValues();
  	}
  
  	@Override
  	public BitSet getLabelStates(final String name)
  	{
  		return model.getLabelStates(name);
  	}
  
  	@Override
  	public Set<String> getLabels()
  	{
  		return model.getLabels();
  	}
  
  	@Override
  	public boolean hasLabel(String name)
  	{
  		return model.hasLabel(name);
  	}
  
  
  	//--- NondetModel ---
  
  	@Override
  	public int getNumChoices(final int state)
  	{
  		return model.getNumChoices(state) + getNumAdditionalChoices(state);
  	}
  
  	@Override
  	public Object getAction(final int state, final int choice)
  	{
  		final int numOriginalChoices = model.getNumChoices(state);
  		if (choice < numOriginalChoices) {
  			return model.getAction(state, choice);
  		}
  		if (actions == null) {
  			final int numChoices = numOriginalChoices + getNumAdditionalChoices(state);
  			if (choice < numChoices) {
  				return null;
  			}
  			throw new IndexOutOfBoundsException("choice index out of bounds");
  		}
  
  		final List<Object> additional = actions.apply(state);
  		return (additional == null) ?  null : additional.get(choice - numOriginalChoices);
  	}
  
  	@Override
  	public boolean areAllChoiceActionsUnique()
  	{
  		return model.areAllChoiceActionsUnique() && super.areAllChoiceActionsUnique();
  	}
  
  
  
  	//--- MDP ---
  
  	@Override
  	public Iterator<Entry<Integer, Double>> getTransitionsIterator(final int state, final int choice)
  	{
  		final int numOriginalChoices = model.getNumChoices(state);
  		if (choice < numOriginalChoices) {
  			return model.getTransitionsIterator(state, choice);
  		}
  		try {
  			return choices.apply(state).get(choice - numOriginalChoices);
  		} catch (NullPointerException | IndexOutOfBoundsException e)
  		{
  			// alter message of exception
  			throw new IndexOutOfBoundsException("choice index out of bounds");
  		}
  	}
  
  	//--- MDPView ---
  
  	@Override
  	protected void fixDeadlocks()
  	{
  		assert !fixedDeadlocks : "deadlocks already fixed";
  
  		model = fixDeadlocks((MDP) this.clone());
  		choices = (int element) -> {return (List<Iterator<Entry<Integer, Double>>>)null;};
  		actions = null;
  	}
  
  
  
  	//--- instance methods ---
  
  	private int getNumAdditionalChoices(final int state)
  	{
  		final List<Iterator<Entry<Integer, Double>>> additional = choices.apply(state);
  		return (additional == null) ? 0 : additional.size();
  	}
  
  
  
  	//--- static methods ---
  
  	public static MDPView fixDeadlocks(final MDP model)
  	{
  		final BitSet deadlockStates = new BitSet();
  		model.getDeadlockStates().forEach(deadlockStates::set);
  		final MDPView fixed = addSelfLoops(model, deadlockStates);
  		fixed.deadlockStates = deadlockStates;
  		fixed.fixedDeadlocks = true;
  		return fixed;
  	}
  
  	public static MDPView addSelfLoops(final MDP model, final BitSet states)
  	{
  		return addSelfLoops(model, states::get);
  	}
  
  	public static MDPView addSelfLoops(final MDP model, final IntPredicate states)
  	{
  		final IntFunction<List<Iterator<Entry<Integer, Double>>>> addSelfLoops = new IntFunction<List<Iterator<Entry<Integer, Double>>>>()
  		{
  			@Override
  			public List<Iterator<Entry<Integer, Double>>> apply(final int state)
  			{
  				if (states.test(state)) {
  					Entry<Integer,Double> transition = new AbstractMap.SimpleImmutableEntry<>(state, 1.0);
  					return Collections.singletonList(new SingletonIterator.Of<>(transition));
  				}
  				return null;
  			}
  		};
  		return new MDPAdditionalChoices(model, addSelfLoops, null);
  	}
  }
  