"""Run a DFA and NFA; compare them with numerical/regex specifications. Python 3.10+."""
from __future__ import annotations
from itertools import product


def divisible_by_three(bits: str) -> bool:
    """Pre: binary string; empty denotes zero. Post: value is divisible by three."""
    if set(bits) - {"0", "1"}:
        raise ValueError("input must contain only 0 and 1")
    state = 0
    # Invariant: state is the remainder of the consumed prefix. Bound: remaining bits.
    for bit in bits:
        state = (2 * state + int(bit)) % 3
    return state == 0


def ending_01(bits: str, *, trace=False):
    """Pre: binary string. Post: NFA accepts exactly strings ending in 01."""
    if set(bits) - {"0", "1"}:
        raise ValueError("input must contain only 0 and 1")
    transitions = {("A", "0"): {"A", "B"}, ("A", "1"): {"A"}, ("B", "1"): {"C"}}
    live, steps = {"A"}, [("start", ["A"])]
    # Invariant: live contains all states reachable after the consumed prefix.
    # Bound: remaining input symbols.
    for bit in bits:
        live = set().union(*(transitions.get((state, bit), set()) for state in live))
        steps.append((bit, sorted(live)))
    return ("C" in live, steps) if trace else "C" in live


def demo():
    print("Divisible by 3:", [(s, divisible_by_three(s)) for s in ("", "110", "111")])
    accepted, steps = ending_01("1101", trace=True)
    print("NFA trace for 1101:", steps)
    print("Accepted:", accepted)


def self_test():
    count = 0
    # Invariant: previous binary words agree with arithmetic and string suffix predicates.
    # Bound: remaining words of length <=7.
    for length in range(8):
        for symbols in product("01", repeat=length):
            word = "".join(symbols)
            check(divisible_by_three(word), int(word or "0", 2) % 3 == 0)
            check(ending_01(word), word.endswith("01"))
            count += 2
    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()