import z3

from common.axioms import of_size, total_strict_order
from common.relations import relation
from problems.z3_utils import consts, fresh_consts

Value = z3.DeclareSort("Value")
ValueType = z3.ExprRef
Acceptor = z3.DeclareSort("Acceptor")
AcceptorType = z3.ExprRef
Quorum = z3.DeclareSort("Quorum")
QuorumType = z3.ExprRef
Ballot = z3.DeclareSort("Ballot")  # Should be Nat U {-1}
BallotType = z3.ExprRef
belongs_to = relation("belongs_to", Acceptor, Quorum)
_1 = z3.Const("_1", Ballot)  # the "-1" element
less_than = relation("less_than", Ballot, Ballot)

# Induction
pred = z3.Function("pred", Ballot, Ballot)


def quorum_axioms() -> z3.BoolRef:
    Q1, Q2, A = consts(("Q1", Quorum), ("Q2", Quorum), ("A", Acceptor))
    return z3.ForAll(
        [Q1, Q2],
        z3.Exists(
            A,
            z3.And(belongs_to(A, Q1), belongs_to(A, Q2)),
            qid="quorums_q1",
        ),
        qid="quorums_q2",
    )


def order_axioms() -> z3.BoolRef:
    (B,) = fresh_consts(Ballot)
    return z3.And(total_strict_order(less_than), z3.ForAll(B, lte(_1, B)))


def gte(b1: BallotType, b2: BallotType) -> z3.BoolRef:
    return z3.Not(less_than(b1, b2))


def lte(b1: BallotType, b2: BallotType) -> z3.BoolRef:
    return z3.Or(b1 == b2, less_than(b1, b2))


def gt(b1: BallotType, b2: BallotType) -> z3.BoolRef:
    return less_than(b2, b1)


def size_constraints() -> z3.BoolRef:
    return z3.And(of_size(Value, 2), of_size(Acceptor, 3), of_size(Quorum, 3))
