"""A capture-avoiding normal-order lambda reducer. Python 3.10+, no dependencies."""
from __future__ import annotations
from dataclasses import dataclass


@dataclass(frozen=True)
class Var:
    name: str


@dataclass(frozen=True)
class Lam:
    name: str
    body: object


@dataclass(frozen=True)
class App:
    function: object
    argument: object


def free(term):
    """Pre: finite lambda AST. Post: its free variable names."""
    if isinstance(term, Var):
        return {term.name}
    if isinstance(term, Lam):
        return free(term.body)-{term.name}
    return free(term.function) | free(term.argument)


def substitute(term, name, replacement):
    """Pre: finite ASTs. Post: replace free occurrences without capturing variables."""
    if isinstance(term, Var):
        return replacement if term.name == name else term
    if isinstance(term, App):
        return App(substitute(term.function, name, replacement),
                   substitute(term.argument, name, replacement))
    if term.name == name or name not in free(term.body):
        return term
    binder, body = term.name, term.body
    if binder in free(replacement):
        used = free(body) | free(replacement) | {name, binder}
        fresh = binder + "_"
        # Invariant: candidate is valid; bound: at most len(used)+1 name candidates.
        while fresh in used:
            fresh += "_"
        body, binder = substitute(body, binder, Var(fresh)), fresh
    return Lam(binder, substitute(body, name, replacement))


def step(term):
    """Pre: finite AST. Post: one leftmost-outermost beta step, or None for normal form."""
    if isinstance(term, Var):
        return None
    if isinstance(term, Lam):
        body = step(term.body)
        return None if body is None else Lam(term.name, body)
    if isinstance(term.function, Lam):
        return substitute(term.function.body, term.function.name, term.argument)
    function = step(term.function)
    if function is not None:
        return App(function, term.argument)
    argument = step(term.argument)
    return None if argument is None else App(term.function, argument)


def normalize(term, *, limit=1000):
    """Pre: limit>=0. Post: normal form within budget, else inconclusive RuntimeError."""
    if limit < 0:
        raise ValueError("limit must be nonnegative")
    # Invariant: current term is beta-equivalent to input. Bound: remaining allowed steps.
    for count in range(limit+1):
        following = step(term)
        if following is None:
            return term
        if count == limit:
            raise RuntimeError("reduction budget exhausted; termination remains unknown")
        term = following


def show(term):
    """Pre: finite AST. Post: fully parenthesized readable lambda term."""
    if isinstance(term, Var):
        return term.name
    if isinstance(term, Lam):
        return f"(lambda {term.name}. {show(term.body)})"
    return f"({show(term.function)} {show(term.argument)})"


def demo():
    identity = Lam("x", Var("x"))
    capture = App(Lam("x", Lam("y", Var("x"))), Var("y"))
    print("Identity application:", show(normalize(App(identity, Var("a")))))
    print("Capture-avoiding result:", show(normalize(capture)))
    omega_body = Lam("x", App(Var("x"), Var("x")))
    omega = App(omega_body, omega_body)
    print("Discard divergent argument:", show(normalize(App(Lam("x", Var("z")), omega))))
    try:
        normalize(omega, limit=10)
    except RuntimeError as error:
        print(error)


def self_test():
    check(normalize(App(Lam("x", Var("x")), Var("a"))), Var("a"))
    check(normalize(App(Lam("x", Lam("x", Var("x"))), Var("a"))), Lam("x", Var("x")))
    captured = normalize(App(Lam("x", Lam("y", Var("x"))), Var("y")))
    check(free(captured), {"y"})
    check(captured.name != "y", True)
    check(captured.body, Var("y"))
    check(normalize(Lam("x", Var("x")), limit=0), Lam("x", Var("x")))
    return 6

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()