import z3

import problems
from common.relations import Relation
from problems.transition_systems import (
    state_function,
    state_relation,
    unchanged,
    update_state,
)
from problems.z3_utils import TRUE, fresh_const, fresh_consts

from . import (
    Acceptor,
    AcceptorType,
    Ballot,
    BallotType,
    Quorum,
    QuorumType,
    Value,
    ValueType,
    _1,
    belongs_to,
    gt,
    gte,
    less_than,
    lte,
    order_axioms,
    pred,
    quorum_axioms,
)

none = z3.Const("none", Value)

# State
msg1a_pre, msg1a_post = state_relation("msg1a", Ballot)
"""
The global proposer says: let us consider this Ballot
"""

msg1b_pre, msg1b_post = state_relation("msg1b", Acceptor, Ballot)
"""
This acceptor promises never to vote on Ballot < B1, 
and says: I last voted at B2, and for Value
"""

msg2a_pre, msg2a_post = state_relation("msg2a", Ballot, Value)
"""
The global proposer says: let us vote for Value at Ballot
"""

msg2b_pre, msg2b_post = state_relation("msg2b", Acceptor, Ballot, Value)
"""
This acceptor votes for Value at Ballot
"""

max_bal_pre, max_bal_post = state_function("max_bal", Acceptor, Ballot)
"""
This acceptor promised to never vote on 'smaller' Ballots
"""


def chosen_at(b: BallotType, v: ValueType, msg2b: Relation) -> z3.BoolRef:
    Q, A = fresh_consts(Quorum, Acceptor)
    return z3.Exists(Q, z3.ForAll(A, z3.Implies(belongs_to(A, Q), msg2b(A, b, v))))


def chosen(v: ValueType, msg2b: Relation) -> z3.BoolRef:
    B = fresh_const(Ballot)

    return z3.Exists(B, chosen_at(B, v, msg2b))


def shows_safe_at_paxos(
    q: QuorumType, b: BallotType, v: ValueType, msg1b: Relation
) -> z3.BoolRef:
    A, Mb, Mv, Mb2, Mv2 = fresh_consts(Acceptor, Ballot, Value, Ballot, Value)
    return z3.And(
        z3.ForAll(
            A, z3.Implies(belongs_to(A, q), z3.Exists([Mb, Mv], msg1b(A, b, Mb, Mv)))
        ),
        z3.Or(
            z3.ForAll(
                [A, Mb, Mv],
                z3.Not(z3.And(belongs_to(A, q), msg1b(A, b, Mb, Mv), Mb != _1)),
            ),
            z3.Exists(
                # this exists should be skolemized and turned into
                # a "predecessor" function, and then we get that
                # the counter-example will use an element which is
                # smaller than everyone, but nobody's predecessor (-inf)
                Mb,
                z3.And(
                    Mb != _1,
                    z3.Exists(
                        A,
                        z3.And(
                            belongs_to(A, q),
                            msg1b(A, b, Mb, v),
                        ),
                    ),
                    z3.ForAll(
                        [A, Mb2, Mv2],
                        z3.Implies(
                            z3.And(belongs_to(A, q), msg1b(A, b, Mb2, Mv2), Mb2 != _1),
                            lte(Mb2, Mb),
                        ),
                    ),
                ),
            ),
        ),
    )


def shows_safe_at_paxos_skolemized(
    q: QuorumType, b: BallotType, v: ValueType, msg1b: Relation, msg2b: Relation
) -> z3.BoolRef:
    A, Mb, Mv, Mb2, Mv2 = fresh_consts(Acceptor, Ballot, Value, Ballot, Value)
    return z3.And(
        z3.ForAll(
            A,
            z3.Implies(
                belongs_to(A, q),
                msg1b(
                    A,
                    b,
                ),
            ),
        ),
        z3.And(
            z3.Or(
                z3.ForAll(
                    [A, Mb, Mv],
                    z3.Not(
                        z3.And(
                            belongs_to(A, q),
                            msg1b(A, b),
                            msg2b(A, Mb, Mv),
                            lte(Mb, b),
                            Mv != v,
                        ),
                    ),
                ),
                z3.And(
                    # pred(b) != _1,
                    less_than(pred(b), b),
                    z3.Exists(
                        A, z3.And(belongs_to(A, q), msg1b(A, b), msg2b(A, pred(b), v))
                    ),
                    z3.ForAll(
                        [A, Mb2, Mv2],
                        z3.Implies(
                            z3.And(
                                belongs_to(A, q),
                                msg1b(A, b),
                                msg2b(A, Mb2, Mv2),
                                Mv2 != v,
                            ),
                            lte(Mb2, pred(b)),
                        ),
                    ),
                ),
            ),
        ),
    )


