-- The sphere absorbs any countable subset: S² and S² ∖ D are
-- SO3-equidecomposable for every countable D ⊆ S²
-- 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.Sphere.Absorb
import Transitive
import Conjugate
import Geometry.Absorption.Final

-- The values of a row of three reals, and rows of unit norm on the sphere.
theorem row3.zero (a b c : Real)
  shows row3(a, b, c)(0) = a
proof
  rw row3.graph
  exact triple.zero at ℝ, a, b, c with *

theorem row3.one (a b c : Real)
  shows row3(a, b, c)(1) = b
proof
  rw row3.graph
  exact triple.one at ℝ, a, b, c with *

theorem row3.two (a b c : Real)
  shows row3(a, b, c)(2) = c
proof
  rw row3.graph
  exact triple.two at ℝ, a, b, c with *

instance row3.function (a b c : Real)
  shows row3(a, b, c) : 3 → ℝ
proof
  have member := row3.member at a, b, c
  exact functionSpace.mem_iff.forward with member

theorem row3.sphere (a b c : Real)
  requires unit : a * a + b * b + c * c = 1
  shows row3(a, b, c) ∈ S²
proof
  have function := row3.function at a, b, c
  have inSpace := ℝ³.of_function with function
  have v0 := row3.zero at a, b, c
  have v1 := row3.one at a, b, c
  have v2 := row3.two at a, b, c
  have norm : normSq(row3(a, b, c)) = (1:R)
    rw [normSq.value, v0, v1, v2]
    exact unit
  exact S².membership.backward with ⟨inSpace, norm⟩

-- The north pole as a row.
theorem basis.two
  shows basis(2) = row3((0:R), (0:R), (1:R))
proof
  have pole := basis.vector at 2 with index.two_mem
  have row := row3.function at (0:R), (0:R), (1:R)
  apply FunctionBetween.ext with pole, row
  intros i hi
  have alternatives := index.cases with hi
  rcases alternatives with at0 | at1 | at2
  . have other := basis.other at 0, 2
      with index.zero_mem, index.two_mem, index.zero_ne_two
    rw [at0, other, row3.zero]
    rfl
  . have other := basis.other at 1, 2
      with index.one_mem, index.two_mem, index.one_ne_two
    rw [at1, other, row3.one]
    rfl
  . have same := basis.same at 2 with index.two_mem
    rw [at2, same, row3.two]
    rfl

-- A point of the sphere on the z-axis is one of the two poles.
theorem sphere.axis_cases (v : Set)
  requires sphere : v ∈ S²
  requires axis : vectorPoint(v) = complexZero
  shows v = row3((0:R), (0:R), (1:R)) ∨ v = row3((0:R), (0:R), -(1:R))
proof
  have vector := S².function with sphere
  have v0 := vector.component_real at v, 0 with vector, index.zero_mem
  have v1 := vector.component_real at v, 1 with vector, index.one_mem
  have v2 := vector.component_real at v, 2 with vector, index.two_mem
  obtain ⟨inSpace, unit⟩ := S².membership.forward with sphere
  rw normSq.value at unit
  rw vectorPoint.graph at axis
  unfold complexZero at axis
  obtain ⟨first, second⟩ := orderedPair.injective with axis
  rw [first, second] at unit
  have square : (v(2):R) * (v(2):R) = 1
    linear_combination unit
  have alternatives : (v(2):R) = (1:R) ∨ (v(2):R) ≠ (1:R) by tauto
  cases alternatives with north | south
  . left
    have row := row3.function at (0:R), (0:R), (1:R)
    apply FunctionBetween.ext with vector, row
    intros i hi
    have cases := index.cases with hi
    rcases cases with at0 | at1 | at2
    . rw [at0, first, row3.zero]
      rfl
    . rw [at1, second, row3.one]
      rfl
    . rw [at2, north, row3.two]
      rfl
  . right
    have nonzero : (v(2):R) - 1 ≠ 0
      intro zero
      apply south
      linarith
    have factored : ((v(2):R) + 1) * ((v(2):R) - 1) = 0 * ((v(2):R) - 1)
      linear_combination square
    have sum := realMulValue.right_cancel_of_nonzero at (v(2):R) + 1, (0:R),
      (v(2):R) - 1 with nonzero, factored
    have third : (v(2):R) = -(1:R) by linarith
    have row := row3.function at (0:R), (0:R), -(1:R)
    apply FunctionBetween.ext with vector, row
    intros i hi
    have cases := index.cases with hi
    rcases cases with at0 | at1 | at2
    . rw [at0, first, row3.zero]
      rfl
    . rw [at1, second, row3.one]
      rfl
    . rw [at2, third, row3.two]
      rfl

