uni/fmi/fmi23block2/exercise8.py

from z3 import *

def all_smt(s, initial_terms):
    def block_term(s, m, t):
        s.add(t != m.eval(t, model_completion=True))
    def fix_term(s, m, t):
        s.add(t == m.eval(t, model_completion=True))
    def all_smt_rec(terms):
        if sat == s.check():
           m = s.model()
           yield m
           for i in range(len(terms)):
               s.push()
               block_term(s, m, terms[i])
               for j in range(i):
                   fix_term(s, m, terms[j])
               yield from all_smt_rec(terms[i:])
               s.pop()   
    yield from all_smt_rec(list(initial_terms))

def main():
    s = Solver()

    # we use ints by normalizing all weights
    # (scaling them by 100, which does not change the result)
    products = [Int(f'g_{i}') for i in range(6)]
    max_weight = 1600
    weights = {
        0: 240,
        1: 275,
        2: 335,
        3: 355,
        4: 420,
        5: 580,
    }

    for i in range(6):
        # all products can be taken at least 0 times
        s.add(products[i] >= 0)
    # finally add the total weight constraint for problem
    s.add(sum([products[i] * weights[i] for i in range(6)]) == max_weight)

    # generate all models using all_smt for our products
    models = list(all_smt(s, products))
    print(f'there are {len(models)} models:')

    # finally we display our results
    for model in models:
        print(model)
        total = sum([model[products[i]] * weights[i] for i in range(6)])
        print(f'where {total == max_weight}')

if __name__ == '__main__':
    main()