from ..AST import Eq, Not

class SmtlibCapability:

    ### serialize constraints with z3 compatibility
    def to_smt2(self):
        res = []
        res.append("; generated by vega")
        res.append("(set-info :status unknown)")

        ### Collect sorts
        sorts = set()
        for v in self.variables.keys():
            sorts.add(v.sort)

        ### Define domain
        res.append("(declare-datatypes () (({} {})))".format(self.domain.name, ' '.join(["({})".format(x) for x in self.domain.values])))

        ### Define sorts
        for sort in sorts: # TODO
            res.append("(declare-datatypes () (({} {})))".format(sort.name, ' '.join(["({})".format(x) for x in sort.values])))
        for v in self.domain.values: # reffers domain.values to avoid redefine symbols (e.g. a of sort s, a of sort t)
            res.append("(declare-fun {} () {})".format(v.name, self.domain.name))

        ### Define all variables with same domain (sort)
        for v in self.variables.keys():
            res.append("(declare-fun {} () {})".format(v.name, self.domain.name))
        
        for expr in self.constraints:
            res.append("(assert\n {})".format(expr.to_smt2()))

        ### Restrict values
        for v in self.variables.keys():
            if v.sort != self.domain: # e.g. sort FileContent
                for value in self.domain.values - v.sort.values:
                    res.append("(assert\n {})".format(Not(Eq(v, value)).to_smt2()))

        res.append("(check-sat)")
        
        return '\n'.join(res)