import sys
import itertools
from satispi import *

def symabs(smtlib, exprs, types, cutoffs):
    if len(exprs) == 0: return []

    if Solver.singleton().model(smtlib) is None:
        return [frozenset() for _ in exprs]

    smtlib = smtlib.copy()

    vars_ = [Variable(f"____symabs{i}", types[i]) for i in range(len(exprs))]
    smtlib += [Term("=", var, expr) for var, expr in zip(vars_, exprs)]

    possible = [frozenset() for _ in exprs]
    aim_for = max(len(possible) // 8, 1)
    for i in itertools.count():
        distinction = [bvize(Term("distinct", v, *map(wordlit, sorted(p))))
                       for v, p in zip(vars_, possible)
                       if p and p is not True]

        if i > 0 and not distinction:
            # everything is Top
            return possible

        if len(distinction) == 1:
            n_deltas = Term("bvadd", *distinction, wordlit(0))
        else:
            n_deltas = Term("bvadd", *distinction)

        # aim_for <= n distincts.
        # if aim_for = 1 and it's unsat, then impossible to get any distincts
        goal = []
        if distinction:
            goal = [Term("bvult", wordlit(aim_for - 1), n_deltas)]

        use_vars = [v for v, p in zip(vars_, possible) if p is not True]
        model = Solver.singleton().model(smtlib + goal, use_vars)

        if model is None:
            assert aim_for >= 0
            if aim_for > 1:
                aim_for = max(aim_for // 2, 1)
                continue
            return possible

        prior = possible.copy()
        for i, (v, c) in enumerate(zip(vars_, cutoffs)):
            if possible[i] is True: continue
            if model[v.name] is None: continue
            possible[i] = possible[i] | frozenset({model[v.name]})
            if c is not None and len(possible[i]) > c:
                possible[i] = True
        assert possible != prior, \
            "Solver did not return a complete enough model to make progress"

def wordlit(val):
    return Literal(f"(_ bv{val} 32)")

def bvize(term):
    return Term("ite", term, wordlit(1), wordlit(0))