"""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.grammar import (
    Symbol, Concat, Alt, Plus, Optional, Star, Epsilon, Empty,
    parse_sore, count_words, lang_size, model_cost, alphabet,
)
from bex.mdl import (
    model_cost as mdl_model_cost, data_cost, lang_size_score,
    mdl_score, score_grammar,
)
from bex.ensemble import infer_ensemble
from bex.crx import CRX


# ── Helper to parse SORE string to AST ──
def p(s):
    return parse_sore(s)


# ── 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 = p('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
        specific = lang_size(p('a.b.c'), n)
        generic = lang_size(p('(a+b+c)+'), n)
        assert specific < generic

    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(p('a.b.c'), n)
        overfit = lang_size(p('a.a.a'), n)
        assert correct == overfit == 1

    def test_paper_per_candidate_n_generic_wins_unfairly(self):
        """Per-candidate n lets generic patterns win unfairly."""
        generic_n = 2 * model_cost(p('info+')) + 1
        specific_n = 2 * model_cost(p('info.file.template.shell.service+')) + 1
        generic_ls = lang_size(p('info+'), generic_n)
        specific_ls = lang_size(p('info.file.template.shell.service+'), specific_n)
        assert generic_ls < specific_ls

    def test_paper_alphabet_size_5_mdL_vs_langsize(self):
        """Paper's result: Language Size 98% vs MDL 21% on alphabet size 5."""
        n = 7
        specific = lang_size(p('a.b.c'), n)
        generic = lang_size(p('(a+b+c)+'), n)
        medium = lang_size(p('a.(b+c)?'), n)
        assert specific < medium < 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):
        seqs = [['a', 'b', 'c'], ['a', 'b'], ['a', 'c'], ['b', 'c']]
        specific = lang_size_score(p('a.b.c'), seqs)
        generic = lang_size_score(p('(a+b+c)+'), seqs)
        assert specific < generic

    def test_info_plus_vs_specific_identical_lengths(self):
        seqs = [['info', 'file', 'template', 'shell', 'service']] * 5
        generic = lang_size_score(p('info+'), seqs)
        specific = lang_size_score(p('info.file.template.shell.service+'), seqs)
        assert generic == specific == 5

    def test_generic_vs_more_generic(self):
        seqs = [['a', 'b', 'c']] * 3
        less_generic = lang_size_score(p('a+'), seqs)
        more_generic = lang_size_score(p('(a+b+c)+'), seqs)
        assert less_generic < more_generic

    def test_single_sequence(self):
        seqs = [['a', 'b', 'c']]
        specific = lang_size_score(p('a.b.c'), seqs)
        generic = lang_size_score(p('(a+b+c)+'), seqs)
        assert specific < generic

    def test_long_sequences(self):
        seqs = [['a', 'b', 'c', 'd', 'e']] * 3
        specific = lang_size_score(p('a.b.c.d.e'), seqs)
        generic = lang_size_score(p('(a+b+c+d+e)+'), seqs)
        assert specific < generic

    def test_empty_sequences(self):
        score = lang_size_score(p('a.b.c'), [])
        expected = lang_size(p('a.b.c'), 2 * model_cost(p('a.b.c')) + 1)
        assert score == expected

    def test_ordered_vs_unordered(self):
        seqs = [['a', 'b', 'c'], ['a', 'b'], ['a', 'c']]
        ordered = lang_size_score(p('a.b.c'), seqs)
        unordered = lang_size_score(p('(a+b+c)+'), seqs)
        assert ordered < unordered

    def test_optional_beats_generic(self):
        seqs = [['a', 'b'], ['a', 'c'], ['a']]
        optional = lang_size_score(p('a.(b+c)?'), seqs)
        generic = lang_size_score(p('(a+b+c)+'), seqs)
        assert optional < generic

    def test_repeat_beats_concat(self):
        seqs = [['a'], ['a', 'a'], ['a', 'a', 'a']]
        repeat = lang_size_score(p('a+'), seqs)
        concat = lang_size_score(p('a.a.a'), seqs)
        assert concat < repeat


# ── MDL Fallback ──

class TestMDLFallback:
    def test_mdl_basic(self):
        score = mdl_score(p('a.b.c'), [['a', 'b', 'c']])
        assert score == mdl_model_cost(p('a.b.c')) + data_cost(p('a.b.c'), [['a', 'b', 'c']])

    def test_mdl_prefers_short_expressions(self):
        seqs = [['info', 'file', 'template', 'shell', 'service']] * 5
        generic = mdl_score(p('info+'), seqs)
        specific = mdl_score(p('info.file.template.shell.service+'), seqs)
        assert generic < specific

    def test_score_grammar_method_switch(self):
        seqs = [['a', 'b', 'c']]
        ls = score_grammar(p('a.b.c'), seqs, method='langsize')
        mdl = score_grammar(p('a.b.c'), seqs, method='mdl')
        assert isinstance(ls, (int, float))
        assert isinstance(mdl, (int, float))

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

    def test_langsize_beats_mdl_on_info_plus(self):
        seqs = [['info', 'file', 'template', 'shell', 'service']] * 5
        ls_generic = score_grammar(p('info+'), seqs, method='langsize')
        ls_specific = score_grammar(p('info.file.template.shell.service+'), seqs, method='langsize')
        mdl_generic = score_grammar(p('info+'), seqs, method='mdl')
        mdl_specific = score_grammar(p('info.file.template.shell.service+'), seqs, method='mdl')
        assert ls_generic == ls_specific, "Language Size should tie"
        assert mdl_generic < mdl_specific, "MDL should pick generic (the bug)"


