-- Infinite cardinal arithmetic: an infinite set is equipotent with its own
-- square, hence with its union or product with any set it dominates. Among
-- the bijections from the square of an infinite subset onto that subset, a
-- maximal one, which exists by the maximal principle since such bijections
-- are closed under nested unions, already covers the whole set.
-- 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.Cardinal.Square
import Algebra
import Set.MaximalPrinciple

-- Bijections from the square of an infinite subset of A onto that subset.
def squareBijections (A : Set) : Set :=
  { f ∈ 𝒫((A × A) × A) | ∃ X ∈ 𝒫(A), Infinite(X) ∧ BijectiveBetween(f, X × X, X) }

theorem squareBijections.of_bijective (A f X : Set)
  requires subset : X ⊆ A
  requires infinite : Infinite(X)
  requires bijective : BijectiveBetween(f, X × X, X)
  shows f ∈ squareBijections(A)
proof
  apply squareBijections.membership.backward
  have relation := FunctionBetween.relation with (Injective.function with (Bijective.injective with bijective))
  have fSubset : f ⊆ (X × X) × X := relation
  have squareInside := product.monotone with subset, subset
  have bound := product.monotone with squareInside, subset
  constructor
  . exact powerset.membership_iff.backward with (Subset.transitive with fSubset, bound)
  . refine ⟨X, (powerset.membership_iff.backward with subset), infinite, bijective⟩

-- The subset is the range of the bijection.
theorem squareBijections.range (f A : Set)
  requires hf : f ∈ squareBijections(A)
  shows Infinite(range(f)) ∧ BijectiveBetween(f, range(f) × range(f), range(f)) ∧ range(f) ⊆ A
proof
  obtain ⟨bound, X, hX, infiniteX, bijective⟩ := squareBijections.membership.forward with hf
  have function := Injective.function with (Bijective.injective with bijective)
  have surjective := Bijective.surjective with bijective
  have equal : range(f) = X
    apply Subset.antisymmetric
    . exact range.subset with (FunctionBetween.relation with function)
    . intros x hx
      obtain ⟨p, hp, related⟩ := surjective.right with x, hx
      have member : (p, x) ∈ f := related
      exact range.contains with member
  rw equal
  have inside : X ⊆ A := powerset.membership_iff.forward with hX
  exact ⟨infiniteX, bijective, inside⟩

-- A countable infinite subset seeds the family.
theorem squareBijections.nonempty (A : Set)
  requires infinite : Infinite(A)
  shows ∃ f, f ∈ squareBijections(A)
proof
  obtain ⟨g, injective⟩ := Infinite.omega_dominated with infinite
  have equipotent := Equipotent.image with injective
  have squareEquipotent := Equipotent.product with equipotent, equipotent
  have chain := Equipotent.transitive with (Equipotent.symmetric with squareEquipotent), omega.square
  obtain ⟨f, bijective⟩ := Equipotent.transitive with chain, equipotent
  witness f
  have inside := functionImage.subset at g, ℕ, A
  have infiniteImage : Infinite(functionImage(g, ℕ, A))
    apply Infinite.iff_omega_dominated.backward
    exact Dominated.of_equipotent with equipotent
  exact squareBijections.of_bijective with inside, infiniteImage, bijective

-- Nested unions stay in the family: the union is a bijection from the
-- square of the union of the ranges onto that union.
theorem squareBijections.union_closed (A S : Set)
  requires inhabited : Nonempty(S)
  requires inside : S ⊆ squareBijections(A)
  requires nested : ∀ C ∈ S, ∀ D ∈ S, C ⊆ D ∨ D ⊆ C
  shows ⋃(S) ∈ squareBijections(A)
