import operator
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from functools import reduce
from itertools import product
from typing import Callable, Self

import z3

from axer.induction_utils import find_induction_relation
from common import config
from common.relations import Relation
from common.z3_utils.expr_utils import (
    ground_terms_in_formula_by_sort,
    uninterpreted_relations,
)


class UnaryPredicate(ABC):
    @property
    @abstractmethod
    def free_variables(self) -> set[z3.Const]: ...

    @property
    @abstractmethod
    def sort(self) -> z3.SortRef: ...

    @abstractmethod
    def __call__(self, e: z3.ExprRef) -> z3.BoolRef: ...

    def __invert__(self) -> "UnaryPredicate":
        return NegatedUnaryPredicate.create(self)

    def __and__(self, other: "UnaryPredicate") -> "UnaryPredicate":
        return ConjunctionUnaryPredicate.create(self, other)

    def negate_if(self, negate: bool) -> "UnaryPredicate":
        return ~self if negate else self

    def __str__(self) -> str:
        x = z3.Const("?", self.sort)
        return self(x).sexpr()


@dataclass(frozen=True)
class CallbackUnaryPredicate(UnaryPredicate):
    _sort: z3.SortRef
    fun: Callable[[z3.ExprRef], z3.BoolRef]
    _free_variables: set[z3.Const] = field(default_factory=set)

    @property
    def sort(self) -> z3.SortRef:
        return self._sort

    @property
    def free_variables(self) -> set[z3.Const]:
        return self._free_variables

    def __call__(self, e: z3.ExprRef) -> z3.BoolRef:
        return (self.fun)(e)


@dataclass
class RelationUnaryPredicate(UnaryPredicate):
    relation: Relation
    args: list[z3.ExprRef] = field(default_factory=list)
    index: int = 0
    _free_variables: set[z3.Const] = field(default_factory=set)

    @classmethod
    def forall_except(cls, relation: Relation, index: int) -> Self:
        domains = [z3.FreshConst(relation.domain(i)) for i in range(relation.arity())]
        args: list[z3.ExprRef] = list(domains)
        domains.pop(index)

        return cls(relation, args, index, set(domains))

    @property
    def free_variables(self) -> set[z3.Const]:
        return self._free_variables

    @property
    def sort(self) -> z3.SortRef:
        return self.relation.domain(self.index)

    def __call__(self, e: z3.ExprRef) -> z3.BoolRef:
        self.args[self.index] = e
        return (self.relation)(*self.args)


@dataclass(frozen=True)
class NegatedUnaryPredicate(UnaryPredicate):
    predicate: UnaryPredicate

    @classmethod
    def create(cls, predicate: UnaryPredicate) -> UnaryPredicate:
        if isinstance(predicate, NegatedUnaryPredicate):
            return predicate.predicate
        return cls(predicate)

    @property
    def free_variables(self) -> set[z3.Const]:
        return self.predicate.free_variables

    @property
    def sort(self) -> z3.SortRef:
        return self.predicate.sort

    def __call__(self, e: z3.ExprRef) -> z3.BoolRef:
        return z3.Not(self.predicate(e))


@dataclass(frozen=True)
class ConjunctionUnaryPredicate(UnaryPredicate):
    predicates: tuple[UnaryPredicate, ...]

    @classmethod
    def create(cls, *predicates: UnaryPredicate) -> UnaryPredicate:
        flattened_predicates: list[UnaryPredicate] = []
        for predicate in predicates:
            if isinstance(predicate, ConjunctionUnaryPredicate):
                flattened_predicates.extend(predicate.predicates)
            else:
                flattened_predicates.append(predicate)

        return cls(tuple(flattened_predicates))

    @property
    def free_variables(self) -> set[z3.Const]:
        return reduce(
            operator.or_, (predicate.free_variables for predicate in self.predicates)
        )

    @property
    def sort(self) -> z3.SortRef:
        if len(self.predicates) > 0:
            return self.predicates[0].sort
        else:
            return z3.BoolSort()

    def __call__(self, e: z3.ExprRef) -> z3.BoolRef:
        return z3.And(*(predicate(e) for predicate in self.predicates))


