from random import shuffle

import z3

from axer.axioms import (
    Semantics,
    build_axioms,
    build_characterization_axioms,
    build_pair_axioms,
)
from axer.induction_utils import find_induction_relation
from axer.predicates import (
    CallbackUnaryPredicate,
    UnaryPredicate,
    non_ground_unary_predicates,
    unary_predicates_from_ground_terms,
)
from common import config
from common.z3_utils import check_with_timeout


def bruteforce_axiomatize(
    formula: z3.BoolRef, *, semantics: Semantics = Semantics.FINITE
) -> z3.BoolRef | None:
    if config.axer.predicates.ground_terms(bool):
        predicates = unary_predicates_from_ground_terms(formula)
    else:
        predicates = non_ground_unary_predicates(formula)

    if config.axer.predicates.dummy(bool):
        predicates.append(
            CallbackUnaryPredicate(z3.BoolSort(), lambda e: z3.BoolVal(True))
        )

    return axiomatize_from_predicates(formula, predicates, semantics)


def axiomatize_from_predicates(
    formula: z3.BoolRef, predicates: list[UnaryPredicate], semantics: Semantics
) -> z3.BoolRef | None:
    order = find_induction_relation(formula)
    assert order is not None, "Formula does not contain an induction relation"

    axioms = build_axioms(order, predicates, semantics)
    shuffle(axioms)

    if config.axer.try_each_axiom(bool):
        for axiom in axioms:
            if (
                check_with_timeout(
                    z3.And(formula, axiom), config.axer.each_axiom_timeout_in_ms(int)
                )
                == z3.unsat
            ):
                return axiom

    axiom = z3.And(*axioms)
    if (
        check_with_timeout(
            z3.And(formula, axiom), config.axer.conjunction_timeout_in_ms(int)
        )
        == z3.unsat
    ):
        return axiom

    if not config.axer.predicates.pairs(bool):
        return None

    pair_axioms = build_pair_axioms(order, predicates, semantics)
    shuffle(pair_axioms)

    if config.axer.try_each_axiom(bool):
        for axiom in pair_axioms:
            if (
                check_with_timeout(
                    z3.And(formula, axiom),
                    config.axer.each_pair_axiom_timeout_in_ms(int),
                )
                == z3.unsat
            ):
                return axiom

    axiom = z3.And(*axioms, *pair_axioms)
    if (
        check_with_timeout(
            z3.And(formula, axiom), config.axer.pair_conjunction_timeout_in_ms(int)
        )
        == z3.unsat
    ):
        return axiom

    if not config.axer.predicates.characterization(bool):
        return None

    characterization_axioms = build_characterization_axioms(
        order, predicates, semantics
    )
    shuffle(characterization_axioms)

    if config.axer.try_each_axiom(bool):
        for axiom in characterization_axioms:
            if (
                check_with_timeout(
                    z3.And(formula, axiom), config.axer.each_axiom_timeout_in_ms(int)
                )
                == z3.unsat
            ):
                return axiom

    return None
