import z3

import problems
from common.relations import Relation
from problems.transition_systems import state_function, state_relation, update_state
from problems.z3_utils import TRUE, fresh_consts

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

# State
votes_pre, votes_post = state_relation("votes", Acceptor, Ballot, Value)
max_bal_pre, max_bal_post = state_function("max_bal", Acceptor, Ballot)

witness = z3.Function("witness", Ballot, Acceptor)

unsafe_v1, unsafe_v2 = z3.Consts("unsafe_v1 unsafe_v2", Value)
unsafe_b1, unsafe_b2 = z3.Consts("unsafe_b1 unsafe_b2", Ballot)


def chosen_at(b: BallotType, v: ValueType, votes: Relation) -> z3.BoolRef:
    (A,) = fresh_consts(Acceptor)
    return z3.ForAll(A, votes(A, b, v))


def chosen(v: ValueType, votes: Relation) -> z3.BoolRef:
    (B,) = fresh_consts(Ballot)
    return z3.Exists(B, chosen_at(B, v, votes))


def shows_safe_at_skolemized(
    b: BallotType, v: ValueType, max_bal: z3.FuncDeclRef, votes: Relation
) -> z3.BoolRef:
    A, D, V = fresh_consts(Acceptor, Ballot, Value)
    return z3.And(
        z3.ForAll(A, gte(max_bal(A), b)),
        z3.And(
            less_than(pred(b), b),
            z3.Implies(
                pred(b) != _1,
                votes(witness(b), pred(b), v),
            ),
            z3.ForAll(
                D,
                z3.Implies(
                    z3.And(less_than(pred(b), D), less_than(D, b)),
                    z3.ForAll([A, V], z3.Not(votes(A, D, V))),
                ),
            ),
        ),
    )


def is_safe_at(
    b: BallotType, v: ValueType, max_bal: z3.FuncDeclRef, votes: Relation
) -> z3.BoolRef:
    return shows_safe_at_skolemized(b, v, max_bal, votes)


# Actions
def increase_max_bal(a: AcceptorType, b: BallotType) -> z3.BoolRef:
    A, B, V = fresh_consts(Acceptor, Ballot, Value)
    return z3.And(
        b != _1,
        gt(b, max_bal_pre(a)),
        update_state(votes_pre, votes_post, {}),
        update_state(max_bal_pre, max_bal_post, {(a,): b}),
    )


def vote_for(a: AcceptorType, b: BallotType, v: ValueType) -> z3.BoolRef:
    A, V = fresh_consts(Acceptor, Value)
    return z3.And(
        b != _1,
        lte(max_bal_pre(a), b),
        z3.ForAll(V, z3.Not(votes_pre(a, b, V))),
        z3.ForAll(
            A, z3.Implies(A != a, z3.ForAll(V, z3.Implies(votes_pre(A, b, V), V == v)))
        ),
        is_safe_at(b, v, max_bal_pre, votes_pre),
        update_state(votes_pre, votes_post, {(a, b, v): TRUE}),
        update_state(max_bal_pre, max_bal_post, {(a,): b}),
    )


def init() -> z3.BoolRef:
    A, B, V = fresh_consts(Acceptor, Ballot, Value)
    return z3.And(
        z3.ForAll([A, B, V], z3.Not(votes_pre(A, B, V))),
        z3.ForAll(A, max_bal_pre(A) == _1),
    )


def specific_transition() -> z3.BoolRef:
    A, B, V = fresh_consts(Acceptor, Ballot, Value)
    return z3.Exists(
        [A, B],
        # [A, B],
        # z3.Or(
        # increase_max_bal(A, B),
        vote_for(A, B, unsafe_v1),
        # ),
    )


def transition() -> z3.BoolRef:
    A, B, V = fresh_consts(Acceptor, Ballot, Value)
    return z3.Exists(
        [A, B, V],
        z3.Or(
            increase_max_bal(A, B),
            vote_for(A, B, V),
        ),
    )


def safety(votes: Relation) -> z3.BoolRef:
    V1, V2 = fresh_consts(Value, Value)
    V1, V2 = z3.Consts("safe_V1 safe_V2", Value)
    return z3.ForAll(
        [V1, V2], z3.Implies(z3.And(chosen(V1, votes), chosen(V2, votes)), V1 == V2)
    )


def skolemized_unsafety(votes: Relation) -> z3.BoolRef:
    return z3.And(
        chosen_at(unsafe_b1, unsafe_v1, votes),
        chosen_at(unsafe_b2, unsafe_v2, votes),
        unsafe_v1 != unsafe_v2,
    )


def inv1(max_bal: z3.FuncDeclRef, votes: Relation) -> z3.BoolRef:
    A, B, V = fresh_consts(Acceptor, Ballot, Value)
    return z3.ForAll(
        [A, B, V], z3.Implies(votes(A, B, V), is_safe_at(B, V, max_bal, votes))
    )


