/tests/datalab/test_datalab.py

# Copyright (C) 2017-2022  Cleanlab Inc.
# This file is part of cleanlab.
#
# cleanlab is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# cleanlab is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with cleanlab.  If not, see <https://www.gnu.org/licenses/>.


import contextlib
import io
import os
import pickle
import timeit
from pathlib import Path
from unittest.mock import MagicMock, Mock, patch

import numpy as np
import pandas as pd
import pytest
from datasets.dataset_dict import DatasetDict
from scipy.sparse import csr_matrix
from sklearn.neighbors import NearestNeighbors
from sklearn.datasets import make_blobs

import cleanlab
from cleanlab.datalab.datalab import Datalab
from cleanlab.datalab.internal.report import Reporter

SEED = 42


def test_datalab_invalid_datasetdict(dataset, label_name):
    with pytest.raises(ValueError) as e:
        datadict = DatasetDict({"train": dataset, "test": dataset})
        Datalab(datadict, label_name)  # type: ignore
        assert "Please pass a single dataset, not a DatasetDict." in str(e)


@pytest.fixture(scope="function")
def list_possible_issue_types(monkeypatch, request):
    return lambda *_: request.param


class TestDatalab:
    """Tests for the Datalab class."""

    @pytest.fixture
    def lab(self, dataset, label_name):
        return Datalab(data=dataset, label_name=label_name)

    def test_print(self, lab, capsys):
        # Can print the object
        print(lab)
        captured = capsys.readouterr()
        expected_output = (
            "Datalab:\n"
            "Checks run: No\n"
            "Number of examples: 5\n"
            "Number of classes: 3\n"
            "Issues identified: Not checked\n"
        )
        assert expected_output == captured.out

    def test_class_names(self):
        y = ["a", "3", "2", "3"]
        lab = Datalab({"y": y}, label_name="y")
        assert lab.class_names == ["2", "3", "a"]

        y = [-1, 4, 0.5, 0, 4, -1]
        lab = Datalab({"y": y}, label_name="y")
        assert lab.class_names == [-1, 0, 0.5, 4]

    def test_list_default_issue_types(self):
        y = ["a", "3", "2", "3"]
        lab = Datalab({"y": y}, label_name="y")
        assert lab.list_default_issue_types() == [
            "label",
            "outlier",
            "near_duplicate",
            "non_iid",
            "class_imbalance",
        ]

    def tmp_path(self):
        # A path for temporarily saving the instance during tests.
        # This is a workaround for the fact that the Datalab class
        # does not have a save method.
        return Path(__file__).parent / "tmp.pkl"

    def test_attributes(self, lab):
        # Has the right attributes
        for attr in ["data", "label_name", "_labels", "info", "issues"]:
            assert hasattr(lab, attr), f"Missing attribute {attr}"

        assert all(lab.labels == np.array([1, 1, 2, 0, 2]))
        assert isinstance(lab.issues, pd.DataFrame), "Issues should by in a dataframe"
        assert isinstance(lab.issue_summary, pd.DataFrame), "Issue summary should be a dataframe"

    def test_get_info(self, lab):
        mock_info: dict = {
            "label": {
                "given_label": [1, 0, 1, 0, 2],
                "predicted_label": [1, 1, 2, 0, 2],
                # get_info("label") adds `class_names` from statistics
            },
            "near_duplicate": {
                "nearest_neighbor": [1, 0, 0, 4, 3],
            },
        }
        mock_info = {**lab.info, **mock_info}
        lab.info = mock_info

        label_info = lab.get_info("label")
        assert label_info["given_label"].tolist() == [4, 3, 4, 3, 5]
        assert label_info["predicted_label"].tolist() == [4, 4, 5, 3, 5]
        assert label_info["class_names"] == [3, 4, 5]

        near_duplicate_info = lab.get_info("near_duplicate")
        assert near_duplicate_info["nearest_neighbor"] == [1, 0, 0, 4, 3]

        assert lab.get_info() == lab.info == mock_info

    def test_get_issue_summary(self, lab, monkeypatch):
        mock_summary: pd.DataFrame = pd.DataFrame(
            {
                "issue_type": ["label", "outlier"],
                "score": [0.5, 0.3],
                "num_issues": [1, 2],
            }
        )
        monkeypatch.setattr(lab, "issue_summary", mock_summary)

        label_summary = lab.get_issue_summary(issue_name="label")
        pd.testing.assert_frame_equal(label_summary, mock_summary.iloc[[0]])

        outlier_summary = lab.get_issue_summary(issue_name="outlier")
        pd.testing.assert_frame_equal(
            outlier_summary, mock_summary.iloc[[1]].reset_index(drop=True)
        )

        summary = lab.get_issue_summary()
        pd.testing.assert_frame_equal(summary, mock_summary)

    def test_get_issues(self, lab, monkeypatch):
        mock_issues: pd.DataFrame = pd.DataFrame(
            {
                "is_label_issue": [True, False, False, True, False],
                "label_score": [0.2, 0.4, 0.6, 0.1, 0.8],
                "is_near_duplicate_issue": [False, True, True, False, True],
                "near_duplicate_score": [0.5, 0.3, 0.1, 0.7, 0.2],
            },
        )
        monkeypatch.setattr(lab, "issues", mock_issues)

        mock_predicted_labels = np.array([0, 1, 2, 1, 2])

        mock_distance_to_nearest_neighbor = [0.1, 0.2, 0.3, 0.4, 0.5]

        lab.info.update(
            {
                "label": {
                    "given_label": lab.labels,
                    "predicted_label": mock_predicted_labels,
                },
                "near_duplicate": {
                    "distance_to_nearest_neighbor": mock_distance_to_nearest_neighbor,
                },
            }
        )

        label_issues = lab.get_issues(issue_name="label")

        expected_label_issues = pd.DataFrame(
            {
                **{key: mock_issues[key] for key in ["is_label_issue", "label_score"]},
                "given_label": [4, 4, 5, 3, 5],
                "predicted_label": [3, 4, 5, 4, 5],
            },
        )

        pd.testing.assert_frame_equal(label_issues, expected_label_issues, check_dtype=False)

        near_duplicate_issues = lab.get_issues(issue_name="near_duplicate")

        expected_near_duplicate_issues = pd.DataFrame(
            {
                **{
                    key: mock_issues[key]
                    for key in ["is_near_duplicate_issue", "near_duplicate_score"]
                },
                "distance_to_nearest_neighbor": mock_distance_to_nearest_neighbor,
            },
        )
        pd.testing.assert_frame_equal(
            near_duplicate_issues, expected_near_duplicate_issues, check_dtype=False
        )

        issues = lab.get_issues()
        pd.testing.assert_frame_equal(issues, mock_issues, check_dtype=False)

    @pytest.mark.parametrize(
        "issue_types",
        [None, {"label": {}}],
        ids=["Default issues", "Only label issues"],
    )
    def test_find_issues_with_pred_probs(self, lab, pred_probs, issue_types):
        assert lab.issues.empty, "Issues should be empty before calling find_issues"
        assert lab.issue_summary.empty, "Issue summary should be empty before calling find_issues"
        assert lab.info["statistics"]["health_score"] is None
        lab.find_issues(pred_probs=pred_probs, issue_types=issue_types)
        assert not lab.issues.empty, "Issues weren't updated"
        assert not lab.issue_summary.empty, "Issue summary wasn't updated"
        assert (
            lab.info["statistics"]["health_score"] == lab.issue_summary["score"].mean()
        )  # TODO: Avoid re-implementing logic in test

        if issue_types is None:
            # Test default issue types
            columns = lab.issues.columns
            for issue_type in ["label", "outlier"]:
                assert f"is_{issue_type}_issue" in columns
                assert f"{issue_type}_score" in columns

    def test_find_issues_without_values_in_issue_types_raises_warning(self, lab, pred_probs):
        issue_types = {}
        with pytest.warns(UserWarning) as record:
            lab.find_issues(pred_probs=pred_probs, issue_types=issue_types)
        warning_message = record[0].message.args[0]
        assert (
            "No issue types were specified so no issues will be found in the dataset. Set `issue_types` as None to consider a default set of issues."
            in warning_message
        )

    @pytest.mark.parametrize(
        "issue_types",
        [
            None,
            {"label": {}},
            {"outlier": {}},
            {"near_duplicate": {}},
            {"non_iid": {}},
            {"outlier": {}, "near_duplicate": {}},
        ],
        ids=[
            "Defaults",
            "Only label issues",
            "Only outlier issues",
            "Only near_duplicate issues",
            "Only non_iid issues",
            "Both outlier and near_duplicate issues",
        ],
    )
    @pytest.mark.parametrize(
        "use_features",
        [True, False],
        ids=["Use features", "Don't use features"],
    )
    @pytest.mark.parametrize(
        "use_pred_probs",
        [True, False],
        ids=["Use pred_probs", "Don't use pred_probs"],
    )
    @pytest.mark.parametrize(
        "use_knn_graph",
        [True, False],
        ids=["Use knn_graph", "Don't use knn_graph"],
    )
    def test_repeat_find_issues_then_report_with_defaults(
        self,
        large_lab,
        issue_types,
        use_features,
        use_pred_probs,
        use_knn_graph,
    ):
        """Test "all" combinations of inputs to find_issues() and make sure repeated calls to it won't change any results. Same applies to report().

        This test does NOT test the correctness of the inputs, so some test cases may lead to missing arguments errors that are silently ignored.
        """

        # Extract features and pred_probs from Datalab object
        features, pred_probs = (
            np.array(large_lab.data[k]) if v else None
            for k, v in zip(["features", "pred_probs"], [use_features, use_pred_probs])
        )

        # Extract sparse knn_graph from Datalab object's info dictionary
        knn_graph = None
        if use_knn_graph:
            knn_graph = large_lab.info["statistics"]["unit_test_knn_graph"]

        # Run find_issues and report() once
        large_lab.find_issues(
            features=features, pred_probs=pred_probs, knn_graph=knn_graph, issue_types=issue_types
        )
        with contextlib.redirect_stdout(io.StringIO()) as f:
            large_lab.report()
        first_report = f.getvalue()
        issues = large_lab.issues.copy()
        issue_summary = large_lab.issue_summary.copy()

        # Rerunning find_issues() and report() with the same default parameters should not change the number of issues
        large_lab.find_issues(
            features=features, pred_probs=pred_probs, knn_graph=knn_graph, issue_types=issue_types
        )
        with contextlib.redirect_stdout(io.StringIO()) as f:
            large_lab.report()
        second_report = f.getvalue()
        pd.testing.assert_frame_equal(large_lab.issues, issues)
        pd.testing.assert_frame_equal(large_lab.issue_summary, issue_summary)
        assert first_report == second_report

    @pytest.mark.parametrize("k", [2, 3])
    @pytest.mark.parametrize("metric", ["euclidean", "cosine"])
    def test_find_issues_with_custom_hyperparams(self, lab, pred_probs, k, metric):
        dataset_size = lab.get_info("statistics")["num_examples"]
        embedding_size = 2
        mock_embeddings = np.random.rand(dataset_size, embedding_size)

        knn = NearestNeighbors(n_neighbors=k, metric=metric)
        issue_types = {"outlier": {"knn": knn}}
        assert lab.get_info("statistics").get("weighted_knn_graph") is None
        lab.find_issues(
            pred_probs=pred_probs,
            features=mock_embeddings,
            issue_types=issue_types,
        )
        assert lab.info["outlier"]["k"] == k
        statistics = lab.get_info("statistics")
        assert statistics["knn_metric"] == metric
        knn_graph = statistics["weighted_knn_graph"]
        assert isinstance(knn_graph, csr_matrix)
        assert knn_graph.shape == (dataset_size, dataset_size)
        assert knn_graph.nnz == dataset_size * k

    # Mock the lab.issues dataframe to have some pre-existing issues
    def test_update_issues(self, lab, pred_probs, monkeypatch):
        """If there are pre-existing issues in the lab,
        find_issues should add columns to the issues dataframe for each example.
        """
        mock_issues = pd.DataFrame(
            {
                "is_foo_issue": [False, True, False, False, False],
                "foo_score": [0.6, 0.8, 0.7, 0.7, 0.8],
            }
        )
        monkeypatch.setattr(lab, "issues", mock_issues)
        mock_issue_summary = pd.DataFrame(
            {
                "issue_type": ["foo"],
                "score": [0.72],
                "num_issues": [1],
            }
        )
        monkeypatch.setattr(lab, "issue_summary", mock_issue_summary)

        lab.find_issues(pred_probs=pred_probs, issue_types={"label": {}})
        # Check that the issues dataframe has the right columns
        expected_issues_df = pd.DataFrame(
            {
                "is_foo_issue": mock_issues.is_foo_issue,
                "foo_score": mock_issues.foo_score,
                "is_label_issue": [False, False, False, False, False],
                "label_score": [0.95071431, 0.15601864, 0.60111501, 0.70807258, 0.18182497],
            }
        )
        pd.testing.assert_frame_equal(lab.issues, expected_issues_df, check_exact=False)

        expected_issue_summary_df = pd.DataFrame(
            {
                "issue_type": ["foo", "label"],
                "score": [0.72, 0.4],
                "num_issues": [1, 0],
            }
        )
        pd.testing.assert_frame_equal(
            lab.issue_summary, expected_issue_summary_df, check_exact=False
        )

    def test_save(self, lab, tmp_path, monkeypatch):
        """Test that the save and load methods work."""
        lab.save(tmp_path, force=True)
        assert tmp_path.exists(), "Save directory was not created"
        assert (tmp_path / "data").is_dir(), "Data directory was not saved"
        assert (tmp_path / "issues.csv").exists(), "Issues file was not saved"
        assert (tmp_path / "summary.csv").exists(), "Issue summary file was not saved"
        assert (tmp_path / "datalab.pkl").exists(), "Datalab file was not saved"

        # Mock the issues dataframe
        mock_issues = pd.DataFrame(
            {
                "is_foo_issue": [False, True, False, False, False],
                "foo_score": [0.6, 0.8, 0.7, 0.7, 0.8],
            }
        )
        monkeypatch.setattr(lab, "issues", mock_issues)

        # Mock the issue summary dataframe
        mock_issue_summary = pd.DataFrame(
            {
                "issue_type": ["foo"],
                "score": [0.72],
            }
        )
        monkeypatch.setattr(lab, "issue_summary", mock_issue_summary)
        lab.save(tmp_path, force=True)
        assert (tmp_path / "issues.csv").exists(), "Issues file was not saved"
        assert (tmp_path / "summary.csv").exists(), "Issue summary file was not saved"

        # Save works in an arbitrary directory, that should be created if it doesn't exist
        new_dir = tmp_path / "subdir"
        assert not new_dir.exists(), "Directory should not exist"
        lab.save(new_dir)
        assert new_dir.exists(), "Directory was not created"

    def test_pickle(self, lab, tmp_path):
        """Test that the class can be pickled."""
        pickle_file = os.path.join(tmp_path, "lab.pkl")
        with open(pickle_file, "wb") as f:
            pickle.dump(lab, f)
        with open(pickle_file, "rb") as f:
            lab2 = pickle.load(f)

        assert lab2.label_name == "star"

    def test_load(self, lab, tmp_path, dataset, monkeypatch):
        """Test that the save and load methods work."""

        # Mock the issues dataframe
        mock_issues = pd.DataFrame(
            {
                "is_foo_issue": [False, True, False, False, False],
                "foo_score": [0.6, 0.8, 0.7, 0.7, 0.8],
            }
        )
        monkeypatch.setattr(lab, "issues", mock_issues)

        # Mock the issue summary dataframe
        mock_issue_summary = pd.DataFrame(
            {
                "issue_type": ["foo"],
                "score": [0.72],
            }
        )
        monkeypatch.setattr(lab, "issue_summary", mock_issue_summary)

        lab.save(tmp_path, force=True)

        loaded_lab = Datalab.load(tmp_path)
        data = lab._data
        loaded_data = loaded_lab._data
        assert loaded_data == data
        assert loaded_lab.info == lab.info
        pd.testing.assert_frame_equal(loaded_lab.issues, mock_issues)
        pd.testing.assert_frame_equal(loaded_lab.issue_summary, mock_issue_summary)

        # Load accepts a `Dataset`.
        loaded_lab = Datalab.load(tmp_path, data=dataset)
        assert loaded_lab.data._data == dataset.data

        # Misaligned dataset raises a ValueError
        with pytest.raises(ValueError) as excinfo:
            Datalab.load(tmp_path, data=dataset.shard(2, 0))
            expected_error_msg = "Length of data (2) does not match length of labels (5)"
            assert expected_error_msg == str(excinfo.value)

        with pytest.raises(ValueError) as excinfo:
            Datalab.load(tmp_path, data=dataset.shuffle())
            expected_error_msg = (
                "Data has been modified since Lab was saved. Cannot load Lab with modified data."
            )
            assert expected_error_msg == str(excinfo.value)

    @pytest.mark.parametrize("list_possible_issue_types", [["erroneous_issue_type"]], indirect=True)
    def test_failed_issue_managers(self, lab, monkeypatch, list_possible_issue_types):
        """Test that a failed issue manager will not be added to the Datalab instance after
        the call to `find_issues`."""
        mock_issue_types = {"erroneous_issue_type": {}}

        mock_issue_manager = Mock()
        mock_issue_manager.issue_name = "erroneous_issue_type"
        mock_issue_manager.find_issues.side_effect = ValueError("Some error")

        class MockIssueManagerFactory:
            @staticmethod
            def from_list(*args, **kwargs):
                return [mock_issue_manager]

        monkeypatch.setattr(
            "cleanlab.datalab.internal.issue_finder._IssueManagerFactory", MockIssueManagerFactory
        )

        assert lab.issues.empty
        with patch("builtins.print") as mock_print:
            lab.find_issues(issue_types=mock_issue_types)
            for expected_msg_substr in [
                "Error in",
                "Audit complete",
                "Failed to check for these issue types: ",
            ]:
                assert any(expected_msg_substr in call[0][0] for call in mock_print.call_args_list)

        assert lab.issues.empty

    def test_report(self, lab, monkeypatch, capsys):
        class MockReporter:
            def __init__(self, *args, **kwargs):
                self.verbosity = kwargs.get("verbosity", None)
                assert self.verbosity is not None, "Reporter should be initialized with verbosity"

            def report(self, *args, **kwargs) -> None:
                print(
                    f"Report with verbosity={self.verbosity} and k={kwargs.get('num_examples', 5)}"
                )

        monkeypatch.setattr(cleanlab.datalab.internal.helper_factory, "Reporter", MockReporter)
        monkeypatch.setattr(
            lab.data_issues,
            "issue_summary",
            pd.DataFrame(np.random.randint(0, 100, size=(100, 4)), columns=list("ABCD")),
        )
        lab.report(verbosity=0)
        captured = capsys.readouterr()
        assert "Report with verbosity=0 and k=5" in captured.out

        lab.report(num_examples=10, verbosity=3)
        captured = capsys.readouterr()
        assert "Report with verbosity=3 and k=10" in captured.out

        lab.report()
        captured = capsys.readouterr()
        assert "Report with verbosity=1 and k=5" in captured.out


