-- Ordinal arithmetic: addition, multiplication, and exponentiation by
-- transfinite recursion, with their recursion laws and closure
-- Part of Cumulus, the Grace17 mathematics library.
-- Copyright (c) 2026 Grant Passmore.
-- Released under the Apache License 2.0. See LICENSE in the repository root.
-- Authors: Grant Passmore

module Grace.Cumulus.Set.Ordinal.Arithmetic
import Basic
import Relation.Algebra

noncomputable transfinite def ordinalAdd (α : Ordinal, β : Ordinal) : Set
  | 0 => α
  | succ γ => next(ordinalAdd(α, γ))
  | limit λ => ⋃(Relation.range(history))

noncomputable transfinite def ordinalMul (α : Ordinal, β : Ordinal) : Set
  | 0 => ∅
  | succ γ => ordinalAdd(ordinalMul(α, γ), α)
  | limit λ => ⋃(Relation.range(history))

noncomputable transfinite def ordinalPow (α : Ordinal, β : Ordinal) : Set
  | 0 => 1
  | succ γ => ordinalMul(ordinalPow(α, γ), α)
  | limit λ => ⋃(Relation.range(history))

notation infixl 60 "+" => ordinalAdd for Ordinal
notation infixl 70 "*" => ordinalMul for Ordinal
notation infixr 80 "^" => ordinalPow for Ordinal

-- Addition

theorem ordinalAdd.zero (α : Ordinal)
  shows α + 0 = α
proof
  simp [ordinalAdd, nat.zero]

theorem ordinalAdd.next (α β : Ordinal)
  shows α + next(β) = next(α + β)
proof
  simp [ordinalAdd.equation.successor]

theorem ordinalAdd.limit (α λ : Ordinal)
  requires h : Limit(λ)
  shows α + λ = ⋃(Relation.range(ordinalAdd.history(α, λ)))
proof
  simp [ordinalAdd.equation.limit]

theorem ordinalAdd.history_membership (p α λ : Set)
  shows p ∈ ordinalAdd.history(α, λ) ↔ ∃ β ∈ λ, p = (β ↦ ordinalAdd(α, β))
proof
  simp [ordinalAdd.history.membership]

theorem ordinalAdd.history_range (y α λ : Set)
  shows y ∈ Relation.range(ordinalAdd.history(α, λ)) ↔ ∃ β ∈ λ, y = ordinalAdd(α, β)
proof
  constructor
  . intro hy
    obtain ⟨β, hpair⟩ := range.membership_iff.forward with hy
    have member : (β, y) ∈ ordinalAdd.history(α, λ) := hpair
    simp [ordinalAdd.history_membership] at member
    obtain ⟨γ, hγ, shape⟩ := member
    have coordinates := orderedPair.injective with shape
    witness γ
    constructor
    . exact hγ
    . exact coordinates.right
  . intro h
    obtain ⟨β, hβ, shape⟩ := h
    have member : (β ↦ ordinalAdd(α, β)) ∈ ordinalAdd.history(α, λ) by simp [ordinalAdd.history_membership]
    rw shape
    exact range.contains at β, ordinalAdd(α, β), ordinalAdd.history(α, λ) with member

theorem ordinalAdd.limit_membership (x : Set, α λ : Ordinal)
  requires limit : Limit(λ)
  shows x ∈ α + λ ↔ ∃ β ∈ λ, x ∈ ordinalAdd(α, β)
proof
  have equation := ordinalAdd.limit at α, λ with limit
  rw equation
  constructor
  . intro hx
    simp [set.union] at hx
    obtain ⟨stage, hStage, hxStage⟩ := hx
    obtain ⟨β, hβ, shape⟩ := ordinalAdd.history_range.forward with hStage
    witness β
    constructor
    . exact hβ
    . rw ← shape
      exact hxStage
  . intro h
    obtain ⟨β, hβ, hx⟩ := h
    have inRange : ordinalAdd(α, β) ∈ Relation.range(ordinalAdd.history(α, λ))
      apply ordinalAdd.history_range.backward
      exact ⟨β, hβ, rfl⟩
    exact union.contains with inRange, hx

