import numpy as np

from neurovault.cognitive.clustering import threshold_cluster


def test_threshold_cluster_groups_similar_vectors():
    vectors = np.array(
        [
            [1.0, 0.0, 0.0],
            [0.99, 0.01, 0.0],
            [0.98, 0.02, 0.0],
            [0.0, 1.0, 0.0],
        ],
        dtype=np.float32,
    )
    clusters = threshold_cluster(vectors, threshold=0.95, min_cluster_size=3)
    assert len(clusters) == 1
    assert len(clusters[0].indices) == 3


def test_threshold_cluster_empty_for_small_input():
    vectors = np.array([[1.0, 0.0]], dtype=np.float32)
    clusters = threshold_cluster(vectors, threshold=0.6, min_cluster_size=3)
    assert clusters == []
