"""
The slicefinder module implements the Slicefinder class.
"""

from __future__ import annotations

import logging
import warnings
from typing import Any

import numpy as np
import numpy.typing as npt
from scipy import sparse as sp
from scipy.stats import rankdata
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.preprocessing import OneHotEncoder
from sklearn.utils.validation import _check_feature_names, check_is_fitted

from sliceline.validation import check_array, check_X_e

ArrayLike = npt.ArrayLike
NDArray = npt.NDArray[Any]

logger = logging.getLogger(__name__)

# Numba availability detection
try:
    from sliceline._numba_ops import (
        compute_slice_ids_numba,
        score_slices_numba,
        score_ub_batch_numba,
    )

    NUMBA_AVAILABLE = True
except ImportError:
    NUMBA_AVAILABLE = False
    score_slices_numba = None
    score_ub_batch_numba = None
    compute_slice_ids_numba = None


def is_numba_available() -> bool:
    """Check if numba is available for acceleration.

    Returns
    -------
    bool
        True if numba is installed and can be used for acceleration.
    """
    return NUMBA_AVAILABLE


def _warn_numba_not_available() -> None:
    """Issue a warning if numba is not available."""
    warnings.warn(
        "Numba not available. Install with: pip install numba\n"
        "Or: pip install sliceline[optimized]\n"
        "Performance will be 5-50x slower without Numba optimization.",
        UserWarning,
        stacklevel=3,
    )