-- A rotation is undone by its transpose, and carries opposite vectors
-- to opposite vectors.
theorem matrixVector.transpose_cancel (σ v : Set)
  requires rotation : σ ∈ so3Matrices
  requires vector : v : 3 → ℝ
  shows transpose(σ) · (σ · v) = v
proof
  have matrix := so3Matrices.matrix with rotation
  have transposed := transpose.matrix at σ with matrix
  have orthogonal := so3Matrices.orthogonal with rotation
  have left := Orthogonal.left with orthogonal
  have merged := matrixVector.mul at transpose(σ), σ, v with vector
  rw [← merged, left]
  exact matrixVector.one with vector

theorem matrixVector.opposite (M : realRing ^ 3 × 3) (u v : Set)
  requires vectorU : u : 3 → ℝ
  requires vectorV : v : 3 → ℝ
  requires opposite : ∀ i ∈ 3, (v(i):R) = -(u(i):R)
  shows ∀ i ∈ 3, ((M · v)(i):R) = -((M · u)(i):R)
proof
  intros i hi
  have entryU := matrixVector.entry at M, u, i with *
  have entryV := matrixVector.entry at M, v, i with *
  have o0 := opposite with 0, index.zero_mem
  have o1 := opposite with 1, index.one_mem
  have o2 := opposite with 2, index.two_mem
  obtain ⟨m0, m1, m2⟩ := Matrix.row_real at M, i with hi
  have u0 := vector.component_real at u, 0 with vectorU, index.zero_mem
  have u1 := vector.component_real at u, 1 with vectorU, index.one_mem
  have u2 := vector.component_real at u, 2 with vectorU, index.two_mem
  rw [entryU, entryV, o0, o1, o2]
  ring

-- The equator, parametrized rationally by the turn, and its antipodes.
defgraph equatorPoint (t ∈ ℝ) : S² := row3(turnCos(t), turnSin(t), (0:R))
proof
  have typed : Real(t) by typecheck
  have circle := turn.circle at t
  have unit : turnCos(t) * turnCos(t) + turnSin(t) * turnSin(t) +
      (0:R) * (0:R) = 1
    linear_combination circle
  exact row3.sphere with unit

defgraph equatorAntipode (t ∈ ℝ) : S² :=
  row3(-turnCos(t), -turnSin(t), (0:R))
proof
  have typed : Real(t) by typecheck
  have circle := turn.circle at t
  have unit : -turnCos(t) * -turnCos(t) + -turnSin(t) * -turnSin(t) +
      (0:R) * (0:R) = 1
    linear_combination circle
  exact row3.sphere with unit

theorem equatorPoint.cancel (s t : Real)
  requires same : equatorPoint(s) = equatorPoint(t)
  shows s = t
proof
  have realS : s ∈ ℝ by typecheck
  have realT : t ∈ ℝ by typecheck
  have valueS := equatorPoint.value at s with *
  have valueT := equatorPoint.value at t with *
  rw [valueS, valueT] at same
  have cosS := row3.zero at turnCos(s), turnSin(s), (0:R)
  have cosT := row3.zero at turnCos(t), turnSin(t), (0:R)
  have sinS := row3.one at turnCos(s), turnSin(s), (0:R)
  have sinT := row3.one at turnCos(t), turnSin(t), (0:R)
  have cosEqual : turnCos(s) = turnCos(t) by cc
  have sinEqual : turnSin(s) = turnSin(t) by cc
  exact turn.injective_of_coordinates with cosEqual, sinEqual