def shows_safe_at_paxos_skolemized_for_inv(
    q: QuorumType, b: BallotType, v: ValueType, msg1b: Relation, msg2b: Relation
) -> z3.BoolRef:
    A, Mb, Mv, Mb2, Mv2 = fresh_consts(Acceptor, Ballot, Value, Ballot, Value)
    return z3.And(
        z3.ForAll(
            A,
            z3.Implies(
                belongs_to(A, q),
                msg1b(
                    A,
                    b,
                ),
            ),
        ),
        z3.And(
            z3.Or(
                z3.ForAll(
                    [A, Mb, Mv],
                    z3.Not(
                        z3.And(
                            belongs_to(A, q),
                            msg1b(A, b),
                            msg2b(A, Mb, Mv),
                            Mb != b,  # todo: lte(Mb, b)?
                            Mv != v,
                        ),
                    ),
                ),
                z3.And(
                    # pred(b) != _1,
                    less_than(pred(b), b),
                    z3.Exists(
                        A, z3.And(belongs_to(A, q), msg1b(A, b), msg2b(A, pred(b), v))
                    ),
                    z3.ForAll(
                        [A, Mb2, Mv2],
                        z3.Implies(
                            z3.And(
                                belongs_to(A, q),
                                msg1b(A, b),
                                msg2b(A, Mb2, Mv2),
                                Mv2 != v,
                            ),
                            lte(Mb2, pred(b)),
                        ),
                    ),
                ),
            ),
        ),
    )


def is_safe_at_paxos(
    b: BallotType, v: ValueType, msg1b: Relation, msg2b: Relation
) -> z3.BoolRef:
    Q = fresh_const(Quorum)
    return z3.Exists(Q, shows_safe_at_paxos_skolemized(Q, b, v, msg1b, msg2b))


def is_safe_at_paxos_for_inv(
    b: BallotType, v: ValueType, msg1b: Relation, msg2b: Relation
) -> z3.BoolRef:
    Q = fresh_const(Quorum)
    return z3.Exists(Q, shows_safe_at_paxos_skolemized(Q, b, v, msg1b, msg2b))


def phase1a(b: BallotType) -> z3.BoolRef:
    return z3.And(
        b != _1,
        update_state(msg1a_pre, msg1a_post, {(b,): TRUE}),
        unchanged(msg1b_pre, msg1b_post),
        unchanged(msg2a_pre, msg2a_post),
        unchanged(msg2b_pre, msg2b_post),
        unchanged(max_bal_pre, max_bal_post),
    )


def phase1b(a: AcceptorType, b: BallotType) -> z3.BoolRef:
    return z3.And(
        b != _1,
        msg1a_pre(b),
        gt(b, max_bal_pre(a)),
        update_state(max_bal_pre, max_bal_post, {(a,): b}),
        update_state(msg1b_pre, msg1b_post, {(a, b): TRUE}),
        unchanged(msg1a_pre, msg1a_post),
        unchanged(msg2a_pre, msg2a_post),
        unchanged(msg2b_pre, msg2b_post),
    )


def phase2a(b: BallotType, v: ValueType) -> z3.BoolRef:
    V = fresh_const(Value)
    return z3.And(
        b != _1,
        v != none,
        z3.Not(z3.Exists(V, msg2a_pre(b, V))),
        is_safe_at_paxos(b, v, msg1b_pre, msg2b_pre),
        update_state(msg2a_pre, msg2a_post, {(b, v): TRUE}),
        unchanged(msg1a_pre, msg1a_post),
        unchanged(msg1b_pre, msg1b_post),
        unchanged(msg2b_pre, msg2b_post),
        unchanged(max_bal_pre, max_bal_post),
    )


