import abc
import heapq
import typing

import numpy

from . import cluster
from ..utils import helpers

logger = helpers.make_logger(__name__)


class GraphCriterion(abc.ABC):
    """ A `GraphCriterion` can be called on the `root` cluster of a tree to
    select a set of clusters from the tree. These clusters can be used to
    create a `Graph`. Subclasses much implement the `select` method. The
    abstract `GraphCriterion` will verify that the set of selected clusters
    obey the graph invariant. See the documentation for `Graph` for details on
    the graph invariant.
    """

    @property
    @abc.abstractmethod
    def name(self) -> str:
        pass

    @abc.abstractmethod
    def select(self, root: cluster.Cluster) -> set[cluster.Cluster]:
        pass

    @staticmethod
    def assert_invariant(root: cluster.Cluster, selected: set[cluster.Cluster]):
        for c in selected:
            if any((c.is_ancestor_of(other) for other in selected)):
                raise ValueError(f'A cluster and its ancestor were both selected.')

        indices = {i for c in selected for i in c.indices}

        if len(indices) != root.cardinality:
            raise ValueError(
                f'There was a mis-match in the number of instances that were '
                f'selected. The selected clusters have {len(indices)} instance '
                f'but the root has {root.cardinality}.'
            )

        return

    def __call__(self, root: cluster.Cluster) -> set[cluster.Cluster]:
        selected = self.select(root)
        self.assert_invariant(root, selected)
        return selected


class Layer(GraphCriterion):
    """ Selects the layer at the specified depth, -1 means go to leaves.
    """
    def __init__(self, depth: int):
        if depth < -1:
            raise ValueError(f'expected a \'-1\' or a non-negative depth. got: {depth}')
        self.depth = depth

    @property
    def name(self) -> str:
        return f'Layer_{self.depth}'

    def select(self, root: cluster.Cluster) -> set[cluster.Cluster]:

        if self.depth == -1:
            return {c for layer in root.subtree for c in layer if c.is_leaf}

        else:
            selected = {c for layer in root.subtree[:self.depth] for c in layer if c.is_leaf}
            selected.update(root.subtree[self.depth])

            return selected


class PropertyThreshold(GraphCriterion):
    """ Selects clusters when the given property crosses the given
    percentile.
    """

    def __init__(
            self,
            value: typing.Literal['cardinality', 'radius', 'lfd'],
            percentile: float,
            mode: typing.Literal['above', 'below'],
    ):
        """
        During a BFT of the tree, this keeps track of the given cluster
        property. A cluster is selected if any of the following is true:
            - it is a leaf, or
            - it qualifies by `mode`, or
            - both children qualify by `mode`.

        If `mode` is 'above', a cluster qualifies if the property goes above the
        `percentile`. If `mode` is 'below', a cluster qualifies if the property
        goes below the `percentile`.

        Args:
            value: The cluster property to keep track of. Must be one of
             'cardinality', 'radius', or 'lfd'.
            percentile: the percentile, in the (0, 100) range, to be crossed.
            mode: select a cluster when `value` is 'above' or 'below' the
             `percentile`.
        """
        if 0. < percentile < 100.:
            self.percentile: float = percentile
        else:
            raise ValueError(f'percentile must be in the (0, 100) range. Got {percentile:.2f} instead.')

        self.value = value

        if mode == 'above':
            self.qualifies = lambda c, v: getattr(c, self.value) > v
        elif mode == 'below':
            self.qualifies = lambda c, v: getattr(c, self.value) < v
        else:
            raise ValueError(f'mode must be \'above\' or \'below\'. Got {mode} instead.')

    @property
    def name(self) -> str:
        return f'PropertyThreshold_{self.value}_{self.name}_{self.percentile:.2f}'

    def select(self, root: cluster.Cluster) -> set[cluster.Cluster]:
        threshold = float(numpy.percentile(
            [getattr(c, self.value) for layer in root.subtree for c in layer if c.cardinality > 1],
            self.percentile,
        ))
        selected: set[cluster.Cluster] = set()
        frontier: set[cluster.Cluster] = {root}
        while frontier:
            c = frontier.pop()
            if (
                    c.is_leaf  # select the leaves
                    or self.qualifies(c, threshold)  # if cluster qualifies
                    or (self.qualifies(c.left_child, threshold) and self.qualifies(c.right_child, threshold))
            ):
                selected.add(c)
            else:  # add children to frontier
                frontier.update(c.children)
        return selected


class MetaMLSelect(GraphCriterion):
    """ Uses the scoring function from a trained meta-ml model to select
    clusters.
    """

    def __init__(self, scorer: typing.Callable[[numpy.ndarray], float], name: typing.Optional[str] = None, min_depth: int = 4):
        """
        Args:
            scorer: A function that takes the ratios of a cluster and returns
             its predicted score. Higher scoring clusters are considered
             better than lower scoring clusters.
            min_depth: The minimum depth in the tree for a cluster to be
             selected.
        """
        if min_depth < 1:
            raise ValueError(f'min-depth must be a positive integer.')

        self.__name = scorer.__name__ if name is None else name
        self.scorer = lambda c: -scorer(numpy.asarray(c.ratios, dtype=numpy.float32))
        self.min_depth = min_depth

    @property
    def name(self) -> str:
        return self.__name

    def select(self, root: cluster.Cluster) -> set[cluster.Cluster]:

        tree = [c for layer in root.subtree for c in layer]

        candidate_clusters = [c for c in tree if c.depth >= self.min_depth]
        normalized_scores = list(map(float, helpers.normalize(
            numpy.asarray(list(map(self.scorer, candidate_clusters))),
            mode='gaussian',
        )))

        heap = list(zip(normalized_scores, candidate_clusters))
        heapq.heapify(heap)

        selected = {c for c in tree if c.depth < self.min_depth and c.is_leaf}
        selected_indices = {i for c in selected for i in c.indices}

        while len(heap) > 0:
            _, c = heapq.heappop(heap)

            if len(selected_indices.intersection(set(c.indices))) > 0:
                continue
            else:
                selected.add(c)
                selected_indices.update(set(c.indices))

        return selected


__all__ = [
    'GraphCriterion',
    'Layer',
    'PropertyThreshold',
    'MetaMLSelect',
]
