package explicit;

import java.util.BitSet;
import java.util.HashMap;

import explicit.rewards.Rewards;
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, AccumulationContext ctx) {
		super(originalModel, ctx);
	}
	
	@SuppressWarnings("unchecked")
	public static <T extends Model, R extends Rewards> AccumulationProductCounting<T> generate(final Model graph, AccumulationContext ctx, BitSet statesOfInterest) throws PrismException {
		switch(graph.getModelType()) {
		case DTMC:
			return (AccumulationProductCounting<T>)generate((DTMC) graph, ctx, statesOfInterest);
		case MDP:
			return (AccumulationProductCounting<T>)generate((MDP)  graph, ctx, statesOfInterest);
		default:
			throw new PrismException("Can't handle accumulation product for " + graph.getModelType());
		}
	}
	
	public static AccumulationProductCounting<DTMC> generate(final DTMC graph, AccumulationContext ctx, BitSet statesOfInterest) throws PrismException {
		final AccumulationProductCounting<DTMC> result = new AccumulationProductCounting<DTMC>(graph, ctx);
		// Create auxiliary data
		result.createAuxData(graph);
		
		// 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
				HashMap<Object, Double> weights = ctx.getWeights(from_state.getFirstState());
				
				// Update accumulation product state, store it and get its ID.
				AccumulationState<Integer> to_accproduct = result.updateAccumulationState(from_state.getFirstState(), from_accstate, weights);
				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);
			}

			@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, AccumulationContext ctx, BitSet statesOfInterest) throws PrismException {
		// This is basically the same thing as for DTMCs
		final AccumulationProductCounting<MDP> result = new AccumulationProductCounting<MDP>(graph, ctx);
		
		// Create auxiliary data
		result.createAuxData(graph);

		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
				HashMap<Object, Double> weights = ctx.getWeights(from_state.getFirstState(), choice_i);

				// Update accumulation product state, store it and get its ID.
				AccumulationState<Integer> to_accproduct = result.updateAccumulationState(from_state.getFirstState(), from_accstate, weights);
				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);
			}

			@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) 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(ctx.accexp.getBoundExpression(), ctx.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) {
		int currentStep = track.getComponent();
		
		int maxStep = 0;
		try {
			maxStep = IntegerBound.fromTemporalOperatorBound(ctx.accexp.getBoundExpression(), ctx.mc.getConstantValues(), true).getHighestInteger();
		} catch(PrismException e) {
			throw new RuntimeException("...");
		}
		
		int nextStep = currentStep+1;
		if(nextStep > maxStep) { return null; }
		return nextStep;
	}
	
	protected void createAuxData(final Model graph) throws PrismException {
		numberOfTracks  = IntegerBound.fromTemporalOperatorBound(ctx.accexp.getBoundExpression(), ctx.mc.getConstantValues(), true).getHighestInteger()+1;
		numberOfWeights = ctx.rewards.size();
		
		if(ctx.simpleMethod) {
			goalStates.add(new BitSet());
		} else {
			for(int i=0; i<numberOfTracks; i++) {
				initStates.add(new BitSet());
				runStates.add(new BitSet());
				goalStates.add(new BitSet());
			}
		}
	}
}