-- Countable unions: a countable family of countable sets has a countable
-- union, coded through choice functions and the pairing of ℕ × ℕ
-- 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.Geometry.Absorption.CountableUnion
import Cardinal.CountableClosure
import Cardinal.Infinite
import Set.ChoiceFunctions

theorem Countable.of_injection (f A B : Set)
  requires injective : InjectiveBetween(f, A, B)
  requires countable : Countable(B)
  shows Countable(A)
proof
  obtain ⟨g, code⟩ := countable
  have composite := compose.injective with injective, code
  exact ⟨compose(g, f, A, B, ℕ), composite⟩

theorem Countable.of_finite (A : Set)
  requires finite : Finite(A)
  shows Countable(A)
proof
  obtain ⟨n, natural, equivalent⟩ := finite
  obtain ⟨f, bijective⟩ := equivalent
  have injective := Bijective.injective with bijective
  have function := Injective.function with injective
  have values : ∀ x ∈ A, f(x) ∈ ℕ
    intros x hx
    have valueIn := FunctionBetween.apply_closed with function, hx
    have typed := natural.member at n with f(x), valueIn
    exact nat.omega.forward with typed
  have widened := FunctionBetween.codomain_of_values with function, values
  have cancel : ∀ x ∈ A, ∀ y ∈ A, f(x) = f(y) → x = y
    intros x hx y hy same
    exact Injective.apply_cancel with injective, hx, hy, same
  have intoNaturals := Injective.of_pointwise with widened, cancel
  exact ⟨f, intoNaturals⟩

theorem Countable.image (f A B : Set)
  requires function : f : A → B
  requires countable : Countable(A)
  shows Countable(functionImage(f, A, B))
proof
  have surjection := functionImage.surjective with function
  exact Countable.of_surjection with countable, surjection

theorem Finite.preimage (f A B Y : Set)
  requires injective : InjectiveBetween(f, A, B)
  requires finite : Finite(Y)
  shows Finite(functionPreimage(f, A, Y))
proof
  have function := Injective.function with injective
  have subset := functionPreimage.subset at f, A, Y
  have restricted := restrict.injective with injective, subset
  have values : ∀ x ∈ functionPreimage(f, A, Y),
      restrict(f, functionPreimage(f, A, Y))(x) ∈ Y
    intros x hx
    have inY := functionPreimage.membership.forward with hx
    have value := restrict.domain_apply with function, subset, hx
    rw value
    exact inY.right
  have restrictedFunction := Injective.function with restricted
  have intoY := FunctionBetween.codomain_of_values
    with restrictedFunction, values
  have cancel : ∀ x ∈ functionPreimage(f, A, Y),
      ∀ y ∈ functionPreimage(f, A, Y),
      restrict(f, functionPreimage(f, A, Y))(x) =
        restrict(f, functionPreimage(f, A, Y))(y) → x = y
    intros x hx y hy same
    exact Injective.apply_cancel with restricted, hx, hy, same
  have intoYInjective := Injective.of_pointwise with intoY, cancel
  have dominated : functionPreimage(f, A, Y) ≼ Y :=
    Dominated.of_injective with intoYInjective
  exact Finite.of_dominated with finite, dominated

-- The injections of a set into ℕ: for A ⊆ X they are subsets of X × ℕ,
-- so a choice function on the subsets of 𝒫(X × ℕ) picks one for every
-- countable A ⊆ X.
def naturalInjections (A : Set) : Set :=
  { g ∈ 𝒫(A × ℕ) | InjectiveBetween(g, A, ℕ) }

theorem naturalInjections.nonempty (A : Set)
  requires countable : Countable(A)
  shows Nonempty(naturalInjections(A))
proof
  obtain ⟨g, injective⟩ := countable
  have function := Injective.function with injective
  have relation := FunctionBetween.relation with function
  unfold RelationBetween at relation
  have inPower := powerset.membership_iff.backward with relation
  witness g
  apply naturalInjections.membership.backward
  exact ⟨inPower, injective⟩

theorem naturalInjections.subset (A X : Set)
  requires subset : A ⊆ X
  shows naturalInjections(A) ⊆ 𝒫(X × ℕ)
proof
  intros g hg
  have data := naturalInjections.membership.forward with hg
  have same := Subset.reflexive at ℕ
  have pairs : A × ℕ ⊆ X × ℕ := product.monotone with subset, same
  have wider := powerset.monotone with pairs
  exact wider with g, data.left

-- The injection of A into ℕ that the choice function c picks.
def naturalCode (c A : Set) : Set := c(naturalInjections(A))

theorem naturalCode.injective (c A X : Set)
  requires choice : SubsetChoice(c, 𝒫(X × ℕ))
  requires subset : A ⊆ X
  requires countable : Countable(A)
  shows InjectiveBetween(naturalCode(c, A), A, ℕ)
proof
  have nonempty := naturalInjections.nonempty with countable
  have bounded := naturalInjections.subset with subset
  have member := SubsetChoice.member with choice, bounded, nonempty
  have data := naturalInjections.membership.forward with member
  unfold naturalCode
  exact data.right

-- The members of the family containing a point of the union, and the
-- one that a choice function d on the subsets of the family picks.
def containers (S x : Set) : Set := { A ∈ S | x ∈ A }
def container (d S x : Set) : Set := d(containers(S, x))

theorem containers.subset (S x : Set)
  shows containers(S, x) ⊆ S
proof
  intros A hA
  have data := containers.membership.forward with hA
  exact data.left

