import z3

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

Person = z3.DeclareSort("Person")

bob = z3.Const("bob", Person)

child_of = z3.Function("child_of", Person, Person)
is_descendant = relation("is_descendant", Person, Person)


@problems.mark
def tptp_descendants() -> z3.BoolRef:
    a, c, d, g, p = fresh_consts(Person, Person, Person, Person, Person)
    return z3.And(
        z3.ForAll([a, d], z3.Implies(is_descendant(a, d), a != d)),
        z3.ForAll(
            [a, c, g],
            z3.Implies(
                z3.And(is_descendant(a, c), is_descendant(c, g)),
                is_descendant(a, g),
            ),
        ),
        z3.ForAll(p, is_descendant(p, child_of(p))),
        z3.ForAll(p, z3.Exists(c, c == child_of(p))),
    )