theorem equatorAntipode.cancel (s t : Real)
  requires same : equatorAntipode(s) = equatorAntipode(t)
  shows s = t
proof
  have realS : s ∈ ℝ by typecheck
  have realT : t ∈ ℝ by typecheck
  have valueS := equatorAntipode.value at s with *
  have valueT := equatorAntipode.value at t with *
  rw [valueS, valueT] at same
  have cosS := row3.zero at -turnCos(s), -turnSin(s), (0:R)
  have cosT := row3.zero at -turnCos(t), -turnSin(t), (0:R)
  have sinS := row3.one at -turnCos(s), -turnSin(s), (0:R)
  have sinT := row3.one at -turnCos(t), -turnSin(t), (0:R)
  have cosOpposite : -turnCos(s) = -turnCos(t) by cc
  have sinOpposite : -turnSin(s) = -turnSin(t) by cc
  have cosEqual : turnCos(s) = turnCos(t) by linarith
  have sinEqual : turnSin(s) = turnSin(t) by linarith
  exact turn.injective_of_coordinates with cosEqual, sinEqual

theorem equatorAntipode.opposite (t : Real)
  shows ∀ i ∈ 3, (equatorAntipode(t)(i):R) = -(equatorPoint(t)(i):R)
proof
  have real : t ∈ ℝ by typecheck
  have valueP := equatorPoint.value at t with *
  have valueA := equatorAntipode.value at t with *
  intros i hi
  rw [valueP, valueA]
  have alternatives := index.cases with hi
  rcases alternatives with at0 | at1 | at2
  . rw [at0, row3.zero, row3.zero]
    rfl
  . rw [at1, row3.one, row3.one]
    rfl
  . rw [at2, row3.two, row3.two]
    ring

-- The parameters whose point, or whose antipode, lies in D.
def equatorBad (D : Set) : Set := functionPreimage(equatorPoint, ℝ, D)
def antipodeBad (D : Set) : Set := functionPreimage(equatorAntipode, ℝ, D)

theorem equatorBad.membership (D t : Set)
  shows t ∈ equatorBad(D) ↔ t ∈ ℝ ∧ equatorPoint(t) ∈ D
proof
  simp [equatorBad, functionPreimage.membership]

theorem antipodeBad.membership (D t : Set)
  shows t ∈ antipodeBad(D) ↔ t ∈ ℝ ∧ equatorAntipode(t) ∈ D
proof
  simp [antipodeBad, functionPreimage.membership]

theorem equatorBad.countable (D : Set)
  requires countable : Countable(D)
  shows Countable(equatorBad(D))
proof
  have function := equatorPoint.function
  have subset : equatorBad(D) ⊆ ℝ
    intros t ht
    obtain ⟨real, inD⟩ := equatorBad.membership.forward with ht
    exact real
  have closed : ∀ t ∈ equatorBad(D), equatorPoint(t) ∈ D
    intros t ht
    obtain ⟨real, inD⟩ := equatorBad.membership.forward with ht
    exact inD
  have restricted := restrict.function at equatorPoint, ℝ, S²,
    equatorBad(D), D with function, subset, closed
  have injective : InjectiveBetween(restrict(equatorPoint, equatorBad(D)),
      equatorBad(D), D)
    apply Injective.of_pointwise with restricted
    intros s hs t ht same
    have realS := subset with s, hs
    have realT := subset with t, ht
    have typedS : Real(s) by typecheck
    have typedT : Real(t) by typecheck
    have stepS := restrict.apply at equatorPoint, ℝ, S², equatorBad(D), D, s
      with function, subset, closed, hs
    have stepT := restrict.apply at equatorPoint, ℝ, S², equatorBad(D), D, t
      with function, subset, closed, ht
    rw [stepS, stepT] at same
    exact equatorPoint.cancel with same
  exact Countable.of_injection with injective, countable