theorem container.graph (d S x : Set)
  shows container(d, S, x) = d(containers(S, x))
proof
  simp [container]

theorem container.member (d S x : Set)
  requires choice : SubsetChoice(d, S)
  requires hx : x ∈ ⋃(S)
  shows container(d, S, x) ∈ S ∧ x ∈ container(d, S, x)
proof
  obtain ⟨A, hA, inA⟩ := union.membership.forward with hx
  have inContainers := containers.membership.backward at S, x, A with ⟨hA, inA⟩
  have nonempty : Nonempty(containers(S, x)) := ⟨A, inContainers⟩
  have bounded := containers.subset at S, x
  have chosen := SubsetChoice.member with choice, bounded, nonempty
  rw container.graph
  exact containers.membership.forward with chosen

-- The code of a point of the union: the code of its container under an
-- injection g of the family into ℕ, paired with its code inside the
-- container under the injection that c picks.
defgraph unionCode (S c d g : Set) (x ∈ ⋃(S)) : ℕ × ℕ :=
  (g(container(d, S, x)), naturalCode(c, container(d, S, x))(x))
requires choiceC : SubsetChoice(c, 𝒫(⋃(S) × ℕ))
requires choiceD : SubsetChoice(d, S)
requires code : InjectiveBetween(g, S, ℕ)
requires members : ∀ A ∈ S, Countable(A)
proof
  obtain ⟨containerIn, xIn⟩ := container.member with choiceD, member
  have outer := Injective.function with code
  have bounded := union.subset_of_member with containerIn
  have countable := members with container(d, S, x), containerIn
  have innerInjective := naturalCode.injective with choiceC, bounded, countable
  have inner := Injective.function with innerInjective
  have first := FunctionBetween.apply_closed with outer, containerIn
  have second := FunctionBetween.apply_closed with inner, xIn
  exact product.pair with first, second

theorem unionCode.injective (S c d g : Set)
  requires choiceC : SubsetChoice(c, 𝒫(⋃(S) × ℕ))
  requires choiceD : SubsetChoice(d, S)
  requires code : InjectiveBetween(g, S, ℕ)
  requires members : ∀ A ∈ S, Countable(A)
  shows InjectiveBetween(unionCode(S, c, d, g), ⋃(S), ℕ × ℕ)
proof
  have function := unionCode.function with choiceC, choiceD, code, members
  apply Injective.of_pointwise with function
  intros x hx y hy same
  have valueX := unionCode.value with choiceC, choiceD, code, members, hx
  have valueY := unionCode.value with choiceC, choiceD, code, members, hy
  rw [valueX, valueY] at same
  obtain ⟨sameContainerCode, sameCode⟩ := orderedPair.injective with same
  obtain ⟨containerXIn, xIn⟩ := container.member with choiceD, hx
  obtain ⟨containerYIn, yIn⟩ := container.member with choiceD, hy
  have sameContainer := Injective.apply_cancel with
    code, containerXIn, containerYIn, sameContainerCode
  rw ← sameContainer at sameCode
  rw ← sameContainer at yIn
  have bounded := union.subset_of_member with containerXIn
  have countable := members with container(d, S, x), containerXIn
  have inner := naturalCode.injective with choiceC, bounded, countable
  exact Injective.apply_cancel with inner, xIn, yIn, sameCode

-- A countable union of countable sets is countable: an injection of the
-- family into ℕ and two choice functions code every point of the union.
theorem Countable.union (S : Set)
  requires countable : Countable(S)
  requires members : ∀ A ∈ S, Countable(A)
  shows Countable(⋃(S))
proof
  obtain ⟨g, code⟩ := countable
  obtain ⟨c, choiceC⟩ := SubsetChoice.exists at 𝒫(⋃(S) × ℕ)
  obtain ⟨d, choiceD⟩ := SubsetChoice.exists at S
  have injective := unionCode.injective with choiceC, choiceD, code, members
  have pairs : Countable(ℕ × ℕ) :=
    ⟨naturalPairing, naturalPairing.injective⟩
  exact Countable.of_injection with injective, pairs

theorem Countable.union2 (A B : Set)
  requires countableA : Countable(A)
  requires countableB : Countable(B)
  shows Countable(A ∪ B)
proof
  have singletonFinite := Finite.singleton at B
  have insertFinite := Finite.insert at {B}, A with singletonFinite
  have same : insert(A, {B}) = {A, B}
    ext y
    simp
  rw same at insertFinite
  have countablePair := Countable.of_finite with insertFinite
  have members : ∀ C ∈ {A, B}, Countable(C)
    intros C hC
    have cases : C = A ∨ C = B
      simp at hC
      exact hC
    cases cases with isA | isB
    . rw isA
      exact countableA
    . rw isB
      exact countableB
  have union := Countable.union with countablePair, members
  rw union.pair at union
  exact union

-- The union of a countable family indexed by a function.
theorem Countable.indexed_union (F I X : Set)
  requires function : F : I → 𝒫(X)
  requires countable : Countable(I)
  requires members : ∀ i ∈ I, Countable(F(i))
  shows Countable(⋃(functionImage(F, I, 𝒫(X))))
proof
  have familyCountable := Countable.image with function, countable
  have each : ∀ A ∈ functionImage(F, I, 𝒫(X)), Countable(A)
    intros A hA
    obtain ⟨bounded, i, hi, shape⟩ :=
      functionImage.membership.forward with hA
    rw shape
    exact members with i, hi
  exact Countable.union with familyCountable, each