"""Explore weakest preconditions on a finite set of states. Python 3.10+, stdlib."""
from __future__ import annotations


def assignment(name, expression):
    """Pre: expression terminates on valid states. Post: returns a state transformer."""
    def transform(state):
        result = dict(state)
        result[name] = expression(state)
        return result
    return transform


def wp_at(transform, goal, state):
    """Pre: transform terminates. Post: True iff it completes normally and meets goal."""
    try:
        return bool(goal(transform(dict(state))))
    except (ZeroDivisionError, IndexError, KeyError):
        return False


def sequence(first, second):
    """Pre: two terminating state transformers. Post: their ordered composition."""
    return lambda state: second(first(state))


def demo():
    """Enumerate examples; finite enumeration is not a general WP solver."""
    withdraw = assignment("balance", lambda s: s["balance"] - 100)
    safe = [x for x in range(95, 106)
            if wp_at(withdraw, lambda s: s["balance"] >= 0, {"balance": x})]
    print("Safe starting balances in 95..105:", safe)
    quotient = assignment("y", lambda s: 12 // s["x"])
    print("12//x==3, safe x in -5..8:",
          [x for x in range(-5, 9) if wp_at(quotient, lambda s: s["y"] == 3, {"x": x})])
    read_last = assignment("y", lambda s: s["a"][s["i"]])
    print("Python index -1 is defined:",
          wp_at(read_last, lambda s: s["y"] == 7, {"a": [3, 7], "i": -1}))


def self_test():
    """Compare finite-state calculations with independently derived predicates."""
    count = 0
    withdraw = assignment("balance", lambda s: s["balance"] - 100)
    # Invariant: previous balances satisfy the derived equivalence. Bound: remaining states.
    for x in range(-200, 301):
        check(wp_at(withdraw, lambda s: s["balance"] >= 0, {"balance": x}), x >= 100)
        count += 1
    first = assignment("x", lambda s: 2*s["x"])
    second = assignment("y", lambda s: s["x"] + 1)
    for x in range(-10, 11):
        check(wp_at(sequence(first, second), lambda s: s["y"] > 7, {"x": x}), x > 3)
        count += 1
    read = assignment("v", lambda s: s["a"][s["i"]])
    for i in range(-5, 6):
        check(wp_at(read, lambda s: True, {"a": [2, 7], "i": i}), -2 <= i < 2)
        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()