"""Tests for QKMeans algorithm."""

import numpy as np
import pytest
from sklearn.datasets import make_blobs
from sklearn.utils.estimator_checks import check_estimator
from sklearn.base import clone


# Import will fail if package isn't installed - skip tests in that case
pytest.importorskip("qkmeans")
from qkmeans import QKMeansSeeding, ANNSType, get_simd_level, is_anns_available


class TestQKMeansSeeding:
    """Tests for QKMeansSeeding class."""

    @pytest.fixture
    def sample_data(self):
        """Generate sample data for tests."""
        X, _ = make_blobs(n_samples=1000, n_features=10, centers=5, random_state=42)
        return X.astype(np.float64)

    @pytest.fixture
    def small_data(self):
        """Generate small data for quick tests."""
        X, _ = make_blobs(n_samples=100, n_features=5, centers=3, random_state=42)
        return X.astype(np.float64)

    def test_basic_fit(self, sample_data):
        """Test basic fit functionality."""
        seeder = QKMeansSeeding(n_clusters=5, random_state=42)
        seeder.fit(sample_data)

        assert hasattr(seeder, 'cluster_centers_')
        assert seeder.cluster_centers_.shape == (5, 10)
        assert len(seeder.center_indices_) == 5
        assert seeder.n_features_in_ == 10
        assert seeder.seeding_time_ >= 0

    def test_predict(self, sample_data):
        """Test predict functionality."""
        seeder = QKMeansSeeding(n_clusters=5, random_state=42)
        seeder.fit(sample_data)
        labels = seeder.predict(sample_data)

        assert labels.shape == (1000,)
        assert set(labels).issubset(set(range(5)))

    def test_fit_predict(self, sample_data):
        """Test fit_predict functionality."""
        seeder = QKMeansSeeding(n_clusters=5, random_state=42)
        labels = seeder.fit_predict(sample_data)

        assert labels.shape == (1000,)

    def test_transform(self, sample_data):
        """Test transform functionality."""
        seeder = QKMeansSeeding(n_clusters=5, random_state=42)
        seeder.fit(sample_data)
        distances = seeder.transform(sample_data)

        assert distances.shape == (1000, 5)
        assert np.all(distances >= 0)

    def test_score(self, sample_data):
        """Test score functionality."""
        seeder = QKMeansSeeding(n_clusters=5, random_state=42)
        seeder.fit(sample_data)
        score = seeder.score(sample_data)

        # Score should be negative inertia
        assert score <= 0
        assert seeder.inertia_ >= 0
        assert np.isclose(score, -seeder.inertia_, rtol=1e-5)

    def test_reproducibility(self, sample_data):
        """Test reproducibility with random_state."""
        seeder1 = QKMeansSeeding(n_clusters=5, random_state=42)
        seeder1.fit(sample_data)

        seeder2 = QKMeansSeeding(n_clusters=5, random_state=42)
        seeder2.fit(sample_data)

        np.testing.assert_array_equal(
            seeder1.center_indices_, seeder2.center_indices_
        )
        np.testing.assert_array_almost_equal(
            seeder1.cluster_centers_, seeder2.cluster_centers_
        )

    def test_different_seeds(self, sample_data):
        """Test that different seeds give different results."""
        seeder1 = QKMeansSeeding(n_clusters=5, random_state=42)
        seeder1.fit(sample_data)

        seeder2 = QKMeansSeeding(n_clusters=5, random_state=123)
        seeder2.fit(sample_data)

        # Very unlikely to get the same centers
        assert not np.array_equal(seeder1.center_indices_, seeder2.center_indices_)

    def test_clone(self, sample_data):
        """Test sklearn clone functionality."""
        seeder = QKMeansSeeding(n_clusters=5, random_state=42)
        seeder_cloned = clone(seeder)

        assert seeder.get_params() == seeder_cloned.get_params()
        assert not hasattr(seeder_cloned, 'cluster_centers_')

    def test_get_set_params(self):
        """Test get_params and set_params."""
        seeder = QKMeansSeeding(n_clusters=10, anns='lsh', random_state=42)

        params = seeder.get_params()
        assert params['n_clusters'] == 10
        assert params['anns'] == 'lsh'
        assert params['random_state'] == 42

        seeder.set_params(n_clusters=20)
        assert seeder.n_clusters == 20

    def test_brute_force_backend(self, small_data):
        """Test brute force ANNS backend."""
        seeder = QKMeansSeeding(
            n_clusters=3, anns='brute_force', random_state=42
        )
        seeder.fit(small_data)

        assert seeder.cluster_centers_.shape == (3, 5)

    def test_lsh_backend(self, sample_data):
        """Test LSH ANNS backend."""
        seeder = QKMeansSeeding(
            n_clusters=5,
            anns='lsh',
            anns_params={'num_tables': 5, 'num_hashes': 4},
            random_state=42
        )
        seeder.fit(sample_data)

        assert seeder.cluster_centers_.shape == (5, 10)

    def test_invalid_n_clusters(self, sample_data):
        """Test error handling for invalid n_clusters."""
        seeder = QKMeansSeeding(n_clusters=0)
        with pytest.raises(ValueError):
            seeder.fit(sample_data)

        seeder = QKMeansSeeding(n_clusters=2000)  # More than n_samples
        with pytest.raises(ValueError):
            seeder.fit(sample_data)

    def test_invalid_anns(self):
        """Test error handling for invalid ANNS backend."""
        with pytest.raises(ValueError):
            seeder = QKMeansSeeding(anns='invalid')
            seeder._validate_params()

    def test_k_equals_1(self, sample_data):
        """Test edge case k=1."""
        seeder = QKMeansSeeding(n_clusters=1, random_state=42)
        seeder.fit(sample_data)

        assert seeder.cluster_centers_.shape == (1, 10)
        labels = seeder.predict(sample_data)
        assert np.all(labels == 0)

    def test_k_equals_n(self, small_data):
        """Test edge case k=n."""
        n = len(small_data)
        seeder = QKMeansSeeding(n_clusters=n, random_state=42)
        seeder.fit(small_data)

        assert seeder.cluster_centers_.shape == (n, 5)

    def test_high_dimensional_data(self):
        """Test with high-dimensional data."""
        X = np.random.randn(200, 500).astype(np.float64)
        seeder = QKMeansSeeding(n_clusters=10, random_state=42)
        seeder.fit(X)

        assert seeder.cluster_centers_.shape == (10, 500)

    def test_input_validation(self):
        """Test input validation."""
        seeder = QKMeansSeeding(n_clusters=5)

        # 1D array should fail
        with pytest.raises(ValueError):
            seeder.fit(np.array([1, 2, 3]))

        # Wrong dtype should be converted
        X = np.random.randn(100, 10).astype(np.float32)
        seeder.fit(X)  # Should work (converted to float64)

    def test_repr(self):
        """Test string representation."""
        seeder = QKMeansSeeding(n_clusters=10, anns='lsh', random_state=42)
        repr_str = repr(seeder)

        assert 'QKMeansSeeding' in repr_str
        assert 'n_clusters=10' in repr_str


class TestUtilities:
    """Tests for utility functions."""

    def test_get_simd_level(self):
        """Test SIMD level detection."""
        level = get_simd_level()
        assert level in ['SCALAR', 'SSE4', 'AVX2']

    def test_is_anns_available(self):
        """Test ANNS availability check."""
        # LSH and brute force should always be available
        assert is_anns_available(ANNSType.LSH)
        assert is_anns_available(ANNSType.BRUTE_FORCE)

        # FAISS and HNSW depend on compilation options
        # Just check that the function runs without error
        is_anns_available(ANNSType.FAISS)
        is_anns_available(ANNSType.HNSW)


# Uncomment to run sklearn estimator checks (may be slow)
# @pytest.mark.slow
# def test_sklearn_estimator_checks():
#     """Run sklearn estimator checks."""
#     check_estimator(QKMeansSeeding(n_clusters=3))


if __name__ == '__main__':
    pytest.main([__file__, '-v'])