proof
  have bases : ∀ f ∈ S, Infinite(range(f)) ∧ BijectiveBetween(f, range(f) × range(f), range(f)) ∧ range(f) ⊆ A
    intros f hf
    exact squareBijections.range with (inside with f, hf)
  have common : ∀ f ∈ S, ∀ g ∈ S, ∃ h ∈ S, f ⊆ h ∧ g ⊆ h
    intros f hf g hg
    have split := nested with f, hf, g, hg
    cases split with fg | gf
    . refine ⟨g, hg, fg, (Subset.reflexive at g)⟩
    . refine ⟨f, hf, (Subset.reflexive at f), gf⟩
  have rangeInside : ∀ f ∈ S, range(f) ⊆ range(⋃(S))
    intros f hf
    exact range.monotone with (union.subset_of_member with hf)
  have bounds : ∀ p ∈ ⋃(S), ∃ f ∈ S, p ∈ f
    intros p hp
    exact union.membership.forward with hp
  have unionRelation : ⋃(S) ⊆ (range(⋃(S)) × range(⋃(S))) × range(⋃(S))
    intros p hp
    obtain ⟨f, hf, pf⟩ := bounds with p, hp
    obtain ⟨infiniteF, bijectiveF, rangeF⟩ := bases with f, hf
    have relationF := FunctionBetween.relation with (Injective.function with (Bijective.injective with bijectiveF))
    have fSubset : f ⊆ (range(f) × range(f)) × range(f) := relationF
    have grow := rangeInside with f, hf
    have squareGrow := product.monotone with grow, grow
    have whole := product.monotone with squareGrow, grow
    exact whole with p, (fSubset with p, pf)
  have function : ⋃(S) : range(⋃(S)) × range(⋃(S)) → range(⋃(S))
    constructor
    . exact unionRelation
    . constructor
      . intros p y hpy z hpz
        have memberY : (p, y) ∈ ⋃(S) := hpy
        have memberZ : (p, z) ∈ ⋃(S) := hpz
        obtain ⟨f, hf, pf⟩ := bounds with (p, y), memberY
        obtain ⟨g, hg, pg⟩ := bounds with (p, z), memberZ
        obtain ⟨h, hh, fh, gh⟩ := common with f, hf, g, hg
        obtain ⟨infiniteH, bijectiveH, rangeH⟩ := bases with h, hh
        have functionalH := FunctionBetween.functional with (Injective.function with (Bijective.injective with bijectiveH))
        have relY : relates(h, p, y) := fh with (p, y), pf
        have relZ : relates(h, p, z) := gh with (p, z), pg
        exact functionalH with p, y, relY, z, relZ
      . intros p hp
        obtain ⟨bound, x, hx, y, hy, shape⟩ := product.membership.forward with hp
        obtain ⟨qx, relatedX⟩ := range.membership_iff.forward with hx
        obtain ⟨qy, relatedY⟩ := range.membership_iff.forward with hy
        have memberX : (qx, x) ∈ ⋃(S) := relatedX
        have memberY : (qy, y) ∈ ⋃(S) := relatedY
        obtain ⟨f, hf, pf⟩ := bounds with (qx, x), memberX
        obtain ⟨g, hg, pg⟩ := bounds with (qy, y), memberY
        obtain ⟨h, hh, fh, gh⟩ := common with f, hf, g, hg
        obtain ⟨infiniteH, bijectiveH, rangeH⟩ := bases with h, hh
        have functionH := Injective.function with (Bijective.injective with bijectiveH)
        have xInH := range.contains with (fh with (qx, x), pf)
        have yInH := range.contains with (gh with (qy, y), pg)
        have pairIn : (x, y) ∈ range(h) × range(h) := product.pair with xInH, yInH
        obtain ⟨z, hz, related⟩ := FunctionBetween.exists_value with functionH, pairIn
        refine ⟨z, (rangeInside with h, hh, z, hz), ?_⟩
        rw shape
        have memberZ : ((x, y), z) ∈ h := related
        exact union.contains with hh, memberZ
  apply squareBijections.of_bijective at A, ⋃(S), range(⋃(S))
  . intros x hx
    obtain ⟨q, related⟩ := range.membership_iff.forward with hx
    have member : (q, x) ∈ ⋃(S) := related
    obtain ⟨f, hf, pf⟩ := bounds with (q, x), member
    obtain ⟨infiniteF, bijectiveF, rangeF⟩ := bases with f, hf
    have inRange := range.contains with pf
    exact rangeF with x, inRange
  . obtain ⟨f, hf⟩ := inhabited
    obtain ⟨infiniteF, bijectiveF, rangeF⟩ := bases with f, hf
    exact Infinite.of_subset with infiniteF, (rangeInside with f, hf)
  . constructor
    . constructor
      . exact function
      . intros p hp q hq z hz hpz hqz
        have memberP : (p, z) ∈ ⋃(S) := hpz
        have memberQ : (q, z) ∈ ⋃(S) := hqz
        obtain ⟨f, hf, pf⟩ := bounds with (p, z), memberP
        obtain ⟨g, hg, pg⟩ := bounds with (q, z), memberQ
        obtain ⟨h, hh, fh, gh⟩ := common with f, hf, g, hg
        obtain ⟨infiniteH, bijectiveH, rangeH⟩ := bases with h, hh
        have injectiveH := Bijective.injective with bijectiveH
        have functionH := Injective.function with injectiveH
        have relP : relates(h, p, z) := fh with (p, z), pf
        have relQ : relates(h, q, z) := gh with (q, z), pg
        have boundsP := FunctionBetween.bounds with functionH, relP
        have boundsQ := FunctionBetween.bounds with functionH, relQ
        exact Injective.cancel with injectiveH, boundsP.left, boundsQ.left, boundsP.right, relP, relQ
    . constructor
      . exact function
      . intros z hz
        obtain ⟨p, related⟩ := range.membership_iff.forward with hz
        have member : (p, z) ∈ ⋃(S) := related
        have pIn := unionRelation with (p, z), member
        have coordinates := product.pair_membership.forward with pIn
        refine ⟨p, coordinates.left, related⟩