theorem antipodeBad.countable (D : Set)
  requires countable : Countable(D)
  shows Countable(antipodeBad(D))
proof
  have function := equatorAntipode.function
  have subset : antipodeBad(D) ⊆ ℝ
    intros t ht
    obtain ⟨real, inD⟩ := antipodeBad.membership.forward with ht
    exact real
  have closed : ∀ t ∈ antipodeBad(D), equatorAntipode(t) ∈ D
    intros t ht
    obtain ⟨real, inD⟩ := antipodeBad.membership.forward with ht
    exact inD
  have restricted := restrict.function at equatorAntipode, ℝ, S²,
    antipodeBad(D), D with function, subset, closed
  have injective : InjectiveBetween(restrict(equatorAntipode, antipodeBad(D)),
      antipodeBad(D), D)
    apply Injective.of_pointwise with restricted
    intros s hs t ht same
    have realS := subset with s, hs
    have realT := subset with t, ht
    have typedS : Real(s) by typecheck
    have typedT : Real(t) by typecheck
    have stepS := restrict.apply at equatorAntipode, ℝ, S², antipodeBad(D),
      D, s with function, subset, closed, hs
    have stepT := restrict.apply at equatorAntipode, ℝ, S², antipodeBad(D),
      D, t with function, subset, closed, ht
    rw [stepS, stepT] at same
    exact equatorAntipode.cancel with same
  exact Countable.of_injection with injective, countable

-- Some point of the equator lies, with its antipode, outside a countable D.
theorem equator.exists_outside (D : Set)
  requires countable : Countable(D)
  shows ∃ t ∈ ℝ, ¬(equatorPoint(t) ∈ D) ∧
    ¬(equatorAntipode(t) ∈ D)
proof
  have badPoints := equatorBad.countable with countable
  have badAntipodes := antipodeBad.countable with countable
  have bad := Countable.union2 with badPoints, badAntipodes
  have uncountable := reals.uncountable
  unfold Uncountable at uncountable
  by_contra
  rename this => none
  have cover : ℝ ⊆ equatorBad(D) ∪ antipodeBad(D)
    intros t ht
    by_contra
    rename this => outside
    apply none
    refine ⟨t, ht, ?_, ?_⟩
    . intro inD
      apply outside
      have member : t ∈ equatorBad(D) :=
        equatorBad.membership.backward with ⟨ht, inD⟩
      exact union2.left with member
    . intro inD
      apply outside
      have member : t ∈ antipodeBad(D) :=
        antipodeBad.membership.backward with ⟨ht, inD⟩
      exact union2.right with member
  have realsCountable := Countable.subset with bad, cover
  exact uncountable with realsCountable

-- The image of a countable subset of the sphere under a rotation is
-- countable: it is the image of D under the rotation as a function on S².
defgraph rotateSphere (σ : Set) (v ∈ S²) : S² := σ · v
requires rotation : σ ∈ so3Matrices
proof
  obtain ⟨subset, closed⟩ := Invariant.sphere
  exact closed with σ, rotation, v, member

theorem actImage.countable (σ D : Set)
  requires rotation : σ ∈ so3Matrices
  requires subset : D ⊆ S²
  requires countable : Countable(D)
  shows Countable(actImage(sphereAction, σ, D))
