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

# 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, Ballot, Value)
"""
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
"""


phase_1b_bal = z3.Const("phase_1b_bal", Ballot)
# phase_1b_bal = z3.Function("phase_1b_bal", Ballot, Ballot)
phase_1b_val = z3.Const("phase_1b_val", Value)


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
) -> 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.And(
                pred(b) != _1,
                z3.Exists(A, z3.And(belongs_to(A, q), msg1b(A, b, pred(b), v))),
                z3.ForAll(
                    [A, Mb2, Mv2],
                    z3.Implies(
                        z3.And(belongs_to(A, q), msg1b(A, b, Mb2, Mv2), Mb2 != _1),
                        lte(Mb2, pred(b)),
                    ),
                ),
            ),
        ),
    )


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


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:
    Mb, Mv, B, V = fresh_consts(Ballot, Value, Ballot, Value)
    return z3.And(
        b != _1,
        msg1a_pre(b),
        gt(b, max_bal_pre(a)),
        update_state(max_bal_pre, max_bal_post, {(a,): b}),
        z3.Exists(
            [Mb, Mv],
            z3.And(
                z3.Or(
                    z3.And(Mb == _1, z3.ForAll([B, V], z3.Not(msg2b_pre(a, B, V)))),
                    z3.And(
                        Mb != _1,
                        msg2b_pre(a, Mb, Mv),
                        z3.ForAll([B, V], z3.Implies(msg2b_pre(a, B, V), lte(B, Mb))),
                    ),
                ),
                update_state(
                    msg1b_pre,
                    msg1b_post,
                    {(a, b, Mb, Mv): TRUE},
                ),
            ),
        ),
        unchanged(msg1a_pre, msg1a_post),
        unchanged(msg2a_pre, msg2a_post),
        unchanged(msg2b_pre, msg2b_post),
    )


def phase1b_half_skolemized(
    a: AcceptorType, b: BallotType, Mb: BallotType, Mv: BallotType
) -> z3.BoolRef:
    B, V = fresh_consts(Ballot, Value)
    return z3.And(
        b != _1,
        msg1a_pre(b),
        gt(b, max_bal_pre(a)),
        z3.Or(
            z3.And(Mb == _1, z3.ForAll([B, V], z3.Not(msg2b_pre(a, B, V)))),
            z3.And(
                Mb != _1,
                msg2b_pre(a, Mb, Mv),
                z3.ForAll([B, V], z3.Implies(msg2b_pre(a, B, V), lte(B, Mb))),
            ),
        ),
        update_state(
            msg1b_pre,
            msg1b_post,
            {(a, b, Mb, Mv): TRUE},
        ),
        update_state(max_bal_pre, max_bal_post, {(a,): b}),
        unchanged(msg1a_pre, msg1a_post),
        unchanged(msg2a_pre, msg2a_post),
        unchanged(msg2b_pre, msg2b_post),
    )


def phase1b_skolemized(a: AcceptorType, b: BallotType) -> z3.BoolRef:
    B, V = fresh_consts(Ballot, Value)
    return z3.And(
        b != _1,
        msg1a_pre(b),
        gt(b, max_bal_pre(a)),
        z3.Or(
            z3.And(phase_1b_bal == _1, z3.ForAll([B, V], z3.Not(msg2b_pre(a, B, V)))),
            z3.And(
                phase_1b_bal != _1,
                msg2b_pre(a, phase_1b_bal, phase_1b_val),
                z3.ForAll([B, V], z3.Implies(msg2b_pre(a, B, V), lte(B, phase_1b_bal))),
            ),
        ),
        update_state(max_bal_pre, max_bal_post, {(a,): b}),
        update_state(
            msg1b_pre,
            msg1b_post,
            {(a, b, phase_1b_bal, phase_1b_val): 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,
        z3.Not(z3.Exists(V, msg2a_pre(b, V))),
        is_safe_at_paxos(b, v, msg1b_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,
        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),
            # phase1b_skolemized(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(B, V, msg1b))
    )


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) -> z3.BoolRef:
    B, V = fresh_consts(Ballot, Value)
    return z3.ForAll([B, V], z3.Implies(msg2a(B, V), is_safe_at_paxos(B, V, msg1b)))


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, Bmax, Vmax = fresh_consts(Acceptor, Ballot, Ballot, Value)

    return z3.ForAll(
        [A, B, Bmax, Vmax], z3.Implies(msg1b(A, B, Bmax, Vmax), 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 inv0(msg2b: Relation, msg1b: Relation) -> z3.BoolRef:
    A, B, V = fresh_consts(Acceptor, Ballot, Value)
    B2 = fresh_const(Ballot)
    return z3.And(
        z3.ForAll([A, B, V], z3.Implies(msg2b(A, B, V), B != _1)),
        z3.ForAll([A, B, B2, V], z3.Implies(msg1b(A, B, B2, V), less_than(B2, B))),
    )


def invariants() -> z3.BoolRef:
    return z3.And(
        inv1(msg1b_pre, msg2b_pre),
        inv2(msg2b_pre),
        inv3(msg2a_pre, msg1b_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),
        inv0(msg2b_pre, msg1b_pre),
    )


def post_invariants() -> z3.BoolRef:
    return z3.And(
        inv1(msg1b_post, msg2b_post),
        inv2(msg2b_post),
        inv3(msg2a_post, msg1b_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, msg1b_post),
    )


def counter_example() -> z3.BoolRef:
    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 implicit_paxos() -> z3.BoolRef:
    return counter_example()
