package de.bmoth.modelchecker.esmc;

import com.microsoft.z3.*;
import de.bmoth.backend.z3.SolutionFinder;
import de.bmoth.backend.z3.Z3SolverFactory;
import de.bmoth.modelchecker.ModelChecker;
import de.bmoth.modelchecker.State;
import de.bmoth.parser.ast.nodes.DeclarationNode;
import de.bmoth.parser.ast.nodes.MachineNode;
import de.bmoth.preferences.BMothPreferences;

import java.util.*;

import static de.bmoth.backend.TranslationOptions.PRIMED_0;

public class ExplicitStateModelChecker extends ModelChecker<ModelCheckingResult> {
    private Solver solver;
    private Solver opSolver;
    private SolutionFinder finder;
    private SolutionFinder opFinder;

    public ExplicitStateModelChecker(MachineNode machine) {
        super(machine);
        this.solver = Z3SolverFactory.getZ3Solver(getContext());
        this.opSolver = Z3SolverFactory.getZ3Solver(getContext());
        this.finder = new SolutionFinder(solver, getContext());
        this.opFinder = new SolutionFinder(opSolver, getContext());
    }

    public static ModelCheckingResult check(MachineNode machine) {
        ExplicitStateModelChecker modelChecker = new ExplicitStateModelChecker(machine);
        return modelChecker.check();
    }

    @Override
    public void abort() {
        super.abort();
        finder.abort();
        opFinder.abort();
    }

    @Override
    protected ModelCheckingResult doModelCheck() {
        Set<State> visited = new HashSet<>();
        Queue<State> queue = new LinkedList<>();
        // prepare initial states
        BoolExpr initialValueConstraint = getMachineTranslator().getInitialValueConstraint();

        Set<Model> models = finder.findSolutions(initialValueConstraint,
            BMothPreferences.getIntPreference(BMothPreferences.IntPreference.MAX_INITIAL_STATE));
        models.stream().map(this::getStateFromModel).filter(s -> !queue.contains(s)).forEach(queue::add);

        final BoolExpr invariant = getMachineTranslator().getInvariantConstraint();
        solver.add(invariant);

        // create joint operations constraint and permanently add to separate
        // solver
        final BoolExpr operationsConstraint = getContext()
            .mkOr(getMachineTranslator().getOperationConstraints().toArray(new BoolExpr[0]));
        opSolver.add(operationsConstraint);

        while (!isAborted() && !queue.isEmpty()) {
            solver.push();
            State current = queue.poll();
            visited.add(current);

            // apply current state - remains stored in solver for loop iteration
            BoolExpr stateConstraint = current.getStateConstraint(getContext());
            solver.add(stateConstraint);

            // check invariant & state
            Status check = solver.check();
            switch (check) {
                case UNKNOWN:
                    return new ModelCheckingResult("check-sat = unknown, reason: " + solver.getReasonUnknown(),
                        visited.size());
                case UNSATISFIABLE:
                    return new ModelCheckingResult(current, visited.size());
                case SATISFIABLE:
                default:
                    // continue
            }

            // compute successors on separate finder
            models = opFinder.findSolutions(stateConstraint,
                BMothPreferences.getIntPreference(BMothPreferences.IntPreference.MAX_TRANSITIONS));
            models.stream().map(model -> getStateFromModel(current, model)).filter(state -> !visited.contains(state))
                .filter(state -> !queue.contains(state)).forEach(queue::add);

            solver.pop();
        }

        if (isAborted()) {
            return new ModelCheckingResult("aborted", visited.size());
        } else {
            return new ModelCheckingResult("correct", visited.size());
        }
    }

    private State getStateFromModel(Model model) {
        return getStateFromModel(null, model);
    }

    private State getStateFromModel(State predecessor, Model model) {
        HashMap<String, Expr> map = new HashMap<>();
        for (DeclarationNode declNode : getMachineTranslator().getVariables()) {
            Expr expr = getMachineTranslator().getPrimedVariable(declNode, PRIMED_0);
            Expr value = model.eval(expr, true);
            map.put(declNode.getName(), value);
        }
        for (DeclarationNode declarationNode : getMachineTranslator().getConstants()) {
            Expr expr = getMachineTranslator().getVariable(declarationNode);
            Expr value = model.eval(expr, true);
            map.put(declarationNode.getName(), value);
        }

        return new State(predecessor, map);
    }
}

@@ -2164,39 +2164,8 @@
 t = 
-getContext()%0A            .mkOr(
 getM
@@ -2186,16 +2186,24 @@
 or().get
+Combined
 Operatio
@@ -2217,36 +2217,9 @@
 aint
-s().toArray(new BoolExpr%5B0%5D)
+(
 );%0A 

refactor esmc to use combined operations