"""kOREInference — Algorithm 4: iDRegEx (arXiv 1004.2372)."""

from .ikoa import ikoa
from .rwrsq import rwr_sq
from .idregex import is_deterministic
from .mdl import mdl_score
from .grammar import Epsilon, Empty, Symbol, Plus, Optional, Star, Concat, Alt, alphabet as ast_alphabet


def validate_k_ore(node, k, alphabet_set=None):
    """Check if a k-ORE satisfies the k-occurrence condition."""
    if node is None or isinstance(node, (Empty, Epsilon)):
        return True, "OK"
    syms = alphabet_set or ast_alphabet(node)
    counts = _count_symbol_occurrences(node)
    violations = [f"{s}:{c}" for s, c in sorted(counts.items()) if c > k]
    if violations:
        return False, f"k={k} violations: {', '.join(violations)}"
    return True, "OK"


def _count_symbol_occurrences(node):
    """Count how many times each symbol appears as a leaf in the AST."""
    if isinstance(node, Symbol):
        return {node.value: 1}
    if isinstance(node, (Epsilon, Empty)):
        return {}
    if isinstance(node, (Plus, Optional, Star)):
        return _count_symbol_occurrences(node.child)
    if isinstance(node, (Concat, Alt)):
        counts = {}
        for p in node.parts:
            for sym, cnt in _count_symbol_occurrences(p).items():
                counts[sym] = counts.get(sym, 0) + cnt
        return counts
    return {}


def _count_nodes(node):
    if isinstance(node, Symbol):
        return 1
    if isinstance(node, (Epsilon, Empty)):
        return 1
    if isinstance(node, (Plus, Optional, Star)):
        return 1 + _count_nodes(node.child)
    if isinstance(node, (Concat, Alt)):
        return 1 + sum(_count_nodes(p) for p in node.parts)
    return 1


class kOREInference:
    def __init__(self, k_max=5, N=5):
        self.k_max = k_max
        self.N = N

    def infer(self, sequences):
        sequences = [s for s in sequences if s]
        if not sequences:
            return None
        candidates = []
        for k in range(1, self.k_max + 1):
            for _ in range(self.N):
                G = ikoa(sequences, k, num_trials=1)
                if G is None:
                    continue
                expr = rwr_sq(G)
                if expr is not None and not isinstance(expr, (Empty, Epsilon)):
                    if is_deterministic(expr):
                        valid, _ = validate_k_ore(expr, k)
                        if valid:
                            candidates.append((G, expr, k))
        if not candidates:
            return None
        return min(candidates, key=lambda c: mdl_score(c[1], sequences))