from enum import Enum, Flag, auto

import z3

from axer.predicates import ConjunctionUnaryPredicate, UnaryPredicate
from common import config
from common.relations import Relation


class Semantics(Flag):
    WELL_FOUNDED = auto()
    UPWARDS_WELL_FOUNDED = auto()

    FINITE = WELL_FOUNDED | UPWARDS_WELL_FOUNDED


class AxiomKind(Enum):
    MIN = auto()
    MAX = auto()


_induction_counter = 0


def induction_axiom(
    order: Relation,
    predicate: UnaryPredicate,
    kind: AxiomKind,
) -> z3.BoolRef:
    global _induction_counter
    sort = order.domain(0)
    n_edge, n = z3.Consts(f"n_{_induction_counter}_{kind.name} n", sort)

    _induction_counter += 1

    order_formula = order(n, n_edge) if kind is AxiomKind.MIN else order(n_edge, n)

    order_formula = z3.And(order_formula, n != n_edge)

    axiom = z3.Implies(
        z3.Exists(n, predicate(n)),
        z3.Exists(
            n_edge,
            z3.And(
                predicate(n_edge),
                z3.ForAll(n, z3.Implies(order_formula, z3.Not(predicate(n)))),
            ),
        ),
    )

    if predicate.free_variables:
        axiom = z3.ForAll(list(predicate.free_variables), axiom)

    return axiom


def build_axioms(
    order: Relation, predicates: list[UnaryPredicate], semantics: Semantics
) -> list[z3.BoolRef]:
    kinds = axiom_kinds_from_semantics(semantics)
    if config.axer.use_negation(bool):
        predicates = predicates + [~predicate for predicate in predicates]

    return map_predicates_to_axioms(kinds, order, predicates)


def build_pair_axioms(
    order: Relation, predicates: list[UnaryPredicate], semantics: Semantics
) -> list[z3.BoolRef]:
    pair_predicates: list[UnaryPredicate] = []

    for i, predicate1 in enumerate(predicates):
        for j in range(i + 1, len(predicates)):
            predicate2 = predicates[j]
            combinations = [(False, False)]

            if config.axer.use_negation(bool):
                combinations = [
                    (False, False),
                    (False, True),
                    (True, False),
                    (True, True),
                ]

            for negated1, negated2 in combinations:
                pair_predicates.append(
                    predicate1.negate_if(negated1) & predicate2.negate_if(negated2)
                )

    return map_predicates_to_axioms(
        axiom_kinds_from_semantics(semantics), order, pair_predicates
    )


def build_characterization_axioms(
    order: Relation, predicates: list[UnaryPredicate], semantics: Semantics
) -> list[z3.BoolRef]:
    characterization_predicates = [
        ConjunctionUnaryPredicate.create(
            *(
                predicate.negate_if(combination & (2**i) > 0)
                for i, predicate in enumerate(predicates)
            )
        )
        for combination in range(2 ** len(predicates))
    ]

    return map_predicates_to_axioms(
        axiom_kinds_from_semantics(semantics), order, characterization_predicates
    )


def map_predicates_to_axioms(
    kinds: list[AxiomKind], order: Relation, predicates: list[UnaryPredicate]
) -> list[z3.BoolRef]:
    return [
        induction_axiom(order, predicate, kind)
        for predicate in predicates
        for kind in kinds
    ]


def axiom_kinds_from_semantics(semantics: Semantics) -> list[AxiomKind]:
    kinds = []
    if Semantics.WELL_FOUNDED in semantics:
        kinds.append(AxiomKind.MIN)
    if Semantics.UPWARDS_WELL_FOUNDED in semantics:
        kinds.append(AxiomKind.MAX)
    return kinds
