"""rwr² — Translate k-OA to k-ORE (Algorithm 3, arXiv 1004.2372).

rwr²(G):
  1: compute a marking H of G
  2: return strip(rwr²₁(H))
"""

from .marking import mark_koa
from .rwr0 import rwr0
from .grammar import (
    Concat, Alt, Plus, Optional, Star, Symbol, Epsilon, Empty,
)


def strip(node):
    """Remove k-ORE markers: Symbol('a_i') → Symbol('a')."""
    if isinstance(node, Symbol):
        value = node.value
        if '_' in value:
            base = value.rsplit('_', 1)[0]
            return Symbol(base)
        return node
    if isinstance(node, (Epsilon, Empty)):
        return node
    if isinstance(node, (Plus, Optional, Star)):
        return type(node)(strip(node.child))
    if isinstance(node, (Concat, Alt)):
        return type(node)([strip(child) for child in node.parts])
    return node


def rwr_sq(G):
    """
    |———— Algorithm 3: rwr² ————|
    Require: k-OA G
    Ensure: k-ORE r with L(G) ⊆ L(r)

    1: H ← marking of G
    2: return strip(rwr²₁(H))
    """
    H = mark_koa(G)
    result = rwr0(H)
    if result is None or isinstance(result, Empty):
        return None
    return strip(result)