"""Comprehensive tests for Language Size scoring (Bex et al. arXiv:1004.2372).

Tests cover:
1. Paper's Language Size measure (Section 4.3.1)
2. Our adaptation (counting at exact sequence lengths)
3. Edge cases: ties, empty sequences, single sequences, long sequences
4. MDL fallback behavior
5. Ensemble integration with method parameter
6. The concrete info+ problem from our codebase
"""

import pytest
from bex.mdl import (
    model_cost, data_cost, lang_size, lang_size_score,
    mdl_score, score_grammar, _count_words_fast,
)
from bex.ensemble import infer_ensemble
from bex.crx import CRX
from bex.idregex import idregex


# ── Paper's Language Size: cumulative |L(r)≤n| ──

class TestPaperLanguageSize:
    """Tests for the paper's original cumulative measure."""

    def test_paper_example_a_dot_a_c_plus(self):
        """Paper's example: a.(a+c+)? has m=3, n=7, |L≤7|=3."""
        expr = 'a.(a+c+)?'
        m = model_cost(expr)
        n = 2 * m + 1
        assert m == 3, f"model_cost should be 3, got {m}"
        assert n == 7, f"n should be 7, got {n}"
        ls = lang_size(expr, n)
        assert ls == 3, f"|L≤7| should be 3, got {ls}"

    def test_paper_same_n_specific_wins(self):
        """At same n, specific grammar beats generic."""
        n = 7  # target n for a.b.c
        specific = lang_size('a.b.c', n)
        generic = lang_size('(a+b+c)+', n)
        assert specific < generic, (
            f"Specific ({specific}) should beat generic ({generic}) at n={n}"
        )

    def test_paper_same_n_correct_beats_overfit(self):
        """At same n, correct grammar and overfit tie (both accept 1 word)."""
        n = 7
        correct = lang_size('a.b.c', n)
        overfit = lang_size('a.a.a', n)
        assert correct == overfit == 1, (
            f"Both should accept 1 word at n={n}, got {correct} and {overfit}"
        )

    def test_paper_per_candidate_n_generic_wins_unfairly(self):
        """Per-candidate n lets generic patterns win unfairly."""
        # info+ has m=1, n=3 → counts words at lengths 0,1,2,3
        # specific has m=5, n=11 → counts words at lengths 0..11
        generic_n = 2 * model_cost('info+') + 1  # = 3
        specific_n = 2 * model_cost('info.file.template.shell.service+') + 1  # = 11

        generic_ls = lang_size('info+', generic_n)
        specific_ls = lang_size('info.file.template.shell.service+', specific_n)

        # Generic wins on paper (3 < 7) but this is wrong
        assert generic_ls < specific_ls, (
            f"Per-candidate n: generic ({generic_ls}) beats specific ({specific_ls}) — this is the bug"
        )

    def test_paper_alphabet_size_5_mdL_vs_langsize(self):
        """Paper's result: Language Size 98% vs MDL 21% on alphabet size 5."""
        # At the same n, language size correctly differentiates
        n = 7
        specific = lang_size('a.b.c', n)           # 1 word
        generic = lang_size('(a+b+c)+', n)         # 3,279 words
        medium = lang_size('a.(b+c)?', n)           # 3 words

        assert specific < medium < generic, (
            f"Order should be specific({specific}) < medium({medium}) < generic({generic})"
        )


# ── Our Adaptation: words at exact sequence lengths ──

