import z3
from pysat.formula import CNF
from pysat.solvers import Solver


class PysatSolver:

    # add formula to Z3 goal
    def add(self, formula):
        self.goal.add(formula)

    # transform goal to CNF, then to DIMACS, then call the PySat solver
    def check(self):
        z3cnf = z3.Then("simplify", "pb2bv", "bit-blast", "tseitin-cnf")(self.goal)
        H = z3.Goal()
        for x in z3cnf[0]:
            if x == False:
                return z3.unsat
            else:
                # print(z3.simplify(x))
                H.add(x)
        print(f"({len(H)} clauses)", end=" ", flush=True)
        self.dimacs = H.dimacs()
        cnf = CNF(from_string=self.dimacs)
        self.PysatSolver = Solver(name=self.satsolver, bootstrap_with=cnf)
        if self.PysatSolver.solve() == True:
            return z3.sat
        else:
            return z3.unsat

    # Get a Z3 model from the PySat model
    def compute_model(self, PysatModel):
        # First read the mapping var -> id from the comments in DIMACS
        mapping = dict()
        for line in self.dimacs.splitlines():
            if line[0] == "c":
                (_, id, var) = line.split(" ")
                if "!" not in var:  # excluding Tseitin variables
                    mapping[var] = int(id)
        # Convert to a Z3 model
        model = dict()
        for x in mapping:
            if PysatModel[mapping[x] - 1] > 0:
                model[z3.Bool(x)] = z3.simplify(z3.And())  # true
            else:
                model[z3.Bool(x)] = z3.simplify(z3.Or())  # false
        return model

    def model(self):
        if self.z3model == None:
            self.z3model = self.compute_model(self.PysatSolver.get_model())
        return self.z3model

    def reset(self):
        self.z3model = None
        self.dimacs = None
        self.goal = z3.Goal()

    def __init__(self, satsolver):
        self.goal = z3.Goal()
        self.z3model = None
        self.dimacs = None
        self.satsolver = satsolver
