#!/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]

extrafuns = []
preds = []

# s0_j is the current stock of the truck
# s1_j is the current stock of village 1
extrafuns += ["(s Int Int Int)"]
extrafuns += ["(l Int Int)"]

# 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 += ["(= (l 0) 0)"]
preds += ["(= (s 0 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 = "(- (s 0 {jprev}) (s 0 {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([
                "(= (s 0 {j}) {truck_capacity})",
            ], vars())
        altpreds += ["(and %s)" % " ".join(cpreds)]
    preds += ["(or %s)" % "\n".join(altpreds)]

if is_b:
    # If stock is higher or equal to initial state, then we found a loop.
    stuff = []
    for j in range(1, m + 1):
        jprev = j - 1
        highers = " ".join([
            "(>= (s {i} {j}) (s {i} 0))".format(i=i, j=j)
            for i in range(1, len(states))
        ])
        stuff += fillin([
            "(and {highers})"
        ], vars())

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

# Begin generator
s = """(benchmark test.smt
:logic QF_UFLIA
:extrafuns (
"""
s += "\n".join(extrafuns)
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