import z3

import problems
from common.axioms import total_strict_order
from common.relations import relation
from problems.z3_utils import fresh_consts

from . import Value

Round = z3.DeclareSort("Round")

lt = relation("lt", Round, Round)
safe = relation("safe", Value)
manually_set = relation("manually_set", Value)
proposal = relation("proposal", Round, Value)

r1, r2 = z3.Consts("r1 r2", Round)
v1, v2 = z3.Consts("v1 v2", Value)


def safety_violation() -> z3.BoolRef:
    return z3.And(proposal(r1, v1), proposal(r2, v2), v1 != v2)


def single_proposal_per_round() -> z3.BoolRef:
    R, V1, V2 = fresh_consts(Round, Value, Value)
    return z3.ForAll(
        [R, V1, V2],
        z3.Implies(z3.And(proposal(R, V1), proposal(R, V2)), V1 == V2),
    )


def single_safe() -> z3.BoolRef:
    V1, V2 = fresh_consts(Value, Value)
    return z3.ForAll([V1, V2], z3.Implies(z3.And(safe(V1), safe(V2)), V1 == V2))


def single_safe_or_manually_set() -> z3.BoolRef:
    V1, V2 = fresh_consts(Value, Value)
    return z3.ForAll(
        [V1, V2],
        z3.Implies(
            z3.And(
                z3.Or(safe(V1), manually_set(V1)),
                z3.Or(safe(V2), manually_set(V2)),
            ),
            V1 == V2,
        ),
    )


def invariant() -> z3.BoolRef:
    R, V, Rp = fresh_consts(Round, Value, Round)
    return z3.ForAll(
        [R, V],
        z3.Implies(
            proposal(R, V),
            z3.Or(safe(V), z3.Exists(Rp, z3.And(lt(Rp, R), proposal(Rp, V)))),
        ),
    )


def bidi_invariant() -> z3.BoolRef:
    R, V, Rp = fresh_consts(Round, Value, Round)
    return z3.ForAll(
        V,
        z3.Or(
            z3.ForAll(
                R,
                z3.Implies(
                    proposal(R, V),
                    z3.Or(
                        safe(V),
                        z3.Exists(Rp, z3.And(lt(Rp, R), proposal(Rp, V))),
                    ),
                ),
            ),
            z3.ForAll(
                R,
                z3.Implies(
                    proposal(R, V),
                    z3.Or(
                        manually_set(V),
                        z3.Exists(Rp, z3.And(lt(R, Rp), proposal(Rp, V))),
                    ),
                ),
            ),
        ),
    )


@problems.mark(
    model=True,
    osc=True,
    bruteforce_axiom=True,
    guided_axiom=True,
    popl25=True,
)
def paxos_simplified_vc() -> z3.BoolRef:
    return z3.And(
        total_strict_order(lt),
        invariant(),
        single_safe(),
        single_proposal_per_round(),
        safety_violation(),
    )


@problems.mark(
    model=True,
    osc=True,
    bruteforce_axiom=True,
    guided_axiom=True,
    popl25=True,
)
def paxos_simplified_vc_bidi() -> z3.BoolRef:
    """We have a _future_ manually_set predicate, ensuring it cannot contradict safe."""
    return z3.And(
        total_strict_order(lt),
        bidi_invariant(),
        single_safe_or_manually_set(),
        single_proposal_per_round(),
        safety_violation(),
    )