@config(
    use_order_predicates=config.axer.predicates.order,
)
def unary_predicates_from_ground_terms(
    formula: z3.BoolRef,
    *,
    use_order_predicates: bool = False,
) -> list[UnaryPredicate]:
    order = find_induction_relation(formula)
    assert order is not None, "Formula does not contain an induction relation"

    sort = order.domain(0)

    predicates: list[UnaryPredicate] = []
    ground_terms = ground_terms_in_formula_by_sort(formula)

    if use_order_predicates:
        for ground in ground_terms[sort]:
            dummy = z3.FreshConst(sort)
            predicates.append(RelationUnaryPredicate(order, [dummy, ground], 0))
            predicates.append(RelationUnaryPredicate(order, [ground, dummy], 1))

    monadic_relations = monadic_in_sort_relations(sort, formula)

    for relation, sort_index in monadic_relations:
        domains: list[list[z3.ExprRef]] = []
        for i in range(relation.arity()):
            domain = relation.domain(i)
            if domain.eq(sort):
                domains.append([z3.FreshConst(domain)])
            else:
                domains.append(ground_terms[domain])

        possibilities = list(product(*domains))

        for possibility in possibilities:
            predicates.append(
                RelationUnaryPredicate(relation, list(possibility), sort_index)
            )

    return predicates


@config(
    use_order_predicates=config.axer.predicates.order,
)
def non_ground_unary_predicates(
    formula: z3.BoolRef,
    *,
    use_order_predicates: bool = False,
) -> list[UnaryPredicate]:
    order = find_induction_relation(formula)
    assert order is not None, "Formula does not contain an induction relation"

    sort = order.domain(0)
    predicates: list[UnaryPredicate] = []
    for relation, sort_index in monadic_in_sort_relations(sort, formula):
        relation_predicate = RelationUnaryPredicate.forall_except(relation, sort_index)

        predicates.append(relation_predicate)

        if (
            config.axer.predicates.inner_exists(bool)
            and len(relation_predicate.free_variables) >= 1
        ):
            predicates.append(_inner_exists_predicate(sort, relation_predicate))

        if (
            config.axer.predicates.inner_forall(bool)
            and len(relation_predicate.free_variables) >= 1
        ):
            predicates.append(_inner_forall_predicate(sort, relation_predicate))

    if use_order_predicates:
        predicates.extend(
            [
                RelationUnaryPredicate.forall_except(order, 0),
                RelationUnaryPredicate.forall_except(order, 1),
            ]
        )

    return predicates


def monadic_in_sort(sort: z3.SortRef, relation: Relation) -> int | None:
    indices = [i for i in range(relation.arity()) if relation.domain(i).eq(sort)]
    if len(indices) != 1:
        return None
    return indices[0]


def monadic_in_sort_relations(
    sort: z3.SortRef, formula: z3.BoolRef
) -> set[tuple[Relation, int]]:
    return {
        (relation, index)
        for relation in uninterpreted_relations(formula)
        if (index := monadic_in_sort(sort, relation)) is not None
    }


def _inner_exists_predicate(
    sort: z3.SortRef, predicate: UnaryPredicate
) -> UnaryPredicate:
    free_variables = list(predicate.free_variables)
    return CallbackUnaryPredicate(
        sort, lambda e: z3.Exists(free_variables, predicate(e))
    )


def _inner_forall_predicate(
    sort: z3.SortRef, predicate: UnaryPredicate
) -> UnaryPredicate:
    free_variables = list(predicate.free_variables)
    return CallbackUnaryPredicate(
        sort, lambda e: z3.ForAll(free_variables, predicate(e))
    )
