package explicit;

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

import parser.ast.AccumulationFactor;
import parser.ast.ExpressionAccumulation;
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
{
	final StoragePool<AccumulationTracker<Component>> trackers;
	final StoragePool<AccumulationState<Component>> accStates;
	
	final BitSet goodStates;

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

	protected abstract boolean isFinalTrack(final AccumulationTrack<Component> track, final ExpressionAccumulation accexp, final ProbModelChecker mc)
			throws PrismException;
	
	protected boolean isGoodTrack(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;
		
		// 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.
		switch(accexp.getSymbol()) {
			case ACCBOXMINUS:
			case ACCBOXPLUS:
				if (!rhs.isInBounds(lhs)) {
					isGood = true;
				}
				break;
			case ACCDIAMINUS:
			case ACCDIAPLUS:
				if (rhs.isInBounds(lhs)) {
					isGood = true;
				}
				break;
			default:
				throw new RuntimeException("Oh boy!");
			}
		//if(isGood) {mc.getLog().print("+");} else {mc.getLog().print("-");}
		return isGood;
	}
	
	protected boolean isGoodAccState(final AccumulationState<Component> state, final ExpressionAccumulation accexp, final ProbModelChecker mc)
			throws PrismException {
		return state.hasGoodTrack();
	}

	protected abstract Component getInitialComponent();
	
	protected abstract AccumulationTrack<Component> updateTrack(Integer modelFromStateId, final AccumulationTrack<Component> track,
			final ExpressionAccumulation accexp, final double[] weights, final StateModelChecker mc);
	
	protected AccumulationState<Component> updateAccumulationState(final int modelFromStateId,
			final AccumulationState<Component> accstate, final ExpressionAccumulation accexp,
			final double[] weights, final ProbModelChecker mc) throws PrismException {
		// We have the current accumulation state, the current model id and the accumulation expression.

		// Get the old tracker and tracks.
		AccumulationTracker<Component> oldTracker = accstate.getTracker(trackers);
		ArrayList<AccumulationTrack<Component>> oldTracks = oldTracker.getTracks();
		
		BitSet oldGoodTracks = accstate.getGoodTracks();
		BitSet newGoodTracks = new BitSet();
		
		// 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;
			boolean oldTrackGood = oldGoodTracks.get(trackNr);
			boolean newTrackGood;
			
			// restart or advance
			if(trackNr == newLastRestartNr) {
				// If we have a restart, we produce an initial track and reset its goodness.
				newTrack = new AccumulationTrack<Component>(numberOfWeights, getInitialComponent());
				newTrackGood = false;
			} else if (oldTrackGood || oldTrack == null) {
				// If the old track is already good or undefined, the new track is as well. Goodness indicators stay.
				newTrack = null;
				newTrackGood = oldTrackGood;
			} else {
				// Otherwise, the track is defined and not good yet.
				assert oldTrack != null;
				newTrack = updateTrack(modelFromStateId, oldTrack, accexp, weights, mc);
				newTrackGood = isGoodTrack(newTrack, accexp, mc);
			}
			
			// if the new track is good, we can forget about its contents and set the new goodness
			if(newTrackGood) {
				newTrack = null;
				newGoodTracks.set(trackNr);
			};
			
			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, newGoodTracks);
	}
	
	@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();
	}
}