# ── Ensemble Integration ──

class TestEnsembleIntegration:
    def test_ensemble_accepts_method(self):
        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):
        seqs = [['a', 'b'], ['a', 'b', 'c']]
        r = infer_ensemble(seqs)
        assert r['best'] is not None

    def test_ensemble_langsize_prefers_specific(self):
        seqs = [['a', 'b', 'c'], ['a', 'b'], ['a', 'c'], ['b', 'c']]
        r = infer_ensemble(seqs, method='langsize')
        best = r['best']['grammar']
        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):
        seqs = [['a', 'b', 'c'], ['a', 'b']]
        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


# ── count_words Correctness ──

class TestCountWordsFast:
    def test_single_symbol(self):
        assert count_words(p('a'), 1) == 1
        assert count_words(p('a'), 0) == 0
        assert count_words(p('a'), 2) == 0

    def test_concatenation(self):
        assert count_words(p('a.b.c'), 3) == 1
        assert count_words(p('a.b.c'), 2) == 0
        assert count_words(p('a.b.c'), 4) == 0

    def test_plus_quantifier(self):
        for l in range(1, 6):
            assert count_words(p('a+'), l) == 1
        assert count_words(p('a+'), 0) == 0

    def test_disjunction(self):
        assert count_words(p('(a+b+c)'), 1) == 3
        assert count_words(p('(a+b+c)'), 0) == 0
        assert count_words(p('(a+b+c)'), 2) == 0

    def test_disjunction_plus(self):
        assert count_words(p('(a+b+c)+'), 1) == 3
        assert count_words(p('(a+b+c)+'), 2) == 9
        assert count_words(p('(a+b+c)+'), 3) == 27

    def test_optional(self):
        assert count_words(p('a?.(b+c)'), 0) == 0
        assert count_words(p('a?.(b+c)'), 1) == 2
        assert count_words(p('a?.(b+c)'), 2) == 2

    def test_epsilon(self):
        assert count_words(Epsilon(), 0) == 1
        assert count_words(Epsilon(), 1) == 0

    def test_empty(self):
        assert count_words(Empty(), 0) == 0
        assert count_words(Empty(), 1) == 0

    def test_info_plus(self):
        for l in range(1, 8):
            assert count_words(p('info+'), l) == 1

    def test_info_dot_concat(self):
        assert count_words(p('info.file.template'), 3) == 1
        assert count_words(p('info.file.template'), 2) == 0
        assert count_words(p('info.file.template'), 4) == 0

    def test_mixed_disj_concat(self):
        assert count_words(p('a.(b+c)+'), 2) == 2
        assert count_words(p('a.(b+c)+'), 3) == 4

    def test_optional_concat(self):
        assert count_words(p('a?.b.(c+d)'), 0) == 0
        assert count_words(p('a?.b.(c+d)'), 2) == 2
        assert count_words(p('a?.b.(c+d)'), 3) == 2


# ── Regression: info+ Problem ──

class TestInfoPlusRegression:
    def test_info_plus_not_preferred_over_specific(self):
        seqs = [['info', 'file', 'template', 'shell', 'service']] * 5
        generic_score = lang_size_score(p('info+'), seqs)
        specific_score = lang_size_score(p('info.file.template.shell.service+'), seqs)
        assert generic_score == specific_score

    def test_info_plus_loses_on_diverse_data(self):
        seqs = [
            ['info', 'file'],
            ['info', 'file', 'template'],
            ['info', 'file', 'template', 'shell'],
        ]
        generic = lang_size_score(p('info+'), seqs)
        specific = lang_size_score(p('info.file.template+'), seqs)
        assert generic > specific

    def test_crx_does_not_produce_info_plus(self):
        seqs = [['info', 'file', 'template', 'shell', 'service']] * 5
        g = CRX().infer(seqs)
        alpha = alphabet(g)
        assert not (len(alpha) == 1 and Symbol('info') in alpha)

    def test_ensemble_does_not_pick_info_plus(self):
        seqs = [['info', 'file', 'template', 'shell', 'service']] * 5
        r = infer_ensemble(seqs)
        best = r['best']['grammar']
        alpha = alphabet(best)
        assert not (len(alpha) == 1 and Symbol('info') in alpha)