class TestAdaptedLanguageSize:
    """Tests for our adaptation (counting at exact sequence lengths)."""

    def test_specific_vs_generic_diverse_lengths(self):
        """With diverse lengths, specific grammar wins clearly."""
        seqs = [['a', 'b', 'c'], ['a', 'b'], ['a', 'c'], ['b', 'c']]
        specific = lang_size_score('a.b.c', seqs)
        generic = lang_size_score('(a+b+c)+', seqs)
        assert specific < generic, (
            f"Specific ({specific}) should beat generic ({generic})"
        )

    def test_info_plus_vs_specific_identical_lengths(self):
        """With identical lengths, both accept 1 word — honest tie."""
        seqs = [['info', 'file', 'template', 'shell', 'service']] * 5
        generic = lang_size_score('info+', seqs)
        specific = lang_size_score('info.file.template.shell.service+', seqs)
        assert generic == specific == 5, (
            f"Both should score 5 (1 word × 5 seqs), got generic={generic}, specific={specific}"
        )

    def test_generic_vs_more_generic(self):
        """(a+b+c)+ accepts more words than a+ at each length."""
        seqs = [['a', 'b', 'c']] * 3
        less_generic = lang_size_score('a+', seqs)
        more_generic = lang_size_score('(a+b+c)+', seqs)
        # a+ accepts 1 word at each length; (a+b+c)+ accepts 3^L
        assert less_generic < more_generic, (
            f"a+ ({less_generic}) should beat (a+b+c)+ ({more_generic})"
        )

    def test_single_sequence(self):
        """Single sequence — specific grammar wins."""
        seqs = [['a', 'b', 'c']]
        specific = lang_size_score('a.b.c', seqs)
        generic = lang_size_score('(a+b+c)+', seqs)
        assert specific < generic

    def test_long_sequences(self):
        """Long sequences — specific grammar still wins."""
        seqs = [['a', 'b', 'c', 'd', 'e']] * 3
        specific = lang_size_score('a.b.c.d.e', seqs)
        generic = lang_size_score('(a+b+c+d+e)+', seqs)
        assert specific < generic

    def test_empty_sequences(self):
        """Empty sequences — falls back to paper formula."""
        score = lang_size_score('a.b.c', [])
        expected = lang_size('a.b.c', 2 * model_cost('a.b.c') + 1)
        assert score == expected

    def test_ordered_vs_unordered(self):
        """Ordered a.b.c beats unordered (a+b+c)+ on ordered data."""
        seqs = [['a', 'b', 'c'], ['a', 'b'], ['a', 'c']]
        ordered = lang_size_score('a.b.c', seqs)
        unordered = lang_size_score('(a+b+c)+', seqs)
        assert ordered < unordered

    def test_optional_beats_generic(self):
        """a.(b+c)? beats (a+b+c)+ on data where a is always first."""
        seqs = [['a', 'b'], ['a', 'c'], ['a']]
        optional = lang_size_score('a.(b+c)?', seqs)
        generic = lang_size_score('(a+b+c)+', seqs)
        assert optional < generic

    def test_repeat_beats_concat(self):
        """a+ beats a.a.a on data with varying lengths."""
        seqs = [['a'], ['a', 'a'], ['a', 'a', 'a']]
        repeat = lang_size_score('a+', seqs)
        concat = lang_size_score('a.a.a', seqs)
        # a+ accepts 1 word at each length; a.a.a accepts 0 at lengths 1,2 and 1 at length 3
        # Total: a+ = 3, a.a.a = 0+0+1 = 1
        # a.a.a actually wins because it rejects shorter sequences!
        assert concat < repeat, (
            f"a.a.a ({concat}) should beat a+ ({repeat}) — a.a.a rejects short seqs"
        )


# ── MDL Fallback ──

