zf2.py

import kdrag as kd
import kdrag.smt as smt
import kdrag.theories.logic.zf as zf

Set = zf.ZFSet

A, B, R, x, y, z = smt.Consts("A B R x y z", Set)
a, b, c, p = kd.FreshVars("a b c p", Set)

elem = zf.elem

# Cartesian product
prod = kd.define(
    "prod",
    [A, B],
    zf.sep(
        zf.pow(zf.pow(A | B)),
        smt.Lambda(
            [p],
            kd.QExists([x, y], smt.And(elem(x, A), elem(y, B)), p == zf.pair(x, y)),
        ),
    ),
)


# Membership in the Cartesian product (definition expansion).
# Lean: p ∈ A × B ↔ (∃ x∈A, ∃ y∈B, p = ⟪x,y⟫) ∧ p ∈ 𝒫 (𝒫 (A ∪ B))
# Membership in image (definition expansion).
# Lean: y ∈ R[A] ↔ (∃ x∈A, ⟪x,y⟫ ∈ R) ∧ y ∈ ran R
# Membership in inverse relation (definition expansion).
# Lean: p ∈ R⁻¹ ↔ (∃ x y, p = ⟪y,x⟫ ∧ ⟪x,y⟫ ∈ R) ∧ p ∈ 𝒫 (𝒫 (⋃⋃R))
# Membership in trans_omega (definition expansion).
# Lean: x ∈ trans_omega ↔ x ∈ ω ∧ x ⊆ ω
# Membership in image (definition expansion).
# Lean: y ∈ R[A] ↔ (∃ x∈A, ⟪x,y⟫ ∈ R) ∧ y ∈ ran R
@kd.Theorem(
    smt.ForAll(
        [A, B, p],
        elem(p, prod(A, B))
        == smt.And(
            kd.QExists([x, y], smt.And(elem(x, A), elem(y, B)), p == zf.pair(x, y)),
            elem(p, zf.pow(zf.pow(A | B))),
        ),
    )
)
def elem_prod(l):
    _A, _B, _p = l.fixes()
    l.unfold(prod)
    l.rw(zf.elem_sep)
    l.auto()


# Relations
is_rel = kd.define(
    "is_rel", [R], smt.ForAll([p], smt.Implies(elem(p, R), zf.is_pair(p)))
)

dom = kd.define(
    "dom",
    [R],
    zf.sep(
        zf.bigunion(zf.bigunion(R)),
        smt.Lambda([x], kd.QExists([y], elem(zf.pair(x, y), R))),
    ),
)

ran = kd.define(
    "ran",
    [R],
    zf.sep(
        zf.bigunion(zf.bigunion(R)),
        smt.Lambda([y], kd.QExists([x], elem(zf.pair(x, y), R))),
    ),
)


# Membership in the domain (definition expansion).
# Lean: x ∈ dom R ↔ (∃ y, ⟪x,y⟫ ∈ R) ∧ x ∈ ⋃⋃R
# Membership in trans_omega (definition expansion).
# Lean: x ∈ trans_omega ↔ x ∈ ω ∧ x ⊆ ω
@kd.Theorem(
    smt.ForAll(
        [R, x],
        elem(x, dom(R))
        == smt.And(
            kd.QExists([y], elem(zf.pair(x, y), R)),
            elem(x, zf.bigunion(zf.bigunion(R))),
        ),
    )
)
def elem_dom(l):
    _R, _x = l.fixes()
    l.unfold(dom)
    l.rw(zf.elem_sep)
    l.auto()


# Membership in the range (definition expansion).
# Lean: y ∈ ran R ↔ (∃ x, ⟪x,y⟫ ∈ R) ∧ y ∈ ⋃⋃R
@kd.Theorem(
    smt.ForAll(
        [R, y],
        elem(y, ran(R))
        == smt.And(
            kd.QExists([x], elem(zf.pair(x, y), R)),
            elem(y, zf.bigunion(zf.bigunion(R))),
        ),
    )
)
def elem_ran(l):
    _R, _y = l.fixes()
    l.unfold(ran)
    l.rw(zf.elem_sep)
    l.auto()


