import cansig.cluster.leiden as le  # pytype: disable=import-error
import pytest  # pytype: disable=import-error
from cansig.cluster.leiden import _binary_search_leiden_resolution  # pytype: disable=import-error
from sklearn import datasets  # pytype: disable=import-error

_binary_search = _binary_search_leiden_resolution


def binary_search_matching_random_seed(*args, **kwargs):
    if kwargs["random_state"] == 0:
        return None
    else:
        assert kwargs["random_state"] == 20_000
        return _binary_search(*args, **kwargs)


def binary_search_no_matching_random_seed(*args, **kwargs):
    return None


@pytest.fixture
def X():
    X, _ = datasets.make_blobs(n_samples=100, n_features=2)
    return X


def test_leiden_resolution_smoke(X) -> None:
    config = le.LeidenResolutionConfig()

    labels = le.LeidenResolution(config).fit_predict(X)
    assert labels.shape == (len(labels),)
    assert min(labels) == 0
    assert max(labels) > 0


class TestLeidenNClusters:
    @pytest.mark.parametrize("clusters", (3, 4, 5))
    def test_leiden_n_clusters(self, X, clusters: int) -> None:
        # Three Gaussian clusters are generated by default. We will partition them
        # into smaller clusters, matching the specified argument.
        config = le.LeidenNClusterConfig(clusters=clusters)

        labels = le.LeidenNCluster(config).fit_predict(X)
        assert labels.shape == (len(labels),)
        assert min(labels) == 0
        assert max(labels) == clusters - 1

    def test_no_matching_resolution(self, X, monkeypatch):
        config = le.LeidenNClusterConfig(clusters=2)
        monkeypatch.setattr(
            "cansig.cluster.leiden._binary_search_leiden_resolution", binary_search_matching_random_seed
        )
        labels = le.LeidenNCluster(config).fit_predict(X)
        assert labels.shape == (len(labels),)
        assert min(labels) == 0
        assert max(labels) == 1

    def test_no_matching_random_seed(self, X, monkeypatch):
        with pytest.raises(ValueError):
            config = le.LeidenNClusterConfig(clusters=2)
            monkeypatch.setattr(
                "cansig.cluster.leiden._binary_search_leiden_resolution", binary_search_no_matching_random_seed
            )
            le.LeidenNCluster(config).fit_predict(X)
