from pysmt.shortcuts import Solver, Symbol, Not, And, Or, Equals, NotEquals, Int
from pysmt.typing import INT


# As a precondition, init and prop contain symbols in variables and trans
# contains those plus their syntactic primes.
def bmc(variables, init, trans, prop):
    primes = [Symbol(x.symbol_name() + "'", x.symbol_type()) for x in variables]

    # Change the formula so that it talks about variables X_step instead of X.
    def marked_state_formula(formula, step):
        return formula.substitute(
            {x: Symbol(x.symbol_name() + str(step), x.symbol_type()) for x in variables}
        )

    # Return the transition formula with X substituted for X_step and
    # X' substituted for X'_{step + 1}.
    def marked_transition_formula(step):
        return marked_state_formula(trans, step).substitute(
            {x: Symbol(x.symbol_name()[:-1] + str(step + 1), x.symbol_type()) for x in primes}
        )

    with Solver() as s:
        s.add_assertion(marked_state_formula(init, 0))

        # Check the initial states.
        s.push()
        s.add_assertion(Not(marked_state_formula(prop, 0)))

        if s.solve():
            return s.get_model()

        s.pop()

        # Unroll the transition relation iteratively - k is the current
        # unrolling length.
        k = 1

        while True:
            s.add_assertion(marked_transition_formula(k - 1))

            s.push()
            s.add_assertion(Not(marked_state_formula(prop, k)))

            if s.solve():
                return s.get_model()

            s.pop()

            k += 1


if __name__ == '__main__':
    x = Symbol('x', INT)
    y = Symbol('y', INT)

    xp = Symbol("x'", INT)
    yp = Symbol("y'", INT)

    variables = [x, y]
    init = And(Equals(x, Int(0)), Equals(y, Int(10)))
    trans = And(
        Or(
            And(NotEquals(x, Int(3)), Equals(xp, x + Int(1))),
            And(Equals(x, Int(3)), Equals(xp, Int(0)))
        ),
        Or(
            And(NotEquals(x, Int(2)), Equals(yp, y + Int(4))),
            And(Equals(x, Int(2)), Equals(yp, y - Int(13)))
        )
    )
    prop = y >= 0

    print(bmc(variables, init, trans, prop))