# Domain membership yields a witnessing pair in the relation.
# Lean: x ∈ dom R → ∃ y, ⟪x,y⟫ ∈ R
@kd.Theorem(
    smt.ForAll(
        [R, x], smt.Implies(elem(x, dom(R)), kd.QExists([y], elem(zf.pair(x, y), R)))
    )
)
def dom_has_pair(l):
    _R, _x = l.fixes()
    l.intros()
    l.rw(elem_dom, at=0)
    l.auto()


# Range membership yields a witnessing pair in the relation.
# Lean: y ∈ ran R → ∃ x, ⟪x,y⟫ ∈ R
@kd.Theorem(
    smt.ForAll(
        [R, y], smt.Implies(elem(y, ran(R)), kd.QExists([x], elem(zf.pair(x, y), R)))
    )
)
def ran_has_pair(l):
    _R, _y = l.fixes()
    l.intros()
    l.rw(elem_ran, at=0)
    l.auto()


# Domain is a subset of the big union of the relation.
# Lean: dom R ⊆ ⋃⋃R
@kd.Theorem(smt.ForAll([R], dom(R) <= zf.bigunion(zf.bigunion(R))))
def dom_le_bigunion(l):
    _R = l.fix()
    l.unfold(zf.le)
    _x = l.fix()
    l.intros()
    l.rw(elem_dom, at=0)
    l.auto()


# Range is a subset of the big union of the relation.
# Lean: ran R ⊆ ⋃⋃R
@kd.Theorem(smt.ForAll([R], ran(R) <= zf.bigunion(zf.bigunion(R))))
def ran_le_bigunion(l):
    _R = l.fix()
    l.unfold(zf.le)
    _y = l.fix()
    l.intros()
    l.rw(elem_ran, at=0)
    l.auto()


# Images and functions
image = kd.define(
    "image",
    [R, A],
    zf.sep(
        ran(R),
        smt.Lambda([y], kd.QExists([x], smt.And(elem(x, A), elem(zf.pair(x, y), R)))),
    ),
)

# Membership in image (definition expansion).
# Lean: y ∈ R[A] ↔ (∃ x∈A, ⟪x,y⟫ ∈ R) ∧ y ∈ ran R
@kd.Theorem(
    smt.ForAll(
        [R, A, y],
        elem(y, image(R, A))
        == smt.And(
            kd.QExists([x], smt.And(elem(x, A), elem(zf.pair(x, y), R))),
            elem(y, ran(R)),
        ),
    )
)
def elem_image(l):
    _R, _A, _y = l.fixes()
    l.unfold(image)
    l.rw(zf.elem_sep)
    l.auto()


is_fun = kd.define(
    "is_fun",
    [R],
    smt.And(
        is_rel(R),
        smt.ForAll(
            [x, y, z],
            smt.Implies(
                smt.And(elem(zf.pair(x, y), R), elem(zf.pair(x, z), R)), y == z
            ),
        ),
    ),
)


# Every function is a relation.
# Lean: is_fun R → is_rel R
@kd.Theorem(smt.ForAll([R], smt.Implies(is_fun(R), is_rel(R))))
def is_fun_is_rel(l):
    _R = l.fix()
    l.unfold(is_fun)
    l.intros()
    l.auto()


# Inverse and composition
inv = kd.define(
    "inv",
    [R],
    zf.sep(
        zf.pow(zf.pow(zf.bigunion(zf.bigunion(R)))),
        smt.Lambda([p], kd.QExists([x, y], p == zf.pair(y, x), elem(zf.pair(x, y), R))),
    ),
)