-- The square of a union splits into the square of the first set and the
-- three mixed products.
theorem product.union2_square (X Y : Set)
  shows (X ∪ Y) × (X ∪ Y) = (X × X) ∪ ((X × Y) ∪ ((Y × X) ∪ (Y × Y)))
proof
  rw [product.union2_left, product.union2_right, product.union2_right, union2.associative]
  rfl

-- With X infinite, X × X ≃ X, and Y a disjoint copy of X, the three mixed
-- products together are a copy of Y.
theorem squareBijections.rest_equipotent (X Y : Set)
  requires infinite : Infinite(X)
  requires square : (X × X) ≃ X
  requires equipotent : X ≃ Y
  requires disjoint : Disjoint(X, Y)
  shows ((X × Y) ∪ ((Y × X) ∪ (Y × Y))) ≃ Y
proof
  obtain ⟨a, ha, b, hb, distinct⟩ := Infinite.two_points with infinite
  have yDisjoint := Disjoint.symmetric with disjoint
  have squareDominated := Dominated.of_equipotent with square
  have yToX := Dominated.of_equipotent with (Equipotent.symmetric with equipotent)
  have xToY := Dominated.of_equipotent with equipotent
  have xSelf := Dominated.reflexive at X
  have xy := Dominated.transitive with (Dominated.product with xSelf, yToX), squareDominated
  have yx := Dominated.transitive with (Dominated.product with yToX, xSelf), squareDominated
  have yy := Dominated.transitive with (Dominated.product with yToX, yToX), squareDominated
  have innerDisjoint : Disjoint(Y × X, Y × Y) := product.disjoint_of_right with disjoint
  have inner := Dominated.union_of_two_points with yx, yy, ha, hb, distinct, innerDisjoint
  have innerX := Dominated.transitive with inner, squareDominated
  have outerDisjoint : Disjoint(X × Y, (Y × X) ∪ (Y × Y))
    apply Disjoint.symmetric
    apply Disjoint.union2_iff.backward
    constructor
    . exact product.disjoint_of_left with yDisjoint
    . exact product.disjoint_of_left with yDisjoint
  have rest := Dominated.union_of_two_points with xy, innerX, ha, hb, distinct, outerDisjoint
  have restY := Dominated.transitive with (Dominated.transitive with rest, squareDominated), xToY
  obtain ⟨g, gBijective⟩ := equipotent
  have gFunction := Injective.function with (Bijective.injective with gBijective)
  have point := FunctionBetween.apply_closed with gFunction, ha
  have inhabited : Nonempty(Y) := ⟨g(a), point⟩
  have up := Dominated.left_factor at Y, Y with inhabited
  have step1 := union2.subset_right at (Y × X), (Y × Y)
  have step2 := union2.subset_right at (X × Y), ((Y × X) ∪ (Y × Y))
  have inside := Subset.transitive with step1, step2
  have yRest := Dominated.transitive with up, (Dominated.of_subset with inside)
  exact Dominated.antisymmetric with restY, yRest

