abd-clam 0.10.0-dev0

import abc
import heapq
import typing

import numpy

from ..utils import helpers
from . import cluster

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):
                msg = "A cluster and its ancestor were both selected."
                raise ValueError(msg)

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

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

    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) -> None:
        if depth < -1:
            msg = f"expected a '-1' or a non-negative depth. got: {depth}"
            raise ValueError(msg)
        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"],
    ) -> None:
        """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.0 < percentile < 100.0:
            self.percentile: float = percentile
        else:
            msg = f"percentile must be in the (0, 100) range. Got {percentile:.2f} instead."
            raise ValueError(
                msg,
            )

        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:
            msg = f"mode must be 'above' or 'below'. Got {mode} instead."
            raise ValueError(msg)

    @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,
    ) -> None:
        """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:
            msg = "min-depth must be a positive integer."
            raise ValueError(msg)

        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",
]