"""Tests for decomposition forest."""

import pytest
from bex.decompose import (
    decompose_sequence,
    decompose_all,
    decompose_with_coverage,
    get_decomposition_stats,
)


class TestDecomposeSequence:
    """Test sequence decomposition."""

    def test_short_sequence(self):
        seq = ["if", "return"]
        result = decompose_sequence(seq, max_length=5)
        # Short sequences returned as-is
        assert result == [seq]

    def test_empty_sequence(self):
        result = decompose_sequence([], max_length=5)
        assert result == []

    def test_long_sequence(self):
        seq = ["if", "return", "if", "return", "if", "return"]
        result = decompose_sequence(seq, max_length=3)
        
        # Should have prefixes, suffixes, and windows
        assert len(result) > 0
        
        # All fragments should be <= max_length
        for frag in result:
            assert len(frag) <= 3

    def test_prefixes(self):
        seq = ["a", "b", "c", "d"]
        result = decompose_sequence(seq, max_length=2)
        
        # Should include prefixes: [a], [a,b]
        assert ["a"] in result
        assert ["a", "b"] in result

    def test_suffixes(self):
        seq = ["a", "b", "c", "d"]
        result = decompose_sequence(seq, max_length=2)
        
        # Should include suffixes: [d], [c,d]
        assert ["d"] in result
        assert ["c", "d"] in result

    def test_windows(self):
        seq = ["a", "b", "c", "d", "e"]
        result = decompose_sequence(seq, max_length=3)
        
        # Should include windows: [a,b,c], [b,c,d], [c,d,e]
        assert ["a", "b", "c"] in result
        assert ["b", "c", "d"] in result
        assert ["c", "d", "e"] in result


class TestDecomposeAll:
    """Test decomposing multiple sequences."""

    def test_basic(self):
        seqs = [
            ["if", "return", "if"],
            ["while", "return"],
        ]
        result = decompose_all(seqs, max_length=2)
        
        # Should have fragments from both sequences
        assert len(result) > len(seqs)

    def test_empty(self):
        result = decompose_all([], max_length=5)
        assert result == []


class TestDecomposeWithCoverage:
    """Test decomposition with coverage filtering."""

    def test_filter_rare(self):
        seqs = [
            ["if", "return", "if", "return"],
            ["if", "return", "if", "return"],
            ["if", "return", "if", "return"],
            ["rare", "other", "rare", "other"],  # Only 1/4 have "rare"
        ]
        result = decompose_with_coverage(seqs, max_length=2, min_coverage=0.5)
        
        # "rare" fragments should be filtered out
        result_strs = [str(f) for f in result]
        assert not any("rare" in s for s in result_strs)

    def test_keep_common(self):
        seqs = [
            ["if", "return", "if"],
            ["if", "return", "if"],
            ["if", "return", "if"],
        ]
        result = decompose_with_coverage(seqs, max_length=2, min_coverage=0.5)
        
        # "if" and "return" fragments should be kept
        result_strs = [str(f) for f in result]
        assert any("if" in s for s in result_strs)
        assert any("return" in s for s in result_strs)


class TestDecompositionStats:
    """Test decomposition statistics."""

    def test_basic(self):
        seqs = [
            ["if", "return", "if", "return"],
            ["while", "return"],
        ]
        stats = get_decomposition_stats(seqs, max_length=2)
        
        assert stats['n_original'] == 2
        assert stats['n_fragments'] > 2
        assert stats['expansion_ratio'] > 1

    def test_empty(self):
        stats = get_decomposition_stats([], max_length=5)
        assert stats['n_original'] == 0
        assert stats['n_fragments'] == 0