package explicit;

import java.util.BitSet;
import java.util.Vector;

import explicit.rewards.MCRewards;
import explicit.rewards.MDPRewards;
import explicit.rewards.Rewards;
import parser.ast.ExpressionAccumulation;
import prism.IntegerBound;
import prism.PrismException;

/**
 * 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 class AccumulationProductCounting<M extends Model> extends AccumulationProduct<M,Integer>
{
	
	public AccumulationProductCounting(M originalModel) {
		super(originalModel);
	}
	
	public static AccumulationProductCounting<DTMC> generate(final DTMC graph, final ExpressionAccumulation accexp, final Vector<MCRewards> rewards, final ProbModelChecker mc, BitSet statesOfInterest) throws PrismException {
		final AccumulationProductCounting<DTMC> result = new AccumulationProductCounting<DTMC>(graph);
		// Create auxiliary data
		result.createAuxData(graph, accexp, rewards, mc);
		
		// Build an operator
		class AccumulationDTMCProductOperator implements DTMCProductOperator
		{
			@Override
			public ProductState getInitialState(Integer dtmc_state)
					throws PrismException {
				int initialAccStateId = result.createInitialStateId();
				return new ProductState(dtmc_state, initialAccStateId);
			}

			@Override
			public ProductState getSuccessor(ProductState from_state,
					Integer dtmc_to_state) throws PrismException {
				// Get the current accumulation state
				AccumulationState<Integer> from_accstate = result.accStates.getById(from_state.getSecondState());
				
				// Get step weights
				double[] weights = new double[rewards.size()];
				
				for (int i=0; i < rewards.size(); i++) {
					weights[i] = rewards.get(i).getStateReward(from_state.getFirstState());
				}
				
				// Update accumulation product state, store it and get its ID.
				AccumulationState<Integer> to_accproduct = result.updateAccumulationState(from_state.getFirstState(), from_accstate, accexp, weights, mc);
				int to_accproduct_id = result.accStates.findOrAdd(to_accproduct);
				
				return new ProductState(dtmc_to_state, to_accproduct_id);
			}

			@Override
			public void notify(ProductState state, Integer index)
					throws PrismException {
				result.generateTrackInfo(state, index, accexp, mc);
			}

			@Override
			public void finish() throws PrismException {
				// Do nothing
			}

			@Override
			public DTMC getGraph() {
				return graph;
			}
		}
		
		// Apply the operator
		AccumulationDTMCProductOperator op = new AccumulationDTMCProductOperator();
		ProductWithProductStates.generate(op, result, statesOfInterest);
		
		return result;
	}
	
	public static AccumulationProductCounting<MDP> generate(final MDP graph, final ExpressionAccumulation accexp, final Vector<MDPRewards> rewards, final ProbModelChecker mc, BitSet statesOfInterest) throws PrismException {
		// This is basically the same thing as for DTMCs
		final AccumulationProductCounting<MDP> result = new AccumulationProductCounting<MDP>(graph);
		
		// Create auxiliary data
		result.createAuxData(graph, accexp, rewards, mc);

		class AccumulationMDPProductOperator implements MDPProductOperator
		{

			@Override
			public ProductState getInitialState(final Integer MDP_state)
					throws PrismException {
				int initialAccStateId = result.createInitialStateId();
				return new ProductState(MDP_state, initialAccStateId);
			}

			@Override
			public ProductState getSuccessor(final ProductState from_state,
					final int choice_i, final Integer mdp_to_state) throws PrismException {
				// Get the current accumulation state
				AccumulationState<Integer> from_accstate = result.accStates.getById(from_state.getSecondState());
				
				// Get step weights
				// THIS IS DIFFERENT FROM ABOVE!
				double[] weights = new double[rewards.size()];
				
				for (int i=0; i < rewards.size(); i++) {
					double currentWeight = rewards.get(i).getStateReward(from_state.getFirstState());
					currentWeight += rewards.get(i).getTransitionReward(from_state.getFirstState(), choice_i);
					weights[i] = currentWeight;
				}

				// Update accumulation product state, store it and get its ID.
				AccumulationState<Integer> to_accproduct = result.updateAccumulationState(from_state.getFirstState(), from_accstate, accexp, weights, mc);
				int to_accproduct_id = result.accStates.findOrAdd(to_accproduct);
				return new ProductState(mdp_to_state, to_accproduct_id);
			}

			@Override
			public void notify(final ProductState state, final Integer index)
					throws PrismException {
				result.generateTrackInfo(state, index, accexp, mc);
			}

			@Override
			public void finish() throws PrismException {
				// Do nothing
				//mc.getLog().println(".");
			}

			@Override
			public MDP getGraph() {
				return graph;
			}
		}
		
		AccumulationMDPProductOperator op = new AccumulationMDPProductOperator();
		ProductWithProductStates.generate(op, result, statesOfInterest);
		return result;
	}
	
	@Override
	protected boolean isFinalTrack(final AccumulationTrack<Integer> track, final ExpressionAccumulation accexp, final ProbModelChecker mc)
			throws PrismException {
		boolean isFinal = false;
		if ( track != null ) {
			Integer step = track.getComponent();
			if ( step > 0 ) { // if the step is 0, we cannot have a goal state.
				IntegerBound stepBound = IntegerBound.fromTemporalOperatorBound(accexp.getBoundExpression(), mc.getConstantValues(), true);
				isFinal = stepBound.isInBounds(step);
			}
		}
		return isFinal;
	}
	
	@Override
	protected Integer getInitialComponent() {
		return 0;
	}
	
	@Override
	protected Integer updateComponent(Integer modelFromStateId, final AccumulationTrack<Integer> track,
			final ExpressionAccumulation accexp, final StateModelChecker mc) {
		int currentStep = track.getComponent();
		
		int maxStep = 0;
		try {
			maxStep = IntegerBound.fromTemporalOperatorBound(accexp.getBoundExpression(), mc.getConstantValues(), true).getHighestInteger();
		} catch(PrismException e) {
			throw new RuntimeException("...");
		}
		
		int nextStep = currentStep+1;
		if(nextStep > maxStep) { return null; }
		return nextStep;
	}
	
	/**
	 * Populates fields:
	 * - numberOfTracks  with the highest relevant integer plus one, and
	 * - numberOfWeights with the size of the reward vector.
	 * @param graph
	 * @param accexp
	 * @param rewards
	 * @param mc
	 * @throws PrismException
	 */
	protected void createAuxData(final Model graph, final ExpressionAccumulation accexp, final Vector<? extends Rewards> rewards, final ProbModelChecker mc) throws PrismException {
		numberOfTracks  = IntegerBound.fromTemporalOperatorBound(accexp.getBoundExpression(), mc.getConstantValues(), true).getHighestInteger()+2;
		numberOfWeights = rewards.size();
		
		for(int i=0; i<numberOfTracks; i++) {
			initStates.add(new BitSet());
			runStates.add(new BitSet());
			goalStates.add(new BitSet());
		}
	}
}