theorem ordinalAdd.ordinal (α : Ordinal, β : Set)
  shows β ∈ Ordinal → ordinalAdd(α, β) ∈ Ordinal
proof
  induction β using membership
  intros current ih hCurrent
  have cases := Ordinal.cases with hCurrent
  cases cases with zero | rest
  . rw zero
    rw ← nat.zero
    rw ordinalAdd.zero
    assumption
  . cases rest with successor | limit
    . obtain ⟨γ, hγ, shape⟩ := successor
      rw ← shape
      have γOrdinal := Ordinal.member with hCurrent, hγ
      rw ordinalAdd.next
      exact Ordinal.next with (ih with hγ, γOrdinal)
    . have equation := ordinalAdd.limit at α, current with limit
      rw equation
      apply Ordinal.union
      intros y hy
      obtain ⟨stage, hStage, shape⟩ := ordinalAdd.history_range.forward with hy
      rw shape
      exact ih with hStage, (Ordinal.member with hCurrent, hStage)

theorem ordinalAdd.one (α : Ordinal)
  shows α + 1 = next(α)
proof
  have one : 1 = next(0) by simp [next.natural]
  rw [one, ordinalAdd.next, ordinalAdd.zero]
  rfl

theorem ordinalAdd.zero_left (β : Set)
  shows β ∈ Ordinal → ordinalAdd(0, β) = β
proof
  induction β using membership
  intros current ih hCurrent
  have cases := Ordinal.cases with hCurrent
  cases cases with zero | rest
  . rw zero
    rw ← nat.zero
    exact ordinalAdd.zero at 0
  . cases rest with successor | limit
    . obtain ⟨γ, hγ, shape⟩ := successor
      rw ← shape
      have γOrdinal := Ordinal.member with hCurrent, hγ
      have step := ih with hγ, γOrdinal
      rw ordinalAdd.next
      rw step
      rfl
    . apply Subset.antisymmetric
      . intros x hx
        obtain ⟨stage, hStage, hxStage⟩ := ordinalAdd.limit_membership.forward at x, 0, current with limit, hx
        have step := ih with hStage, (Ordinal.member with hCurrent, hStage)
        rw step at hxStage
        exact Ordinal.member_subset at stage, current with hCurrent, hStage, x, hxStage
      . intros x hx
        apply ordinalAdd.limit_membership.backward at x, 0, current with limit
        witness next(x)
        have xOrdinal := Ordinal.member with hCurrent, hx
        constructor
        . exact limit.right.right with x, hx
        . have step := ih with (limit.right.right with x, hx), (Ordinal.next with xOrdinal)
          rw step
          exact next.contains at x

theorem ordinalAdd.subset_left (α : Ordinal, β : Set)
  shows β ∈ Ordinal → α ⊆ ordinalAdd(α, β)
proof
  induction β using membership
  intros current ih hCurrent
  have cases := Ordinal.cases with hCurrent
  cases cases with zero | rest
  . rw zero
    rw ← nat.zero
    rw ordinalAdd.zero
    exact Subset.reflexive at α
  . cases rest with successor | limit
    . obtain ⟨γ, hγ, shape⟩ := successor
      rw ← shape
      have γOrdinal := Ordinal.member with hCurrent, hγ
      rw ordinalAdd.next
      have step := ih with hγ, γOrdinal
      intros x hx
      apply next.membership.backward
      left
      exact step with x, hx
    . intros x hx
      apply ordinalAdd.limit_membership.backward at x, α, current with limit
      have nonempty := Nonempty.of_ne_empty with limit.right.left
      simp [Nonempty] at nonempty
      obtain ⟨stage, hStage⟩ := nonempty
      witness stage
      constructor
      . exact hStage
      . exact ih with hStage, (Ordinal.member with hCurrent, hStage), x, hx

theorem ordinalAdd.monotone (α β : Ordinal, γ : Set)
  shows γ ∈ Ordinal → β ∈ γ → ordinalAdd(α, β) ∈ ordinalAdd(α, γ)
