#!/usr/bin/env python3
# Usage:
#   $0 [number-of-steps [truck-capacity [b]]] | z3 -smt -in
# Examples:
#   $0 64 240
#   $0 30 260 b

import sys

# Number of steps
m = 64 if len(sys.argv) <= 1 else int(sys.argv[1])
truck_capacity = 240 if len(sys.argv) <= 2 else int(sys.argv[2])
is_b = len(sys.argv) >= 4 and sys.argv[3] == 'b'

states = "SABCD"

capacities = [
None,
120,
160,
100,
160
]

# Possible transitions from x to y and the travel cost c.
# trans[x, y] = c
# trans[x, y] = None  if there is no such transition.
trans = {}
trans[0, 1] = 15
trans[0, 3] = 15
trans[1, 2] = 17
trans[1, 3] = 12
trans[2, 3] = 10
trans[2, 4] = 20
trans[3, 4] = 20
# and the other direction
for (x, y), c in dict(trans).items():
    trans[y, x] = c

if False:
    # Print table for debugging
    for x in range(len(states)):
        for y in range(len(states)):
            c = trans.get((x, y))
            if c:
                print(" %2d " % c, end="")
            else:
                print(" -- ", end="")
        print("")

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 = []

# s0_j is the current stock of the truck
# s1_j is the current stock of village 1
for j in range(m + 1):
    literals += ["s%d_%d" % (i, j) for i in range(len(states))]
    literals += ["l%d" % j]

# Initially the villages have 80 food,
# the truck is at the supply (l=0),
# and the truck is fully loaded
for i in range(1, len(states)):
    preds += ["(= s%d_0 80)" % i]
preds += ["(= l0 0)"]
preds += ["(= s0_0 %d)" % truck_capacity]

# The truck and village capacity should never drop below zero in any state.
for j in range(m + 1):
    for i, capacity in enumerate(capacities):
        preds += ["(>= s%d_%d 0)" % (i, j)]
        # The village also has a maximum capacity
        if i > 0:
            preds += ["(<= s%d_%d %d)" % (i, j, capacity)]

# For every step, any of the edge transitions can be done.
for j in range(1, m + 1):
    altpreds = []
    jprev = j - 1
    for (x, y), cost in trans.items():
        # Preds for this choice
        cpreds = fillin([
            "(= l{jprev} {x})",
            "(= l{j} {y})",
        ], vars())

        # All villages must have sufficient time to make the transition.
        for i in range(1, len(states)):
            cpreds += fillin([
                "(>= s{i}_{jprev} {cost})",
            ], vars())

        # Calculate new food stock per village.  All village lose the cost, but
        # the next village can be stocked from the truck.
        for i in range(1, len(states)):
            newval = "(- s{i}_{jprev} {cost})"
            if i == y:
                # This is next village, so restock village from truck.
                drop = "(- s0_{jprev} s0_{j})"
                newval = "(+ %s %s)" % (newval, drop)
                cpreds += fillin([
                    "(>= %s 0)" % drop,
                ], vars())
            cpreds += fillin([
                "(= s{i}_{j} %s)" % newval,
            ], vars())
        # If next is truck, restock.
        if y == 0:
            cpreds += fillin([
                "(= s0_{j} {truck_capacity})",
            ], vars())
        altpreds += ["(and %s)" % " ".join(cpreds)]
    preds += ["(or %s)" % "\n".join(altpreds)]

if is_b:
    # If any stock is higher or equal to another state in the same location,
    # then we found a loop.
    stuff = []
    literals += ["k0", "k1"]
    for k in range(1, m + 1):
        for j in range(1, m + 1):
            if j >= k:
                continue
            highers = " ".join([
                "(<= s{i}_{j} s{i}_{k})".format(i=i, j=j, k=k)
                for i in range(len(states))
            ])
            stuff += fillin([
                "(and (= k0 {j}) (= k1 {k}) (= l{j} l{k}) {highers})"
            ], vars())

    preds += ["(or %s)" % "\n".join(stuff)]

# 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