class TestDatalabUsingKNNGraph:
    """The Datalab class can accept a `knn_graph` argument to `find_issues` that should
    be used instead of computing a new one from the `features` argument."""

    @pytest.fixture
    def data_tuple(self):
        # from cleanlab.datalab.datalab import Datalab
        np.random.seed(SEED)
        N = 10
        data = {"label": np.random.randint(0, 2, size=N)}
        features = np.random.rand(N, 5)
        knn_graph = (
            NearestNeighbors(n_neighbors=3, metric="cosine")
            .fit(features)
            .kneighbors_graph(mode="distance")
        )
        return Datalab(data=data, label_name="label"), knn_graph, features

    def test_knn_graph(self, data_tuple):
        """Test that the `knn_graph` argument to `find_issues` is used instead of computing a new
        one from the `features` argument."""
        lab, knn_graph, _ = data_tuple
        assert lab.get_info("statistics").get("weighted_knn_graph") is None
        lab.find_issues(knn_graph=knn_graph)
        knn_graph_stats = lab.get_info("statistics").get("weighted_knn_graph")
        np.testing.assert_array_equal(knn_graph_stats.toarray(), knn_graph.toarray())

        assert lab.get_info("statistics").get("knn_metric") is None

    def test_features_and_knn_graph(self, data_tuple):
        """Test that the `knn_graph` argument to `find_issues` is used instead of computing a new
        one from the `features` argument."""
        lab, knn_graph, features = data_tuple
        k = 4
        lab.find_issues(knn_graph=knn_graph, features=features, issue_types={"outlier": {"k": k}})
        knn_graph_stats = lab.get_info("statistics").get("weighted_knn_graph")
        assert knn_graph_stats.nnz == k * len(
            lab.data
        ), f"Expected {k * len(lab.data)} nnz, got {knn_graph_stats.nnz}"
        three_nn_dists = knn_graph_stats.data.reshape(len(lab.data), k)[:, :3]
        knn_graph_three_nn_dists = knn_graph.data.reshape(len(lab.data), k - 1)
        np.testing.assert_array_equal(three_nn_dists, knn_graph_three_nn_dists)
        assert lab.get_info("statistics").get("knn_metric") == "cosine"

    def test_without_features_or_knn_graph(self, data_tuple):
        """Test that only the class_imbalance issue is run
        when no features, knn_graph or pred_probs are passed."""
        lab, _, _ = data_tuple

        # Test that a warning is raised
        lab.find_issues()
        # Only class_imbalance issue columns should be present
        assert list(lab.issues.columns) == ["is_class_imbalance_issue", "class_imbalance_score"]

    def test_data_valuation_issue_with_knn_graph(self, data_tuple):
        lab, knn_graph, features = data_tuple
        assert lab.get_info("statistics").get("weighted_knn_graph") is None
        lab.find_issues(knn_graph=knn_graph, issue_types={"data_valuation": {}})
        score = lab.get_issues().get(["data_valuation_score"])
        assert isinstance(score, pd.DataFrame)
        assert len(score) == len(lab.data)

    def test_data_valuation_issue_with_existing_knn_graph(self, data_tuple):
        lab, knn_graph, features = data_tuple
        lab.find_issues(features=features, issue_types={"outlier": {"k": 3}})
        lab.find_issues(issue_types={"data_valuation": {}})
        score = lab.get_issues().get(["data_valuation_score"])
        assert isinstance(score, pd.DataFrame)
        assert len(score) == len(lab.data)

        # Compare this with directly passing in a knn_graph
        lab_2 = Datalab(data=lab.data, label_name=lab.label_name)
        lab_2.find_issues(knn_graph=knn_graph, issue_types={"data_valuation": {}})
        score_2 = lab_2.get_issues().get(["data_valuation_score"])
        pd.testing.assert_frame_equal(score, score_2)

    def test_data_valuation_issue_without_knn_graph(self, data_tuple):
        lab, _, features = data_tuple
        lab.find_issues(features=features, issue_types={"data_valuation": {}})
        assert (
            lab.issues.empty
        ), "The issues dataframe should be empty as the issue manager expects an existing knn_graph"