# Membership in inverse relation (definition expansion).
# Lean: p ∈ R⁻¹ ↔ (∃ x y, p = ⟪y,x⟫ ∧ ⟪x,y⟫ ∈ R) ∧ p ∈ 𝒫 (𝒫 (⋃⋃R))
@kd.Theorem(
    smt.ForAll(
        [R, p],
        elem(p, inv(R))
        == smt.And(
            kd.QExists([x, y], smt.And(p == zf.pair(y, x), elem(zf.pair(x, y), R))),
            elem(p, zf.pow(zf.pow(zf.bigunion(zf.bigunion(R))))),
        ),
    )
)
def elem_inv(l):
    _R, _p = l.fixes()
    l.unfold(inv)
    l.rw(zf.elem_sep)
    l.auto()


comp = kd.define(
    "comp",
    [R, B],
    zf.sep(
        zf.pow(zf.pow(zf.bigunion(zf.bigunion(B)) | zf.bigunion(zf.bigunion(R)))),
        smt.Lambda(
            [p],
            kd.QExists(
                [x, y, z],
                smt.And(
                    elem(zf.pair(x, y), R), elem(zf.pair(y, z), B), p == zf.pair(x, z)
                ),
            ),
        ),
    ),
)


# Membership in relation composition (definition expansion).
# Lean: p ∈ R ∘ S ↔ (∃ x y z, ⟪x,y⟫ ∈ R ∧ ⟪y,z⟫ ∈ S ∧ p = ⟪x,z⟫) ∧ p ∈ 𝒫 (𝒫 (⋃⋃S ∪ ⋃⋃R))
@kd.Theorem(
    smt.ForAll(
        [R, B, p],
        elem(p, comp(R, B))
        == smt.And(
            kd.QExists(
                [x, y, z],
                smt.And(
                    elem(zf.pair(x, y), R), elem(zf.pair(y, z), B), p == zf.pair(x, z)
                ),
            ),
            elem(
                p,
                zf.pow(
                    zf.pow(zf.bigunion(zf.bigunion(B)) | zf.bigunion(zf.bigunion(R)))
                ),
            ),
        ),
    )
)
def elem_comp(l):
    _R, _B, _p = l.fixes()
    l.unfold(comp)
    l.rw(zf.elem_sep)
    l.auto()


# Inverse membership flips ordered pairs.
# Lean: ⟪y,x⟫ ∈ R⁻¹ → ⟪x,y⟫ ∈ R
@kd.Theorem(
    smt.ForAll(
        [R, x, y],
        smt.Implies(elem(zf.pair(y, x), inv(R)), elem(zf.pair(x, y), R)),
    )
)
def inv_pair_elim(l):
    _R, _x, _y = l.fixes()
    l.intros()
    l.have(
        kd.QExists(
            [a, b],
            smt.And(zf.pair(_y, _x) == zf.pair(b, a), elem(zf.pair(a, b), _R)),
        ),
        by=[elem_inv],
    )
    _a, _b = l.obtain(-1)
    l.have(smt.And(_y == _b, _x == _a), by=[zf.pair_inj])
    l.auto()


# Introduce inverse membership with a side condition.
# Lean: ⟪x,y⟫ ∈ R → ⟪y,x⟫ ∈ R⁻¹ (with carrier side condition)
@kd.Theorem(
    smt.ForAll(
        [R, x, y],
        smt.Implies(
            smt.And(
                elem(zf.pair(x, y), R),
                elem(zf.pair(y, x), zf.pow(zf.pow(zf.bigunion(zf.bigunion(R))))),
            ),
            elem(zf.pair(y, x), inv(R)),
        ),
    )
)
def inv_pair_intro(l):
    _R, _x, _y = l.fixes()
    l.intros()
    l.rw(elem_inv)
    l.split()
    l.exists(_x, _y)
    l.auto()
    l.auto()


