"""MDL scoring for iDRegEx (Algorithm 4, arXiv 1004.2372)."""

import math
from .grammar import (
    Concat, Alt, Plus, Optional, Star, Symbol, Epsilon, Empty,
    alphabet, count_words, lang_size, model_cost, match as grammar_match,
    _COUNT_CAP,
)


def data_cost(expr, sequences):
    """MDL data cost: Σ_i log₂(|L_i(r)|) where |L_i(r)| is the number
    of words of length len(seq_i) accepted by the grammar.

    Lower cost = more specific grammar that still covers the data.
    Exact computation is capped at max_len=50 to prevent combinatorial
    explosion. Longer sequences use an alphabet-size upper bound.
    """
    MAX_EXACT = 50
    n = 2 * model_cost(expr) + 1
    runtime_n = min(max(n, max((len(s) for s in sequences), default=0)), MAX_EXACT)

    lang_sizes = [count_words(expr, l) for l in range(runtime_n + 1)]

    alpha_size = len(alphabet(expr))

    total_cost = 0.0
    for seq in sequences:
        length = len(seq)
        if length <= runtime_n:
            ls = lang_sizes[length]
            if ls > 0:
                total_cost += math.log2(ls)
            else:
                total_cost += length * math.log2(max(alpha_size, 1))
        else:
            total_cost += length * math.log2(max(alpha_size, 1))
    return total_cost


def lang_size_score(expr, sequences):
    """Language Size: Σ |L(r)|_len(seq) — sum of words at each sequence length.

    From Bex et al. (arXiv:1004.2372), Section 4.3.1, adapted for our setting
    where candidates have different n values.

    Counts words at exactly the lengths present in the input sequences.
    Lower is better — the grammar that accepts the fewest words at the
    observed lengths wins. Generic grammars like `info+` accept many words
    at each length; specific grammars like `a.b.c.d.e+` accept exactly one.
    """
    if not sequences:
        return lang_size(expr, 2 * model_cost(expr) + 1)

    total = 0
    for seq in sequences:
        length = len(seq)
        total += count_words(expr, length)
        if total >= _COUNT_CAP:
            return _COUNT_CAP
    return total


def mdl_score(expr, sequences):
    """MDL = model cost + data cost. (Fallback, Bex et al. Section 4.3.2.)"""
    model = model_cost(expr)
    data = data_cost(expr, sequences)
    return model + data


_SCORERS = {
    'langsize': lang_size_score,
    'mdl': mdl_score,
}


def score_grammar(expr, sequences, method='langsize'):
    """Score a grammar using the specified method.

    Args:
        expr: Grammar AST node.
        sequences: List of sequences (each a list of strings).
        method: 'langsize' (default, Bex et al.) or 'mdl' (fallback).

    Returns:
        Numeric score (lower is better).
    """
    fn = _SCORERS.get(method)
    if fn is None:
        raise ValueError(f"Unknown scoring method '{method}'. Choose from: {list(_SCORERS)}")
    return fn(expr, sequences)


# For backward compatibility
class MDLScorer:
    def score(self, expr, sequences):
        return mdl_score(expr, sequences)