class TestDatalabIssueManagerInteraction:
    """The Datalab class should integrate with the IssueManager class correctly.

    Tests include:
    - Make sure a custom manager needs to be registered to work with Datalab
    - Make sure that `find_issues()` with different affects the outcome (e.g. `Datalab.issues`)
        differently depending on the issue manager.
    """

    def test_custom_issue_manager_not_registered(self, lab):
        """Test that a custom issue manager that is not registered will not be used."""
        # Mock registry dictionary
        mock_registry = MagicMock()
        mock_registry.__getitem__.side_effect = KeyError("issue type not registered")

        with patch("cleanlab.datalab.internal.issue_manager_factory.REGISTRY", mock_registry):
            with pytest.raises(ValueError) as excinfo:
                lab.find_issues(issue_types={"custom_issue": {}})

                assert "issue type not registered" in str(excinfo.value)

            assert mock_registry.__getitem__.called_once_with("custom_issue")

            assert lab.issues.empty
            assert lab.issue_summary.empty

    def test_custom_issue_manager_registered(self, lab, custom_issue_manager):
        """Test that a custom issue manager that is registered will be used."""
        from cleanlab.datalab.internal.issue_manager_factory import register

        register(custom_issue_manager)

        assert lab.issues.empty
        assert lab.issue_summary.empty

        lab.find_issues(issue_types={"custom_issue": {}})

        expected_is_custom_issue_issue = [False, True] + [False] * 3
        expected_custom_issue_score = [1 / 1, 0 / 2, 1 / 3, 2 / 4, 3 / 5]
        expected_issues = pd.DataFrame(
            {
                "is_custom_issue_issue": expected_is_custom_issue_issue,
                "custom_issue_score": expected_custom_issue_score,
            }
        )
        assert pd.testing.assert_frame_equal(lab.issues, expected_issues) is None

    @pytest.mark.parametrize("list_possible_issue_types", [["custom_issue"]], indirect=True)
    def test_find_issues_for_custom_issue_manager_with_custom_kwarg(
        self, lab, custom_issue_manager, list_possible_issue_types
    ):
        """Test that a custom issue manager that is registered will be used."""
        from cleanlab.datalab.internal.issue_manager_factory import register

        register(custom_issue_manager)

        assert lab.issues.empty
        assert lab.issue_summary.empty

        lab.find_issues(issue_types={"custom_issue": {"custom_argument": 3}})

        expected_is_custom_issue_issue = [False, False, False, True, False]
        expected_custom_issue_score = [3 / 3, 2 / 4, 1 / 5, 0 / 6, 1 / 7]
        expected_issues = pd.DataFrame(
            {
                "is_custom_issue_issue": expected_is_custom_issue_issue,
                "custom_issue_score": expected_custom_issue_score,
            }
        )
        assert pd.testing.assert_frame_equal(lab.issues, expected_issues) is None

        # Clean up registry
        from cleanlab.datalab.internal.issue_manager_factory import REGISTRY

        # Find the custom issue manager in the registry and remove it
        REGISTRY["classification"].pop(custom_issue_manager.issue_name)