proof
  have function := rotateSphere.function at σ with rotation
  have closed : ∀ d ∈ D, rotateSphere(σ)(d) ∈ S²
    intros d hd
    exact FunctionBetween.apply_closed with function, (subset with d, hd)
  have restricted := restrict.function at rotateSphere(σ), S², S², D, S²
    with function, subset, closed
  have image := Countable.image at restrict(rotateSphere(σ), D), D, S²
    with restricted, countable
  have bounded : actImage(sphereAction, σ, D) ⊆
      functionImage(restrict(rotateSphere(σ), D), D, S²)
    intros y hy
    obtain ⟨d, hd, moved⟩ := actImage.membership.forward with hy
    have dIn := subset with d, hd
    have applied := sphereAction.apply with rotation, dIn
    have value := rotateSphere.value at σ, d with *
    have step := restrict.apply at rotateSphere(σ), S², S², D, S², d
      with function, subset, closed, hd
    have member := functionImage.value at restrict(rotateSphere(σ), D), D,
      S², d with restricted, hd
    rw [step, value] at member
    rw [← moved, applied]
    exact member
  exact Countable.subset with image, bounded

-- A rotation carrying the equator point q to the north pole carries its
-- antipode to the south pole, so a point of D off q and -q leaves the axis.
theorem equatorAntipode.image (σ t : Set)
  requires rotation : σ ∈ so3Matrices
  requires real : t ∈ ℝ
  requires pole : σ · equatorPoint(t) = basis(2)
  shows σ · equatorAntipode(t) = row3((0:R), (0:R), -(1:R))
proof
  have typed : Real(t) by typecheck
  have matrix := so3Matrices.matrix with rotation
  have pointIn := FunctionBetween.apply_closed with equatorPoint.function, real
  have antipodeIn := FunctionBetween.apply_closed
    with equatorAntipode.function, real
  have vectorP := S².function with pointIn
  have vectorA := S².function with antipodeIn
  have opposite := equatorAntipode.opposite at t
  have coordinates := matrixVector.opposite at σ, equatorPoint(t),
    equatorAntipode(t) with matrix, vectorP, vectorA, opposite
  have image := matrixVector.vector at σ, equatorAntipode(t)
    with matrix, vectorA
  have row := row3.function at (0:R), (0:R), -(1:R)
  apply FunctionBetween.ext with image, row
  intros i hi
  have coordinate := coordinates with i, hi
  rw [coordinate, pole]
  have alternatives := index.cases with hi
  rcases alternatives with at0 | at1 | at2
  . have other := basis.other at 0, 2
      with index.zero_mem, index.two_mem, index.zero_ne_two
    rw [at0, other, row3.zero]
    ring
  . have other := basis.other at 1, 2
      with index.one_mem, index.two_mem, index.one_ne_two
    rw [at1, other, row3.one]
    ring
  . have same := basis.same at 2 with index.two_mem
    rw [at2, same, row3.two]
    rfl

theorem actImage.off_axis (σ D t : Set)
  requires rotation : σ ∈ so3Matrices
  requires subset : D ⊆ S²
  requires real : t ∈ ℝ
  requires pole : σ · equatorPoint(t) = basis(2)
  requires offPoint : ¬(equatorPoint(t) ∈ D)
  requires offAntipode : ¬(equatorAntipode(t) ∈ D)
  shows ∀ y ∈ actImage(sphereAction, σ, D), vectorPoint(y) ≠ complexZero
proof
  intros y hy axis
  obtain ⟨d, hd, moved⟩ := actImage.membership.forward with hy
  have dIn := subset with d, hd
  have applied := sphereAction.apply with rotation, dIn
  rw applied at moved
  have vectorD := S².function with dIn
  obtain ⟨sphereSubset, closed⟩ := Invariant.sphere
  have yIn : y ∈ S²
    rw ← moved
    exact closed with σ, rotation, d, dIn
  have poles := sphere.axis_cases with yIn, axis
  have back := matrixVector.transpose_cancel with rotation, vectorD
  rw moved at back
  have pointIn := FunctionBetween.apply_closed with equatorPoint.function, real
  have antipodeIn := FunctionBetween.apply_closed
    with equatorAntipode.function, real
  have vectorP := S².function with pointIn
  have vectorA := S².function with antipodeIn
  cases poles with north | south
  . have backP := matrixVector.transpose_cancel with rotation, vectorP
    rw [pole, basis.two] at backP
    rw north at back
    apply offPoint
    rw [← backP, back]
    exact hd
  . have backA := matrixVector.transpose_cancel with rotation, vectorA
    have antipode := equatorAntipode.image with rotation, real, pole
    rw antipode at backA
    rw south at back
    apply offAntipode
    rw [← backA, back]
    exact hd

