from collections.abc import Callable

import click
import z3

from axer.bruteforce import bruteforce_axiomatize
from axer.guided import guided_axiomatize
from common import config
from common.cli_utils import error_reporting
from problems import Problem
from problems.benchmark_utils import average_time
from problems.cli_utils import problem_options


def try_auto(formula: z3.BoolRef) -> Callable[[], None]:
    def attempt() -> None:
        axiom = bruteforce_axiomatize(formula)
        if axiom is None:
            raise RuntimeError()

    return attempt


def try_guided(formula: z3.BoolRef) -> Callable[[], None]:
    def attempt() -> None:
        axiom = guided_axiomatize(formula)
        if axiom is None:
            raise RuntimeError()

    return attempt


@click.command(
    help="Benchmark for the POPL 2025 paper "
    "`Axe 'Em: Eliminating Spurious States with Induction Axioms`",
    context_settings=dict(
        show_default=True,
        help_option_names=["-h", "--help"],
    ),
)
@click.option("--bruteforce/--no-bruteforce", "-b / -B", default=True)
@click.option("--guided/--no-guided", "-g / -G", default=True)
@click.option(
    "-r",
    "--runs",
    type=int,
    help="Number of runs to average runtime",
    default=config.popl25.runs[int],
)
@error_reporting
@problem_options(popl25=True)
def main(
    problems: set[Problem],
    bruteforce: bool,
    guided: bool,
    runs: int,
) -> None:
    for problem in problems:
        print(f"(problem {problem.name}")

        if bruteforce:
            print(f"\t(bruteforce '{average_time(try_auto(problem.formula), runs)}')")

        if guided:
            print(f"\t(guided '{average_time(try_guided(problem.formula), runs)}')")

        print(")")


if __name__ == "__main__":
    main()