def phase2b(a: AcceptorType, b: BallotType, v: ValueType) -> z3.BoolRef:
    return z3.And(
        b != _1,
        v != none,
        msg2a_pre(b, v),
        gte(b, max_bal_pre(a)),
        update_state(max_bal_pre, max_bal_post, {(a,): b}),
        update_state(msg2b_pre, msg2b_post, {(a, b, v): TRUE}),
        unchanged(msg1a_pre, msg1a_post),
        unchanged(msg1b_pre, msg1b_post),
        unchanged(msg2a_pre, msg2a_post),
    )


def transition() -> z3.BoolRef:
    A, B, V = fresh_consts(Acceptor, Ballot, Value)
    return z3.Exists(
        [A, B, V],
        z3.Or(
            # phase1a(B),
            # phase1b(A, B),
            # phase2a(B, V),
            phase2b(A, B, V),
        ),
    )


def simplified_transition() -> z3.BoolRef:
    A, B, V = fresh_consts(Acceptor, Ballot, Value)
    return z3.Exists(
        [A, B, V],
        phase2b(A, B, V),
    )


def safety(msg2b: Relation) -> z3.BoolRef:
    V1, V2 = fresh_consts(Value, Value)

    return z3.ForAll(
        [V1, V2], z3.Implies(z3.And(chosen(V1, msg2b), chosen(V2, msg2b)), V1 == V2)
    )


# Invariants
def inv1(msg1b: Relation, msg2b: Relation) -> z3.BoolRef:
    A, B, V = fresh_consts(Acceptor, Ballot, Value)
    return z3.ForAll(
        [A, B, V],
        z3.Implies(msg2b(A, B, V), is_safe_at_paxos_for_inv(B, V, msg1b, msg2b)),
    )


def inv2(msg2b: Relation) -> z3.BoolRef:
    A, B, V1, V2 = fresh_consts(Acceptor, Ballot, Value, Value)
    return z3.ForAll(
        [A, B, V1, V2],
        z3.Implies(z3.And(chosen_at(B, V1, msg2b), msg2b(A, B, V2)), V1 == V2),
    )


def inv3(msg2a: Relation, msg1b: Relation, msg2b: Relation) -> z3.BoolRef:
    B, V = fresh_consts(Ballot, Value)
    return z3.ForAll(
        [B, V], z3.Implies(msg2a(B, V), is_safe_at_paxos_for_inv(B, V, msg1b, msg2b))
    )


def inv4(msg2a: Relation) -> z3.BoolRef:
    B, V1, V2 = fresh_consts(Ballot, Value, Value)
    return z3.ForAll(
        [B, V1, V2], z3.Implies(z3.And(msg2a(B, V1), msg2a(B, V2)), V1 == V2)
    )


def inv5(msg2a: Relation, msg2b: Relation) -> z3.BoolRef:
    A, B, V = fresh_consts(Acceptor, Ballot, Value)
    return z3.ForAll([A, B, V], z3.Implies(msg2b(A, B, V), msg2a(B, V)))


def inv6(msg1b: Relation, max_bal: z3.FuncDeclRef) -> z3.BoolRef:
    A, B = fresh_consts(Acceptor, Ballot)

    return z3.ForAll([A, B], z3.Implies(msg1b(A, B), gte(max_bal(A), B)))


