#!/usr/bin/env python
# Usage for 1a (where -1 is the number of prittles, -1 means unrestricted)
#   $0 -1 | z3 -smt -in
# Usage for 1b:
#   $0 -1 b | z3 -smt -in

import sys

number_of_prittles = int(sys.argv[1])
prtl = None if number_of_prittles == -1 else number_of_prittles
is_b = len(sys.argv) >= 3 and sys.argv[2] == 'b'

# List of m types ("m types"):
# - exact number of items needed (or None for no restriction)
# - Name (first letter is used for variable)
# - Weight (in kg)
stuff = [
    (4,     "Nuzzles",  700),
    (prtl,  "Prittles", 400),
    (8,     "Skipples", 1000),
    (10,    "Crottles", 2500),
    (20,    "Dupples",  200),
]

# T trucks ("n trucks")
n_trucks = 8
# Capacity C_T
truck_capacity = 8000
# Maximum pallets p_T
max_pallets = 8

def fillin(cases, template_vars):
    if type(cases) == str:
        return cases.format(**template_vars)
    return [template.format(**template_vars) for template in cases]

literals = []
preds = []

# For each truck,
# define the number of pallets of each type that will be contained in.
for i in range(1, n_trucks + 1):
    for needed, name, weight in stuff:
        name = "{name[0]}{i}".format(name=name, i=i)
        literals.append(name)
        # Cannot have negative number of items...
        preds += fillin([
            "(>= {name} 0)"
        ], vars())

# For each item, require exact as many items spread over all trucks
for needed, name, weight in stuff:
    if needed is None:
        continue
    # total_items = item_in_truck0 + ... + item_in_truckN
    items_in_truckX = " ".join(
        "{name[0]}{i}".format(name=name, i=i)
        for i in range(1, n_trucks + 1)
    )
    preds += fillin([
        "(= {needed} (+ {items_in_truckX}))",
    ], vars())

# Each truck may hold a maximum number of pallets.
# and a maximum weight.
for i in range(1, n_trucks + 1):
    # max_pallets >= item0_in_truck + ... + itemM_in_truck
    itemX_in_truck = " ".join(
        "{name[0]}{i}".format(name=name, i=i)
        for needed, name, weight in stuff
    )
    total_weight_in_truck = " ".join(
        "(* {name[0]}{i} {weight})".format(name=name, i=i, weight=weight)
        for needed, name, weight in stuff
    )
    preds += fillin([
        "(>= {max_pallets} (+ {itemX_in_truck}))",
        "(>= {truck_capacity} (+ {total_weight_in_truck}))",
    ], vars())


# Skipples must be distributed over at most three trucks.
# So, Any combination of n-3 trucks must have zero skipples.
number_of_skipples = None
for needed, name, weight in stuff:
    if name[0] == 'S':
        number_of_skipples = needed
assert number_of_skipples is not None
skipple_preds = []
for x in range(1, n_trucks + 1):
    for y in range(1, n_trucks + 1):
        if x == y:
            # Require unique pairs (x, y, 0, ..., 0) as optimization.
            # This is weaker than (x, 0, ... 0)
            continue
        for z in range(1, n_trucks + 1):
            if x == z or y == z:
                # Same as above, require unique x, y, z such that
                # (x, y, z, 0, ..., 0)
                continue

            # "These threee trucks hold all available skipples"
            #items_all_in_these_trucks = [
            #    "S{i}".format(i=i)
            #    for i in (x, y, z)
            #]
            #skipple_preds += fillin([
            #    "(= {number_of_skipples} (+ %s))" % " ".join(items_all_in_these_trucks)
            #], vars())

            # "All other n-3" trucks have no skipples
            items_in_other_trucks_zero = [
                "(= 0 S{i})".format(i=i)
                for i in range(1, n_trucks + 1)
                if i not in (x, y, z)
            ]
            preds += [
                "(or %s)" % " ".join(items_in_other_trucks_zero)
            ]
#preds += [
#    "(or\n%s\n)" % "\n".join(skipple_preds)
#]

# Nuzzles: max 1 per truck
for i in range(1, n_trucks + 1):
    preds += fillin([
        "(>= 1 N{i})"
    ], vars())

if is_b:
    # Prittles and crottles cannot both be on the same truck. Either one of them
    # should be zero for every truck.
    for i in range(1, n_trucks + 1):
        preds += fillin([
            "(or (= 0 P{i}) (= 0 C{i}))"
        ], vars())

# Begin generator
s = """(benchmark test.smt
:logic QF_UFLIA
:extrafuns
(
"""
s += "\n".join("(%s Int)" % x for x in literals)
s += """
)
:formula
(and
"""
s += "\n".join(preds)
s += """
))
"""

try:
    import sys
    print(s)
    # Hack to avoid printing a ugly error in case stdin of next pipe is closed.
    sys.stdout.flush()
    sys.stdout.close()
except BrokenPipeError:
    pass