"""
A database to search for tree.
"""

import abc
import itertools
from collections import defaultdict
from typing import (
    Any,
    Dict,
    Iterable,
    Iterator,
    List,
    MutableMapping,
    Optional,
    Set,
    Tuple,
)

import sympy
import tabulate
from logzero import logger

from comb_spec_searcher.equiv_db import EquivalenceDB
from comb_spec_searcher.exception import InvalidOperationError
from comb_spec_searcher.rule_db.abstract import RuleDBAbstract, ensure_specification
from comb_spec_searcher.specification_extrator import SpecificationRuleExtractor
from comb_spec_searcher.strategies import AbstractStrategy, VerificationRule
from comb_spec_searcher.strategies.rule import AbstractRule
from comb_spec_searcher.tree_searcher import (
    Node,
    iterative_proof_tree_finder,
    iterative_prune,
    proof_tree_generator_dfs,
    prune,
    smallish_random_proof_tree,
)
from comb_spec_searcher.typing import CombinatorialClassType, RuleKey, RulesDict

__all__ = ["RuleDBBase", "RuleDB"]


class RuleDBBase(RuleDBAbstract):
    """
    A database to find specification using the pruning method.
    """

    def __init__(self) -> None:
        super().__init__()
        self.equivdb = EquivalenceDB()
        # Store the pruned dict. Should be set back to None every time the ruledb is
        # edited. Use the propertu for clean access.
        self._pruned_dict: Optional[RulesDict] = None

    @property
    def iterative(self) -> bool:
        """
        Indicate whether we are searching for an iterative specification.
        """
        return self.strategy_pack.iterative

    @property
    def pruned_dict(self) -> RulesDict:
        """
        Returned the prune dict of all the rules that are in a specification.

        The dict needs to be recomputed every time the database changes and it
        is costly to so. Call with moderation.
        """
        if self._pruned_dict is None:
            rules_dict = self.rules_up_to_equivalence()
            if self.iterative:
                rules_dict = iterative_prune(rules_dict, root=self.root_label)
            else:
                prune(rules_dict)
            self._pruned_dict = rules_dict
            for ver_label in rules_dict.keys():
                self.equivdb.set_verified(ver_label)
        return self._pruned_dict

    @property
    @abc.abstractmethod
    def rule_to_strategy(self) -> MutableMapping[RuleKey, AbstractStrategy]:
        pass

    @property
    @abc.abstractmethod
    def eqv_rule_to_strategy(self) -> MutableMapping[RuleKey, AbstractStrategy]:
        pass

    def __eq__(self, other: object) -> bool:
        """Check if all stored information is the same."""
        if not isinstance(other, self.__class__):
            return NotImplemented
        return (
            self.rule_to_strategy == other.rule_to_strategy
            and self.eqv_rule_to_strategy == other.eqv_rule_to_strategy
        )

    def add(self, start: int, ends: Tuple[int, ...], rule: AbstractRule) -> None:
        self._pruned_dict = None
        ends = self._clean_labels(ends, rule)
        if ends == [start]:
            return
        if isinstance(rule, VerificationRule):
            self.equivdb.set_verified(start)
        if len(ends) == 1:
            if rule.is_two_way():
                self.equivdb.add_two_way_edge(start, ends[0])
                self.eqv_rule_to_strategy[(start, ends)] = rule.strategy
                self.rule_to_strategy.pop((start, ends), None)
                self.rule_to_strategy.pop((ends[0], (start,)), None)
            else:
                self.equivdb.add_one_way_edge(start, ends[0])
                self.rule_to_strategy[(start, ends)] = rule.strategy
        else:
            self.rule_to_strategy[(start, ends)] = rule.strategy

    def _clean_labels(
        self, ends: Tuple[int, ...], rule: AbstractRule
    ) -> Tuple[int, ...]:
        """
        Remove the empty label and sort the remaining one.
        """
        cleaned_ends = []
        for comb_class, child_label in zip(rule.children, ends):
            if rule.possibly_empty and self.classdb.is_empty(comb_class, child_label):
                logger.debug("Label %s is empty.", child_label)
                self.searcher.classqueue.set_stop_yielding(child_label)
                continue
            cleaned_ends.append(child_label)
        return tuple(sorted(cleaned_ends))

    def is_verified(self, label: int) -> bool:
        """Return True if label has been verified."""
        return self.equivdb.is_verified(label)

    def are_equivalent(self, label: int, other: int) -> bool:
        """Return true if label and other are equivalent."""
        return self.equivdb.equivalent(label, other)

    def rules_up_to_equivalence(self) -> Dict[int, Set[Tuple[int, ...]]]:
        """Return a defaultdict containing all rules up to the equivalence."""
        self.equivdb.connect_cycles()
        rules_dict: Dict[int, Set[Tuple[int, ...]]] = defaultdict(set)
        for start, ends in self:
            if len(ends) == 1 and self.are_equivalent(start, ends[0]):
                continue
            rules_dict[self.equivdb[start]].add(
                tuple(sorted(self.equivdb[e] for e in ends))
            )
        return rules_dict

    def __iter__(self) -> Iterator[Tuple[int, Tuple[int, ...]]]:
        """Iterate through rules as the pairs (start, ends)."""
        return itertools.chain(self.rule_to_strategy, self.eqv_rule_to_strategy)

    def contains(self, start: int, ends: Tuple[int, ...]) -> bool:
        """Return true if the rule start -> ends is in the database."""
        key = (start, sorted(ends))
        return key in self.rule_to_strategy or key in self.eqv_rule_to_strategy

    def status(self, elaborate: bool) -> str:
        """Return a string describing the status of the rule database."""
        status = "RuleDB status:\n"
        table: List[Tuple[str, str]] = []
        table.append(("Combinatorial rules", f"{len(self.rule_to_strategy):,d}"))
        table.append(("Equivalence rules", f"{len(self.eqv_rule_to_strategy):,d}"))
        if elaborate:
            eqv_labels = set()
            strategy_verified_labels = set()
            verified_labels = set()
            eqv_rules = set()
            for start, ends in self.rule_to_strategy:
                if not ends:
                    strategy_verified_labels.add(start)
                if self.is_verified(start):
                    verified_labels.add(start)
                eqv_start = self.equivdb[start]
                eqv_ends = tuple(sorted(self.equivdb[label] for label in ends))
                eqv_labels.add(eqv_start)
                eqv_labels.update(eqv_ends)
                eqv_rules.add((eqv_start, eqv_ends))
            table.append(
                ("Combintorial rules up to equivalence", f"{len(eqv_rules):,d}")
            )
            table.append(
                (
                    "Strategy verified combinatorial classes",
                    f"{len(strategy_verified_labels):,d}",
                )
            )
            table.append(
                (
                    "Verified combinatorial classes",
                    f"{len(verified_labels):,d}",
                )
            )
            table.append(
                (
                    "combinatorial classes up to equivalence",
                    f"{len(eqv_labels):,d}",
                )
            )
        status += "    "
        status += tabulate.tabulate(
            table, headers=("", "Total number"), colalign=("left", "right")
        ).replace("\n", "\n    ")
        time_taken = round(self.equivdb.func_times["find_paths"], 2)
        status += (
            f"\n\tCalled find equiv path {self.equivdb.func_calls['find_paths']}"
            + f" times, for total time of {time_taken}"
            + " seconds.\n"
        )
        return status

    ################################################################
    # Below are methods for finding a combinatorial specification. #
    ################################################################

    def has_specification(self) -> bool:
        """Return True if a specification has been found, false otherwise."""
        return self.equivdb[self.root_label] in self.pruned_dict

    def rule_from_equivalence_rule(
        self, eqv_start: int, eqv_ends: Iterable[int]
    ) -> Optional[Tuple[int, Tuple[int, ...]]]:
        """
        Return a rule that satisfies the equivalence rule.

        Returns None if no such rule exists.
        """
        eqv_start = self.equivdb[eqv_start]
        eqv_ends = tuple(sorted(map(self.equivdb.__getitem__, eqv_ends)))
        for rule in self:
            start, ends = rule
            temp_start = self.equivdb[start]
            temp_ends = tuple(sorted(map(self.equivdb.__getitem__, ends)))
            if eqv_start == temp_start and eqv_ends == temp_ends:
                return start, ends
        return None

    def rule_from_equivalence_rule_dict(
        self, eqv_rules: Iterable[RuleKey]
    ) -> Dict[RuleKey, RuleKey]:
        """
        Return a dictionary pointing from an equivalence rule to an actual rule.
        """
        eqv_rules = set(eqv_rules)
        res: Dict[RuleKey, RuleKey] = {}
        for start, ends in self.rule_to_strategy:
            eqv_start = self.equivdb[start]
            eqv_ends = tuple(sorted(map(self.equivdb.__getitem__, ends)))
            eqv_key = (eqv_start, eqv_ends)
            if eqv_key in eqv_rules:
                res[eqv_key] = (start, ends)
        return res

    @ensure_specification
    def get_specification_rules(  # pylint: disable=arguments-differ
        self,
        *,
        minimization_time_limit: float = 10,
        smallest: bool = False,
        **kwargs,
    ) -> Iterator[AbstractRule]:
        node = self._get_specification_node(minimization_time_limit, smallest)
        logger.info("Found specification with %s rules.", len(node.labels()))
        spec_extractor = SpecificationRuleExtractor(
            self.root_label, node, self, self.classdb
        )
        return spec_extractor.rules()

    @ensure_specification
    def _get_specification_node(
        self, minimization_time_limit: float, smallest: bool
    ) -> Node:
        """
        Return a specification node for the given label.
        """
        if self.iterative:
            if smallest:
                raise InvalidOperationError("can't use iterative and smallest")
            node = self._get_iterative_node()
        else:
            if smallest:
                node = self._get_smallest_node(minimization_time_limit)
            else:
                node = self._get_smallish_node(
                    minimization_time_limit=minimization_time_limit,
                )
        return node

    @ensure_specification
    def _get_iterative_node(self) -> Node:
        if not self.iterative:
            raise InvalidOperationError("Not supported in non-iterative mode")
        return iterative_proof_tree_finder(
            self.pruned_dict, root=self.equivdb[self.root_label]
        )

    @ensure_specification
    def _get_smallish_node(self, minimization_time_limit: float) -> Node:
        """Search for a smallish node for the given minimization time."""
        if self.iterative:
            raise InvalidOperationError("Not supported in iterative mode.")
        logger.info("Minimizing for %s seconds.", round(minimization_time_limit))
        return smallish_random_proof_tree(
            self.pruned_dict, self.equivdb[self.root_label], minimization_time_limit
        )

    @ensure_specification
    def _get_smallest_node(self, minimization_time_limit: float) -> Node:
        """
        Return the smallest specification node in the universe.

        Will look for a smallish spec for minimization_time_limit and then use
        binary search to find it to find the smallest one.

        This doesn't consider the length of the equivalence paths.
        """
        if self.iterative:
            raise InvalidOperationError("Not supported in iterative mode.")
        node = self._get_smallish_node(minimization_time_limit)
        minimum = 1
        maximum = len(node)
        logger.info(
            "Found a specification of size %s. Looking for the smallest.", len(node)
        )
        # Binary search to find a smallest proof tree.
        while minimum < maximum:
            middle = (minimum + maximum) // 2
            logger.info("Looking for specification of size %s", middle)
            try:
                node = next(
                    proof_tree_generator_dfs(
                        self.pruned_dict,
                        root=self.equivdb[self.root_label],
                        maximum=middle,
                    )
                )
                maximum = min(middle, len(node))
            except StopIteration:
                minimum = middle + 1
        logger.info("The smallest specification is of size %s.", len(node))
        return node

    @ensure_specification
    def _all_nodes(self, iterative: bool = False) -> Iterator[Node]:
        """
        A generator that yields all specifications in the universe for
        the given label.
        """
        if self.iterative:
            raise NotImplementedError(
                "There is no method for yielding all iterative proof trees."
            )
        yield from proof_tree_generator_dfs(
            self.pruned_dict, root=self.equivdb[self.root_label]
        )


