import z3

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

Sort = z3.DeclareSort("Sort")
SortType = z3.ExprRef

less_than = relation("less_than", Sort, Sort)

pred = z3.Function("pred", Sort, Sort)
succ = z3.Function("succ", Sort, Sort)
is_process = relation("is_process", Sort)

first = z3.Const("first", Sort)
last = z3.Const("last", Sort)

key_pre, key_post = state_function("key", Sort, Sort)  # z3.Function("key", Sort, Sort)


def lte(x: SortType, y: SortType) -> z3.BoolRef:
    return z3.Or(x == y, less_than(x, y))


def dijkstra_self_stabilizing_system() -> z3.BoolRef:
    x, y = fresh_consts(Sort, Sort)
    return z3.And(
        total_strict_order(less_than),
        z3.ForAll(x, less_than(x, succ(x))),
        z3.ForAll(
            x,
            z3.Or(
                z3.And(x == first, pred(x) == last),
                z3.And(
                    less_than(pred(x), x),
                    z3.ForAll(
                        y, z3.Not(z3.And(less_than(pred(x), y), less_than(y, x)))
                    ),
                ),
            ),
        ),
        z3.ForAll(x, lte(first, x)),
        z3.ForAll(x, lte(x, last) == is_process(x)),
        first != last,
    )


def is_privileged(s: SortType, key: z3.FuncDeclRef = key_pre) -> z3.BoolRef:
    return z3.Or(
        z3.And(s == first, key(s) == key(pred(s))),
        z3.And(s != first, key(s) != key(pred(s))),
    )