# def inv7(msg1b: Relation, msg2b: Relation) -> z3.BoolRef:
#     A, B, Bmax, Vmax = fresh_consts(Acceptor, Ballot, Ballot, Value)
#     return z3.ForAll(
#         [A, B, Bmax, Vmax],
#         z3.Implies(
#             z3.And(
#                 B != _1,
#                 Bmax != _1,
#                 msg1b(A, B, Bmax, Vmax),
#             ),
#             msg2b(A, Bmax, Vmax),
#         ),
#     )
#
#
# def inv8(msg1b: Relation, msg2b: Relation) -> z3.BoolRef:
#     A, B, Bmid, Bmax, V, Vmax = fresh_consts(
#         Acceptor, Ballot, Ballot, Ballot, Value, Value
#     )
#     return z3.ForAll(
#         [A, B, Bmid, Bmax, V, Vmax],
#         z3.Implies(
#             z3.And(
#                 gt(B, Bmid),
#                 gt(Bmid, Bmax),
#                 msg1b(A, B, Bmax, Vmax),
#             ),
#             z3.Not(msg2b(A, Bmid, V)),
#         ),
#     )
#
#
# def inv9(max_vbal: z3.FuncDeclRef, max_bal: z3.FuncDeclRef) -> z3.BoolRef:
#     A = fresh_const(Acceptor)
#     return z3.ForAll(A, lte(max_vbal(A), max_bal(A)))
#
#
# def inv10(max_vbal: z3.FuncDeclRef, msg2b: Relation) -> z3.BoolRef:
#     A, B, V = fresh_consts(Acceptor, Ballot, Value)
#     return z3.ForAll([A, B, V], z3.Implies(msg2b(A, B, V), gte(max_vbal(A), B)))
#
#
# def inv11(
#     max_vbal: z3.FuncDeclRef, max_val: z3.FuncDeclRef, msg2b: Relation
# ) -> z3.BoolRef:
#     A = fresh_const(Acceptor)
#     return z3.ForAll(
#         A, z3.Implies(max_vbal(A) != _1, msg2b(A, max_vbal(A), max_val(A)))
#     )


def inv0(msg2b: Relation) -> z3.BoolRef:
    A, B, V = fresh_consts(Acceptor, Ballot, Value)
    return z3.And(
        z3.ForAll([A, B, V], z3.Implies(msg2b(A, B, V), B != _1)),
        # z3.ForAll([A, B, V], z3.Implies(msg2b(A, B, V), V != none)),
    )


def invariants() -> z3.BoolRef:
    return z3.And(
        inv1(msg1b_pre, msg2b_pre),
        inv2(msg2b_pre),
        inv3(msg2a_pre, msg1b_pre, msg2b_pre),
        inv4(msg2a_pre),
        inv5(msg2a_pre, msg2b_pre),
        inv6(msg1b_pre, max_bal_pre),
        # inv7(msg1b_pre, msg2b_pre),
        # inv8(msg1b_pre, msg2b_pre),
        # inv9(max_vbal_pre, max_bal_pre),
        # inv10(max_vbal_pre, msg2b_pre),
        # inv11(max_vbal_pre, max_val_pre, msg2b_pre),
        inv0(msg2b_pre),
    )


def post_invariants() -> z3.BoolRef:
    return z3.And(
        inv1(msg1b_post, msg2b_post),
        inv2(msg2b_post),
        inv3(msg2a_post, msg1b_post, msg2b_post),
        inv4(msg2a_post),
        inv5(msg2a_post, msg2b_post),
        inv6(msg1b_post, max_bal_post),
        # inv7(msg1b_post, msg2b_post),
        # inv8(msg1b_post, msg2b_post),
        inv0(msg2b_post),
    )


def counter_example() -> z3.BoolRef:
    """
    Long ago, it was proposed to consider voting at -∞.                     (1a)
    Two brave knights (a1, a2) accepted the challenge and declared
    they have never voted before, and will vote at -∞                       (1b)
    It was then proposed to vote for v1 at -∞                               (2a)
    Alas only knight voted (a1), and the other (a2) disappeared.            (2b)
    *
    An infinity passed since then,
    and no one seems to remember the before-times.
    Instead, a new tradition has emerged:
    Every day it is proposed to vote,                                       (1a)
    and two brave knights (a1, a3) accept the challenge and declare
    we have voted yesterday for v2, and we will vote today.                 (1b)
    Every day it is then proposed to vote for v2,                           (2a)
    and every day the two brave knights (a1, a3) vote.                      (2b)
    *
    But a strange thing happened today,
    a long-forgotten knight (a2) has re-appeared,
    and he remembers the before-times,
    he remembers what he vowed to do,
    and he still thinks this is the before-times,
    and he votes for v1                                                     (2b)
    """
    return z3.And(
        quorum_axioms(),
        order_axioms(),
        safety(msg2b_pre),
        invariants(),
        transition(),
        z3.Not(safety(msg2b_post)),
        # or some invariant is broken?
    )


@problems.mark(
    model=True,
    slow=True,
    popl24=True,
)
def simple_paxos() -> z3.BoolRef:
    return counter_example()
