import itertools
import typing
import warnings

import networkx as nx

from kele.control.registry import register
from kele.syntax.base_classes import FACT_TYPE, Assertion, CompoundTerm, Constant, Formula, Operator, Rule, Variable

from .strategy_protocol import Feedback, RuleSelectionStrategy


class SccSortStrategyBase(RuleSelectionStrategy):
    _reverse_order: bool

    def __init__(self) -> None:
        warnings.warn("SccSort and ReverseSccSort are not stable features. Use them at your own risk.", stacklevel=3)
        self._rules: list[Rule]
        self._graph: nx.DiGraph
        """Under current implementation, this is a bipartite graph. There may
        be further performance optimizations with this property?"""
        self._order: list[Rule]
        self._stateful_gen: typing.Generator[Rule]

    @classmethod
    def _all_operators_in_fact_type(cls, f: FACT_TYPE | CompoundTerm | Constant | Variable | None) -> set[Operator]:
        match f:
            case CompoundTerm(operator=operator, arguments=arguments):
                s = set()
                s.add(operator)
                for argument in arguments:
                    s |= cls._all_operators_in_fact_type(argument)
                return s
            case Assertion(lhs=lhs, rhs=rhs):
                return cls._all_operators_in_fact_type(lhs) | cls._all_operators_in_fact_type(rhs)
            case Formula(formula_left=l, formula_right=r):
                return cls._all_operators_in_fact_type(l) | cls._all_operators_in_fact_type(r)
            case _:
                return set()

    def _rules_to_graph(self, rules: list[Rule]) -> nx.DiGraph:
        graph = nx.DiGraph()
        nodes_left = []
        for rule in rules:
            graph.add_node(rule, type_='rule')
            nodes_left.append(rule)
            heads = self._all_operators_in_fact_type(rule.head)
            bodies = self._all_operators_in_fact_type(rule.body)

            for op in bodies:
                graph.add_edge(op, rule)
            for op in heads:
                graph.add_edge(rule, op)
        return graph

    @staticmethod
    def _compile_order(graph: nx.DiGraph) -> typing.Generator[Rule]:
        def _order_within_scc(subgraph: nx.DiGraph) -> typing.Generator[Rule]:
            temp_g = subgraph.copy()
            while temp_g.nodes:
                out_d = dict(temp_g.out_degree)
                in_d = dict(temp_g.in_degree)
                n = max(temp_g.nodes, key=lambda x: out_d[x] - in_d[x])
                yield n
                temp_g.remove_node(n)

        scc_dag = nx.condensation(graph)
        scc_order = nx.topological_sort(scc_dag)
        for scc_id in scc_order:
            members: frozenset[Rule | Operator] = scc_dag.nodes[scc_id]['members']
            subgraph = typing.cast("nx.DiGraph", graph.subgraph(members))
            for m in _order_within_scc(subgraph):
                node = graph.nodes[m]
                # This "isinstance" clause is useless, just to satisfy the LSP
                if node.get('type_', None) == 'rule' and isinstance(m, Rule):
                    yield m

    def _select_next_generator(self) -> typing.Generator[Rule]:
        yield from itertools.cycle(self._order)

    def set_rules(self, rules: typing.Sequence[Rule]) -> None:
        self._rules = list(rules)
        if not self._rules:
            raise ValueError('rules cannot be empty')
        self._graph = self._rules_to_graph(self._rules)
        self._order = list(self._compile_order(self._graph))
        if self._reverse_order:
            self._order.reverse()
        self._stateful_gen = self._select_next_generator()

    def reset(self) -> None:
        self._rules = []
        self._graph = nx.DiGraph()
        self._order = []

    def select_next(self) -> typing.Sequence[Rule]:
        return [next(self._stateful_gen)]

    def on_feedback(self, feedback: Feedback) -> None:  # noqa: PLR6301
        return


@register.rule_selector('SccSort')
class SccSortStrategy(SccSortStrategyBase):
    _reverse_order = False


@register.rule_selector('ReverseSccSort')
class ReverseSccSortStrategy(SccSortStrategyBase):
    _reverse_order = True
