package explicit;

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

import automata.finite.DeterministicFiniteAutomaton;
import automata.finite.EdgeLabel;
import automata.finite.NondeterministicFiniteAutomaton;
import automata.finite.State;
import explicit.rewards.MCRewards;
import explicit.rewards.MDPRewards;
import explicit.rewards.Rewards;
import parser.ast.ExpressionAccumulation;
import parser.ast.ExpressionRegular;
import prism.PrismException;
import prism.PrismSettings;

/**
 * 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 AccumulationProductComplexRegular<M extends Model> extends AccumulationProductComplex<M,State>
{
	protected DeterministicFiniteAutomaton<String> automaton;
	
	public AccumulationProductComplexRegular(M originalModel) {
		super(originalModel);
	}
	
	@SuppressWarnings("unchecked")
	public static <T extends Model, R extends Rewards> AccumulationProductComplexRegular<T> generate(final Model graph, final ExpressionAccumulation accexp, final Vector<R> rewards, final ProbModelChecker mc, BitSet statesOfInterest) throws PrismException {
		switch(graph.getModelType()) {
		case DTMC:
			return (AccumulationProductComplexRegular<T>)generate((DTMC) graph, accexp, (Vector<MCRewards>)  rewards, mc, statesOfInterest);
		case MDP:
			return (AccumulationProductComplexRegular<T>)generate((MDP)  graph, accexp, (Vector<MDPRewards>) rewards, mc, statesOfInterest);
		default:
			throw new PrismException("Can't handle accumulation product for " + graph.getModelType());
		}
	}
	
	public static AccumulationProductComplexRegular<DTMC> generate(final DTMC graph, final ExpressionAccumulation accexp, final Vector<MCRewards> rewards, final ProbModelChecker mc, BitSet statesOfInterest) throws PrismException {
		final AccumulationProductComplexRegular<DTMC> result = new AccumulationProductComplexRegular<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<State> 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<State> 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 AccumulationProductComplexRegular<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 AccumulationProductComplexRegular<MDP> result = new AccumulationProductComplexRegular<MDP>(graph);
		// Create auxiliary data
		result.createAuxData(graph, accexp, rewards, mc);

		class AccumulationMDPProductOperator implements MDPProductOperator
		{

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

			@Override
			public ProductState getSuccessor(ProductState from_state,
					int choice_i, Integer mdp_to_state) throws PrismException {
				// Get the current accumulation state
				AccumulationState<State> 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<State> 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(ProductState state, Integer index)
					throws PrismException {
				result.generateTrackInfo(state, index, accexp, mc);
			}

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

			@Override
			public MDP getGraph() {
				return graph;
			}
		}
		
		AccumulationMDPProductOperator op = new AccumulationMDPProductOperator();
		ProductWithProductStates.generate(op, result, statesOfInterest);
		return result;
	}
	
	@Override
	protected boolean isFinalTrack(AccumulationTrack<State> track, ExpressionAccumulation accexp, ProbModelChecker mc) throws PrismException {
		boolean isFinal = false;
		if ( track != null ) {
				isFinal = automaton.isAcceptingState(track.getComponent());
		}
		
		return isFinal;
	}
	
	@Override
	protected State getInitialComponent() {
		return automaton.getInitialState();
	}
	
	@Override
	protected State updateComponent(Integer modelFromStateId, final AccumulationTrack<State> track,
			final ExpressionAccumulation accexp, final StateModelChecker mc) {
		State currentState = track.getComponent();
		
		// Build EdgeLabel from labels.
		// labels is a BitSet with labels L0,...,Ln
		ArrayList<String> edgeSymbols = new ArrayList<>();
		BitSet edgeValues = new BitSet();
		
		for (int i=0; i < labels.size(); i++) {
			// Each bitLabel contains one L
			edgeSymbols.add(i, "L"+i);
			if(labels.get(i).get(modelFromStateId)) {
				edgeValues.set(i);
			}
		}

		EdgeLabel<String> edgeLabel = new EdgeLabel<>(edgeSymbols, edgeValues);
		
		State nextState = automaton.getSuccessor(currentState,edgeLabel);
		if(nextState == null) { return null; }
		return nextState;
	}
	
	/**
	 * Populates fields:
	 * - automaton with a trimmed deterministic finite automaton,
	 * - numberOfTracks  with the length of its longest path plus one, and
	 * - numberOfWeights with the size of the reward vector.
	 * 
	 * @param graph
	 * @param accexp
	 * @param mc
	 * @throws PrismException
	 */
	protected void createAuxData(final Model graph, final ExpressionAccumulation accexp, final Vector<? extends Rewards> rewards, final ProbModelChecker mc) throws PrismException {
		// Build labels and DFA
		AccumulationModelChecker accMc = new AccumulationModelChecker();
		ExpressionRegular reg = (ExpressionRegular)accMc.checkMaximalStateFormulas(mc, graph, accexp.getRegularExpression(), labels);
		
		// Create and store the actual DFA
		NondeterministicFiniteAutomaton<String> nfa = NondeterministicFiniteAutomaton.fromExpressionRegular(reg);
		//System.out.println(nfa.toDot());
		
		automaton = nfa.determinize();
		automaton.trim(); // This should remove cycles.
		
		if(mc.getSettings().getBoolean(PrismSettings.ACC_GENERATE_DOTS)) {
		  nfa.exportToDotFile("DEBUG-automaton-nfa.dot");
		  automaton.exportToDotFile("DEBUG-automaton-dfa.dot");
		}
		
		if (!automaton.isAcyclic()) {
			throw new PrismException("Cannot handle cyclic automata!");
		}
		
		numberOfTracks = automaton.getLongestPathLength()+1;
		numberOfWeights= rewards.size();
		
		for(int i=0; i<numberOfTracks; i++) {
			initStates.add(new BitSet());
			runStates.add(new BitSet());
			goalStates.add(new BitSet());
		}
	}
}