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


# This has mostly the same structure as bmc in ex03. It returns a model with
# a counterexample, or None in case the system is safe.
def kind(variables, init, trans, prop):
    primes = [Symbol(x.symbol_name() + "'", x.symbol_type()) for x in variables]

    def marked_state_formula(formula, step):
        return formula.substitute(
            {x: Symbol(x.symbol_name() + str(step), x.symbol_type()) for x in variables}
        )

    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}
        )

    # The base solver can prove unsafety, the inductive one safety.
    with Solver() as s_base, Solver() as s_ind:
        s_base.add_assertion(marked_state_formula(init, 0))

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

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

        s_base.pop()

        s_ind.add_assertion(marked_transition_formula(0))
        s_ind.add_assertion(marked_state_formula(prop, 0))

        s_ind.push()
        s_ind.add_assertion(Not(marked_state_formula(prop, 1)))

        if not s_ind.solve():
            return None

        s_ind.pop()

        # Unroll the transition relation iteratively.
        k = 1

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

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

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

            s_base.pop()

            s_ind.add_assertion(marked_transition_formula(k))
            s_ind.add_assertion(marked_state_formula(prop, k))

            s_ind.push()
            s_ind.add_assertion(Not(marked_state_formula(prop, k + 1)))

            if not s_ind.solve():
                return None

            s_ind.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(11)))
        )
    )
    prop = y >= 0

    res = kind(variables, init, trans, prop)

    if res is None:
        print("The system is safe")
    else:
        print(res)