import z3

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

Sort = z3.DeclareSort("Sort")

lt = relation("lt", Sort, Sort)
f = z3.Function("f", Sort, Sort)
P = relation("P", Sort)


@problems.mark
def simple_order() -> z3.BoolRef:
    x, y = z3.Consts("x y", Sort)
    return z3.And(
        total_strict_order(lt),
        z3.ForAll(x, z3.Exists(y, lt(x, y))),
    )


@problems.mark
def simple_order_f() -> z3.BoolRef:
    x = z3.Const("x", Sort)
    return z3.And(
        total_strict_order(lt),
        z3.ForAll(x, lt(x, f(x))),
    )


@problems.mark
def simple_succ_f() -> z3.BoolRef:
    x, y = z3.Consts("x y", Sort)
    return z3.And(
        simple_order_f(),
        z3.ForAll(
            [x, y],
            z3.Or(
                x == y,
                f(x) == y,
                lt(y, x),
                lt(f(x), y),
            ),
        ),
    )


@problems.mark
def simple_flip_p() -> z3.BoolRef:
    x = z3.Const("x", Sort)
    return z3.And(simple_succ_f(), z3.ForAll(x, P(x) != P(f(x))))
