"""Balanced parentheses and a CYK parser for a^n b^n (n>=1). Python 3.10+, stdlib."""
from __future__ import annotations
from itertools import product


def balanced(text: str) -> bool:
    """Pre: only parentheses. Post: True iff every close matches an earlier open."""
    if set(text) - {"(", ")"}:
        raise ValueError("use only '(' and ')'")
    stack = []
    # Invariant: stack contains the unmatched opening parentheses of the prefix.
    # Bound: remaining characters.
    for character in text:
        if character == "(":
            stack.append(character)
        elif not stack:
            return False
        else:
            stack.pop()
    return not stack


def cyk_ab(word: str) -> tuple[bool, list]:
    """Pre: finite word. Post: membership/table for CNF S->AB|AC, C->SB, A->a, B->b."""
    if not word:
        return False, []
    n = len(word)
    table = [[set() for _ in range(n + 1)] for _ in range(n)]
    terminal = {"a": {"A"}, "b": {"B"}}
    rules = (("S", "A", "B"), ("S", "A", "C"), ("C", "S", "B"))
    # Invariant: all earlier length-one spans are correct. Bound: remaining positions.
    for i, symbol in enumerate(word):
        table[i][1] = set(terminal.get(symbol, set()))
    # Invariant: all shorter spans are correct; every split is considered exactly once.
    # Bound: remaining lengths, starts, splits, and rules.
    for length in range(2, n + 1):
        for start in range(n - length + 1):
            for split in range(1, length):
                for parent, left, right in rules:
                    if left in table[start][split] and right in table[start+split][length-split]:
                        table[start][length].add(parent)
    return "S" in table[0][n], table


def demo():
    print("Balanced:", [(s, balanced(s)) for s in ("", "(()())", "())(", "(")])
    accepted, table = cyk_ab("aabb")
    print("CYK aabb accepted:", accepted)
    print("Full-span variables:", sorted(table[0][4]))
    print("Length-2 middle span ab:", sorted(table[1][2]))


def self_test():
    count = 0
    # Invariant: previous finite words agree with direct counting specifications.
    # Bound: remaining words of length <=6.
    for length in range(7):
        for chars in product("()", repeat=length):
            word = "".join(chars)
            levels = [word[:i].count("(")-word[:i].count(")") for i in range(length+1)]
            check(balanced(word), min(levels) >= 0 and levels[-1] == 0)
            count += 1
        for chars in product("ab", repeat=length):
            word = "".join(chars)
            reference = length >= 2 and length % 2 == 0 and word == "a"*(length//2)+"b"*(length//2)
            check(cyk_ab(word)[0], reference)
            count += 1
    return count

def check(actual, expected):
    """Require an expected result; this check also runs under python -O."""
    if actual != expected:
        raise AssertionError(f"expected {expected!r}, got {actual!r}")


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--self-test", action="store_true")
    args = parser.parse_args()
    if args.self_test:
        print(f"{self_test()} checks passed")
    else:
        demo()