"""Tests for CRX Refined (cluster-then-infer)."""

import pytest
from bex.crx_refined import crx_refined, crx_with_confidence, _cluster_by_structure
from bex.crx import CRX
from bex.grammar import (
    Symbol, Concat, Alt, Plus, Optional, Star, Epsilon, Empty,
    match, alphabet,
)


class TestClusterByStructure:
    def test_empty(self):
        assert _cluster_by_structure([]) == {}

    def test_single_sequence(self):
        result = _cluster_by_structure([['a', 'b', 'c']])
        assert len(result) == 1
        assert ('a', 'c', 'short') in result

    def test_same_start_end_same_length(self):
        seqs = [['a', 'b', 'c'], ['a', 'd', 'c']]
        result = _cluster_by_structure(seqs)
        assert len(result) == 1

    def test_different_start_end(self):
        seqs = [['a', 'b', 'c'], ['x', 'y', 'z']]
        result = _cluster_by_structure(seqs)
        assert len(result) == 2

    def test_length_buckets(self):
        seqs = [
            ['a', 'b', 'a'],
            ['a', 'b', 'c', 'a'],
            ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'a'],
            ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'a'],
        ]
        result = _cluster_by_structure(seqs)
        keys = set(result.keys())
        assert ('a', 'a', 'short') in keys
        assert ('a', 'a', 'med') in keys
        assert ('a', 'a', 'long') in keys


class TestCrxRefined:
    def test_empty(self):
        assert isinstance(crx_refined([]), Epsilon)

    def test_single_sequence(self):
        result = crx_refined([['a', 'b', 'c']])
        assert result is not None
        assert match(result, ['a', 'b', 'c'])

    def test_identical_sequences(self):
        seqs = [['a', 'b', 'c']] * 5
        result = crx_refined(seqs)
        alpha = alphabet(result)
        assert 'a' in alpha
        assert 'b' in alpha

    def test_linear_pattern(self):
        seqs = [['a', 'b', 'c'], ['a', 'b', 'c'], ['a', 'b', 'c']]
        result = crx_refined(seqs)
        assert result == CRX().infer(seqs)

    def test_branching_pattern(self):
        seqs = [['a', 'b', 'c'], ['a', 'd', 'e'], ['a', 'b', 'e']]
        result = crx_refined(seqs)
        assert result is not None
        alpha = alphabet(result)
        assert 'a' in alpha

    def test_falls_back_to_standard(self):
        seqs = [['a'], ['b'], ['c'], ['d']]
        result = crx_refined(seqs, min_cluster=2)
        assert result == CRX().infer(seqs)


class TestCrxWithConfidence:
    def test_empty(self):
        result = crx_with_confidence([])
        assert isinstance(result['grammar'], Epsilon)
        assert result['confidence'] == 1.0
        assert result['n_clusters'] == 0

    def test_returns_all_fields(self):
        seqs = [['a', 'b', 'c'], ['a', 'b', 'c']]
        result = crx_with_confidence(seqs)
        assert 'grammar' in result
        assert 'confidence' in result
        assert 'n_clusters' in result
        assert 'largest_cluster' in result
        assert 'n_sequences' in result

    def test_confidence_range(self):
        seqs = [['a', 'b', 'c'], ['a', 'd', 'e']]
        result = crx_with_confidence(seqs)
        assert 0.0 <= result['confidence'] <= 1.0

    def test_confident_when_tight(self):
        seqs = [['a', 'b', 'c']] * 10
        result = crx_with_confidence(seqs)
        assert result['confidence'] >= 0.9

    def test_less_confident_when_diverse(self):
        seqs = [
            ['a', 'b', 'c', 'z'],
            ['a', 'x', 'y', 'z'],
            ['a', 'p', 'q', 'z'],
            ['a', 'r', 's', 'z'],
        ]
        result = crx_with_confidence(seqs)
        assert result['n_clusters'] >= 1


class TestComparisonWithStandard:
    """Compare refined vs standard CRX on various inputs."""

    def test_linear_same_result(self):
        seqs = [['a', 'b', 'c']] * 5
        assert crx_refined(seqs) == CRX().infer(seqs)

    def test_single_symbol(self):
        seqs = [['a']] * 5
        result = crx_refined(seqs)
        assert isinstance(result, Symbol) and result.value == 'a'

    def test_two_symbols(self):
        seqs = [['a', 'b']] * 5
        result = crx_refined(seqs)
        assert isinstance(result, Concat) and len(result.parts) == 2

    def test_disjunction(self):
        seqs = [['a', 'b'], ['a', 'c']]
        result = crx_refined(seqs)
        assert isinstance(result, Alt) or (isinstance(result, Concat) and any(isinstance(p, Alt) for p in result.parts))