"""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,
    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


# ── AST builders (replacing old parse_sore helper) ──

_a = Symbol('a')
_b = Symbol('b')
_c = Symbol('c')
_d = Symbol('d')
_e = Symbol('e')
_info = Symbol('info')
_file = Symbol('file')
_template = Symbol('template')
_shell = Symbol('shell')
_service = Symbol('service')


# ── 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+)? → a concat optional(a|c), m=3, n=7, |L≤7|=3."""
        expr = Concat([_a, Optional(Alt([_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(Concat([_a, _b, _c]), n)
        generic = lang_size(Plus(Alt([_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(Concat([_a, _b, _c]), n)
        overfit = lang_size(Concat([_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(Plus(_info)) + 1
        specific_n = 2 * model_cost(Concat([Plus(_info), _file, _template, _shell, Plus(_service)])) + 1
        generic_ls = lang_size(Plus(_info), generic_n)
        specific_ls = lang_size(Concat([Plus(_info), _file, _template, _shell, Plus(_service)]), specific_n)
        assert generic_ls < specific_ls

    def test_paper_alphabet_size_5_langsize_vs_mdl(self):
        """Paper's result: Language Size 98% vs MDL 21% on alphabet size 5."""
        n = 7
        specific = lang_size(Concat([_a, _b, _c]), n)
        generic = lang_size(Plus(Alt([_a, _b, _c])), n)
        medium = lang_size(Concat([_a, Optional(Alt([_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(Concat([_a, _b, _c]), seqs)
        generic = lang_size_score(Plus(Alt([_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(Plus(_info), seqs)
        specific = lang_size_score(Concat([_info, _file, _template, _shell, Plus(_service)]), seqs)
        assert generic == specific == 5

    def test_generic_vs_more_generic(self):
        seqs = [['a', 'b', 'c']] * 3
        less_generic = lang_size_score(Plus(_a), seqs)
        more_generic = lang_size_score(Plus(Alt([_a, _b, _c])), seqs)
        assert less_generic < more_generic

    def test_single_sequence(self):
        seqs = [['a', 'b', 'c']]
        specific = lang_size_score(Concat([_a, _b, _c]), seqs)
        generic = lang_size_score(Plus(Alt([_a, _b, _c])), seqs)
        assert specific < generic

    def test_long_sequences(self):
        seqs = [['a', 'b', 'c', 'd', 'e']] * 3
        specific = lang_size_score(Concat([_a, _b, _c, _d, _e]), seqs)
        generic = lang_size_score(Plus(Alt([_a, _b, _c, _d, _e])), seqs)
        assert specific < generic

    def test_empty_sequences(self):
        g = Concat([_a, _b, _c])
        score = lang_size_score(g, [])
        expected = lang_size(g, 2 * model_cost(g) + 1)
        assert score == expected

    def test_ordered_vs_unordered(self):
        seqs = [['a', 'b', 'c'], ['a', 'b'], ['a', 'c']]
        ordered = lang_size_score(Concat([_a, _b, _c]), seqs)
        unordered = lang_size_score(Plus(Alt([_a, _b, _c])), seqs)
        assert ordered < unordered

    def test_optional_beats_generic(self):
        seqs = [['a', 'b'], ['a', 'c'], ['a']]
        optional = lang_size_score(Concat([_a, Optional(Alt([_b, _c]))]), seqs)
        generic = lang_size_score(Plus(Alt([_a, _b, _c])), seqs)
        assert optional < generic

    def test_repeat_beats_concat(self):
        seqs = [['a'], ['a', 'a'], ['a', 'a', 'a']]
        repeat = lang_size_score(Plus(_a), seqs)
        concat = lang_size_score(Concat([_a, _a, _a]), seqs)
        assert concat < repeat


# ── MDL Fallback ──

class TestMDLFallback:
    def test_mdl_basic(self):
        g = Concat([_a, _b, _c])
        score = mdl_score(g, [['a', 'b', 'c']])
        assert score == mdl_model_cost(g) + data_cost(g, [['a', 'b', 'c']])

    def test_mdl_prefers_short_expressions(self):
        seqs = [['info', 'file', 'template', 'shell', 'service']] * 5
        generic = mdl_score(Plus(_info), seqs)
        specific = mdl_score(Concat([_info, _file, _template, _shell, Plus(_service)]), seqs)
        assert generic < specific

    def test_score_grammar_method_switch(self):
        g = Concat([_a, _b, _c])
        ls = score_grammar(g, [['a', 'b', 'c']], method='langsize')
        mdl = score_grammar(g, [['a', 'b', 'c']], 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(_a, [['a']], method='bogus')

    def test_langsize_beats_mdl_on_info_plus(self):
        seqs = [['info', 'file', 'template', 'shell', 'service']] * 5
        g_generic = Plus(_info)
        g_specific = Concat([_info, _file, _template, _shell, Plus(_service)])
        ls_generic = score_grammar(g_generic, seqs, method='langsize')
        ls_specific = score_grammar(g_specific, seqs, method='langsize')
        mdl_generic = score_grammar(g_generic, seqs, method='mdl')
        mdl_specific = score_grammar(g_specific, 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')
        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(_a, 1) == 1
        assert count_words(_a, 0) == 0
        assert count_words(_a, 2) == 0

    def test_concatenation(self):
        assert count_words(Concat([_a, _b, _c]), 3) == 1
        assert count_words(Concat([_a, _b, _c]), 2) == 0
        assert count_words(Concat([_a, _b, _c]), 4) == 0

    def test_plus_quantifier(self):
        for l in range(1, 6):
            assert count_words(Plus(_a), l) == 1
        assert count_words(Plus(_a), 0) == 0

    def test_disjunction(self):
        assert count_words(Alt([_a, _b, _c]), 1) == 3
        assert count_words(Alt([_a, _b, _c]), 0) == 0
        assert count_words(Alt([_a, _b, _c]), 2) == 0

    def test_disjunction_plus(self):
        assert count_words(Plus(Alt([_a, _b, _c])), 1) == 3
        assert count_words(Plus(Alt([_a, _b, _c])), 2) == 9
        assert count_words(Plus(Alt([_a, _b, _c])), 3) == 27

    def test_optional(self):
        assert count_words(Concat([Optional(_a), Alt([_b, _c])]), 0) == 0
        assert count_words(Concat([Optional(_a), Alt([_b, _c])]), 1) == 2
        assert count_words(Concat([Optional(_a), Alt([_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(Plus(_info), l) == 1

    def test_info_dot_concat(self):
        assert count_words(Concat([_info, _file, _template]), 3) == 1
        assert count_words(Concat([_info, _file, _template]), 2) == 0
        assert count_words(Concat([_info, _file, _template]), 4) == 0

    def test_mixed_disj_concat(self):
        assert count_words(Concat([_a, Plus(Alt([_b, _c]))]), 2) == 2
        assert count_words(Concat([_a, Plus(Alt([_b, _c]))]), 3) == 4

    def test_optional_concat(self):
        assert count_words(Concat([Optional(_a), _b, Alt([_c, _d])]), 0) == 0
        assert count_words(Concat([Optional(_a), _b, Alt([_c, _d])]), 2) == 2
        assert count_words(Concat([Optional(_a), _b, Alt([_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
        g_generic = Plus(_info)
        g_specific = Concat([_info, _file, _template, _shell, Plus(_service)])
        generic_score = lang_size_score(g_generic, seqs)
        specific_score = lang_size_score(g_specific, 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(Plus(_info), seqs)
        specific = lang_size_score(Concat([_info, _file, Plus(_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 _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 _info in alpha)