"""Tests for distributional clustering."""

import pytest
from bex.distributional import (
    extract_contexts,
    build_distribution_matrix,
    cosine_similarity,
    jaccard_similarity,
    cluster_symbols,
    split_by_cluster,
    split_by_distributional,
    distributional_split,
)


class TestExtractContexts:
    """Test context extraction."""

    def test_simple_sequence(self):
        seqs = [["if", "return", "if"]]
        ctxs = extract_contexts(seqs)
        
        # "if" at position 0: (None, "return")
        # "return" at position 1: ("if", "if")
        # "if" at position 2: ("return", None)
        assert ctxs["if"] == [(None, "return"), ("return", None)]
        assert ctxs["return"] == [("if", "if")]

    def test_empty_sequence(self):
        seqs = [[]]
        ctxs = extract_contexts(seqs)
        assert ctxs == {}

    def test_single_symbol(self):
        seqs = [["return"]]
        ctxs = extract_contexts(seqs)
        assert ctxs["return"] == [(None, None)]

    def test_multiple_sequences(self):
        seqs = [
            ["if", "return"],  # if at pos 0: (None, return), return at pos 1: (if, None)
            ["return", "if"],  # return at pos 0: (None, if), if at pos 1: (return, None)
        ]
        ctxs = extract_contexts(seqs)
        
        # if appears at: pos 0 in seq1, pos 1 in seq2
        assert (None, "return") in ctxs["if"]  # from seq1
        assert ("return", None) in ctxs["if"]  # from seq2
        
        # return appears at: pos 1 in seq1, pos 0 in seq2
        assert ("if", None) in ctxs["return"]  # from seq1
        assert (None, "if") in ctxs["return"]  # from seq2


class TestBuildDistributionMatrix:
    """Test distribution matrix construction."""

    def test_simple(self):
        seqs = [
            ["if", "return", "if", "return"],
            ["return", "if", "return", "if"],
        ]
        syms, ctxs, mat = build_distribution_matrix(seqs, min_occurrences=1)
        
        assert len(syms) == 2
        assert "if" in syms
        assert "return" in syms
        assert len(ctxs) > 0
        assert len(mat) == 2

    def test_filter_rare(self):
        seqs = [
            ["if", "return"],
            ["rare"],  # rare symbol
        ]
        syms, ctxs, mat = build_distribution_matrix(seqs, min_occurrences=2)
        
        # "rare" should be filtered out
        assert "rare" not in syms

    def test_empty(self):
        syms, ctxs, mat = build_distribution_matrix([], min_occurrences=1)
        assert syms == []


class TestSimilarity:
    """Test similarity measures."""

    def test_cosine_identical(self):
        assert cosine_similarity([1, 2, 3], [1, 2, 3]) == pytest.approx(1.0)

    def test_cosine_orthogonal(self):
        assert cosine_similarity([1, 0], [0, 1]) == pytest.approx(0.0)

    def test_cosine_similar(self):
        sim = cosine_similarity([1, 2, 3], [1, 2, 4])
        assert sim > 0.9

    def test_jaccard_identical(self):
        assert jaccard_similarity({1, 2}, {1, 2}) == 1.0

    def test_jaccard_disjoint(self):
        assert jaccard_similarity({1}, {2}) == 0.0

    def test_jaccard_partial(self):
        assert jaccard_similarity({1, 2}, {2, 3}) == pytest.approx(1/3)


class TestClusterSymbols:
    """Test symbol clustering."""

    def test_similar_symbols(self):
        # Two symbols with identical context distributions
        symbols = ["a", "b"]
        matrix = [
            [1, 2, 0],  # a
            [1, 2, 0],  # b (same as a)
        ]
        clusters = cluster_symbols(symbols, matrix, threshold=0.5)
        
        # Should be in same cluster
        assert clusters["a"] == clusters["b"]

    def test_dissimilar_symbols(self):
        # Two symbols with different contexts
        symbols = ["a", "b"]
        matrix = [
            [1, 0, 0],  # a
            [0, 0, 1],  # b (different from a)
        ]
        clusters = cluster_symbols(symbols, matrix, threshold=0.5)
        
        # Should be in different clusters
        assert clusters["a"] != clusters["b"]

    def test_empty(self):
        clusters = cluster_symbols([], [], threshold=0.5)
        assert clusters == {}


class TestSplitByCluster:
    """Test sequence splitting by cluster."""

    def test_basic(self):
        seqs = [
            ["if", "return"],
            ["if", "return"],
            ["return", "if"],
            ["return", "if"],
        ]
        clusters = {"if": 0, "return": 1}
        groups = split_by_cluster(seqs, clusters, min_cluster_size=1)
        
        assert 0 in groups
        assert 1 in groups
        assert len(groups[0]) == 2
        assert len(groups[1]) == 2

    def test_empty_sequences(self):
        seqs = [
            [],
            ["if", "return"],
        ]
        clusters = {"if": 0}
        groups = split_by_cluster(seqs, clusters, min_cluster_size=1)
        
        assert "__empty__" in groups
        assert len(groups["__empty__"]) == 1

    def test_filter_small(self):
        seqs = [
            ["if", "return"],
            ["return", "if"],
        ]
        clusters = {"if": 0, "return": 1}
        groups = split_by_cluster(seqs, clusters, min_cluster_size=2)
        
        # Both clusters have only 1 sequence, should be filtered
        assert 0 not in groups
        assert 1 not in groups


class TestDistributionalSplit:
    """Test high-level distributional split."""

    def test_similar_first_symbols(self):
        # "if" and "while" both appear before "return"
        # They should cluster together
        seqs = [
            ["if", "return", "if", "return"],
            ["if", "return"],
            ["while", "return", "while", "return"],
            ["while", "return"],
            ["return", "if"],
            ["return", "if"],
        ]
        groups = split_by_distributional(seqs, threshold=0.3, min_cluster_size=2)
        
        # Should have fewer groups than first-symbol split
        # (if and while should merge)
        assert len(groups) < 6

    def test_different_first_symbols(self):
        # "if" and "return" have different contexts
        seqs = [
            ["if", "return", "if"],
            ["if", "return"],
            ["return", "if", "return"],
            ["return", "if"],
        ]
        groups = split_by_distributional(seqs, threshold=0.5, min_cluster_size=2)
        
        # Should still have 2 groups
        assert len(groups) >= 2

    def test_empty(self):
        groups = split_by_distributional([], threshold=0.5)
        assert groups == {}


class TestDropInReplacement:
    """Test that distributional_split is a drop-in replacement."""

    def test_same_signature(self):
        # Should work like _split_by_first_symbol but smarter
        seqs = [
            ["if", "return", "if"],
            ["if", "return"],
            ["return", "if"],
            ["return", "if"],
        ]
        groups = distributional_split(seqs, threshold=0.5)
        
        # Should return dict of groups
        assert isinstance(groups, dict)
        assert len(groups) > 0
        
        # All sequences should be in some group
        total = sum(len(v) for v in groups.values())
        assert total == len(seqs)