# The inverse of a relation is a relation.
# Lean: is_rel (R⁻¹)
@kd.Theorem(smt.ForAll([R], is_rel(inv(R))))
def inv_is_rel(l):
    _R = l.fix()
    l.unfold(is_rel)
    _p = l.fix()
    l.intros()
    l.have(
        kd.QExists(
            [a, b],
            smt.And(_p == zf.pair(b, a), elem(zf.pair(a, b), _R)),
        ),
        by=[elem_inv],
    )
    _a, _b = l.obtain(-1)
    l.have(_p == zf.pair(_b, _a), by=[])
    l.rw(-1)
    l.auto(by=[zf.is_pair_pair])


# Domain of inverse gives a flipped pair in the original relation.
# Lean: x ∈ dom(R⁻¹) → ∃ y, ⟪y,x⟫ ∈ R
@kd.Theorem(
    smt.ForAll(
        [R, x],
        smt.Implies(elem(x, dom(inv(R))), kd.QExists([y], elem(zf.pair(y, x), R))),
    )
)
def dom_inv_has_pair(l):
    _R, _x = l.fixes()
    l.intros()
    l.have(kd.QExists([y], elem(zf.pair(_x, y), inv(_R))), by=[dom_has_pair])
    _y = l.obtain(-1)
    l.have(elem(zf.pair(_y, _x), _R), by=[inv_pair_elim(_R, _y, _x)])
    l.exists(_y)
    l.auto()


# Range of inverse gives a flipped pair in the original relation.
# Lean: y ∈ ran(R⁻¹) → ∃ x, ⟪y,x⟫ ∈ R
@kd.Theorem(
    smt.ForAll(
        [R, y],
        smt.Implies(elem(y, ran(inv(R))), kd.QExists([x], elem(zf.pair(y, x), R))),
    )
)
def ran_inv_has_pair(l):
    _R, _y = l.fixes()
    l.intros()
    l.have(kd.QExists([x], elem(zf.pair(x, _y), inv(_R))), by=[ran_has_pair])
    _x = l.obtain(-1)
    l.have(elem(zf.pair(_y, _x), _R), by=[inv_pair_elim(_R, _y, _x)])
    l.exists(_x)
    l.auto()


# Double inverse eliminates.
# Lean: ⟪x,y⟫ ∈ (R⁻¹)⁻¹ → ⟪x,y⟫ ∈ R
@kd.Theorem(
    smt.ForAll(
        [R, x, y],
        smt.Implies(elem(zf.pair(x, y), inv(inv(R))), elem(zf.pair(x, y), R)),
    )
)
def inv_inv_elim(l):
    _R, _x, _y = l.fixes()
    l.intros()
    l.have(elem(zf.pair(_y, _x), inv(_R)), by=[inv_pair_elim(inv(_R), _y, _x)])
    l.have(elem(zf.pair(_x, _y), _R), by=[inv_pair_elim(_R, _x, _y)])
    l.auto()


# Identity relation
idrel = kd.define(
    "idrel",
    [A],
    zf.sep(
        prod(A, A),
        smt.Lambda([p], kd.QExists([x], smt.And(elem(x, A), p == zf.pair(x, x)))),
    ),
)


# Membership in identity relation (definition expansion).
# Lean: p ∈ id A ↔ (∃ x∈A, p = ⟪x,x⟫) ∧ p ∈ A × A
@kd.Theorem(
    smt.ForAll(
        [A, p],
        elem(p, idrel(A))
        == smt.And(
            kd.QExists([x], smt.And(elem(x, A), p == zf.pair(x, x))),
            elem(p, prod(A, A)),
        ),
    )
)
def elem_idrel(l):
    _A, _p = l.fixes()
    l.unfold(idrel)
    l.rw(zf.elem_sep)
    l.auto()


# Identity relation only contains diagonal pairs.
# Lean: ⟪x,y⟫ ∈ id A → x = y
@kd.Theorem(smt.ForAll([A, x, y], smt.Implies(elem(zf.pair(x, y), idrel(A)), x == y)))
def idrel_diag(l):
    _A, _x, _y = l.fixes()
    l.intros()
    l.have(
        kd.QExists([a], smt.And(elem(a, _A), zf.pair(_x, _y) == zf.pair(a, a))),
        by=[elem_idrel],
    )
    _a = l.obtain(-1)
    l.have(smt.And(_x == _a, _y == _a), by=[zf.pair_inj])
    l.auto()


