from itertools import product

import z3

from axer.axioms import Semantics
from axer.bruteforce import axiomatize_from_predicates
from axer.induction_utils import find_induction_relation
from axer.predicates import (
    RelationUnaryPredicate,
    UnaryPredicate,
    monadic_in_sort_relations,
)
from common import config
from fest.heuristic import find_sym_model_by_heuristic
from fest.templates.stdlib.arith import TemplateChoice


def guided_axiomatize(
    formula: z3.BoolRef,
    *,
    template: TemplateChoice = TemplateChoice.EXTENDED,
    semantics: Semantics = Semantics.FINITE,
) -> z3.BoolRef | None:
    order = find_induction_relation(formula)
    assert order is not None, "Formula does not contain an induction relation"

    infinite_sort = order.domain(0)

    model = find_sym_model_by_heuristic(formula, template=template)

    if model is None:
        return None

    sym_model = model

    monadic_relations = monadic_in_sort_relations(infinite_sort, formula)

    summary_nodes = [
        element for element in model.elements[infinite_sort] if element.is_summary()
    ]

    predicates = []
    for relation, sort_index in monadic_relations:
        args_options = [
            model.elements[relation.domain(i)] for i in range(relation.arity())
        ]

        predicate: UnaryPredicate = RelationUnaryPredicate.forall_except(
            relation, sort_index
        )

        literals = set()
        for node in summary_nodes:
            args_options[sort_index] = {node}
            evaluations = [
                not z3.is_true(sym_model.eval_relation(relation, *args))
                for args in product(*args_options)
            ]

            if all(evaluations):
                literals.add(True)
            else:
                literals.add(False)

        for literal in literals:
            predicates.append(predicate.negate_if(literal))

    with config.axer.use_negation.override(False):
        return axiomatize_from_predicates(formula, predicates, semantics)
