"""Tests for analysis module."""

import numpy as np

from sltk.analysis.datasets import (
    compare_datasets,
    lexical_statistics,
    pos_analysis,
    vocabulary_growth,
)
from sltk.data.core import Sample


class TestLexicalStatistics:
    """Tests for lexical statistics computation."""

    def test_basic_statistics(self):
        """Test basic vocabulary statistics."""
        glosses = ["HELLO", "WORLD", "HELLO", "GOODBYE", "HELLO"]

        stats = lexical_statistics(glosses)

        assert stats.vocabulary_size == 3
        assert stats.total_tokens == 5
        assert stats.hapax_legomena == 2  # WORLD, GOODBYE appear once
        assert ("HELLO", 3) in stats.most_common

    def test_type_token_ratio(self):
        """Test type-token ratio calculation."""
        glosses = ["A", "B", "C", "D", "E"]  # 5 types, 5 tokens

        stats = lexical_statistics(glosses)

        assert stats.type_token_ratio == 1.0

    def test_hapax_legomena(self):
        """Test hapax legomena (words appearing once)."""
        glosses = ["A"] * 10 + ["B"] * 5 + ["C", "D", "E"]

        stats = lexical_statistics(glosses)

        # C, D, E appear once
        assert stats.hapax_legomena == 3
        assert stats.hapax_ratio == 3 / 5  # 3 out of 5 types (A, B, C, D, E)

    def test_dis_legomena(self):
        """Test dis legomena (words appearing twice)."""
        glosses = ["A", "A", "B", "B", "C"]

        stats = lexical_statistics(glosses)

        assert stats.dis_legomena == 2  # A and B appear twice

    def test_zipf_coefficient(self):
        """Test Zipf coefficient calculation."""
        # Generate Zipf-like distribution
        glosses = []
        for i, word in enumerate(["A", "B", "C", "D", "E", "F", "G", "H", "I", "J"]):
            count = int(100 / (i + 1))
            glosses.extend([word] * count)

        stats = lexical_statistics(glosses)

        # Should have a Zipf coefficient close to 1.0
        assert stats.zipf_coefficient is not None
        assert 0.5 < stats.zipf_coefficient < 2.0

    def test_with_samples(self):
        """Test with list of Sample objects."""
        samples = [
            Sample(id="1", glosses=["HELLO", "WORLD"]),
            Sample(id="2", glosses=["HELLO", "GOODBYE"]),
            Sample(id="3", glosses=["WORLD"]),
        ]

        stats = lexical_statistics(samples)

        assert stats.vocabulary_size == 3
        assert stats.total_tokens == 5

    def test_empty_input(self):
        """Test with empty input."""
        stats = lexical_statistics([])

        assert stats.vocabulary_size == 0
        assert stats.total_tokens == 0


class TestCompareDatasets:
    """Tests for dataset comparison."""

    def test_vocabulary_overlap(self):
        """Test vocabulary overlap calculation."""
        samples1 = [
            Sample(id="1", glosses=["A", "B", "C"]),
        ]
        samples2 = [
            Sample(id="2", glosses=["B", "C", "D"]),
        ]

        comparison = compare_datasets(
            {
                "Dataset1": samples1,
                "Dataset2": samples2,
            }
        )

        # Overlap: B, C (2 items)
        # Union: A, B, C, D (4 items)
        # Jaccard: 2/4 = 0.5
        assert comparison.vocabulary_overlap == 0.5
        assert set(comparison.shared_glosses) == {"B", "C"}

    def test_unique_glosses(self):
        """Test identification of unique glosses per dataset."""
        samples1 = [Sample(id="1", glosses=["A", "B", "C"])]
        samples2 = [Sample(id="2", glosses=["B", "C", "D"])]

        comparison = compare_datasets(
            {
                "Dataset1": samples1,
                "Dataset2": samples2,
            }
        )

        assert "A" in comparison.unique_glosses["Dataset1"]
        assert "D" in comparison.unique_glosses["Dataset2"]

    def test_sample_counts(self):
        """Test sample count reporting."""
        samples1 = [Sample(id=str(i)) for i in range(10)]
        samples2 = [Sample(id=str(i)) for i in range(20)]

        comparison = compare_datasets(
            {
                "Small": samples1,
                "Large": samples2,
            }
        )

        assert comparison.sample_counts["Small"] == 10
        assert comparison.sample_counts["Large"] == 20


class TestPOSAnalysis:
    """Tests for parts of speech analysis."""

    def test_basic_pos_analysis(self):
        """Test basic POS analysis with heuristics."""
        glosses = ["I", "WHAT", "NOT", "RUN", "HELLO"]

        result = pos_analysis(glosses)

        assert "distribution" in result
        assert "total_tokens" in result
        assert result["total_tokens"] == 5

    def test_custom_pos_mapping(self):
        """Test POS analysis with custom mapping."""
        glosses = ["HELLO", "WORLD", "HELLO"]
        pos_mapping = {"HELLO": "NOUN", "WORLD": "NOUN"}

        result = pos_analysis(glosses, pos_mapping=pos_mapping)

        assert result["distribution"]["NOUN"]["count"] == 3


class TestVocabularyGrowth:
    """Tests for vocabulary growth analysis."""

    def test_growth_tracking(self):
        """Test vocabulary growth over samples."""
        samples = []
        for i in range(100):
            # Add new word every 10 samples
            words = ["COMMON"] + ([f"WORD_{i//10}"] if i % 10 == 0 else [])
            samples.append(Sample(id=str(i), glosses=words))

        growth = vocabulary_growth(samples, window_size=20)

        # Vocabulary should grow over time
        assert len(growth["vocabulary_sizes"]) > 0
        assert growth["vocabulary_sizes"][-1] >= growth["vocabulary_sizes"][0]


class TestClusteringResult:
    """Tests for ClusteringResult dataclass."""

    def test_cluster_extraction(self):
        """Test extracting samples from a specific cluster."""
        from sltk.analysis.clustering import ClusteringResult

        embeddings = np.random.randn(100, 32)
        projections = np.random.randn(100, 2)
        cluster_ids = np.array([0] * 50 + [1] * 30 + [-1] * 20)
        sample_ids = [f"sample_{i}" for i in range(100)]

        result = ClusteringResult(
            embeddings=embeddings,
            projections=projections,
            cluster_ids=cluster_ids,
            sample_ids=sample_ids,
        )

        assert result.num_clusters == 2  # Excludes -1 (noise)

        cluster_0_embs, cluster_0_ids = result.get_cluster(0)
        assert len(cluster_0_ids) == 50

        noise_embs, noise_ids = result.get_noise_points()
        assert len(noise_ids) == 20

    def test_cluster_sizes(self):
        """Test cluster size reporting."""
        from sltk.analysis.clustering import ClusteringResult

        embeddings = np.random.randn(100, 32)
        projections = np.random.randn(100, 2)
        cluster_ids = np.array([0] * 40 + [1] * 35 + [2] * 25)
        sample_ids = [f"sample_{i}" for i in range(100)]

        result = ClusteringResult(
            embeddings=embeddings,
            projections=projections,
            cluster_ids=cluster_ids,
            sample_ids=sample_ids,
        )

        sizes = result.cluster_sizes()

        assert sizes[0] == 40
        assert sizes[1] == 35
        assert sizes[2] == 25