class TestMDLFallback:
    """Tests for the old MDL scoring method."""

    def test_mdl_basic(self):
        """MDL = model_cost + data_cost."""
        score = mdl_score('a.b.c', [['a', 'b', 'c']])
        assert score == model_cost('a.b.c') + data_cost('a.b.c', [['a', 'b', 'c']])

    def test_mdl_prefers_short_expressions(self):
        """MDL rewards short expressions — the info+ bug."""
        seqs = [['info', 'file', 'template', 'shell', 'service']] * 5
        generic = mdl_score('info+', seqs)
        specific = mdl_score('info.file.template.shell.service+', seqs)
        assert generic < specific, (
            f"MDL should pick info+ ({generic}) over specific ({specific}) — this is the bug"
        )

    def test_score_grammar_method_switch(self):
        """score_grammar dispatches to the correct scorer."""
        seqs = [['a', 'b', 'c']]
        ls = score_grammar('a.b.c', seqs, method='langsize')
        mdl = score_grammar('a.b.c', seqs, method='mdl')
        assert isinstance(ls, (int, float))
        assert isinstance(mdl, (int, float))

    def test_score_grammar_invalid_method(self):
        """Invalid method raises ValueError."""
        with pytest.raises(ValueError, match="Unknown scoring method"):
            score_grammar('a.b.c', [['a']], method='bogus')

    def test_langsize_beats_mdl_on_info_plus(self):
        """Language Size ties on info+ scenario; MDL picks info+."""
        seqs = [['info', 'file', 'template', 'shell', 'service']] * 5
        ls_generic = score_grammar('info+', seqs, method='langsize')
        ls_specific = score_grammar('info.file.template.shell.service+', seqs, method='langsize')
        mdl_generic = score_grammar('info+', seqs, method='mdl')
        mdl_specific = score_grammar('info.file.template.shell.service+', seqs, method='mdl')

        # Language Size: tie (honest)
        assert ls_generic == ls_specific, "Language Size should tie"
        # MDL: generic wins (the bug)
        assert mdl_generic < mdl_specific, "MDL should pick generic (the bug)"


# ── Ensemble Integration ──

class TestEnsembleIntegration:
    """Tests for the ensemble with method parameter."""

    def test_ensemble_accepts_method(self):
        """Ensemble accepts method= parameter."""
        seqs = [['a', 'b'], ['a', 'b', 'c']]
        r_ls = infer_ensemble(seqs, method='langsize')
        r_mdl = infer_ensemble(seqs, method='mdl')
        assert r_ls['best'] is not None
        assert r_mdl['best'] is not None

    def test_ensemble_default_is_langsize(self):
        """Default method is langsize."""
        seqs = [['a', 'b'], ['a', 'b', 'c']]
        r = infer_ensemble(seqs)
        assert r['best'] is not None

    def test_ensemble_langsize_prefers_specific(self):
        """With diverse sequences, langsize picks the specific grammar."""
        seqs = [['a', 'b', 'c'], ['a', 'b'], ['a', 'c'], ['b', 'c']]
        r = infer_ensemble(seqs, method='langsize')
        # Should pick a.b.c or a.(b+c)? — something specific
        best = r['best']['grammar']
        # The specific grammar should have a low score
        score = r['best']['mdl_score']
        assert score < 100, f"Score should be low for specific grammar, got {score}"

    def test_ensemble_method_threaded_to_algorithms(self):
        """Method parameter is passed through to scoring."""
        seqs = [['a', 'b', 'c'], ['a', 'b']]
        r_ls = infer_ensemble(seqs, method='langsize')
        r_mdl = infer_ensemble(seqs, method='mdl')
        # Both should produce results
        assert r_ls['best'] is not None
        assert r_mdl['best'] is not None
        # Scores may differ
        # (not necessarily — depends on what the algorithms produce)


# ── _count_words_fast Correctness ──

