#!/usr/bin/env python
# Usage:
#   $0 [number-of-states] | z3 -smt -in
# Examples:
#   $0 25

import itertools
import sys

# Number of steps
m = 25 if len(sys.argv) <= 1 else int(sys.argv[1])

# all possible strings
lang = [
    "aa",
    "bb",
    "aba",
    "baa",
    "abab",
    "babb",
    "bbba",
]

# sledgehammer approach: specify all bad words
fulllang = []
for n in range(1, 4 + 1):
    fulllang += [
        "".join(x)
        for x in itertools.product(*["ab"] * n)
    ]
# Remove good words.
badlang = [word for word in fulllang if word not in lang]

literals = []
preds = []
extrafuns = []

# Populate Q = { q0, q1, ..., q_(m-1) } with unique values.
literals += ["q%d" % i for i in range(m)]
preds += ["(= q%d %d)" % (i, i) for i in range(m)]

def mks(word_index, letter_index):
    """Return variable for a state of a word after some transitions."""
    return "s%d_%d" % (word_index, letter_index)

# Assuming a single initial state (which may or may not be the final state),
# add states for every letter j in each word i.
for i, word in enumerate(lang):
    for j, letter in enumerate(word):
        v = mks(i, j)
        literals += [v]
        # Each of the states must be assigned to any of the states in Q
        #preds += ["(>= %s 0)" % v, "(< %s %d)" % (v, m)]

# Transition Δ: Q × Σ -> P(Q). For Σ = {a,b}, we assume this range:
# transitions (0..; 1..2; 0..)
extrafuns += ["(trans Int Int Int)"]
# final states F \subset Q
extrafuns += ["(final Int Bool)"]
# Require a transition between each letter
for i, word in enumerate(lang):
    state_prev = "q0"
    for j, letter in enumerate(word):
        state_next = mks(i, j)
        label = "ab".index(letter)
        preds += ["(= (trans %s %d) %s)" % (state_prev, label, state_next)]
        state_prev = state_next
    preds += ["(final %s)" % state_prev]

# "a" and "b" are not allowed to accept (length 1)
#preds += ["(not (final (trans q0 %d)))" % label for label in range(2)]
for i, word in enumerate(badlang):
    state_prev = "q0"
    move = "q0"
    preconds = []
    for j, letter in enumerate(word):
        label = "ab".index(letter)
        move = "(trans %s %d)" % (move, label)
        # Ensure result is within range.
        preconds += ["(>= %s 0)" % move, "(< %s %d)" % (move, m)]
    preds += ["(=> (and %s) (not (final %s)))" % (" ".join(preconds), move)]
    #preds += ["(not (final %s))" % move]
    #valid = "(>= {move} 0) (< {move} {m})".format(**vars())
    #preds += ["(=> (and %s) (not (final %s)))" % (valid, move)]

# Try to fit any of the letter states with any state from Q
for i, word in enumerate(lang):
    for j, letter in enumerate(word):
        alts = []
        for qi in range(m):
            state = "q%d" % qi
            alts += ["(= %s %s)" % (state, mks(i, j))]
        preds += ["(or %s)" % " ".join(alts)]

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