#!/usr/bin/env python
# Usage for 4a (where 7 is the number of steps to run):
#   $0 7 | z3 -smt -in
# Usage for 4b:
#   $0 17 b | z3 -smt -in
import sys

# i = 1..8
n = 8
steps = int(sys.argv[1])
is_b = len(sys.argv) >= 3 and sys.argv[2] == 'b'

literals = []
preds = []

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

# a. Repeat a number of steps, check in each step
# if a_3 = a_7 is possible in that step.
# Number a_i in step s is encoded as a_i_s.
# The initial state is a_i_0.
for step in range(steps + 1):
    for i in range(1, n + 1):
        literals += ["a{i}_{step}".format(**vars())]

# Initial state, a_i_0 = i
for i in range(1, n + 1):
    preds += ["(= a{i}_0 {i})".format(i=i)]

# Goal: have a_3 = a_7 in any step:
if not is_b:
    goal = "(= a3_{i} a7_{i})"
else:
    goal = "(and (= a3_{i} a5_{i}) (= a5_{i} a7_{i}))"
preds += [
    "(or %s)" % " ".join(goal.format(i=i)
    for i in range(steps + 1))
]

# In each step, can sum neighbors (leaving others unchanged).
for step in range(1, steps + 1):
    # The first and last one cannot be changed. In a step, one can choose to
    # change a_i for any i, leaving all other j (i != j) intact.
    altpreds = []
    for i in range(2, n):
        laststep = step - 1
        prev = i - 1
        next = i + 1

        # Require a'_i = a_{i-1} + a_{i+1}
        statepreds = fillin([
            "(= a{i}_{step} (+ a{prev}_{laststep} a{next}_{laststep}))"
        ], vars())

        # Require a'_j = a_j for i != j
        for j in range(1, n + 1):
            if i == j:
                continue
            statepreds += fillin([
                "(= a{j}_{step} a{j}_{laststep})"
            ], vars())
        # All assignments must apply for this step.
        altpreds += [
            "(and %s)" % "\n".join(statepreds)
        ]
    preds += [
        "(or",
        "%s" % "\n".join(altpreds),
        ")"
    ]


# 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