class TestCountWordsFast:
    """Tests for the word counting function used by Language Size."""

    def test_single_symbol(self):
        """Single symbol: 1 word of length 1, 0 otherwise."""
        assert _count_words_fast('a', 1) == 1
        assert _count_words_fast('a', 0) == 0
        assert _count_words_fast('a', 2) == 0

    def test_concatenation(self):
        """a.b.c: 1 word of length 3, 0 otherwise."""
        assert _count_words_fast('a.b.c', 3) == 1
        assert _count_words_fast('a.b.c', 2) == 0
        assert _count_words_fast('a.b.c', 4) == 0

    def test_plus_quantifier(self):
        """a+: 1 word of each length ≥ 1."""
        for l in range(1, 6):
            assert _count_words_fast('a+', l) == 1
        assert _count_words_fast('a+', 0) == 0

    def test_disjunction(self):
        """(a+b+c): 3 words of length 1, 0 otherwise."""
        assert _count_words_fast('(a+b+c)', 1) == 3
        assert _count_words_fast('(a+b+c)', 0) == 0
        assert _count_words_fast('(a+b+c)', 2) == 0

    def test_disjunction_plus(self):
        """(a+b+c)+: 3^L words of length L."""
        assert _count_words_fast('(a+b+c)+', 1) == 3
        assert _count_words_fast('(a+b+c)+', 2) == 9
        assert _count_words_fast('(a+b+c)+', 3) == 27

    def test_optional(self):
        """a?.(b+c): 2 words of length 2 (ab, ac), 2 words of length 1 (b, c)."""
        assert _count_words_fast('a?.(b+c)', 0) == 0
        assert _count_words_fast('a?.(b+c)', 1) == 2  # b, c (a? absent)
        assert _count_words_fast('a?.(b+c)', 2) == 2  # ab, ac (a? present)

    def test_epsilon(self):
        """ε: 1 word of length 0."""
        assert _count_words_fast('ε', 0) == 1
        assert _count_words_fast('ε', 1) == 0

    def test_empty(self):
        """∅: 0 words at any length."""
        assert _count_words_fast('∅', 0) == 0
        assert _count_words_fast('∅', 1) == 0

    def test_info_plus(self):
        """info+: 1 word of each length ≥ 1 (info repeated L times)."""
        for l in range(1, 8):
            assert _count_words_fast('info+', l) == 1

    def test_info_dot_concat(self):
        """info.file.template: 1 word of length 3, 0 otherwise."""
        assert _count_words_fast('info.file.template', 3) == 1
        assert _count_words_fast('info.file.template', 2) == 0
        assert _count_words_fast('info.file.template', 4) == 0

    def test_mixed_disj_concat(self):
        """a.(b+c)+: a followed by 1+ of b or c."""
        # length 2: ab, ac (2 words)
        assert _count_words_fast('a.(b+c)+', 2) == 2
        # length 3: abb, abc, acb, acc (4 words)
        assert _count_words_fast('a.(b+c)+', 3) == 4

    def test_optional_concat(self):
        """a?.b.(c+d): a optional, then b, then c or d."""
        assert _count_words_fast('a?.b.(c+d)', 0) == 0
        assert _count_words_fast('a?.b.(c+d)', 2) == 2  # bc, bd
        assert _count_words_fast('a?.b.(c+d)', 3) == 2  # abc, abd


# ── Regression: info+ Problem ──

class TestInfoPlusRegression:
    """Regression tests for the concrete info+ problem from our codebase."""

    def test_info_plus_not_preferred_over_specific(self):
        """info+ should not beat the specific grammar on diverse data."""
        seqs = [['info', 'file', 'template', 'shell', 'service']] * 5
        generic_score = lang_size_score('info+', seqs)
        specific_score = lang_size_score('info.file.template.shell.service+', seqs)
        # They tie — which is correct
        assert generic_score == specific_score

    def test_info_plus_loses_on_diverse_data(self):
        """info+ loses when sequences have different lengths."""
        seqs = [
            ['info', 'file'],
            ['info', 'file', 'template'],
            ['info', 'file', 'template', 'shell'],
        ]
        generic = lang_size_score('info+', seqs)
        specific = lang_size_score('info.file.template+', seqs)
        assert generic > specific, (
            f"info+ ({generic}) should lose to specific ({specific}) on diverse data"
        )

    def test_crx_does_not_produce_info_plus(self):
        """CRX does not produce info+ for identical sequences."""
        seqs = [['info', 'file', 'template', 'shell', 'service']] * 5
        g = CRX().infer(seqs)
        assert g != 'info+', f"CRX should not produce info+, got {g}"

    def test_ensemble_does_not_pick_info_plus(self):
        """Ensemble does not pick info+ for5 identical sequences."""
        seqs = [['info', 'file', 'template', 'shell', 'service']] * 5
        r = infer_ensemble(seqs)
        assert r['best']['grammar'] != 'info+', (
            f"Ensemble should not pick info+, got {r['best']['grammar']}"
        )