proof
  induction γ using membership
  intros current ih hCurrent hBelow
  have cases := Ordinal.cases with hCurrent
  cases cases with zero | rest
  . exfalso
    rw zero at hBelow
    exact set.empty at β with hBelow
  . cases rest with successor | limit
    . obtain ⟨δ, hδ, shape⟩ := successor
      rw ← shape
      have δOrdinal := Ordinal.member with hCurrent, hδ
      rw ordinalAdd.next
      rw ← shape at hBelow
      simp [next.membership] at hBelow
      apply next.membership.backward
      cases hBelow with earlier | equal
      . left
        exact ih with hδ, δOrdinal, earlier
      . right
        rw equal
        rfl
    . apply ordinalAdd.limit_membership.backward at ordinalAdd(α, β), α, current with limit
      witness next(β)
      constructor
      . exact limit.right.right with β, hBelow
      . rw ordinalAdd.next
        exact next.contains at α + β

-- Multiplication

theorem ordinalMul.zero (α : Ordinal)
  shows α * 0 = ∅
proof
  simp [ordinalMul, nat.zero]

theorem ordinalMul.next (α β : Ordinal)
  shows α * next(β) = α * β + α
proof
  simp [ordinalMul.equation.successor]

theorem ordinalMul.limit (α λ : Ordinal)
  requires h : Limit(λ)
  shows α * λ = ⋃(Relation.range(ordinalMul.history(α, λ)))
proof
  simp [ordinalMul.equation.limit]

theorem ordinalMul.history_range (y α λ : Set)
  shows y ∈ Relation.range(ordinalMul.history(α, λ)) ↔ ∃ β ∈ λ, y = ordinalMul(α, β)
proof
  constructor
  . intro hy
    obtain ⟨β, hpair⟩ := range.membership_iff.forward with hy
    have member : (β, y) ∈ ordinalMul.history(α, λ) := hpair
    simp [ordinalMul.history.membership] at member
    obtain ⟨γ, hγ, shape⟩ := member
    have coordinates := orderedPair.injective with shape
    witness γ
    constructor
    . exact hγ
    . exact coordinates.right
  . intro h
    obtain ⟨β, hβ, shape⟩ := h
    have member : (β ↦ ordinalMul(α, β)) ∈ ordinalMul.history(α, λ) by simp [ordinalMul.history.membership]
    rw shape
    exact range.contains at β, ordinalMul(α, β), ordinalMul.history(α, λ) with member

theorem ordinalMul.ordinal (α : Ordinal, β : Set)
  shows β ∈ Ordinal → ordinalMul(α, β) ∈ Ordinal
proof
  induction β using membership
  intros current ih hCurrent
  have cases := Ordinal.cases with hCurrent
  cases cases with zero | rest
  . rw zero
    rw ← nat.zero
    rw ordinalMul.zero
    exact Ordinal.empty
  . cases rest with successor | limit
    . obtain ⟨γ, hγ, shape⟩ := successor
      rw ← shape
      have γOrdinal := Ordinal.member with hCurrent, hγ
      rw ordinalMul.next
      have product := ih with hγ, γOrdinal
      exact ordinalAdd.ordinal at α * γ, α with product
    . have equation := ordinalMul.limit at α, current with limit
      rw equation
      apply Ordinal.union
      intros y hy
      obtain ⟨stage, hStage, shape⟩ := ordinalMul.history_range.forward with hy
      rw shape
      exact ih with hStage, (Ordinal.member with hCurrent, hStage)

theorem ordinalMul.one (α : Ordinal)
  shows α * 1 = α
proof
  have one : 1 = next(0) by simp [next.natural]
  rw [one, ordinalMul.next, ordinalMul.zero]
  rw ← nat.zero
  exact ordinalAdd.zero_left at α with assumption

-- Exponentiation

theorem ordinalPow.zero (α : Ordinal)
  shows α ^ 0 = 1
proof
  simp [ordinalPow, nat.zero]

theorem ordinalPow.next (α β : Ordinal)
  shows α ^ next(β) = α ^ β * α
proof
  simp [ordinalPow.equation.successor]

theorem ordinalPow.limit (α λ : Ordinal)
  requires h : Limit(λ)
  shows α ^ λ = ⋃(Relation.range(ordinalPow.history(α, λ)))
proof
  simp [ordinalPow.equation.limit]

