import z3

from common.axioms import total_strict_order
from common.relations import Relation, relation
from problems.z3_utils import fresh_consts

Sort = z3.DeclareSort("Sort")

is_key = relation("is_key", Sort)
less_than = relation("less_than", Sort, Sort)
key = z3.Function("key", Sort, Sort)
card_witness_f = z3.Function("card_witness_f", Sort, Sort)
card_witness_c = z3.Const("card_witness_c", Sort)


def has_source(f: z3.FuncDeclRef) -> z3.BoolRef:
    x, y = fresh_consts(Sort, Sort)
    the_relation = relation(f"has_source_{f}", Sort)
    definition = z3.Exists(y, z3.And(is_node(y), f(y) == x))

    return z3.ForAll(
        x,
        z3.Implies(
            is_key(x),
            z3.And(
                z3.Implies(the_relation(x), definition),
                z3.Implies(definition, the_relation(x)),
            ),
        ),
    )


def is_node(x: z3.ExprRef) -> z3.BoolRef:
    return z3.Not(is_key(x))


def cardinality_problem() -> tuple[z3.BoolRef, Relation]:
    x, y = fresh_consts(Sort, Sort)
    return (
        z3.And(
            total_strict_order(less_than),
            is_key(card_witness_c),
            z3.ForAll(
                x,
                z3.Implies(
                    is_node(x),
                    z3.And(
                        is_key(card_witness_f(x)), card_witness_f(x) != card_witness_c
                    ),
                ),
            ),
            z3.ForAll(
                [x, y],
                z3.Implies(
                    z3.And(
                        is_node(x), is_node(y), card_witness_f(x) == card_witness_f(y)
                    ),
                    x == y,
                ),
            ),
            z3.ForAll(
                x, z3.Implies(is_key(x), z3.Exists(y, z3.And(is_node(y), key(y) == x)))
            ),
            has_source(key),
            has_source(card_witness_f),
        ),
        less_than,
    )
