import z3

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

Node = z3.DeclareSort("Node")

n, n1, n2 = z3.Consts("n n1 n2", Node)

leader = relation("leader", Node)
sent = relation("sent", Node)

pending = relation("pending", Node, Node)

less_than = relation("less_than", Node, Node)

pending_source = z3.Function("pending_source", Node, Node)


def between(left: z3.Const, middle: z3.Const, right: z3.Const) -> z3.BoolRef:
    return z3.Or(
        z3.And(less_than(left, middle), less_than(middle, right)),
        z3.And(less_than(left, middle), less_than(right, left)),
        z3.And(less_than(right, left), less_than(middle, right)),
    )


def ring_topology() -> z3.BoolRef:
    w, x, y, z = z3.Consts("w x y z", Node)
    return z3.And(
        z3.ForAll(
            [w, x, y, z],
            z3.Implies(z3.And(between(w, x, y), between(w, y, z)), between(w, x, z)),
        ),
        z3.ForAll([w, x, y], z3.Implies(between(w, x, y), z3.Not(between(w, y, x)))),
        z3.ForAll(
            [w, x, y], z3.Or(between(w, x, y), between(w, y, x), w == x, w == y, x == y)
        ),
        z3.ForAll([x, y, z], z3.Implies(between(x, y, z), between(y, z, x))),
    )


def safety() -> z3.BoolRef:
    return z3.Or(
        z3.Exists(n, leader(n)),
        z3.Exists([n1, n2], z3.Or(z3.Not(sent(n1)), pending(n1, n2))),
    )


def invariant() -> z3.BoolRef:
    return z3.ForAll(
        n1,
        z3.Implies(
            z3.And(sent(n1), z3.Not(leader(n1))),
            # z3.Exists(n2, z3.Or(pending(n1, n2), less_than(n1, n2))),
            z3.Or(pending(n1, pending_source(n1)), less_than(n1, pending_source(n1))),
        ),
    )


def formula() -> z3.BoolRef:
    return z3.And(invariant(), z3.Not(safety()), total_strict_order(less_than))


@problems.mark(
    model=True,
    popl24=True,
    bruteforce_axiom=True,
    guided_axiom=True,
    popl25=True,
)
def ring_leader() -> z3.BoolRef:
    return formula()