-- A disjoint copy of the range extends a member of the family properly.
theorem squareBijections.extend (A M Y : Set)
  requires hM : M ∈ squareBijections(A)
  requires hY : Y ⊆ A
  requires disjoint : Disjoint(range(M), Y)
  requires equipotent : range(M) ≃ Y
  shows ∃ J ∈ squareBijections(A), M ⊆ J ∧ J ≠ M
proof
  obtain ⟨infiniteX, bijective, rangeInside⟩ := squareBijections.range with hM
  have squareEquipotent : (range(M) × range(M)) ≃ range(M) := ⟨M, bijective⟩
  obtain ⟨k, kBijective⟩ := squareBijections.rest_equipotent with infiniteX, squareEquipotent, equipotent, disjoint
  have yDisjoint := Disjoint.symmetric with disjoint
  have domainsDisjoint : Disjoint(range(M) × range(M), (range(M) × Y) ∪ ((Y × range(M)) ∪ (Y × Y)))
    apply Disjoint.symmetric
    apply Disjoint.union2_iff.backward
    constructor
    . exact product.disjoint_of_right with yDisjoint
    . apply Disjoint.union2_iff.backward
      constructor
      . exact product.disjoint_of_left with yDisjoint
      . exact product.disjoint_of_left with yDisjoint
  have jBijective := Bijective.union_disjoint with bijective, kBijective, domainsDisjoint, disjoint
  have square := product.union2_square at range(M), Y
  rw ← square at jBijective
  have unionInside := union2.subset with rangeInside, hY
  have infiniteUnion := Infinite.of_subset with infiniteX, (union2.subset_left at range(M), Y)
  have hJ := squareBijections.of_bijective with unionInside, infiniteUnion, jBijective
  refine ⟨M ∪ k, hJ, (union2.subset_left at M, k), ?_⟩
  intro equal
  obtain ⟨a, ha, b, hb, distinct⟩ := Infinite.two_points with infiniteX
  obtain ⟨g, gBijective⟩ := equipotent
  have gFunction := Injective.function with (Bijective.injective with gBijective)
  have point := FunctionBetween.apply_closed with gFunction, ha
  have kSurjective := Bijective.surjective with kBijective
  obtain ⟨p, hp, related⟩ := kSurjective.right with g(a), point
  have inK : (p, g(a)) ∈ k := related
  have inJ := union2.right at (p, g(a)), M, k with inK
  rw equal at inJ
  have relationM := FunctionBetween.relation with (Injective.function with (Bijective.injective with bijective))
  have mSubset : M ⊆ (range(M) × range(M)) × range(M) := relationM
  have inProduct := mSubset with (p, g(a)), inJ
  have coordinates := product.pair_membership.forward with inProduct
  have separated := Disjoint.iff.forward with disjoint
  exact separated with g(a), coordinates.right, point

-- A maximal member of the family covers A: either the complement of the
-- range is dominated by the range, and then A is a copy of the range, or
-- the range injects into its complement, which yields a proper extension.
theorem squareBijections.maximal_full (A M : Set)
  requires hM : M ∈ squareBijections(A)
  requires maximal : ∀ J ∈ squareBijections(A), M ⊆ J → J = M
  shows A ≃ range(M)