class RuleDB(RuleDBBase):
    def __init__(self) -> None:
        super().__init__()
        self._rule_to_strategy: Dict[Tuple[int, Tuple[int, ...]], AbstractStrategy] = {}
        self._eqv_rule_to_strategy: Dict[
            Tuple[int, Tuple[int, ...]], AbstractStrategy
        ] = {}

    @property
    def rule_to_strategy(self) -> Dict[Tuple[int, Tuple[int, ...]], AbstractStrategy]:
        return self._rule_to_strategy

    @property
    def eqv_rule_to_strategy(
        self,
    ) -> Dict[Tuple[int, Tuple[int, ...]], AbstractStrategy]:
        return self._eqv_rule_to_strategy

    def all_rules(self) -> Iterator[AbstractRule]:
        """Yield all the rules found so far."""
        for rulekey in self:
            parent = self.classdb.get_class(rulekey[0])
            try:
                strategy = self.rule_to_strategy[rulekey]
            except KeyError:
                strategy = self.eqv_rule_to_strategy[rulekey]
            yield strategy(parent)

    def all_equations(self) -> Set[sympy.Eq]:
        """
        Returns a set of equations for all rules currently found.
        """

        def get_function(comb_class: CombinatorialClassType) -> Any:
            return comb_class.get_function(self.classdb.get_label)

        eqs = set()
        for rule in self.all_rules():
            try:
                eq = rule.get_equation(get_function)
            except NotImplementedError:
                logger.info(
                    "can't find generating function for %s. The comb class is:\n%s",
                    get_function(rule.comb_class),
                    rule.comb_class,
                )
                eq = sympy.Eq(
                    get_function(rule.comb_class),
                    sympy.Function("NOTIMPLEMENTED")(sympy.var("x")),
                )
            eqs.add(eq)
        return eqs