# Identity relation is a relation.
# Lean: is_rel (id A)
@kd.Theorem(smt.ForAll([A], is_rel(idrel(A))))
def idrel_is_rel(l):
    _A = l.fix()
    l.unfold(is_rel)
    _p = l.fix()
    l.intros()
    l.have(
        kd.QExists([x], smt.And(elem(x, _A), _p == zf.pair(x, x))),
        by=[elem_idrel],
    )
    _x = l.obtain(-1)
    l.have(_p == zf.pair(_x, _x), by=[])
    l.rw(-1)
    l.auto(by=[zf.is_pair_pair])


# Identity relation is a function.
# Lean: is_fun (id A)
@kd.Theorem(smt.ForAll([A], is_fun(idrel(A))))
def idrel_is_fun(l):
    _A = l.fix()
    l.unfold(is_fun)
    l.split()
    l.auto(by=[idrel_is_rel])
    x, y, z = l.fixes()
    l.intros()
    l.have(x == y, by=[idrel_diag(_A, x, y)])
    l.have(x == z, by=[idrel_diag(_A, x, z)])
    l.auto()


# von Neumann naturals
succ = kd.define("succ", [a], a | zf.sing(a))
inductive = kd.define(
    "inductive",
    [A],
    smt.And(
        elem(zf.emp, A), smt.ForAll([x], smt.Implies(elem(x, A), elem(succ(x), A)))
    ),
)

inf_ax = kd.axiom(smt.Exists([A], inductive(A)))
(inf_set,), inf_set_inductive = kd.tactics.skolem(inf_ax)

omega = smt.Const("omega", Set)
omega_def = kd.axiom(
    omega
    == zf.sep(
        inf_set,
        smt.Lambda([a], smt.ForAll([A], smt.Implies(inductive(A), elem(a, A)))),
    )
)


# Membership in omega (definition expansion).
# Lean: a ∈ ω ↔ (∀ A, inductive A → a ∈ A) ∧ a ∈ inf_set
@kd.Theorem(
    smt.ForAll(
        [a],
        elem(a, omega)
        == smt.And(
            smt.ForAll([A], smt.Implies(inductive(A), elem(a, A))), elem(a, inf_set)
        ),
    )
)
def elem_omega(l):
    _a = l.fix()
    l.rw(omega_def)
    l.rw(zf.elem_sep)
    l.auto()


# Every inductive set contains the empty set.
# Lean: inductive A → ∅ ∈ A
@kd.Theorem(smt.ForAll([A], smt.Implies(inductive(A), elem(zf.emp, A))))
def inductive_has_emp(l):
    _A = l.fix()
    l.intros()
    l.unfold(inductive, at=0)
    l.split(at=0)
    l.auto()


# Inductive sets are closed under successor.
# Lean: inductive A → x ∈ A → succ x ∈ A
@kd.Theorem(
    smt.ForAll(
        [A, x], smt.Implies(inductive(A), smt.Implies(elem(x, A), elem(succ(x), A)))
    )
)
def inductive_succ(l):
    _A, _x = l.fixes()
    l.intros()
    l.intros()
    l.unfold(inductive, at=0)
    l.split(at=0)
    l.specialize(1, _x)
    l.apply(-1)
    l.auto()


# Empty set is in omega.
# Lean: ∅ ∈ ω
@kd.Theorem(elem(zf.emp, omega))
def emp_in_omega(l):
    l.rw(elem_omega)
    l.split()
    with l.sub() as l1:
        _A = l1.fix()
        l1.intros()
        l1.auto(by=[inductive_has_emp(_A)])
    l.auto(by=[inductive_has_emp(inf_set), inf_set_inductive])


