import z3

from common.relations import Relation, relation

from .z3_utils import fresh_consts


def derive_relation(
    the_relation: Relation, *, exists: bool = True
) -> tuple[z3.BoolRef, Relation]:
    assert (
        the_relation.arity() == 2
    ), "Only binary relations are supported for derived relations"

    derived_sort = the_relation.domain(0)
    quantifier_sort = the_relation.domain(1)

    quantifier_name = "exists" if exists else "forall"

    name = f"{the_relation}_derived_{quantifier_name}"
    derived_relation = relation(name, derived_sort)

    x, y = fresh_consts(derived_sort, quantifier_sort)

    if exists:
        definition = z3.Exists(y, the_relation(x, y))
    else:
        definition = z3.ForAll(y, the_relation(x, y))

    axiom = z3.ForAll(
        x,
        z3.And(
            z3.Implies(derived_relation(x), definition),
            z3.Implies(definition, derived_relation(x)),
        ),
    )

    return axiom, derived_relation