@pytest.mark.parametrize(
    "find_issues_kwargs",
    [
        ({"pred_probs": np.random.rand(3, 2)}),
        ({"features": np.random.rand(3, 2)}),
        ({"pred_probs": np.random.rand(3, 2), "features": np.random.rand(6, 2)}),
    ],
    ids=["pred_probs", "features", "pred_probs and features"],
)
def test_report_for_outlier_issues_via_pred_probs(find_issues_kwargs):
    data = {"labels": [0, 1, 0]}
    lab = Datalab(data=data, label_name="labels")
    find_issues_kwargs["issue_types"] = {"outlier": {"k": 1}}
    lab.find_issues(**find_issues_kwargs)

    reporter = Reporter(
        lab.data_issues, task="classification", verbosity=0, include_description=False
    )
    report = reporter.get_report(num_examples=3)
    assert report, "Report should not be empty"


def test_near_duplicates_reuses_knn_graph():
    """'outlier' and 'near_duplicate' issues both require a KNN graph.
    This test ensures that the KNN graph is only computed once.
    E.g. if outlier is called first, and then near_duplicate can reuse the
    resulting graph.
    """
    N = 3000
    num_features = 1000
    k = 20
    data = {"labels": np.random.randint(0, 2, size=N)}

    np.random.seed(SEED)
    features = np.random.rand(N, num_features)

    # Run 1: only near_duplicate
    lab = Datalab(data=data, label_name="labels")
    find_issues_kwargs = {"issue_types": {"near_duplicate": {"k": k}}}
    time_only_near_duplicates = timeit.timeit(
        lambda: lab.find_issues(features=features, **find_issues_kwargs),
        number=1,
    )

    # Run 2: near_duplicate and outlier with same k
    lab = Datalab(data=data, label_name="labels")
    # Outliers need more neighbors, so this should be slower, so the graph will be computed twice
    find_issues_kwargs = {
        "issue_types": {"near_duplicate": {"k": k}, "outlier": {"k": 2 * k}},
    }
    time_near_duplicates_and_outlier = timeit.timeit(
        lambda: lab.find_issues(features=features, **find_issues_kwargs),
        number=1,
    )

    # Run 3: Same Datalab instance with same issues, but in different order
    find_issues_kwargs = {
        "issue_types": {"outlier": {"k": 2 * k}, "near_duplicate": {"k": k}},
    }
    time_outliers_before_near_duplicates = timeit.timeit(
        lambda: lab.find_issues(features=features, **find_issues_kwargs),
        number=1,
    )

    # Run 2 does an extra check, so it should be slower
    assert time_only_near_duplicates < time_near_duplicates_and_outlier, (
        "Run 2 should be slower because it does an extra check "
        "for outliers, which requires a KNN graph."
    )

    # Run 3 should be faster because it reuses the KNN graph from Run 2
    # in both issue checks
    assert (
        time_outliers_before_near_duplicates < time_near_duplicates_and_outlier
    ), "KNN graph reuse should make this run of find_issues faster."


def pred_probs_from_features(features):
    """
    Converts an array of features into an array of predicted probabilities
    by normalizing each row to sum to 1.

    Note
    ----
    This function is only used for testing purposes
    and may not align with conventional methodologies
    used in real-world applications.
    """
    return features / features.sum(axis=1, keepdims=True)