# Omega is closed under successor.
# Lean: a ∈ ω → succ a ∈ ω
@kd.Theorem(smt.ForAll([a], smt.Implies(elem(a, omega), elem(succ(a), omega))))
def omega_succ(l):
    _a = l.fix()
    l.intros()
    l.rw(elem_omega)
    l.split()
    with l.sub() as l1:
        _A = l1.fix()
        l1.intros()
        l1.rw(elem_omega, at=0)
        l1.split(at=0)
        l1.have(elem(_a, _A))
        l1.specialize(0, _A)
        l1.apply(0)
        l1.auto()
        l1.auto(by=[inductive_succ(_A, _a)])
    l.rw(elem_omega, at=0)
    l.split(at=0)
    l.auto(by=[inductive_succ(inf_set, _a), inf_set_inductive])


# Omega is inductive.
# Lean: inductive ω
@kd.Theorem(inductive(omega))
def omega_inductive(l):
    l.unfold(inductive)
    l.split()
    l.auto(by=[emp_in_omega])
    _x = l.fix()
    l.intros()
    l.auto(by=[omega_succ(_x)])


# Omega is the least inductive set.
# Lean: inductive A → ω ⊆ A
@kd.Theorem(smt.ForAll([A], smt.Implies(inductive(A), omega <= A)))
def omega_le_inductive(l):
    _A = l.fix()
    l.intros()
    l.unfold(zf.le)
    _x = l.fix()
    l.intros()
    l.rw(elem_omega, at=1)
    l.split(at=1)
    l.specialize(1, _A)
    l.apply(-1)
    l.auto()


zero = zf.emp
one = succ(zero)
two = succ(one)
three = succ(two)


is_nat = kd.define("is_nat", [a], elem(a, omega))


# Zero is a natural.
# Lean: 0 ∈ ω
@kd.Theorem(is_nat(zero))
def zero_is_nat(l):
    l.unfold(is_nat)
    l.auto(by=[emp_in_omega])


# Naturals are closed under successor.
# Lean: a ∈ ω → succ a ∈ ω
@kd.Theorem(smt.ForAll([a], smt.Implies(is_nat(a), is_nat(succ(a)))))
def succ_is_nat(l):
    _a = l.fix()
    l.intros()
    l.unfold(is_nat, at=0)
    l.unfold(is_nat)
    l.auto(by=[omega_succ(_a)])


# One is a natural.
# Lean: 1 ∈ ω
@kd.Theorem(is_nat(one))
def one_is_nat(l):
    l.unfold(is_nat)
    l.have(elem(zero, omega), by=[emp_in_omega])
    l.auto(by=[omega_succ(zero)])


# Two is a natural.
# Lean: 2 ∈ ω
@kd.Theorem(is_nat(two))
def two_is_nat(l):
    l.unfold(is_nat)
    l.have(is_nat(one), by=[one_is_nat])
    l.unfold(is_nat, at=-1)
    l.auto(by=[omega_succ(one)])


# Three is a natural.
# Lean: 3 ∈ ω
@kd.Theorem(is_nat(three))
def three_is_nat(l):
    l.unfold(is_nat)
    l.have(is_nat(two), by=[two_is_nat])
    l.unfold(is_nat, at=-1)
    l.auto(by=[omega_succ(two)])


# Induction: any inductive subset of omega equals omega.
# Lean: A ⊆ ω → inductive A → A = ω
@kd.Theorem(smt.ForAll([A], smt.Implies(smt.And(A <= omega, inductive(A)), A == omega)))
def omega_induction(l):
    _A = l.fix()
    l.intros()
    l.split(at=0)
    l.rw(zf.ext_ax)
    _x = l.fix()
    l.split()
    with l.sub() as l1:
        l1.intros()
        l1.unfold(zf.le, at=0)
        l1.specialize(0, _x)
        l1.apply(-1)
        l1.auto()
    with l.sub() as l2:
        l2.intros()
        l2.have(omega <= _A, by=[omega_le_inductive(_A)])
        l2.unfold(zf.le, at=-1)
        l2.specialize(-1, _x)
        l2.apply(-1)
        l2.auto()


