package explicit;

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

import common.StopWatch;
import explicit.rewards.ConstructRewards;
import explicit.rewards.Rewards;
import parser.ast.Expression;
import parser.ast.ExpressionAccumulation;
import parser.ast.ExpressionReward;
import parser.ast.RewardStruct;
import parser.visitor.ReplaceAccumulationExpressionComplex;
import parser.visitor.ReplaceAccumulationExpressionSimple;
import parser.visitor.ASTVisitor;
import parser.visitor.ReplaceAccumulationBox;
import prism.PrismException;
import prism.PrismSettings;

public class AccumulationTransformation<M extends ModelExplicit> implements ModelExpressionTransformation<M, M> {
	final protected Expression originalExpression;
	final protected M originalModel;
	final protected BitSet statesOfInterest;
	final protected StateModelChecker mc;
	
	protected Expression transformedExpression;
	protected AccumulationProduct<M,?> product;
	
	ArrayList<String> initLabels;
	ArrayList<String> runLabels;
	ArrayList<String> goalLabels;
	
	public AccumulationTransformation(
			StateModelChecker mc,
			M originalModel, Expression expr,
			BitSet statesOfInterest, boolean forceComplex) throws PrismException{
		super();
		this.originalExpression = expr;
		this.originalModel = originalModel;
		this.mc = mc;
		this.statesOfInterest = statesOfInterest;
		
		this.initLabels = new ArrayList<>();
		this.runLabels  = new ArrayList<>();
		this.goalLabels = new ArrayList<>();
		
		doTransformation(forceComplex);
	}
	

	public AccumulationTransformation(
			StateModelChecker mc,
			M originalModel, Expression expr,
			BitSet statesOfInterest) throws PrismException{
		this(mc, originalModel, expr, statesOfInterest, false);
	}
	
	@Override
	public M getOriginalModel() {
		return originalModel;
	}
	@Override
	public Expression getTransformedExpression() {
		return transformedExpression;
	}
	@Override
	public Expression getOriginalExpression() {
		return originalExpression;
	}
	
	@Override
	public BitSet getTransformedStatesOfInterest() {
		return product.getTransformedStatesOfInterest();
	}

	@Override
	public M getTransformedModel() {
		return product.getTransformedModel();
	}
	@Override
	public StateValues projectToOriginalModel(StateValues svTransformedModel)
			throws PrismException {
		return product.projectToOriginalModel(svTransformedModel);
	}
	