proof
  obtain ⟨infiniteX, bijective, rangeInside⟩ := squareBijections.range with hM
  have squareEquipotent : (range(M) × range(M)) ≃ range(M) := ⟨M, bijective⟩
  have comparison := Dominated.total at A ∖ range(M), range(M)
  cases comparison with small | large
  . obtain ⟨a, ha, b, hb, distinct⟩ := Infinite.two_points with infiniteX
    have cover := difference.union2_of_subset with rangeInside
    have disjoint := Disjoint.symmetric with (difference.disjoint at A, range(M))
    have unionDominated := Dominated.union_of_two_points with (Dominated.reflexive at range(M)), small, ha, hb, distinct, disjoint
    rw cover at unionDominated
    have down := Dominated.transitive with unionDominated, (Dominated.of_equipotent with squareEquipotent)
    exact Dominated.antisymmetric with down, (Dominated.of_subset with rangeInside)
  . exfalso
    obtain ⟨h, hInjective⟩ := large
    have yInside := functionImage.subset at h, range(M), A ∖ range(M)
    have yInA := Subset.transitive with yInside, (difference.subset_left at A, range(M))
    have disjoint := Disjoint.symmetric with (Disjoint.of_subset_left with (difference.disjoint at A, range(M)), yInside)
    have equipotent := Equipotent.image with hInjective
    obtain ⟨J, hJ, extends, different⟩ := squareBijections.extend with hM, yInA, disjoint, equipotent
    exact different with (maximal with J, hJ, extends)

theorem Infinite.square (A : Set)
  requires infinite : Infinite(A)
  shows (A × A) ≃ A
proof
  obtain ⟨f, seed⟩ := squareBijections.nonempty with infinite
  have closed : ∀ S, Nonempty(S) → S ⊆ squareBijections(A) →
      (∀ C ∈ S, ∀ D ∈ S, C ⊆ D ∨ D ⊆ C) → ⋃(S) ∈ squareBijections(A)
    intros S inhabited inside nested
    exact squareBijections.union_closed with inhabited, inside, nested
  obtain ⟨M, hM, extends, maximal⟩ := maximal_extension with seed, closed
  have full := squareBijections.maximal_full with hM, maximal
  obtain ⟨infiniteX, bijective, rangeInside⟩ := squareBijections.range with hM
  have squareEquipotent : (range(M) × range(M)) ≃ range(M) := ⟨M, bijective⟩
  have squares := Equipotent.product with full, full
  have step := Equipotent.transitive with squares, squareEquipotent
  exact Equipotent.transitive with step, (Equipotent.symmetric with full)

-- Consequences: an infinite set absorbs any set it dominates, by union or
-- by product.
theorem Infinite.union (A B : Set)
  requires infinite : Infinite(A)
  requires dominated : B ≼ A
  shows (A ∪ B) ≃ A
proof
  obtain ⟨a, ha, b, hb, distinct⟩ := Infinite.two_points with infinite
  have restDominated := Dominated.transitive with (Dominated.of_subset with (difference.subset_left at B, A)), dominated
  have disjoint := Disjoint.symmetric with (difference.disjoint at B, A)
  have unionDominated := Dominated.union_of_two_points with (Dominated.reflexive at A), restDominated, ha, hb, distinct, disjoint
  have cover := difference.union2_self at A, B
  rw cover at unionDominated
  have down := Dominated.transitive with unionDominated, (Dominated.of_equipotent with (Infinite.square with infinite))
  exact Dominated.antisymmetric with down, (Dominated.of_subset with (union2.subset_left at A, B))

theorem Infinite.product (A B : Set)
  requires infinite : Infinite(A)
  requires dominated : B ≼ A
  requires inhabited : Nonempty(B)
  shows (A × B) ≃ A
proof
  have down := Dominated.transitive with (Dominated.product with (Dominated.reflexive at A), dominated), (Dominated.of_equipotent with (Infinite.square with infinite))
  have up := Dominated.left_factor at A, B with inhabited
  exact Dominated.antisymmetric with down, up