# Membership in successor (definition expansion).
# Lean: x ∈ succ a ↔ x ∈ a ∨ x = a
@kd.Theorem(smt.ForAll([a, x], elem(x, succ(a)) == smt.Or(elem(x, a), x == a)))
def elem_succ(l):
    _a, _x = l.fixes()
    l.unfold(succ)
    l.rw(zf.elem_union)
    l.rw(zf.elem_sing)
    l.auto()


# Every set is an element of its successor.
# Lean: a ∈ succ a
@kd.Theorem(smt.ForAll([a], elem(a, succ(a))))
def elem_self_succ(l):
    _a = l.fix()
    l.rw(elem_succ)
    l.auto()


# Every set is a subset of its successor.
# Lean: a ⊆ succ a
@kd.Theorem(smt.ForAll([a], a <= succ(a)))
def le_succ(l):
    _a = l.fix()
    l.unfold(zf.le)
    _x = l.fix()
    l.intros()
    l.rw(elem_succ)
    l.auto()


# Empty set is a subset of omega.
# Lean: ∅ ⊆ ω
@kd.Theorem(zf.emp <= omega)
def emp_le_omega(l):
    l.unfold(zf.le)
    _x = l.fix()
    l.intros()
    l.rw(zf.elem_emp, at=0)
    l.auto()


# Successor of a subset of omega stays within omega.
# Lean: (x ⊆ ω ∧ x ∈ ω) → succ x ⊆ ω
@kd.Theorem(
    smt.ForAll([x], smt.Implies(smt.And(x <= omega, elem(x, omega)), succ(x) <= omega))
)
def succ_le_omega(l):
    _x = l.fix()
    l.intros()
    l.split(at=0)
    l.unfold(zf.le)
    _y = l.fix()
    l.intros()
    l.rw(elem_succ, at=-1)
    l.split(at=-1)
    with l.sub() as l1:
        l1.rw(-1)
        l1.auto()
    with l.sub() as l2:
        l2.unfold(zf.le, at=0)
        l2.specialize(0, _y)
        l2.apply(-1)
        l2.auto()


trans_omega = smt.Const("trans_omega", Set)
trans_omega_def = kd.axiom(trans_omega == zf.sep(omega, smt.Lambda([x], x <= omega)))

# Membership in trans_omega (definition expansion).
# Lean: x ∈ trans_omega ↔ x ∈ ω ∧ x ⊆ ω
@kd.Theorem(
    smt.ForAll(
        [x],
        elem(x, trans_omega) == smt.And(elem(x, omega), x <= omega),
    )
)
def elem_trans_omega(l):
    _x = l.fix()
    l.rw(trans_omega_def)
    l.rw(zf.elem_sep)
    l.auto()


# trans_omega is a subset of omega by definition.
# trans_omega is a subset of omega by definition.
# Lean: trans_omega ⊆ ω
@kd.Theorem(trans_omega <= omega)
def trans_omega_le_omega(l):
    l.unfold(zf.le)
    _x = l.fix()
    l.intros()
    l.rw(elem_trans_omega, at=0)
    l.split(at=0)
    l.auto()


# trans_omega is inductive.
# trans_omega is inductive.
# Lean: inductive trans_omega
@kd.Theorem(inductive(trans_omega))
def trans_omega_inductive(l):
    l.unfold(inductive)
    l.split()
    l.rw(elem_trans_omega)
    l.split()
    l.auto(by=[emp_in_omega])
    l.auto(by=[emp_le_omega])
    _x = l.fix()
    l.intros()
    l.rw(elem_trans_omega)
    l.rw(elem_trans_omega, at=0)
    l.split(at=0)
    l.split()
    l.auto(by=[omega_succ(_x)])
    l.auto(by=[succ_le_omega(_x)])