theorem ordinalPow.history_range (y α λ : Set)
  shows y ∈ Relation.range(ordinalPow.history(α, λ)) ↔ ∃ β ∈ λ, y = ordinalPow(α, β)
proof
  constructor
  . intro hy
    obtain ⟨β, hpair⟩ := range.membership_iff.forward with hy
    have member : (β, y) ∈ ordinalPow.history(α, λ) := hpair
    simp [ordinalPow.history.membership] at member
    obtain ⟨γ, hγ, shape⟩ := member
    have coordinates := orderedPair.injective with shape
    witness γ
    constructor
    . exact hγ
    . exact coordinates.right
  . intro h
    obtain ⟨β, hβ, shape⟩ := h
    have member : (β ↦ ordinalPow(α, β)) ∈ ordinalPow.history(α, λ) by simp [ordinalPow.history.membership]
    rw shape
    exact range.contains at β, ordinalPow(α, β), ordinalPow.history(α, λ) with member

theorem ordinalPow.ordinal (α : Ordinal, β : Set)
  shows β ∈ Ordinal → ordinalPow(α, β) ∈ Ordinal
proof
  induction β using membership
  intros current ih hCurrent
  have cases := Ordinal.cases with hCurrent
  cases cases with zero | rest
  . rw zero
    rw ← nat.zero
    rw ordinalPow.zero
    exact Ordinal.natural at 1
  . cases rest with successor | limit
    . obtain ⟨γ, hγ, shape⟩ := successor
      rw ← shape
      have γOrdinal := Ordinal.member with hCurrent, hγ
      rw ordinalPow.next
      have power := ih with hγ, γOrdinal
      exact ordinalMul.ordinal at α ^ γ, α with power
    . have equation := ordinalPow.limit at α, current with limit
      rw equation
      apply Ordinal.union
      intros y hy
      obtain ⟨stage, hStage, shape⟩ := ordinalPow.history_range.forward with hy
      rw shape
      exact ih with hStage, (Ordinal.member with hCurrent, hStage)

theorem ordinalMul.one_left (β : Set)
  shows β ∈ Ordinal → ordinalMul(1, β) = β
proof
  induction β using membership
  intros current ih hCurrent
  have cases := Ordinal.cases with hCurrent
  cases cases with zero | rest
  . rw zero
    have base := ordinalMul.zero at 1
    rw nat.zero at base
    exact base
  . cases rest with successor | limit
    . obtain ⟨γ, hγ, shape⟩ := successor
      rw ← shape
      have γOrdinal := Ordinal.member with hCurrent, hγ
      have step := ih with hγ, γOrdinal
      rw [ordinalMul.next, step]
      exact ordinalAdd.one at γ
    . have oneOrdinal := Ordinal.natural at 1
      apply Subset.antisymmetric
      . intros x hx
        have equation := ordinalMul.limit at 1, current with limit
        rw equation at hx
        simp [set.union] at hx
        obtain ⟨stage, hStage, hxStage⟩ := hx
        obtain ⟨earlier, hEarlier, shape⟩ := ordinalMul.history_range.forward with hStage
        have step := ih with hEarlier, (Ordinal.member with hCurrent, hEarlier)
        rw [shape, step] at hxStage
        exact Ordinal.member_subset at earlier, current with hCurrent, hEarlier, x, hxStage
      . intros x hx
        have equation := ordinalMul.limit at 1, current with limit
        rw equation
        have xOrdinal := Ordinal.member with hCurrent, hx
        have nextIn := limit.right.right with x, hx
        have step := ih with nextIn, (Ordinal.next with xOrdinal)
        have inRange : ordinalMul(1, next(x)) ∈ Relation.range(ordinalMul.history(1, current))
          apply ordinalMul.history_range.backward
          exact ⟨next(x), nextIn, rfl⟩
        rw step at inRange
        exact union.contains with inRange, (next.contains at x)

theorem ordinalPow.one (α : Ordinal)
  shows α ^ 1 = α
proof
  have one : 1 = next(0) by simp [next.natural]
  rw [one, ordinalPow.next, ordinalPow.zero]
  exact ordinalMul.one_left at α