package de.bmoth.backend;

import com.microsoft.z3.BoolExpr;
import com.microsoft.z3.Context;
import com.microsoft.z3.Expr;
import com.microsoft.z3.Sort;
import de.bmoth.parser.ast.nodes.*;

import java.util.*;

public class MachineToZ3Translator {
    private final MachineNode machineNode;
    private final Context z3Context;
    private final BoolExpr initialisationConstraint;
    private final BoolExpr invariantConstraint;
    private final HashMap<String, String> primedVariablesToVariablesMap;
    private final List<BoolExpr> operationConstraints;

    public MachineToZ3Translator(MachineNode machineNode, Context ctx) {
        this.machineNode = machineNode;
        this.z3Context = ctx;
        this.initialisationConstraint = visitSubstitution(machineNode.getInitialisation());
        this.invariantConstraint = (BoolExpr) FormulaToZ3Translator.translatePredicate(machineNode.getInvariant(),
            z3Context);
        this.operationConstraints = visitOperations(machineNode.getOperations());

        {
            primedVariablesToVariablesMap = new HashMap<>();
            for (DeclarationNode node : machineNode.getVariables()) {
                primedVariablesToVariablesMap.put(getPrimedName(node.getName()), node.getName());
            }
        }

    }

    private List<BoolExpr> visitOperations(List<OperationNode> operations) {
        List<BoolExpr> results = new ArrayList<>(operations.size());
        for (OperationNode operationNode : this.machineNode.getOperations()) {
            BoolExpr temp = visitSubstitution(operationNode.getSubstitution());
            // for unassigned variables add a dummy assignment, e.g. x' = x
            Set<DeclarationNode> set = new HashSet<>(this.getVariables());
            set.removeAll(operationNode.getSubstitution().getAssignedVariables());
            for (DeclarationNode node : set) {
                BoolExpr mkEq = z3Context.mkEq(getPrimedVariable(node), getVariableAsZ3Expression(node));
                temp = z3Context.mkAnd(temp, mkEq);
            }
            results.add(temp);
        }
        return results;
    }

    public List<DeclarationNode> getVariables() {
        return machineNode.getVariables();
    }

    public List<DeclarationNode> getConstants() {
        return machineNode.getConstants();
    }

    public Expr getVariableAsZ3Expression(DeclarationNode node) {
        Sort type = FormulaToZ3Translator.bTypeToZ3Sort(z3Context, node.getType());
        Expr expr = z3Context.mkConst(node.getName(), type);
        return expr;
    }

    public Expr getPrimedVariable(DeclarationNode node) {
        String primedName = getPrimedName(node.getName());
        Sort type = FormulaToZ3Translator.bTypeToZ3Sort(z3Context, node.getType());
        Expr expr = z3Context.mkConst(primedName, type);
        return expr;
    }

    public BoolExpr getInitialValueConstraint() {
        PredicateNode properties = machineNode.getProperties();
        if (properties != null) {
            BoolExpr prop = FormulaToZ3Translator.translatePredicate(machineNode.getProperties(), z3Context,
                new TranslationOptions(1));
            return z3Context.mkAnd(initialisationConstraint, prop);
        } else {
            return initialisationConstraint;
        }
    }

    public BoolExpr getInvariantConstraint() {
        return invariantConstraint;
    }

    private BoolExpr visitSubstitution(SubstitutionNode node) {
        if (node instanceof SingleAssignSubstitutionNode) {
            return visitSingleAssignSubstitution((SingleAssignSubstitutionNode) node);
        } else if (node instanceof ParallelSubstitutionNode) {
            return visitParallelSubstitution((ParallelSubstitutionNode) node);
        } else if (node instanceof AnySubstitutionNode) {
            return visitAnySubstitution((AnySubstitutionNode) node);
        } else if (node instanceof SelectSubstitutionNode) {
            return visitSelectSubstitutionNode((SelectSubstitutionNode) node);
        }
        throw new AssertionError("Not implemented" + node.getClass());
    }

    private BoolExpr visitSelectSubstitutionNode(SelectSubstitutionNode node) {
        BoolExpr condition = (BoolExpr) FormulaToZ3Translator.translatePredicate(node.getCondition(), z3Context);
        BoolExpr substitution = visitSubstitution(node.getSubstitution());
        return z3Context.mkAnd(condition, substitution);
    }

    private BoolExpr visitAnySubstitution(AnySubstitutionNode node) {
        throw new AssertionError("Not implemented: " + node.getClass());// TODO
    }

    private BoolExpr visitParallelSubstitution(ParallelSubstitutionNode node) {
        List<SubstitutionNode> substitutions = node.getSubstitutions();
        BoolExpr boolExpr = null;
        for (SubstitutionNode substitutionNode : substitutions) {
            BoolExpr temp = visitSubstitution(substitutionNode);
            if (boolExpr == null) {
                boolExpr = temp;
            } else {
                boolExpr = z3Context.mkAnd(boolExpr, temp);
            }
        }
        return boolExpr;
    }

    private BoolExpr visitSingleAssignSubstitution(SingleAssignSubstitutionNode node) {
        Sort bTypeToZ3Sort = FormulaToZ3Translator.bTypeToZ3Sort(z3Context, node.getIdentifier().getType());
        Expr value = FormulaToZ3Translator.translateExpr(node.getValue(), z3Context);
        String name = getPrimedName(node.getIdentifier().getName());
        Expr variable = z3Context.mkConst(name, bTypeToZ3Sort);
        return this.z3Context.mkEq(variable, value);
    }

    private String getPrimedName(String name) {
        return name + "'";
    }

    public List<BoolExpr> getOperationConstraints() {
        return operationConstraints;
    }
}

@@ -326,38 +326,32 @@
 xt;%0A    private 
-final 
 BoolExpr initial
@@ -359,32 +359,39 @@
 sationConstraint
+ = null
 ;%0A    private fi
@@ -388,22 +388,16 @@
 private 
-final 
 BoolExpr
@@ -412,24 +412,31 @@
 ntConstraint
+ = null
 ;%0A    privat
@@ -697,16 +697,76 @@
  = ctx;%0A
+%0A        if (machineNode.getInitialisation() != null) %7B%0A    
         
@@ -841,32 +841,96 @@
 tialisation());%0A
+        %7D%0A        if (machineNode.getInvariant() != null) %7B%0A    
         this.inv
@@ -1040,16 +1040,20 @@
         
+    
 z3Contex
@@ -1048,32 +1048,127 @@
     z3Context);%0A
+        %7D else %7B%0A            this.invariantConstraint = z3Context.mkTrue();%0A        %7D%0A        %0A
         this.ope
@@ -3170,24 +3170,68 @@
 operties();%0A
+        BoolExpr prop = z3Context.mkTrue();%0A
         if (
@@ -3260,33 +3260,24 @@
             
-BoolExpr 
 prop = Formu
@@ -3396,28 +3396,35 @@
 ns(1));%0A
+%0A
     
+    %7D%0A
         
 return z
@@ -3415,38 +3415,19 @@
         
-return z3Context.mkAnd
+if 
 (initial
@@ -3447,38 +3447,55 @@
 aint
-, prop);%0A        %7D else %7B
+ == null) %7B%0A            return prop;
 %0A    
+    %7D%0A
     
@@ -3505,16 +3505,32 @@
  return 
+z3Context.mkAnd(
 initiali
@@ -3545,27 +3545,24 @@
 nstraint
-;%0A        %7D
+, prop);
 %0A    %7D%0A%0A

handle machines without variables / initialisations / invariants