	protected void doTransformation(boolean forceComplexFlag) throws PrismException {
		StopWatch clock = new StopWatch(mc.getLog());

		mc.getLog().println("Handling maximal state formulas...");
		transformedExpression = mc.handleMaximalStateFormulas(originalModel, originalExpression);		
				
		// Get the first ExpressionAccumulation
		mc.getLog().println("Looking for an accumulation formula...");
		
		boolean isSingleTrack = true;
		ExpressionAccumulation accexp = transformedExpression.getFirstAccumulationExpressionWithoutTemporal();
		if(accexp == null || mc.getSettings().getBoolean(PrismSettings.ACC_FORCE_MULTI)) {
			isSingleTrack = false;
			accexp = transformedExpression.getFirstAccumulationExpression();
		}
		
		mc.getLog().println("... found " + accexp);
	    
	    // Rewrite it, if it is a BOX
	    if (accexp.getSymbol().isBox()) {
	    	ReplaceAccumulationBox replace = new ReplaceAccumulationBox(accexp);
	    	transformedExpression = (Expression)transformedExpression.accept(replace);
	    	accexp = transformedExpression.getFirstAccumulationExpression();
	    	mc.getLog().println("... after unboxing: " + accexp + " (in " + transformedExpression + ")");
	    }
		
		// Get the rewards
	    HashMap<Object, Rewards> rewards = new HashMap<Object, Rewards>();
		
		for (int i=0; i < accexp.getFunctions().size(); i++) {
			Object rewardIndex = accexp.getFunctions().get(i).getRewardIndex();
			
			RewardStruct rewStruct = ExpressionReward.getRewardStructByIndexObject(rewardIndex, mc.modulesFile, originalModel.getConstantValues());
			ConstructRewards constructRewards = new ConstructRewards();
			constructRewards.allowNegativeRewards();
			
			Rewards reward = constructRewards.buildRewardStructure(originalModel, rewStruct, mc.getConstantValues());
			rewards.put(rewardIndex,reward);
		}
		
		// Figure out if we need a complex or a simple trafo...
		//boolean isSimple     = !accexp.hasFireOn() && accexp.isNullary();
		boolean isSimple     = accexp.isNullary();
		boolean isFiltered   = accexp.hasRecordSet();
		boolean isPast       = accexp.getSymbol().isMinus();
		boolean forceComplex = mc.getSettings().getBoolean(PrismSettings.ACC_FORCE_COMPLEX) || forceComplexFlag;
		boolean forceMulti   = mc.getSettings().getBoolean(PrismSettings.ACC_FORCE_MULTI);
		boolean forceUntil   = mc.getSettings().getBoolean(PrismSettings.ACC_FORCE_UNTIL);
		
		if(forceComplex && isPast) {
			throw new PrismException("Unable to handle <->/[-] with -accforcecomplex. Bailing.");
		}
		
		// Build the AccumulationContext
		AccumulationContext ctx = new AccumulationContext(accexp, rewards, mc);
		ctx.singleTrack  = isSingleTrack && !forceMulti;
		ctx.simpleMethod = isSimple      && !forceComplex;
		ctx.reachConstr  = (!isSimple || !isFiltered) && !forceUntil;
		
		mc.getLog().print("Context flags: ");
		if(ctx.singleTrack)  { mc.getLog().print(" singleTrack "); }
		if(ctx.simpleMethod) { mc.getLog().print(" simpleMethod "); }
		if(ctx.reachConstr)  { mc.getLog().print(" reachConstr" ); }
		mc.getLog().println();
		
		// Build the product
		clock.start("accumulation product construction");
		if(ctx.accexp.hasRegularExpression()) {
			product = AccumulationProductRegular.generate(originalModel, ctx, statesOfInterest);
		} else if (ctx.accexp.hasBoundExpression()) {
			product = AccumulationProductCounting.generate(originalModel, ctx, statesOfInterest);
		} else {
			throw new PrismException("Accumulation Expression has no valid monitor!");
		}
		clock.stop(product.getTransformedModel().getNumStates() + " states");
		
		// Attach labels
		clock.start("attaching labels");

		if(ctx.simpleMethod) {
			BitSet goodStates = product.getGoalStates(0);
			String goodLabel = ((ModelExplicit)product.getTransformedModel()).addUniqueLabel("goal", goodStates, product.getTransformedModel().getLabels());
			
			goalLabels.add(goodLabel);
		} else {
			for(int track=0; track < product.getNumberOfTracks(); track++) {
				BitSet initStates = product.getInitStates(track);
				BitSet runStates  = product.getRunStates(track);
				BitSet goalStates = product.getGoalStates(track);
				
				String initLabel = ((ModelExplicit)product.getTransformedModel()).addUniqueLabel("init" + track, initStates, product.getTransformedModel().getLabels());
				initLabels.add(initLabel);
				String runLabel  = ((ModelExplicit)product.getTransformedModel()).addUniqueLabel("run" + track,  runStates, product.getTransformedModel().getLabels());
				runLabels.add(runLabel);
				String goalLabel = ((ModelExplicit)product.getTransformedModel()).addUniqueLabel("goal" + track, goalStates, product.getTransformedModel().getLabels());
				goalLabels.add(goalLabel);
			}
		}
		
		clock.stop(initLabels.size() + "/" + runLabels.size() + "/" + goalLabels.size() + " labels");
		
		// Transform the expression
		clock.start("replacing formula");
		ASTVisitor replace;
		if(ctx.simpleMethod) {
			replace = new ReplaceAccumulationExpressionSimple(ctx.accexp, goalLabels.get(0), product.getAccumulationLength(), product.getNumberOfTracks(), ctx.reachConstr);
		} else {
			replace = new ReplaceAccumulationExpressionComplex(ctx.accexp, initLabels, runLabels, goalLabels, product.getAccumulationLength(), product.getNumberOfTracks());
		}
		transformedExpression = (Expression)transformedExpression.accept(replace);
		clock.stop("got " + transformedExpression);
		//clock.stop();
		
		if(mc.getSettings().getBoolean(PrismSettings.ACC_GENERATE_DOTS)) {
			product.exportToDotFile("DEBUG-product.dot");
		}
	}
}