class TestDatalabFindNonIIDIssues:
    """This class focuses on testing the end-to-end functionality of calling Datalab.find_issues()
    only for non-IID issues. The tests in this class are not meant to test the underlying
    functionality of the non-IID issue finders themselves, but rather to test that the
    Datalab.find_issues() method correctly calls the non-IID issue finders and results are consistent.
    """

    @pytest.fixture
    def random_embeddings(self):
        np.random.seed(SEED)
        return np.random.rand(100, 10)

    @pytest.fixture
    def sorted_embeddings(self):
        np.random.seed(SEED)
        n_samples = 1000

        # Stack features to create a 3D dataset
        x = np.linspace(0, 4 * np.pi, n_samples)
        y = np.sin(x) + np.random.normal(0, 0.1, n_samples)
        z = np.cos(x) + np.random.normal(0, 0.1, n_samples)
        return np.column_stack((x, y, z))

    @pytest.fixture
    def lab(self):
        data = {"labels": [0, 1, 0]}
        lab = Datalab(data=data, label_name="labels")
        return lab

    def test_find_non_iid_issues(self, lab, random_embeddings):
        lab.find_issues(features=random_embeddings, issue_types={"non_iid": {}})
        summary = lab.get_issue_summary()
        assert ["non_iid"] == summary["issue_type"].values
        assert summary["score"].values[0] > 0.05
        assert lab.get_issues()["is_non_iid_issue"].sum() == 0

    def test_find_non_iid_issues_using_pred_probs(self, lab, random_embeddings):
        pred_probs = pred_probs_from_features(random_embeddings)
        lab.find_issues(pred_probs=pred_probs, issue_types={"non_iid": {}})
        summary = lab.get_issue_summary()
        assert ["non_iid"] == summary["issue_type"].values
        assert summary["score"].values[0] > 0.05
        assert lab.get_issues()["is_non_iid_issue"].sum() == 0
        assert "weighted_knn_graph" not in lab.get_info("statistics")

    def test_find_non_iid_issues_sorted(self, lab, sorted_embeddings):
        lab.find_issues(features=sorted_embeddings, issue_types={"non_iid": {}})
        summary = lab.get_issue_summary()
        assert ["non_iid"] == summary["issue_type"].values
        assert summary["score"].values[0] == 0
        assert lab.get_issues()["is_non_iid_issue"].sum() == 1

    def test_find_non_iid_issues_sorted_using_pred_probs(self, lab, sorted_embeddings):
        pred_probs = pred_probs_from_features(sorted_embeddings)
        lab.find_issues(pred_probs=pred_probs, issue_types={"non_iid": {}})
        summary = lab.get_issue_summary()
        assert ["non_iid"] == summary["issue_type"].values
        assert summary["score"].values[0] == 0
        assert lab.get_issues()["is_non_iid_issue"].sum() == 1
        assert "weighted_knn_graph" not in lab.get_info("statistics")

    def test_incremental_search(self, lab, sorted_embeddings):
        lab.find_issues(features=sorted_embeddings)
        summary = lab.get_issue_summary()
        assert len(summary) == 4
        lab.find_issues(features=sorted_embeddings, issue_types={"non_iid": {}})
        summary = lab.get_issue_summary()
        assert len(summary) == 4
        assert "non_iid" in summary["issue_type"].values
        non_iid_summary = lab.get_issue_summary("non_iid")
        assert non_iid_summary["score"].values[0] == 0
        assert non_iid_summary["num_issues"].values[0] == 1

    def test_incremental_search_using_pred_probs(self, lab, sorted_embeddings):
        pred_probs = pred_probs_from_features(sorted_embeddings)
        lab.find_issues(pred_probs=pred_probs, issue_types={"non_iid": {}})
        summary = lab.get_issue_summary()
        assert len(summary) == 1
        lab.find_issues(pred_probs=pred_probs, issue_types={"non_iid": {}})
        summary = lab.get_issue_summary()
        assert len(summary) == 1
        assert "non_iid" in summary["issue_type"].values
        non_iid_summary = lab.get_issue_summary("non_iid")
        assert non_iid_summary["score"].values[0] == 0
        assert non_iid_summary["num_issues"].values[0] == 1

    def test_non_iid_issues_pred_probs_knn_graph_checks(self, lab, random_embeddings):
        """
        Note
        ----
        If the user did provides `features` or there was already a KNN graph
        constructed in Datalab, the results should be returned as they currently
        are, not using the `pred_probs` at all.
        """
        pred_probs = pred_probs_from_features(random_embeddings)
        # knn graph is computed and stored
        lab.find_issues(
            features=random_embeddings, pred_probs=pred_probs, issue_types={"outlier": {}}
        )
        cached_knn_graph = lab.get_info("statistics").get("weighted_knn_graph")
        assert cached_knn_graph is not None
        lab.find_issues(pred_probs=pred_probs, issue_types={"non_iid": {}})
        knn_graph_after_non_iid = lab.get_info("statistics").get("weighted_knn_graph")
        # Check if stored knn graph is same as before
        assert cached_knn_graph is knn_graph_after_non_iid
        issues_1 = lab.get_issues("non_iid")

        lab_2 = Datalab(data={"labels": lab._labels}, label_name="labels")
        lab_2.find_issues(
            pred_probs=pred_probs, knn_graph=cached_knn_graph, issue_types={"non_iid": {}}
        )
        knn_graph_after_non_iid = lab.get_info("statistics").get("weighted_knn_graph")
        # Check if stored knn graph is same as the knn graph passed to find_issues
        assert cached_knn_graph is knn_graph_after_non_iid
        # Check that explicitly passing the cached knn-graph to a new datalab instance
        # leads to the same results.
        issues_2 = lab_2.get_issues("non_iid")
        pd.testing.assert_frame_equal(issues_1, issues_2)


class TestDatalabFindLabelIssues:
    @pytest.fixture
    def random_embeddings(self):
        np.random.seed(SEED)
        return np.random.rand(100, 10)

    @pytest.fixture
    def pred_probs(self):
        np.random.seed(SEED)
        pred_probs_array = np.random.rand(100, 2)
        return pred_probs_array / pred_probs_array.sum(axis=1, keepdims=True)

    def test_incremental_search(self, pred_probs, random_embeddings):
        data = {"labels": np.random.randint(0, 2, 100)}
        lab = Datalab(data=data, label_name="labels")
        lab.find_issues(features=random_embeddings)
        summary = lab.get_issue_summary()
        assert len(summary) == 5
        assert "label" in summary["issue_type"].values
        lab.find_issues(pred_probs=pred_probs, issue_types={"label": {}})
        summary = lab.get_issue_summary()
        assert len(summary) == 5
        assert "label" in summary["issue_type"].values
        label_summary = lab.get_issue_summary("label")
        assert label_summary["num_issues"].values[0] > 0
        # Compare results with low_memory=True
        issues_df = lab.get_issues("label")
        issue_types = {"label": {"clean_learning_kwargs": {"low_memory": True}}}
        lab_lm = Datalab(data=data, label_name="labels")
        lab_lm.find_issues(pred_probs=pred_probs, issue_types=issue_types)
        issues_df_lm = lab_lm.get_issues("label")
        # jaccard similarity
        intersection = len(list(set(issues_df).intersection(set(issues_df_lm))))
        union = len(set(issues_df)) + len(set(issues_df_lm)) - intersection
        assert float(intersection) / union > 0.95

    def test_build_pred_probs_from_features(self, random_embeddings):
        data = {"labels": np.random.randint(0, 2, 100)}
        lab = Datalab(data=data, label_name="labels")
        lab.find_issues(features=random_embeddings, issue_types={"label": {}})
        summary = lab.get_issue_summary()
        assert len(summary) == 1
        assert "label" in summary["issue_type"].values
        lab.find_issues(features=random_embeddings, issue_types={"label": {"k": 5}})
        summary = lab.get_issue_summary()
        assert len(summary) == 1
        assert "label" in summary["issue_type"].values

    def test_pred_probs_precedence(self, pred_probs, random_embeddings):
        data = {"labels": np.random.randint(0, 2, 100)}
        lab = Datalab(data=data, label_name="labels")
        lab.find_issues(pred_probs=pred_probs, issue_types={"label": {}})
        summary = lab.get_issue_summary()
        assert "label" in summary["issue_type"].values
        label_summary_pred_probs = lab.get_issue_summary("label")
        assert label_summary_pred_probs["num_issues"].values[0] > 0
        lab = Datalab(data=data, label_name="labels")
        lab.find_issues(
            features=random_embeddings, pred_probs=pred_probs, issue_types={"label": {}}
        )
        summary = lab.get_issue_summary()
        assert "label" in summary["issue_type"].values
        label_summary_both = lab.get_issue_summary("label")
        assert (
            label_summary_both["num_issues"].values[0]
            == label_summary_pred_probs["num_issues"].values[0]
        )


