package de.bmoth.modelchecker;

import com.microsoft.z3.*;
import de.bmoth.backend.Abortable;
import de.bmoth.backend.z3.MachineToZ3Translator;
import de.bmoth.backend.z3.SolutionFinder;
import de.bmoth.backend.z3.Z3SolverFactory;
import de.bmoth.parser.Parser;
import de.bmoth.parser.ast.nodes.DeclarationNode;
import de.bmoth.parser.ast.nodes.MachineNode;
import de.bmoth.preferences.BMothPreferences;

import java.util.*;

public class ModelChecker implements Abortable {
    private Context ctx;
    private Solver solver;
    private MachineToZ3Translator machineTranslator;
    private SolutionFinder finder;
    private boolean isAborted;

    public ModelChecker(MachineNode machine) {
        this.ctx = new Context();
        this.solver = Z3SolverFactory.getZ3Solver(ctx);
        this.machineTranslator = new MachineToZ3Translator(machine, ctx);
        this.finder = new SolutionFinder(solver, ctx);
    }

    @Override
    public void abort() {
        isAborted = true;
        finder.abort();
    }

    public static ModelCheckingResult doModelCheck(String machineAsString) {
        MachineNode machineAsSemanticAst = Parser.getMachineAsSemanticAst(machineAsString);
        return doModelCheck(machineAsSemanticAst);
    }

    public static ModelCheckingResult doModelCheck(MachineNode machine) {
        ModelChecker modelChecker = new ModelChecker(machine);
        return modelChecker.doModelCheck();
    }

    public ModelCheckingResult doModelCheck() {
        isAborted = false;
        Set<State> visited = new HashSet<>();
        Queue<State> queue = new LinkedList<>();
        // prepare initial states
        BoolExpr initialValueConstraint = machineTranslator.getInitialValueConstraint();

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

        final BoolExpr invariant = machineTranslator.getInvariantConstraint();
        solver.add(invariant);
        while (!isAborted && !queue.isEmpty()) {
            solver.push();
            State current = queue.poll();

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

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

            List<BoolExpr> operationConstraints = machineTranslator.getOperationConstraints();
            for (BoolExpr currentOperationConstraint : operationConstraints) {
                // compute successors
                models = finder.findSolutions(currentOperationConstraint, 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");
        } else {
            return new ModelCheckingResult("correct");
        }
    }

    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 : machineTranslator.getVariables()) {
            Expr expr = machineTranslator.getPrimedVariable(declNode);
            Expr value = model.eval(expr, true);
            map.put(declNode.getName(), value);
        }
        for (DeclarationNode declarationNode : machineTranslator.getConstants()) {
            Expr expr = machineTranslator.getVariable(declarationNode);
            Expr value = model.eval(expr, true);
            map.put(declarationNode.getName(), value);
        }

        return new State(predecessor, map);
    }
}

@@ -516,24 +516,53 @@
 ver solver;%0A
+    private Solver opSolver;%0A
     private 
@@ -633,24 +633,61 @@
 der finder;%0A
+    private SolutionFinder opFinder;%0A
     private 
@@ -839,24 +839,82 @@
 olver(ctx);%0A
+        this.opSolver = Z3SolverFactory.getZ3Solver(ctx);%0A
         this
@@ -1026,24 +1026,83 @@
 lver, ctx);%0A
+        this.opFinder = new SolutionFinder(opSolver, ctx);%0A
     %7D%0A%0A    @
@@ -2255,24 +2255,281 @@
 invariant);%0A
+%0A        // create joint operations constraint and permanently add to separate solver%0A        final BoolExpr operationsConstraint = ctx.mkOr(machineTranslator.getOperationConstraints().toArray(new BoolExpr%5B0%5D));%0A        opSolver.add(operationsConstraint);%0A%0A
         whil
@@ -3337,212 +3337,49 @@
     
-List%3CBoolExpr%3E operationConstraints = machineTranslator.getOperationConstraints();%0A            for (BoolExpr currentOperationConstraint : operationConstraints) %7B%0A                // compute successors%0A    
+// compute successors on separate finder%0A
     
@@ -3387,33 +3387,35 @@
        models = 
-f
+opF
 inder.findSoluti
@@ -3422,24 +3422,13 @@
 ons(
-currentOperation
+state
 Cons
@@ -3527,28 +3527,24 @@
             
-    
 models.strea
@@ -3603,36 +3603,32 @@
                 
-   
- 
 .filter(state -%3E
@@ -3662,36 +3662,32 @@
                 
-    
 .filter(state -%3E
@@ -3719,36 +3719,32 @@
                 
-  
-  
 .forEach(queue::
@@ -3749,29 +3749,16 @@
 ::add);%0A
-            %7D
 %0A       

use separate solver / finder and combined operations constraint for better performance