class Slicefinder(BaseEstimator, TransformerMixin):
    """Slicefinder class.

    SliceLine is a fast, linear-algebra-based slice finding for ML Model Debugging.

    Given an input dataset (`X`) and a model error vector (`errors`), SliceLine finds
    the `k` slices in `X` that identify where the model performs significantly worse.
    A slice is a subspace of `X` defined by one or more predicates.

    The maximal dimension of this subspace is controlled by `max_l`.

    The slice scoring function is the linear combination of two objectives:
        - Find sufficiently large slices, with more than `min_sup` elements
          (high impact on the overall model)
        - With substantial errors
          (high negative impact on sub-group/model)

    The importance of each objective is controlled through a single parameter `alpha`.

    Slice enumeration and pruning techniques are done via sparse linear algebra.

    Parameters
    ----------
    alpha: float, default=0.6
        Weight parameter for the importance of the average slice error.
        0 < `alpha` <= 1.

    k: int, default=1
        Maximum number of slices to return.
        Note: in case of equality between `k`-th slice score and the following ones,
        all those slices are returned, leading to `_n_features_out` slices returned.
        (`_n_features_out` >= `k`)

    max_l: int, default=4
        Maximum lattice level.
        In other words: the maximum number of predicate to define a slice.

    min_sup: int or float, default=10
        Minimum support threshold. Inspired by frequent itemset mining,
        it ensures statistical significance. If `min_sup` is a float (0 < `min_sup` < 1),
        it represents the faction of the input dataset (`X`).

    verbose: bool, default=True
        Controls the verbosity.

    Attributes
    ----------
    top_slices: np.ndarray of shape (_n_features_out, number of columns of the input dataset)
        The `_n_features_out` slices with the highest score.
        `None` values in slices represent unused column in the slice.

    average_error: float
        Mean value of the input error.

    top_slices_statistics: list of dict of length `len(top_slices_)`
        The statistics of the slices found sorted by slice's scores.
        For each slice, the following statistics are stored:
        - slice_score: the score of the slice (defined in `_score` method)
        - sum_slice_error: the sum of all the errors in the slice
        - max_slice_error: the maximum of all errors in the slice
        - slice_size: the number of elements in the slice
        - slice_average_error: the average error in the slice (sum_slice_error / slice_size)

    References
    ----------
    `SliceLine: Fast, Linear-Algebra-based Slice Finding for ML Model Debugging
    <https://mboehm7.github.io/resources/sigmod2021b_sliceline.pdf>`__,
    from *Svetlana Sagadeeva* and *Matthias Boehm* of Graz University of Technology.
    """

    def __init__(
        self,
        alpha: float = 0.6,
        k: int = 1,
        max_l: int = 4,
        min_sup: int | float = 10,
        verbose: bool = True,
    ) -> None:
        self.alpha = alpha
        self.k = k
        self.max_l = max_l
        self.min_sup = min_sup
        self.verbose = verbose

        self._one_hot_encoder = self._top_slices_enc = None
        self.top_slices_ = self.top_slices_statistics_ = None
        self.average_error_ = None
        self._min_sup_actual = min_sup

        if self.verbose:
            logger.setLevel(logging.DEBUG)
        else:
            logger.setLevel(logging.INFO)

        # Warn user once if Numba optimization is not available
        if not NUMBA_AVAILABLE and verbose:
            warnings.warn(
                "Numba JIT optimization not available. "
                "Install with 'pip install sliceline[optimized]' "
                "for 5-50x performance improvements on scoring operations. "
                "See https://github.com/DataDome/sliceline for details.",
                UserWarning,
                stacklevel=2,
            )

    def _check_params(self) -> None:
        """Check transformer parameters."""
        if not 0 < self.alpha <= 1:
            raise ValueError(f"Invalid 'alpha' parameter: {self.alpha}")

        if self.k <= 0:
            raise ValueError(f"Invalid 'k' parameter: {self.k}")

        if self.max_l <= 0:
            raise ValueError(f"Invalid 'max_l' parameter: {self.max_l}")

        if self.min_sup < 0 or (
            isinstance(self.min_sup, float) and self.min_sup >= 1
        ):
            raise ValueError(f"Invalid 'min_sup' parameter: {self.min_sup}")

    def _check_top_slices(self) -> None:
        """Check if slices have been found."""
        # Check if fit has been called
        check_is_fitted(self)

        # Check if a slice has been found
        if self.top_slices_.size == 0:
            raise ValueError("No transform: Sliceline did not find any slice.")

    def fit(self, X: ArrayLike, errors: ArrayLike) -> Slicefinder:
        """Search for slice(s) on `X` based on `errors`.

        Parameters
        ----------
        X: array-like of shape (n_samples, n_features)
            Training data, where `n_samples` is the number of samples
            and `n_features` is the number of features.

        errors: array-like of shape (n_samples, )
            Errors of a machine learning model.

        Returns
        -------
        self: object
            Returns the instance itself.
        """
        self._check_params()

        # Compute actual min_sup value (convert fraction to count if needed)
        if 0 < self.min_sup < 1:
            self._min_sup_actual = int(self.min_sup * len(X))
        else:
            self._min_sup_actual = self.min_sup

        # Check that X and e have correct shape
        X_array, errors = check_X_e(X, errors, y_numeric=True)

        _check_feature_names(self, X, reset=True)

        self._search_slices(X_array, errors)

        return self

    def transform(self, X: ArrayLike) -> NDArray:
        """Generate slices masks for `X`.

        Parameters
        ----------
        X: array-like of shape (n_samples, n_features)
            Training data, where `n_samples` is the number of samples
            and `n_features` is the number of features.

        Returns
        -------
        slices_masks: np.ndarray of shape (n_samples, _n_features_out)
            `slices_masks[i, j] == 1`: the `i`-th sample of `X` is in the `j`-th `top_slices_`.
        """
        self._check_top_slices()

        # Input validation
        X = check_array(X)

        slices_masks = self._get_slices_masks(X)

        return slices_masks.T

    def get_slice(self, X: ArrayLike, slice_index: int) -> NDArray:
        """Filter `X` samples according to the `slice_index`-th slice.

        Parameters
        ----------
        X: array-like of shape (n_samples, n_features)
            Dataset, where `n_samples` is the number of samples
            and `n_features` is the number of features.

        slice_index: int
            Index of the slice to get from `top_slices_`.

        Returns
        -------
        X_slice: np.ndarray of shape (n_samples in the `slice_index`-th slice, n_features)
            Filter `X` samples that are in the `slice_index`-th slice.
        """
        self._check_top_slices()

        # Input validation
        X = check_array(X, force_all_finite=False)

        slices_masks = self._get_slices_masks(X)

        return X[np.where(slices_masks[slice_index])[0], :]

    def get_feature_names_out(self) -> NDArray:
        """Get output feature names for transformation.

        Returns
        -------
        feature_names_out : ndarray of str objects
            The following output feature names are generated:
            `["slice_0", "slice_1", ..., "slice_(_n_features_out)"]`.
        """
        check_is_fitted(self)

        feature_names = [f"slice_{i}" for i in range(self._n_features_out)]

        return np.array(feature_names, dtype=object)

    def _get_slices_masks(self, X: NDArray) -> NDArray:
        """Private utilities function generating slices masks for `X`."""
        X_encoded = self._one_hot_encoder.transform(X)

        # Shape X_encoded: (X.shape[0], total number of modalities in _one_hot_encoder.categories_)
        # Shape _top_slices_enc: (top_slices_.shape[0], X_encoded[1])
        slice_candidates = self._top_slices_enc @ X_encoded.T

        # self._top_slices_enc.sum(axis=1) is the number of predicate(s) for each top_slices_
        slices_masks = (
            slice_candidates == self._top_slices_enc.sum(axis=1)
        ).A.astype(int)

        return slices_masks

    @property
    def _n_features_out(self) -> int:
        """Number of transformed output features."""
        return self.top_slices_.shape[0]

    @staticmethod
    def _dummify(array: NDArray, n_col_x_encoded: int) -> sp.csr_matrix:
        """Dummify `array` with respect to `n_col_x_encoded`.

        Creates a sparse one-hot encoding matrix where each row corresponds
        to an element in array and has a single True value in the column
        specified by that element (adjusted for 1-based indexing).

        Args:
            array: 1-based indices to encode (must not contain 0)
            n_col_x_encoded: Number of columns in output matrix

        Returns:
            Sparse CSR matrix of shape (len(array), n_col_x_encoded)

        Raises:
            ValueError: If array contains 0, which cannot be one-hot encoded.
        """
        if 0 in array:
            raise ValueError(
                "Modality 0 is not expected to be one-hot encoded."
            )

        # Direct CSR construction: 2-3x faster than lil_matrix approach
        n = array.size
        return sp.csr_matrix(
            (np.ones(n, dtype=np.bool_), (np.arange(n), array - 1)),
            shape=(n, n_col_x_encoded),
            dtype=np.bool_,
        )

    def _maintain_top_k(
        self,
        slices: sp.csr_matrix,
        statistics: NDArray,
        top_k_slices: sp.csr_matrix,
        top_k_statistics: NDArray,
    ) -> tuple[sp.csr_matrix, NDArray]:
        """Add new `slices` to `top_k_slices` and update the top-k slices."""
        # prune invalid min_sup and scores
        valid_slices_mask = (statistics[:, 3] >= self._min_sup_actual) & (
            statistics[:, 0] > 0
        )
        if np.sum(valid_slices_mask) != 0:
            slices, statistics = (
                slices[valid_slices_mask],
                statistics[valid_slices_mask],
            )

            if (slices.shape[1] != top_k_slices.shape[1]) and (
                slices.shape[1] == 1
            ):
                slices, statistics = slices.T, statistics.T

            # evaluated candidates and previous top-k
            slices = sp.vstack([top_k_slices, slices])
            statistics = np.concatenate([top_k_statistics, statistics])

            # extract top-k
            top_slices_bool = (
                rankdata(-statistics[:, 0], method="min") <= self.k
            )
            top_k_slices, top_k_statistics = (
                slices[top_slices_bool],
                statistics[top_slices_bool],
            )
            # Sort by score (descending), then lexicographically by slice representation
            # to ensure deterministic ordering when scores are equal
            scores = -top_k_statistics[:, 0]
            slice_keys = tuple(
                top_k_slices.toarray()[:, i]
                for i in range(top_k_slices.shape[1])
            )
            top_slices_indices = np.lexsort(slice_keys[::-1] + (scores,))
            top_k_slices, top_k_statistics = (
                top_k_slices[top_slices_indices],
                top_k_statistics[top_slices_indices],
            )
        return top_k_slices, top_k_statistics

    def _score_ub(
        self,
        slice_sizes_ub: NDArray,
        slice_errors_ub: NDArray,
        max_slice_errors_ub: NDArray,
        n_col_x_encoded: int,
    ) -> NDArray:
        """Compute the upper-bound score for all the slices.

        Uses Numba JIT compilation when available for 5-10x speedup.
        """
        if NUMBA_AVAILABLE and score_ub_batch_numba is not None:
            return score_ub_batch_numba(
                slice_sizes_ub.astype(np.float64),
                slice_errors_ub.astype(np.float64),
                max_slice_errors_ub.astype(np.float64),
                n_col_x_encoded,
                float(self._min_sup_actual),
                self.alpha,
                self.average_error_,
            )

        # Fallback to NumPy implementation
        # Since slice_scores is either monotonically increasing or decreasing, we
        # probe interesting points of slice_scores in the interval [min_sup, ss],
        # and compute the maximum to serve as the upper bound
        potential_solutions = np.column_stack(
            (
                self._min_sup_actual * np.ones(slice_sizes_ub.shape[0]),
                np.maximum(
                    slice_errors_ub / max_slice_errors_ub, self._min_sup_actual
                ),
                slice_sizes_ub,
            )
        )
        slice_scores_ub = np.amax(
            (
                self.alpha
                * (
                    np.minimum(
                        potential_solutions.T * max_slice_errors_ub,
                        slice_errors_ub,
                    ).T
                    / self.average_error_
                    - potential_solutions
                )
                - (1 - self.alpha) * (n_col_x_encoded - potential_solutions)
            )
            / potential_solutions,
            axis=1,
        )
        return slice_scores_ub

    @staticmethod
    def _analyse_top_k(top_k_statistics: NDArray) -> tuple[float, float]:
        """Get the maximum and the minimum slices scores."""
        max_slice_scores = min_slice_scores = -np.inf
        if top_k_statistics.shape[0] > 0:
            max_slice_scores = top_k_statistics[0, 0]
            min_slice_scores = top_k_statistics[
                top_k_statistics.shape[0] - 1, 0
            ]
        return max_slice_scores, min_slice_scores

    def _score(
        self,
        slice_sizes: NDArray,
        slice_errors: NDArray,
        n_row_x_encoded: int,
    ) -> NDArray:
        """Compute the score for all the slices.

        Uses Numba JIT compilation when available for 5-10x speedup.
        """
        if NUMBA_AVAILABLE and score_slices_numba is not None:
            # Ensure inputs are float64 for numba
            sizes = np.asarray(slice_sizes, dtype=np.float64)
            errors = np.asarray(slice_errors, dtype=np.float64)
            return score_slices_numba(
                sizes,
                errors,
                n_row_x_encoded,
                self.alpha,
                self.average_error_,
            )

        # Fallback to NumPy implementation
        with np.errstate(divide="ignore", invalid="ignore"):
            slice_scores = self.alpha * (
                (slice_errors / slice_sizes) / self.average_error_ - 1
            ) - (1 - self.alpha) * (n_row_x_encoded / slice_sizes - 1)
        return np.nan_to_num(slice_scores, nan=-np.inf)

    def _eval_slice(
        self,
        x_encoded: sp.csr_matrix,
        errors: NDArray,
        slices: sp.csr_matrix,
        level: int,
    ) -> NDArray:
        """Compute several statistics for all the slices."""
        slice_candidates = x_encoded @ slices.T == level
        slice_sizes = slice_candidates.sum(axis=0).A[0]
        slice_errors = errors @ slice_candidates
        # Here we can't use the .A shorthand because it is not
        # implemented in all scipy versions for coo_matrix objects
        max_slice_errors = (
            slice_candidates.T.multiply(errors).max(axis=1).toarray()
        )

        # score of relative error and relative size
        slice_scores = self._score(
            slice_sizes, slice_errors, x_encoded.shape[0]
        )
        return np.column_stack(
            [slice_scores, slice_errors, max_slice_errors, slice_sizes]
        )

    def _create_and_score_basic_slices(
        self,
        x_encoded: sp.csr_matrix,
        n_col_x_encoded: int,
        errors: NDArray,
    ) -> tuple[sp.csr_matrix, NDArray]:
        """Initialise 1-slices, i.e. slices with one predicate."""
        slice_sizes = x_encoded.sum(axis=0).A[0]
        slice_errors = errors @ x_encoded
        # Here we can't use the .A shorthand because it is not
        # implemented in all scipy versions for coo_matrix objects
        max_slice_errors = (
            x_encoded.T.multiply(errors).max(axis=1).toarray()[:, 0]
        )

        # working set of active slices (#attr x #slices) and top-k
        valid_slices_mask = (slice_sizes >= self._min_sup_actual) & (
            slice_errors > 0
        )
        attr = np.arange(1, n_col_x_encoded + 1)[valid_slices_mask]
        slice_sizes = slice_sizes[valid_slices_mask]
        slice_errors = slice_errors[valid_slices_mask]
        max_slice_errors = max_slice_errors[valid_slices_mask]
        slices = self._dummify(attr, n_col_x_encoded)

        # score 1-slices and create initial top-k
        slice_scores = self._score(
            slice_sizes, slice_errors, x_encoded.shape[0]
        )
        statistics = np.column_stack(
            (slice_scores, slice_errors, max_slice_errors, slice_sizes)
        )

        n_col_dropped = n_col_x_encoded - sum(valid_slices_mask)
        logger.debug(
            "Dropping %i/%i features below min_sup = %i."
            % (n_col_dropped, n_col_x_encoded, self._min_sup_actual)
        )

        return slices, statistics

    def _get_pruned_s_r(
        self, slices: sp.csr_matrix, statistics: NDArray
    ) -> tuple[sp.csr_matrix, NDArray]:
        """Prune invalid slices.
        Do not affect overall pruning effectiveness due to handling of missing parents.
        """
        valid_slices_mask = (statistics[:, 3] >= self._min_sup_actual) & (
            statistics[:, 1] > 0
        )
        return slices[valid_slices_mask], statistics[valid_slices_mask]

    @staticmethod
    def _join_compatible_slices(
        slices: sp.csr_matrix, level: int
    ) -> sp.csr_matrix:
        """Join compatible slices keeping sparse format when beneficial.

        Returns a sparse boolean matrix where entry (i,j) is True if slices
        i and j are compatible for joining at the given level. Only upper
        triangular entries (i < j) are populated.

        For level==2 (looking for disjoint slices), uses dense format since
        most pairs are compatible. For higher levels, keeps sparse format.
        """
        n_slices = slices.shape[0]
        if n_slices == 0:
            return sp.csr_matrix((0, 0), dtype=np.bool_)

        slices_int = slices.astype(int)
        join_counts = slices_int @ slices_int.T

        if level == 2:
            # For level 2, we're looking for pairs with dot product == 0
            # Most pairs will match, so dense is more efficient
            join_dense = join_counts.toarray() == 0
        else:
            # For higher levels, most pairs won't match, so sparse is better
            # Use dense conversion for smaller matrices to ensure consistent ordering
            # This matches the original behavior and ensures deterministic results
            join_dense = join_counts.toarray() == level - 2

        join_upper = np.triu(join_dense, 1)
        rows, cols = np.where(join_upper)
        return sp.csr_matrix(
            (np.ones(len(rows), dtype=np.bool_), (rows, cols)),
            shape=join_counts.shape,
            dtype=np.bool_,
        )

    @staticmethod
    def _combine_slices(
        slices: sp.csr_matrix,
        statistics: NDArray,
        compatible_slices: sp.csr_matrix,
    ) -> tuple[sp.csr_matrix, NDArray, NDArray, NDArray]:
        """Combine slices by exploiting parents node statistics.

        Works with sparse compatible_slices matrix returned by
        _join_compatible_slices.
        """
        parent_1_idx, parent_2_idx = compatible_slices.nonzero()
        pair_candidates = slices[parent_1_idx] + slices[parent_2_idx]

        slice_errors = np.minimum(
            statistics[parent_1_idx, 1], statistics[parent_2_idx, 1]
        )
        max_slice_errors = np.minimum(
            statistics[parent_1_idx, 2], statistics[parent_2_idx, 2]
        )
        slice_sizes = np.minimum(
            statistics[parent_1_idx, 3], statistics[parent_2_idx, 3]
        )
        return pair_candidates, slice_sizes, slice_errors, max_slice_errors

    @staticmethod
    def _prune_invalid_self_joins(
        feature_offset_start: NDArray,
        feature_offset_end: NDArray,
        pair_candidates: sp.csr_matrix,
        slice_sizes: NDArray,
        slice_errors: NDArray,
        max_slice_errors: NDArray,
    ) -> tuple[sp.csr_matrix, NDArray, NDArray, NDArray]:
        """Prune invalid self joins (>1 bit per feature)."""
        valid_slices_mask = np.full(pair_candidates.shape[0], True)
        for start, end in zip(feature_offset_start, feature_offset_end):
            valid_slices_mask = (
                valid_slices_mask
                * (pair_candidates[:, start:end].sum(axis=1) <= 1).A[:, 0]
            )
        return (
            pair_candidates[valid_slices_mask],
            slice_sizes[valid_slices_mask],
            slice_errors[valid_slices_mask],
            max_slice_errors[valid_slices_mask],
        )

    @staticmethod
    def _prepare_deduplication_and_pruning(
        feature_offset_start: NDArray,
        feature_offset_end: NDArray,
        feature_domains: NDArray,
        pair_candidates: sp.csr_matrix,
    ) -> NDArray:
        """Prepare IDs for deduplication and pruning.

        Uses Numba JIT compilation when available for 10-50x speedup.
        """
        if NUMBA_AVAILABLE and compute_slice_ids_numba is not None:
            return compute_slice_ids_numba(
                pair_candidates.data.astype(np.float64),
                pair_candidates.indices.astype(np.int64),
                pair_candidates.indptr.astype(np.int64),
                feature_offset_start.astype(np.int64),
                feature_offset_end.astype(np.int64),
                feature_domains.astype(np.float64),
            )

        # Fallback to Python implementation
        ids = np.zeros(pair_candidates.shape[0])
        dom = feature_domains + 1
        for j, (start, end) in enumerate(
            zip(feature_offset_start, feature_offset_end)
        ):
            sub_pair_candidates = pair_candidates[:, start:end]
            # sub_p should not contain multiple True on the same line
            i = sub_pair_candidates.argmax(axis=1).T + np.any(
                # Here we can't use the .A shorthand because it is not
                # implemented in all scipy versions for coo_matrix objects
                sub_pair_candidates.toarray(),
                axis=1,
            )
            ids = ids + i.A * np.prod(dom[(j + 1) : dom.shape[0]])
        return ids

    def _get_pair_candidates(
        self,
        slices: sp.csr_matrix,
        statistics: NDArray,
        top_k_statistics: NDArray,
        level: int,
        n_col_x_encoded: int,
        feature_domains: NDArray,
        feature_offset_start: NDArray,
        feature_offset_end: NDArray,
    ) -> sp.csr_matrix:
        """Compute and prune plausible slices candidates."""
        compatible_slices = self._join_compatible_slices(slices, level)

        if compatible_slices.nnz == 0:
            return sp.csr_matrix(np.empty((0, slices.shape[1])))

        (
            pair_candidates,
            slice_sizes,
            slice_errors,
            max_slice_errors,
        ) = self._combine_slices(slices, statistics, compatible_slices)

        (
            pair_candidates,
            slice_sizes,
            slice_errors,
            max_slice_errors,
        ) = self._prune_invalid_self_joins(
            feature_offset_start,
            feature_offset_end,
            pair_candidates,
            slice_sizes,
            slice_errors,
            max_slice_errors,
        )

        if pair_candidates.shape[0] == 0:
            return sp.csr_matrix(np.empty((0, slices.shape[1])))

        ids = self._prepare_deduplication_and_pruning(
            feature_offset_start,
            feature_offset_end,
            feature_domains,
            pair_candidates,
        )

        # remove duplicate candidates and select corresponding statistics
        _, unique_candidate_indices, duplicate_counts = np.unique(
            ids, return_index=True, return_counts=True
        )

        # Slices at level i normally have i parents (cf. section 3.1 in the paper)
        # We want to keep only slices whose parents have not been pruned.
        # If all the parents are present they are going to get combined 2 by 2 in i*(i-1)/2 ways
        # So, we select only candidates which appear with the correct cardinality.
        all_parents_mask = duplicate_counts == level * (level - 1) / 2
        unique_candidate_indices = unique_candidate_indices[all_parents_mask]

        pair_candidates = pair_candidates[unique_candidate_indices]
        slice_sizes = slice_sizes[unique_candidate_indices]
        slice_errors = slice_errors[unique_candidate_indices]
        max_slice_errors = max_slice_errors[unique_candidate_indices]

        slice_scores = self._score_ub(
            slice_sizes,
            slice_errors,
            max_slice_errors,
            n_col_x_encoded,
        )

        # Seems to be always fully True
        # Due to maintain_top_k that apply slice_sizes filter
        pruning_sizes = slice_sizes >= self._min_sup_actual

        _, min_slice_scores = self._analyse_top_k(top_k_statistics)

        pruning_scores = (slice_scores > min_slice_scores) & (slice_scores > 0)

        return pair_candidates[pruning_scores & pruning_sizes]

    def _search_slices(
        self,
        input_x: NDArray,
        errors: NDArray,
    ) -> None:
        """Main function of the SliceLine algorithm."""
        # prepare offset vectors and one-hot encoded input_x
        self._one_hot_encoder = OneHotEncoder(handle_unknown="ignore")
        x_encoded = self._one_hot_encoder.fit_transform(input_x)
        feature_domains: NDArray = np.array(
            [len(sub_array) for sub_array in self._one_hot_encoder.categories_]
        )
        feature_offset_end = np.cumsum(feature_domains)
        feature_offset_start = feature_offset_end - feature_domains

        # initialize statistics and basic slices
        n_col_x_encoded = x_encoded.shape[1]
        self.average_error_ = float(np.mean(errors))
        slices, statistics = self._create_and_score_basic_slices(
            x_encoded,
            n_col_x_encoded,
            errors,
        )

        # initialize top-k
        top_k_slices, top_k_statistics = self._maintain_top_k(
            slices,
            statistics,
            sp.csr_matrix((0, n_col_x_encoded)),
            np.zeros((0, 4)),
        )

        max_slice_scores, min_slice_scores = self._analyse_top_k(
            top_k_statistics
        )
        logger.debug(
            "Initial top-K: count=%i, max=%f, min=%f"
            % (top_k_slices.shape[0], max_slice_scores, min_slice_scores)
        )

        # lattice enumeration w/ size/error pruning, one iteration per level
        # termination condition (max #feature levels)
        level = 1
        min_condition = min(input_x.shape[1], self.max_l)
        while (
            (slices.shape[0] > 0)
            and (slices.sum() > 0)
            and (level < min_condition)
        ):
            level += 1

            # enumerate candidate join pairs, including size/error pruning
            slices, statistics = self._get_pruned_s_r(slices, statistics)
            nr_s = slices.shape[0]
            slices = self._get_pair_candidates(
                slices,
                statistics,
                top_k_statistics,
                level,
                n_col_x_encoded,
                feature_domains,
                feature_offset_start,
                feature_offset_end,
            )

            logger.debug("Level %i:" % level)
            logger.debug(
                " -- generated paired slice candidates: %i -> %i"
                % (nr_s, slices.shape[0])
            )

            # extract and evaluate candidate slices
            statistics = self._eval_slice(x_encoded, errors, slices, level)

            # maintain top-k after evaluation
            top_k_slices, top_k_statistics = self._maintain_top_k(
                slices, statistics, top_k_slices, top_k_statistics
            )

            max_slice_scores, min_slice_scores = self._analyse_top_k(
                top_k_statistics
            )
            valid = np.sum(
                (statistics[:, 3] >= self._min_sup_actual)
                & (statistics[:, 1] > 0)
            )
            logger.debug(
                " -- valid slices after eval: %s/%i" % (valid, slices.shape[0])
            )
            logger.debug(
                " -- top-K: count=%i, max=%f, min=%f"
                % (top_k_slices.shape[0], max_slice_scores, min_slice_scores)
            )

        self._top_slices_enc = top_k_slices.copy()
        if top_k_slices.shape[0] == 0:
            self.top_slices_ = np.empty((0, input_x.shape[1]))
        else:
            self.top_slices_ = self._one_hot_encoder.inverse_transform(
                top_k_slices
            )

        # compute slices' average errors
        top_k_statistics = np.column_stack(
            (
                top_k_statistics,
                np.divide(top_k_statistics[:, 1], top_k_statistics[:, 3]),
            )
        )

        # transform statistics to a list of dict
        statistics_names = [
            "slice_score",
            "sum_slice_error",
            "max_slice_error",
            "slice_size",
            "slice_average_error",
        ]
        self.top_slices_statistics_ = [
            {
                stat_name: float(stat_value)
                for stat_value, stat_name in zip(statistic, statistics_names)
            }
            for statistic in top_k_statistics
        ]

        logger.debug("Terminated at level %i." % level)