class TestDatalabForRegression:
    @pytest.fixture
    def regression_data(self, num_examples=400, num_features=3, error_frac=0.025, error_noise=0.25):
        np.random.seed(SEED)
        X = np.random.random(size=(num_examples, num_features))
        coefficients = np.random.uniform(-5, 5, size=num_features)

        true_y = np.dot(X, coefficients)

        # add extra noisy examples
        num_errors = int(num_examples * error_frac)
        label_noise = np.clip(
            np.random.normal(loc=error_noise, scale=error_noise / 4, size=num_errors),
            0.25 * error_noise,
            4 * error_noise,
        ) * np.random.choice([-1, 1], size=num_errors)
        random_idx = np.random.choice(num_examples, num_errors)
        y = true_y.copy()
        y[random_idx] += label_noise
        error_idx = np.argsort(abs(y - true_y))[-num_errors:]  # get the noisiest examples idx

        # Ensure that the MSE is not too large
        from sklearn.linear_model import LinearRegression

        linear_model = LinearRegression()

        # Validate that the label noise affects the MSE for a linear model
        from sklearn.metrics import mean_squared_error

        y_pred_true = linear_model.fit(X, true_y).predict(X)
        mse_true = mean_squared_error(true_y, y_pred_true)
        assert mse_true < 1e-10

        y_pred = linear_model.fit(X, y).predict(X)
        mse = mean_squared_error(y, y_pred)
        assert 1e-3 < mse < 1e-2

        return {
            "X": X,
            "y": y,
            "true_y": true_y,
            "error_idx": error_idx,
        }

    @pytest.fixture
    def lab(self, regression_data):
        X, y = regression_data["X"], regression_data["y"]
        test_df = pd.DataFrame(X, columns=["c1", "c2", "c3"])
        test_df["y"] = y
        lab = Datalab(data=test_df, label_name="y", task="regression")
        return lab

    def test_available_issue_types(self, lab):
        assert set(lab.list_default_issue_types()) == set(["label"])
        assert set(lab.list_possible_issue_types()) == set(["label"])

    def test_regression_with_features(self, lab, regression_data):
        """Test that the regression issue checks finds 40 label issues, based on the
        numerical features."""
        X = regression_data["X"]
        issue_types = {"label": {"clean_learning_kwargs": {"seed": SEED}}}
        lab.find_issues(features=X, issue_types=issue_types)
        lab.report()

        issues = lab.get_issues("label")
        issue_ids = issues.query("is_label_issue").index
        expected_issue_ids = regression_data["error_idx"]

        # jaccard similarity
        intersection = len(list(set(issue_ids).intersection(set(expected_issue_ids))))
        union = len(set(issue_ids)) + len(set(expected_issue_ids)) - intersection
        assert float(intersection) / union >= 0.7

        # FPR
        fpr = len(list(set(issue_ids).difference(set(expected_issue_ids)))) / len(issue_ids)
        assert fpr < 0.2

    def test_regression_with_predictions(self, lab, regression_data):
        """Test that the regression issue checks find 9 label issues, based on the
        predictions of a model.

        Instead of running a model, we use the ground-truth to emulate a perfect model's predictions.

        Testing the default behavior, we expect to find some label issues with a given mean score.
        Increasing a threshold for flagging issues will flag more issues, but won't change the score.
        """

        # Use ground-truth to emulate a perfect model's predictions
        y_pred = regression_data["true_y"]

        lab.find_issues(pred_probs=y_pred)
        summary = lab.get_issue_summary()

        issues = lab.get_issues("label")
        issue_ids = issues.query("is_label_issue").index
        expected_issue_ids = regression_data["error_idx"]

        # jaccard similarity
        intersection = len(list(set(issue_ids).intersection(set(expected_issue_ids))))
        union = len(set(issue_ids)) + len(set(expected_issue_ids)) - intersection
        assert float(intersection) / union > 0.8

        # FPR
        fpr = len(list(set(issue_ids).difference(set(expected_issue_ids)))) / len(issue_ids)
        assert fpr == 0.0

        # Try running with a different threshold
        lab.find_issues(pred_probs=y_pred, issue_types={"label": {"threshold": 0.4}})
        issues = lab.get_issues("label")
        issue_ids = issues.query("is_label_issue").index

        intersection = len(list(set(issue_ids).intersection(set(expected_issue_ids))))
        union = len(set(issue_ids)) + len(set(expected_issue_ids)) - intersection
        assert float(intersection) / union > 0.3

    def test_regression_with_model_and_features(self, lab, regression_data):
        """Test that the regression issue checks find label issue with another model."""
        from sklearn.linear_model import RANSACRegressor

        model = RANSACRegressor(random_state=SEED)
        X = regression_data["X"]
        issue_types = {"label": {"clean_learning_kwargs": {"model": model, "seed": SEED}}}
        lab.find_issues(features=X, issue_types=issue_types)

        issues = lab.get_issues("label")
        issue_ids = issues.query("is_label_issue").index
        expected_issue_ids = regression_data[
            "error_idx"
        ]  # Set to 5% of the data, but random noise may be too small to detect

        # jaccard similarity
        intersection = len(list(set(issue_ids).intersection(set(expected_issue_ids))))
        union = len(set(issue_ids)) + len(set(expected_issue_ids)) - intersection
        assert float(intersection) / union > 0.3

        # FPR
        fpr = len(list(set(issue_ids).difference(set(expected_issue_ids)))) / len(issue_ids)
        assert fpr < 0.3


class TestDatalabFindOutlierIssues:
    @pytest.fixture
    def random_embeddings(self):
        np.random.seed(SEED)
        X = np.random.rand(100, 10)
        X[-1] += 10 * np.random.rand(10)
        return np.random.rand(100, 10)

    @pytest.fixture
    def pred_probs(self):
        np.random.seed(SEED)
        pred_probs_array = np.random.rand(100, 2)
        return pred_probs_array / pred_probs_array.sum(axis=1, keepdims=True)

    def test_incremental_search(self, pred_probs, random_embeddings):
        data = {"labels": np.random.randint(0, 2, 100)}
        lab = Datalab(data=data, label_name="labels")
        lab.find_issues(pred_probs=pred_probs, issue_types={"label": {}})
        summary = lab.get_issue_summary()
        assert len(summary) == 1
        assert "outlier" not in summary["issue_type"].values
        lab.find_issues(features=random_embeddings, issue_types={"outlier": {}})
        summary = lab.get_issue_summary()
        assert len(summary) == 2
        assert "outlier" in summary["issue_type"].values
        outlier_summary = lab.get_issue_summary("outlier")
        assert outlier_summary["num_issues"].values[0] > 0