-- The separation theorem for every countable subset of the sphere.
theorem sphere.exists_separating_any (D : Set)
  requires countable : Countable(D)
  requires subset : D ⊆ S²
  shows ∃ ρ ∈ Group.carrier(SO3), ∀ n ∈ ℕ,
    Disjoint(actImage(sphereAction, groupPow(SO3, ρ, n + 1), D), D)
proof
  have outside := equator.exists_outside with countable
  obtain ⟨t, real, offPoint, offAntipode⟩ := outside
  have pointIn := FunctionBetween.apply_closed with equatorPoint.function, real
  obtain ⟨σ, rotation, pole⟩ := so3.transitive_pole with pointIn
  have action := sphereAction.action
  have group := SO3.group
  have σIn : σ ∈ Group.carrier(SO3)
    rw SO3.carrierOf
    exact rotation
  have subsetImage : actImage(sphereAction, σ, D) ⊆ S² :=
    actImage.subset at sphereAction, SO3, S², σ, D with action, subset
  have countableImage := actImage.countable with rotation, subset, countable
  have offAxis := actImage.off_axis
    with rotation, subset, real, pole, offPoint, offAntipode
  obtain ⟨τ, turnIn, separated⟩ := sphere.exists_separating
    with countableImage, subsetImage, offAxis
  have τRotation : τ ∈ so3Matrices
    rw ← SO3.carrierOf
    exact turnIn
  have inverseIn : SO3.inv(σ) ∈ Group.carrier(SO3) :=
    Group.inv_mem with group
  have leftIn : SO3.mul(SO3.inv(σ), τ) ∈ Group.carrier(SO3) :=
    Group.mul_mem with group
  have conjugateIn : SO3.mul(SO3.mul(SO3.inv(σ), τ), σ) ∈ SO3.carrier :=
    Group.mul_mem with group
  refine ⟨SO3.mul(SO3.mul(SO3.inv(σ), τ), σ), conjugateIn, ?_⟩
  intros n hn
  have natural := omega.natural with hn
  have apart := separated with n, hn
  have powerIn : groupPow(SO3, τ, n + 1) ∈ Group.carrier(SO3) :=
    Group.pow_mem with group
  have powerRotation : groupPow(SO3, τ, n + 1) ∈ so3Matrices
    rw ← SO3.carrierOf
    exact powerIn
  have powers := so3.pow_conjugate at σ, τ, n + 1 with rotation, τRotation
  have conjugated := actImage.conjugate at σ, groupPow(SO3, τ, n + 1), D
    with rotation, powerRotation, subset
  have movedSubset := actImage.subset at sphereAction, SO3, S²,
    groupPow(SO3, τ, n + 1), actImage(sphereAction, σ, D)
    with action, subsetImage
  have moved := actImage.disjoint at sphereAction, SO3, S², SO3.inv(σ),
    actImage(sphereAction, groupPow(SO3, τ, n + 1),
      actImage(sphereAction, σ, D)),
    actImage(sphereAction, σ, D) with action, movedSubset, subsetImage, apart
  have back := actImage.inv_image at sphereAction, SO3, S², σ, D
    with action, subset
  rw back at moved
  rw [powers, conjugated]
  exact moved

-- The sphere absorbs every countable subset.
theorem sphere.absorb_any (D : Set)
  requires countable : Countable(D)
  requires subset : D ⊆ S²
  shows Equidecomposable(sphereAction, SO3, S², S², S² ∖ D)
proof
  have separating := sphere.exists_separating_any with countable, subset
  obtain ⟨ρ, rotation, separated⟩ := separating
  exact sphere.absorb with rotation, subset, separated