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

from .ikoa import ikoa
from .rwrsq import rwr_sq
from .grammar import (
    Concat, Alt, Plus, Optional, Star, Symbol, Epsilon, Empty,
    alphabet, count_words, lang_size as grammar_lang_size, model_cost as grammar_model_cost,
)


def is_deterministic(node):
    """Check if a k-ORE is deterministic (Glushkov determinism).

    A k-ORE is deterministic iff for every subexpression Alt([r, s, ...]),
    first(r) ∩ first(s) = ∅.
    """
    if node is None or isinstance(node, (Empty, Epsilon)):
        return True
    return _check_det(node)


def _check_det(node):
    """Recursive determinism check on AST nodes."""
    if isinstance(node, Symbol):
        return True
    if isinstance(node, (Epsilon, Empty)):
        return True
    if isinstance(node, (Plus, Optional, Star)):
        return _check_det(node.child)
    if isinstance(node, Alt):
        first_sets = [_first_set(child) for child in node.parts]
        for j, fs1 in enumerate(first_sets):
            for fs2 in first_sets[j + 1:]:
                if fs1 & fs2:
                    return False
        for child in node.parts:
            if not _check_det(child):
                return False
        return True
    if isinstance(node, Concat):
        for child in node.parts:
            if not _check_det(child):
                return False
        return True
    return True


def _first_set(node):
    """Compute first(r) — set of alphabet symbols that can appear at the start of a word in L(r)."""
    if node is None or isinstance(node, Empty):
        return set()
    if isinstance(node, Epsilon):
        return set()
    if isinstance(node, Symbol):
        return {node.value}
    if isinstance(node, (Plus, Optional, Star)):
        return _first_set(node.child)
    if isinstance(node, Concat):
        if node.parts:
            return _first_set(node.parts[0])
        return set()
    if isinstance(node, Alt):
        result = set()
        for child in node.parts:
            result |= _first_set(child)
        return result
    return set()


def idregex(sequences, kmax=4, N=5, criterion='langsize'):
    """
    |———— Algorithm 4: iDRegEx ————|
    Require: sample S
    Ensure: k-ORE r

    1: C ← ∅
    2: for k = 1 to kmax do
    3:   for n = 1 to N do
    4:     G ← iKoa(S, k)
    5:     if rwr²(G) is deterministic then
    6:       add rwr²(G) to C
    7: return best(C)
    """
    C = set()
    for k in range(1, kmax + 1):
        for _ in range(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):
                if is_deterministic(expr):
                    C.add(expr)
    if not C:
        return None
    if criterion == 'langsize':
        return min(C, key=lambda e: (grammar_lang_size(e, 2 * grammar_model_cost(e) + 1), _ast_size(e)))
    return min(C, key=lambda e: _ast_size(e))


def _ast_size(node):
    """Count nodes in AST for length comparison."""
    if isinstance(node, Symbol):
        return 1
    if isinstance(node, (Epsilon, Empty)):
        return 1
    if isinstance(node, (Plus, Optional, Star)):
        return 1 + _ast_size(node.child)
    if isinstance(node, (Concat, Alt)):
        return 1 + sum(_ast_size(child) for child in node.parts)
    return 1