class TestDatalabFindNearDuplicateIssues:
    @pytest.fixture
    def random_embeddings(self):
        np.random.seed(SEED)
        X = np.random.rand(100, 10)
        X[-1] = X[-1] * -1
        X[-2] = X[-1] + 0.0001 * np.random.rand(10)
        return X

    @pytest.fixture
    def fixed_embeddings(self):
        near_duplicate_scale = 0.0001
        non_duplicate_scale = 100
        X = np.array(
            [[0, 0]] * 4  # Points with 3 exact duplicates
            + [[1, 1]] * 2  # Points with 1 exact duplicate
            + [[1, 0]] * 3
            + [[1 + near_duplicate_scale, 0]]
            + [
                [1, 0 + near_duplicate_scale]
            ]  # Points with 2 exact duplicates and 2 near duplicates
            + [
                [-1, -1] + np.random.rand(2) * near_duplicate_scale for _ in range(5)
            ]  # Points with 5 near duplicates
            + [
                [-1, 0] + np.random.rand(2) * near_duplicate_scale for _ in range(2)
            ]  # Points with 1 near duplicate
            + (np.random.rand(20, 2) * non_duplicate_scale).tolist()  # Random points
        )
        return X

    @pytest.fixture
    def pred_probs(self):
        np.random.seed(SEED)
        pred_probs_array = np.random.rand(100, 2)
        return pred_probs_array / pred_probs_array.sum(axis=1, keepdims=True)

    def test_incremental_search(self, pred_probs, random_embeddings):
        data = {"labels": np.random.randint(0, 2, 100)}
        lab = Datalab(data=data, label_name="labels")
        lab.find_issues(pred_probs=pred_probs, issue_types={"label": {}})
        summary = lab.get_issue_summary()
        assert len(summary) == 1
        assert "near_duplicate" not in summary["issue_type"].values
        lab.find_issues(features=random_embeddings, issue_types={"near_duplicate": {}})
        summary = lab.get_issue_summary()
        assert len(summary) == 2
        assert "near_duplicate" in summary["issue_type"].values
        near_duplicate_summary = lab.get_issue_summary("near_duplicate")
        assert near_duplicate_summary["num_issues"].values[0] > 1

    def test_fixed_embeddings_outputs(self, fixed_embeddings):
        lab = Datalab(data={"a": ["" for _ in range(len(fixed_embeddings))]})
        lab.find_issues(features=fixed_embeddings, issue_types={"near_duplicate": {}})
        issues = lab.get_issues("near_duplicate")

        assert issues["is_near_duplicate_issue"].sum() == 18
        assert all(
            issues["is_near_duplicate_issue"].values
            == [True] * 18 + [False] * (len(fixed_embeddings) - 18)
        )

        # Test the first set of near duplicates (only 3 exact duplicates)
        near_duplicate_sets = issues["near_duplicate_sets"].values

        expected_near_duplicate_sets = np.array(
            [
                # 3 exact duplicates
                np.array([3, 1, 2]),
                np.array([0, 3, 2]),
                np.array([0, 3, 1]),
                np.array([0, 1, 2]),
                # 1 exact duplicate
                np.array([5]),
                np.array([4]),
                # 2 exact duplicates and 2 near duplicates
                np.array([8, 7, 9, 10]),
                np.array([8, 6, 9, 10]),
                np.array([6, 7, 9, 10]),
                np.array([8, 6, 7, 10]),
                np.array([7, 8, 6, 9]),
                # 4 near duplicates
                np.array([15, 13, 14, 12]),
                np.array([13, 14, 15, 11]),
                np.array([14, 12, 15, 11]),
                np.array([13, 12, 15, 11]),
                np.array([11, 13, 14, 12]),
                # 1 near duplicate
                np.array([17]),
                np.array([16]),
            ]
            +
            # Random points
            [np.array([])] * 20,
            dtype=object,
        )

        # Exact duplicates may have arbitrary order, so sort the sets before comparing
        equal_sets = [
            np.array_equal(sorted(a), sorted(b))
            for a, b in zip(near_duplicate_sets, expected_near_duplicate_sets)
        ]
        assert all(equal_sets)

        # Assert self-idx is not included in near duplicate sets
        assert all([i not in s for i, s in enumerate(near_duplicate_sets)])

        # Assert near duplicate sets are unique, ignoring empty sets
        unique_non_empty_sets = [tuple(s) for s in near_duplicate_sets if len(s) > 0]
        assert len(set(unique_non_empty_sets)) == 18


class TestDatalabWithoutLabels:
    num_examples = 100
    num_features = 10
    K = 2

    @pytest.fixture
    def features(self):
        np.random.seed(SEED)
        return np.random.rand(self.num_examples, self.num_features)

    @pytest.fixture
    def pred_probs(self):
        np.random.seed(SEED)
        pred_probs_array = np.random.rand(self.num_examples, self.K)
        return pred_probs_array / pred_probs_array.sum(axis=1, keepdims=True)

    @pytest.fixture
    def lab(self, features):
        return Datalab(data={"X": features})

    @pytest.fixture
    def labels(self):
        np.random.seed(SEED)
        return np.random.randint(0, self.K, self.num_examples)

    def test_init(self, lab, features):
        assert np.array_equal(lab.data["X"], features)
        assert np.array_equal(lab.labels, [])

    def test_find_issues(self, lab, features, pred_probs):
        lab = Datalab(data={"X": features})
        lab.find_issues(pred_probs=pred_probs)
        assert set(lab.issues.columns) == {"is_non_iid_issue", "non_iid_score"}

        lab = Datalab(data={"X": features})
        lab.find_issues(features=features)
        assert not lab.issues.empty

    def test_find_issues_features_works_with_and_without_labels(self, features, labels):
        lab_without_labels = Datalab(data={"X": features})
        lab_without_labels.find_issues(features=features)

        lab_with_labels = Datalab(data={"X": features, "labels": labels}, label_name="labels")
        lab_with_labels.find_issues(features=features)

        lab_without_label_name = Datalab(data={"X": features, "labels": labels})
        lab_without_label_name.find_issues(features=features)

        issues_without_labels = lab_without_labels.issues
        issues_with_labels = lab_with_labels.issues
        issues_without_label_name = lab_without_label_name.issues

        # issues_with_labels should have four additional columns, which include label issues
        # and class_imbalance issues
        assert len(issues_without_labels.columns) + 4 == len(issues_with_labels.columns)
        pd.testing.assert_frame_equal(issues_without_labels, issues_without_label_name)


class TestDataLabClassImbalanceIssues:
    K = 3
    N = 100
    num_features = 2

    @pytest.fixture
    def random_embeddings(self):
        np.random.seed(SEED)
        return np.random.rand(self.N, self.num_features)

    @pytest.fixture
    def imbalance_labels(self):
        np.random.seed(SEED)
        labels = np.random.choice(np.arange(self.K - 1), 100, p=[0.5] * (self.K - 1))
        labels[0] = 2
        return labels

    @pytest.fixture
    def pred_probs(self):
        np.random.seed(SEED)
        pred_probs_array = np.random.rand(self.N, self.K)
        return pred_probs_array / pred_probs_array.sum(axis=1, keepdims=True)

    def test_incremental_search(self, pred_probs, random_embeddings, imbalance_labels):
        data = {"labels": imbalance_labels}
        lab = Datalab(data=data, label_name="labels")
        lab.find_issues(pred_probs=pred_probs, issue_types={"label": {}})
        summary = lab.get_issue_summary()
        assert len(summary) == 1
        assert "class_imbalance" not in summary["issue_type"].values
        lab.find_issues(features=random_embeddings, issue_types={"class_imbalance": {}})
        summary = lab.get_issue_summary()
        assert len(summary) == 2
        assert "class_imbalance" in summary["issue_type"].values
        class_imbalance_summary = lab.get_issue_summary("class_imbalance")
        assert class_imbalance_summary["num_issues"].values[0] > 0

    def test_find_imbalance_issues_no_args(self, imbalance_labels):
        data = {"labels": imbalance_labels}
        lab = Datalab(data=data, label_name="labels")
        lab.find_issues(issue_types={"class_imbalance": {}})
        summary = lab.get_issue_summary()
        assert len(summary) == 1
        assert "class_imbalance" in summary["issue_type"].values
        class_imbalance_summary = lab.get_issue_summary("class_imbalance")
        assert class_imbalance_summary["num_issues"].values[0] > 0