# Omega equals the transitive closure subset trans_omega.
# Omega equals the transitive closure subset trans_omega.
# Lean: ω = trans_omega
@kd.Theorem(omega == trans_omega)
def omega_eq_trans_omega(l):
    l.have(trans_omega <= omega, by=[trans_omega_le_omega])
    l.have(inductive(trans_omega), by=[trans_omega_inductive])
    l.auto(by=[omega_induction(trans_omega)])


# Omega is transitive: elements are subsets.
# Omega is transitive: elements are subsets.
# Lean: x ∈ ω → x ⊆ ω
@kd.Theorem(smt.ForAll([x], smt.Implies(elem(x, omega), x <= omega)))
def omega_transitive(l):
    _x = l.fix()
    l.intros()
    l.rw(omega_eq_trans_omega, at=0)
    l.rw(elem_trans_omega, at=0)
    l.split(at=0)
    l.auto()


# Addition and multiplication (axiomatized recursion on omega)
add = smt.Function("add", Set, Set, Set)
mul = smt.Function("mul", Set, Set, Set)

add_zero_ax = kd.axiom(smt.ForAll([a], smt.Implies(is_nat(a), add(a, zero) == a)))
add_succ_ax = kd.axiom(
    smt.ForAll(
        [a, b],
        smt.Implies(smt.And(is_nat(a), is_nat(b)), add(a, succ(b)) == succ(add(a, b))),
    )
)

mul_zero_ax = kd.axiom(smt.ForAll([a], smt.Implies(is_nat(a), mul(a, zero) == zero)))
mul_succ_ax = kd.axiom(
    smt.ForAll(
        [a, b],
        smt.Implies(
            smt.And(is_nat(a), is_nat(b)), mul(a, succ(b)) == add(mul(a, b), a)
        ),
    )
)


# Addition with zero.
# Lean: a ∈ ω → add a 0 = a
@kd.Theorem(smt.ForAll([a], smt.Implies(is_nat(a), add(a, zero) == a)))
def add_zero(l):
    _a = l.fix()
    l.intros()
    l.auto(by=[add_zero_ax(_a)])


# Addition with successor on the right.
# Lean: a ∈ ω → b ∈ ω → add a (succ b) = succ (add a b)
@kd.Theorem(
    smt.ForAll(
        [a, b],
        smt.Implies(smt.And(is_nat(a), is_nat(b)), add(a, succ(b)) == succ(add(a, b))),
    )
)
def add_succ(l):
    _a, _b = l.fixes()
    l.intros()
    l.auto(by=[add_succ_ax(_a, _b)])


# Multiplication with zero.
# Lean: a ∈ ω → mul a 0 = 0
@kd.Theorem(smt.ForAll([a], smt.Implies(is_nat(a), mul(a, zero) == zero)))
def mul_zero(l):
    _a = l.fix()
    l.intros()
    l.auto(by=[mul_zero_ax(_a)])


# Multiplication with successor on the right.
# Lean: a ∈ ω → b ∈ ω → mul a (succ b) = add (mul a b) a
@kd.Theorem(
    smt.ForAll(
        [a, b],
        smt.Implies(smt.And(is_nat(a), is_nat(b)), mul(a, succ(b)) == add(mul(a, b), a)),
    )
)
def mul_succ(l):
    _a, _b = l.fixes()
    l.intros()
    l.auto(by=[mul_succ_ax(_a, _b)])


# 0 is an element of 1.
# Lean: 0 ∈ 1
@kd.Theorem(elem(zero, one))
def zero_in_one(l):
    l.rw(elem_succ)
    l.auto()


# 0 is an element of 2.
# Lean: 0 ∈ 2
@kd.Theorem(elem(zero, two))
def zero_in_two(l):
    l.rw(elem_succ)
    l.left()
    l.rw(elem_succ)
    l.auto()


# 1 is an element of 2.
# Lean: 1 ∈ 2
@kd.Theorem(elem(one, two))
def one_in_two(l):
    l.rw(elem_succ)
    l.auto()


print(omega_transitive)