def inv2(votes: 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, votes), votes(A, B, V2)), V1 == V2),
    )


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


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


def inv_no_double_votes_between_acceptors(votes: Relation) -> z3.BoolRef:
    A1, A2, B, V1, V2 = fresh_consts(Acceptor, Acceptor, Ballot, Value, Value)

    return z3.ForAll(
        [A1, A2, B, V1, V2],
        z3.Implies(z3.And(votes(A1, B, V1), votes(A2, B, V2)), V1 == V2),
    )


def unsafe_values_are_everything() -> z3.BoolRef:
    (V,) = fresh_consts(Value)
    return z3.ForAll(V, z3.Or(V == unsafe_v1, V == unsafe_v2))


def counter_example() -> z3.BoolRef:
    (V,) = fresh_consts(Value)
    return z3.And(
        order_axioms(),
        safety(votes_pre),
        inv1(max_bal_pre, votes_pre),
        inv2(votes_pre),
        inv0(votes_pre),
        # inv_no_double_votes(votes_pre),
        inv_no_double_votes_between_acceptors(votes_pre),
        transition(),
        z3.Not(safety(votes_post)),
        # skolemized_unsafety(votes_post),
        # unsafe_values_are_everything(),
    )


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


def voting_noq_induction_axiom_votes_for_v(
    unsafe_v: ValueType, votes: Relation = votes_pre
) -> z3.BoolRef:
    (A,) = fresh_consts(Acceptor)
    Amin, Bmin, B = fresh_consts(Acceptor, Ballot, Ballot)

    return z3.Implies(
        z3.Exists([A, B], votes(A, B, unsafe_v)),
        z3.Exists(
            [Amin, Bmin],
            z3.And(
                votes(Amin, Bmin, unsafe_v),
                z3.ForAll(
                    [A, B],
                    z3.Implies(less_than(B, Bmin), z3.Not(votes(A, B, unsafe_v))),
                ),
            ),
        ),
    )


def voting_noq_induction_axiom_votes(votes: Relation = votes_pre) -> z3.BoolRef:
    A, V = fresh_consts(Acceptor, Value)
    Amin, Bmin, B = fresh_consts(Acceptor, Ballot, Ballot)

    return z3.ForAll(
        [V],
        z3.Implies(
            z3.Exists([A, B], votes(A, B, V)),
            z3.Exists(
                [Amin, Bmin],
                z3.And(
                    votes(Amin, Bmin, V),
                    z3.ForAll(
                        [A, B],
                        z3.Implies(less_than(B, Bmin), z3.Not(votes(A, B, V))),
                    ),
                ),
            ),
        ),
    )


def voting_noq_induction_axiom_votes_per_acceptor(
    votes: Relation = votes_pre,
) -> z3.BoolRef:
    A, V = fresh_consts(Acceptor, Value)
    Amin, Bmin, B = fresh_consts(Acceptor, Ballot, Ballot)

    return z3.ForAll(
        [V, A],
        z3.Implies(
            z3.Exists([B], votes(A, B, V)),
            z3.Exists(
                [Bmin],
                z3.And(
                    votes(A, Bmin, V),
                    z3.ForAll(
                        [B],
                        z3.Implies(less_than(B, Bmin), z3.Not(votes(A, B, V))),
                    ),
                ),
            ),
        ),
    )


def voting_noq_induction_axiom_safe(
    votes: Relation = votes_pre, max_bal: z3.FuncDeclRef = max_bal_pre
) -> z3.BoolRef:
    (V,) = fresh_consts(Value)
    Bmin, B = fresh_consts(Ballot, Ballot)

    return z3.ForAll(
        V,
        z3.Implies(
            z3.Exists(B, is_safe_at(B, V, max_bal, votes)),
            z3.Exists(
                Bmin,
                z3.And(
                    is_safe_at(Bmin, V, max_bal, votes),
                    z3.ForAll(
                        B,
                        z3.Implies(
                            less_than(B, Bmin),
                            z3.Not(is_safe_at(B, V, max_bal, votes)),
                        ),
                    ),
                ),
            ),
        ),
    )


def voting_noq_induction_axiom_safe_for_v(
    unsafe_v: ValueType,
    votes: Relation = votes_pre,
    max_bal: z3.FuncDeclRef = max_bal_pre,
) -> z3.BoolRef:
    Bmin, B = fresh_consts(Ballot, Ballot)

    return z3.Implies(
        z3.Exists(B, is_safe_at(B, unsafe_v, max_bal, votes)),
        z3.Exists(
            Bmin,
            z3.And(
                is_safe_at(Bmin, unsafe_v, max_bal, votes),
                z3.ForAll(
                    B,
                    z3.Implies(
                        less_than(B, Bmin),
                        z3.Not(is_safe_at(B, unsafe_v, max_bal, votes)),
                    ),
                ),
            ),
        ),
    )
