package explicit;

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

import explicit.rewards.ConstructRewards;
import explicit.rewards.MCRewards;
import explicit.rewards.MDPRewards;
import parser.ast.Expression;
import parser.ast.ExpressionAccumulation;
import parser.ast.ExpressionReward;
import parser.ast.RewardStruct;
import parser.visitor.ReplaceAccumulationExpression;
import prism.PrismException;

public class AccumulationTransformation<M extends Model> implements ModelExpressionTransformation<M, M> {
	final private Expression originalExpression;
	final private M originalModel;
	final private BitSet statesOfInterest;
	final ProbModelChecker mc;
	
	private Expression transformedExpression;
	private AccumulationProduct<M,?> product;
	
	public AccumulationTransformation(
			ProbModelChecker mc,
			M originalModel, Expression expr,
			BitSet statesOfInterest) throws PrismException{
		super();
		this.originalExpression = expr;
		this.originalModel = originalModel;
		this.mc = mc;
		this.statesOfInterest = statesOfInterest;
		doTransformation();
	}
	
	@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);
	}
	
	@SuppressWarnings("unchecked")
	private void doTransformation() throws PrismException {		
		mc.getLog().println("Performing accumulation transformation...");
		// We work on a copy
		transformedExpression = originalExpression.deepCopy();
		
		// Get the first ExpressionAccumulation
		ExpressionAccumulation accexp = transformedExpression.getFirstAccumulationExpression();
		
		// Get the rewards and build the product
		switch(originalModel.getModelType()) {
		case DTMC:
			Vector<MCRewards> dtmc_rewards = new Vector<MCRewards>();
			
			for (int i=0; i < accexp.getConstraint().getFactors().size(); i++) {
				Object rewardIndex = accexp.getConstraint().getFactors().get(i).getFunction().getRewardIndex();
				
				RewardStruct rewStruct = ExpressionReward.getRewardStructByIndexObject(rewardIndex, mc.modulesFile, originalModel.getConstantValues());
				ConstructRewards constructRewards = new ConstructRewards();
				constructRewards.allowNegativeRewards();
				
				MCRewards dtmc_reward = constructRewards.buildMCRewardStructure((DTMC)originalModel, rewStruct, mc.getConstantValues());
				dtmc_rewards.add(i,dtmc_reward);
			}
			mc.getLog().println(" [AT] performing product construction...");
			if(accexp.hasRegularExpression()) {
				product = (AccumulationProductRegular<M>) AccumulationProductRegular.generate((DTMC)originalModel, accexp, dtmc_rewards, mc, statesOfInterest);
			} else if (accexp.hasBoundExpression()) {
				product = (AccumulationProductCounting<M>) AccumulationProductCounting.generate((DTMC)originalModel, accexp, dtmc_rewards, mc, statesOfInterest);
			} else {
				throw new PrismException("Accumulation Expression has no valid monitor!");
			}
			break;
		case MDP:
			Vector<MDPRewards> mdp_rewards = new Vector<MDPRewards>();
			
			for (int i=0; i < accexp.getConstraint().getFactors().size(); i++) {
				Object rewardIndex = accexp.getConstraint().getFactors().get(i).getFunction().getRewardIndex();
				
				RewardStruct rewStruct = ExpressionReward.getRewardStructByIndexObject(rewardIndex, mc.modulesFile, originalModel.getConstantValues());
				ConstructRewards constructRewards = new ConstructRewards();
				constructRewards.allowNegativeRewards();
				
				MDPRewards mdp_reward = constructRewards.buildMDPRewardStructure((MDP)originalModel, rewStruct, mc.getConstantValues());
				mdp_rewards.add(i,mdp_reward);
			}
			mc.getLog().println(" [AT] performing product construction...");
			if(accexp.hasRegularExpression()) {
				product = (AccumulationProductRegular<M>) AccumulationProductRegular.generate((MDP)originalModel, accexp, mdp_rewards, mc, statesOfInterest);
			} else if (accexp.hasBoundExpression()) {
				product = (AccumulationProductCounting<M>) AccumulationProductCounting.generate((MDP)originalModel, accexp, mdp_rewards, mc, statesOfInterest);
			} else {
				throw new PrismException("Accumulation Expression has no valid monitor!");
			}
			
			break;
		default:
			throw new PrismException("Can't handle weight functions for " + originalModel.getModelType());
		}
		
		// Transform the model
		BitSet goodStates = product.getGoodStates();
		String label = gensymLabel("good", product.getTransformedModel());
		
		((ModelExplicit)product.getTransformedModel()).addLabel(label, goodStates);
		
		//System.out.println("Good states " + goodStates);
		
		// Transform the expression
		ReplaceAccumulationExpression replace = new ReplaceAccumulationExpression(accexp, label, product.getNumberOfTracks()-1);
		transformedExpression = (Expression)transformedExpression.accept(replace);
		mc.getLog().println("Transformed " + originalExpression.toString() +
							"\n into " + transformedExpression.toString());
		//DEBUG: output dotfile
		product.exportToDotFile("DEBUG-product.dot");
	}
	
	public String gensymLabel(String prefix, Model model) {
		int suffix = 0;
		String label = prefix + "_" + suffix;
		while(product.getTransformedModel().getLabels().contains(label)) {
			suffix++;
			label = prefix + "_" + suffix;
		}
		return label;
	}

	@Override
	public Integer mapToTransformedModel(int state) {
		return product.mapToTransformedModel(state);
	}

	@Override
	public BitSet mapToTransformedModel(BitSet states) {
		return product.mapToTransformedModel(states);
	}
	
}