class TestDataLabUnderperformingIssue:
    K = 4
    N = 400
    num_features = 2

    @pytest.fixture
    def data(self):
        features, labels = make_blobs(
            n_samples=self.N, centers=self.K, n_features=self.num_features, random_state=SEED
        )
        pred_probs = np.full((self.N, self.K), 0.01)
        pred_probs[np.arange(self.N), labels] = 0.99
        # Generate incorrect prediction for 0th label samples
        zeroth_label_indices = np.nonzero(labels == 0)
        pred_probs[zeroth_label_indices, 0] = 0.01
        pred_probs[zeroth_label_indices, 1] = 0.99
        pred_probs = pred_probs / np.sum(pred_probs, axis=-1, keepdims=True)
        data = {"features": features, "pred_probs": pred_probs, "labels": labels}
        return data

    def test_incremental_search(self, data):
        features, labels, pred_probs = data["features"], data["labels"], data["pred_probs"]
        lab = Datalab(data={"labels": labels}, label_name="labels")
        lab.find_issues(pred_probs=pred_probs, issue_types={"label": {}})
        summary = lab.get_issue_summary()
        assert len(summary) == 1
        assert "underperforming_group" not in summary["issue_type"].values
        lab.find_issues(
            features=features, pred_probs=pred_probs, issue_types={"underperforming_group": {}}
        )
        summary = lab.get_issue_summary()
        assert len(summary) == 2
        assert "underperforming_group" in summary["issue_type"].values
        underperforming_group_summary = lab.get_issue_summary("underperforming_group")
        assert underperforming_group_summary["num_issues"].values[0] > 1

    def test_underperforming_cluster_id(self):
        """
        Test that the issue manager finds the correct underperforming cluster ID.
        """
        np.random.seed(SEED)
        N = 200
        K = 5
        n_clusters = 4
        bad_cluster_id = 2
        flip_rate = 0.95
        cluster_ids = np.random.choice(np.arange(n_clusters), size=N)
        labels = np.random.choice(np.arange(K), size=N)
        pred_probs = np.full((N, K), 0.001)
        pred_probs[np.arange(N), labels] = 0.99  # Set any high value
        pred_probs = pred_probs / np.sum(pred_probs, axis=-1, keepdims=True)
        # Flip prediction probabilities of bad cluster samples
        bad_cluster_indices = np.where(cluster_ids == bad_cluster_id)[0]
        flip_indices = np.random.choice(
            bad_cluster_indices, size=int(flip_rate * len(bad_cluster_indices)), replace=False
        )
        pred_probs[flip_indices] = 1 - pred_probs[flip_indices]  # Incorrect predictions
        features = np.random.rand(len(labels), 2)
        lab = Datalab(data={"labels": labels}, label_name="labels")
        lab.find_issues(
            features=features,
            pred_probs=pred_probs,
            issue_types={"underperforming_group": {"cluster_ids": cluster_ids}},
        )
        info = lab.get_info("underperforming_group")
        assert info["clustering"]["stats"]["underperforming_cluster_id"] == bad_cluster_id
        # Check that no underperforming cluster is found for correct predictions
        pred_probs[flip_indices] = 1 - pred_probs[flip_indices]
        lab.find_issues(
            features=features,
            pred_probs=pred_probs,
            issue_types={"underperforming_group": {"cluster_ids": cluster_ids}},
        )
        info = lab.get_info("underperforming_group")
        assert info["clustering"]["stats"]["underperforming_cluster_id"] not in cluster_ids

    def test_features_and_knn_graph(self, data):
        """
        Test that if the pre-computed knn graph is passed as an argument to `find_issues`,
        it is preferred over the `features` argument.
        """
        np.random.seed(SEED)
        features, labels, pred_probs = data["features"], data["labels"], data["pred_probs"]
        lab = Datalab(data={"labels": labels}, label_name="labels")
        lab.find_issues(
            features=features, pred_probs=pred_probs, issue_types={"underperforming_group": {}}
        )
        issues_1 = lab.get_issues("underperforming_group")
        knn_graph = lab.get_info("statistics")["weighted_knn_graph"]
        # Use another datalab instance to check preference of knn graph over features
        lab = Datalab(data={"labels": labels}, label_name="labels")
        features = np.random.rand(len(labels), 2)
        lab.find_issues(
            features=features,
            pred_probs=pred_probs,
            knn_graph=knn_graph,
            issue_types={"underperforming_group": {}},
        )
        issues_2 = lab.get_issues("underperforming_group")
        pd.testing.assert_frame_equal(issues_1, issues_2)

    def test_precomputed_cluster_ids(self, data):
        features, labels, pred_probs = data["features"], data["labels"], data["pred_probs"]
        lab = Datalab(data={"labels": labels}, label_name="labels")
        lab.find_issues(
            features=features,
            pred_probs=pred_probs,
            issue_types={"underperforming_group": {"cluster_ids": labels}},
        )
        info = lab.get_info("underperforming_group")
        assert info["clustering"]["algorithm"] is None
        underperforming_group_summary = lab.get_issue_summary("underperforming_group")
        assert underperforming_group_summary["num_issues"].values[0] > 1

        # Use new datalab instance to compute KNN graph and perform clustering
        lab = Datalab(data={"labels": labels}, label_name="labels")
        time_with_clustering = timeit.timeit(
            lambda: lab.find_issues(
                features=features, pred_probs=pred_probs, issue_types={"underperforming_group": {}}
            ),
            number=1,
        )
        # Use another datalab instance to emphasize absense of precomputed info
        lab = Datalab(data={"labels": labels}, label_name="labels")
        time_without_clustering = timeit.timeit(
            lambda: lab.find_issues(
                features=features,
                pred_probs=pred_probs,
                issue_types={"underperforming_group": {"cluster_ids": labels}},
            ),
            number=1,
        )
        # find_issues must be slower when clustering needs to be performed.
        assert (
            time_without_clustering < time_with_clustering
        ), "Passing cluster labels should make this run of find_issues faster."

    def test_no_cluster_ids(self, data):
        features, labels, pred_probs = data["features"], data["labels"], data["pred_probs"]
        lab = Datalab(data={"labels": labels}, label_name="labels")
        lab.find_issues(
            features=features,
            pred_probs=pred_probs,
            issue_types={"underperforming_group": {"cluster_ids": np.array([], dtype=int)}},
        )
        assert len(lab.issue_summary["issue_type"].values) == 0


class TestIssueManagersReuseKnnGraph:
    """
    `outlier`, `underperforming_group` and `near_duplicate` issue managers require
    a KNN graph. This test ensures that the KNN graph is only computed once.  E.g. if outlier is called first,
    then underperforming_group can reuse the resulting graph.
    """

    N = 3000
    num_features = 10
    k = 20
    num_classes = 2

    @pytest.fixture
    def features(self):
        np.random.seed(SEED)
        return np.random.uniform(low=0, high=0.2, size=(self.N, self.num_features))

    @pytest.fixture
    def labels(self):
        np.random.seed(SEED)
        return np.random.randint(0, self.num_classes, size=self.N)

    @pytest.fixture
    def pred_probs(self):
        np.random.seed(SEED)
        pred_probs = np.random.rand(self.N, self.num_classes)
        pred_probs = pred_probs / pred_probs.sum(axis=1, keepdims=True)
        return pred_probs

    def test_underperforming_group_reuses_knn_graph(self, features, pred_probs, labels):
        # Run 1: only underperforming_group
        lab = Datalab(data={"labels": labels}, label_name="labels")
        find_issues_kwargs = {"issue_types": {"underperforming_group": {"k": self.k}}}
        time_only_underperforming_group = timeit.timeit(
            lambda: lab.find_issues(features=features, pred_probs=pred_probs, **find_issues_kwargs),
            number=1,
        )

        # Run 2: Run outlier issue first, then underperforming_group
        find_issues_kwargs = {
            "issue_types": {"outlier": {"k": self.k}, "underperforming_group": {"k": self.k}},
        }
        time_underperforming_after_outlier = timeit.timeit(
            lambda: lab.find_issues(features=features, pred_probs=pred_probs, **find_issues_kwargs),
            number=1,
        )
        assert (
            time_underperforming_after_outlier < time_only_underperforming_group
        ), "KNN graph reuse should make this run of find_issues faster."