-- Every point of the unit sphere is carried to the north pole by a
-- rotation, written by coordinates
-- 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.Transitive
import Rotation.Freeness

-- The half turn about the axis through q + e₃ for q = (x, y, z) on the
-- sphere with z ≠ -1 and w = 1 / (1 + z): the matrix
-- w (q + e₃)(q + e₃)ᵀ - 1.
def poleRotation (x y z w : Real) : Set :=
  mat3(x * x * w - 1, x * y * w, x, x * y * w, y * y * w - 1, y, x, y, z)

-- The half turn about the x-axis.
def flipRotation : Set :=
  mat3((1:R), (0:R), (0:R), (0:R), -(1:R), (0:R), (0:R), (0:R), -(1:R))

section PoleEntries
variable (x y z w : Real)

instance poleRotation.matrix
  shows poleRotation(x, y, z, w) : realRing ^ 3 × 3
proof
  unfold poleRotation
  exact mat3.matrix with *

theorem poleRotation.entries
  shows poleRotation(x, y, z, w)[0, 0] = x * x * w - 1 ∧
    poleRotation(x, y, z, w)[0, 1] = x * y * w ∧
    poleRotation(x, y, z, w)[0, 2] = x ∧
    poleRotation(x, y, z, w)[1, 0] = x * y * w ∧
    poleRotation(x, y, z, w)[1, 1] = y * y * w - 1 ∧
    poleRotation(x, y, z, w)[1, 2] = y ∧
    poleRotation(x, y, z, w)[2, 0] = x ∧
    poleRotation(x, y, z, w)[2, 1] = y ∧
    poleRotation(x, y, z, w)[2, 2] = z
proof
  have e00 := mat3.e00 at x * x * w - 1, x * y * w, x, x * y * w, y * y * w - 1,
    y, x, y, z with *
  have e01 := mat3.e01 at x * x * w - 1, x * y * w, x, x * y * w, y * y * w - 1,
    y, x, y, z with *
  have e02 := mat3.e02 at x * x * w - 1, x * y * w, x, x * y * w, y * y * w - 1,
    y, x, y, z with *
  have e10 := mat3.e10 at x * x * w - 1, x * y * w, x, x * y * w, y * y * w - 1,
    y, x, y, z with *
  have e11 := mat3.e11 at x * x * w - 1, x * y * w, x, x * y * w, y * y * w - 1,
    y, x, y, z with *
  have e12 := mat3.e12 at x * x * w - 1, x * y * w, x, x * y * w, y * y * w - 1,
    y, x, y, z with *
  have e20 := mat3.e20 at x * x * w - 1, x * y * w, x, x * y * w, y * y * w - 1,
    y, x, y, z with *
  have e21 := mat3.e21 at x * x * w - 1, x * y * w, x, x * y * w, y * y * w - 1,
    y, x, y, z with *
  have e22 := mat3.e22 at x * x * w - 1, x * y * w, x, x * y * w, y * y * w - 1,
    y, x, y, z with *
  have unfolded : poleRotation(x, y, z, w) =
      mat3(x * x * w - 1, x * y * w, x, x * y * w, y * y * w - 1, y, x, y, z)
    simp [poleRotation]
  rw unfolded
  exact ⟨e00, e01, e02, e10, e11, e12, e20, e21, e22⟩

-- Under x² + y² + z² = 1 and w (1 + z) = 1 the half turn is special
-- orthogonal: each entry condition is a combination of the two relations.
theorem poleRotation.so3
  requires sphere : x * x + y * y + z * z = 1
  requires wz : w * (1 + z) = 1
  shows poleRotation(x, y, z, w) ∈ so3Matrices
proof
  obtain ⟨e00, e01, e02, e10, e11, e12, e20, e21, e22⟩ :=
    poleRotation.entries at x, y, z, w
  have matrix := poleRotation.matrix at x, y, z, w
  have orthogonal : Orthogonal(poleRotation(x, y, z, w))
    apply Orthogonal.of_entries with matrix
    . rw [e00, e10, e20]
      linear_combination (x * x * w * w) * sphere +
        (x * x * (w * (1 - z) - 1)) * wz
    . rw [e01, e11, e21]
      linear_combination (y * y * w * w) * sphere +
        (y * y * (w * (1 - z) - 1)) * wz
    . rw [e02, e12, e22]
      linear_combination sphere
    . rw [e00, e01, e10, e11, e20, e21]
      linear_combination (x * y * w * w) * sphere +
        (x * y * (w * (1 - z) - 1)) * wz
    . rw [e00, e02, e10, e12, e20, e22]
      linear_combination (x * w) * sphere +
        (x * (1 - z)) * wz
    . rw [e01, e02, e11, e12, e21, e22]
      linear_combination (y * w) * sphere +
        (y * (1 - z)) * wz
    . rw [e00, e01, e02]
      linear_combination (x * x * w * w) * sphere +
        (x * x * (w * (1 - z) - 1)) * wz
    . rw [e10, e11, e12]
      linear_combination (y * y * w * w) * sphere +
        (y * y * (w * (1 - z) - 1)) * wz
    . rw [e20, e21, e22]
      linear_combination sphere
    . rw [e00, e01, e02, e10, e11, e12]
      linear_combination (x * y * w * w) * sphere +
        (x * y * (w * (1 - z) - 1)) * wz
    . rw [e00, e01, e02, e20, e21, e22]
      linear_combination (x * w) * sphere +
        (x * (1 - z)) * wz
    . rw [e10, e11, e12, e20, e21, e22]
      linear_combination (y * w) * sphere +
        (y * (1 - z)) * wz
  have unit : determinant3(poleRotation(x, y, z, w)) = (1:R)
    rw [determinant3.value, e00, e01, e02, e10, e11, e12, e20, e21, e22]
    linear_combination (1 - w * z) * sphere - (z * (1 - z)) * wz
  exact so3Matrices.mk with orthogonal, unit

end PoleEntries

-- The half turn carries q to the north pole e₃ = basis(2).
theorem poleRotation.pole (q : Set) (w : Real)
  requires sphere : q ∈ S²
  requires wz : w * ((1:R) + (q(2):R)) = (1:R)
  shows poleRotation(q(0), q(1), q(2), w) · q = basis(2)
proof
  have vector := S².function with sphere
  have q0 := vector.component_real at q, 0 with vector, index.zero_mem
  have q1 := vector.component_real at q, 1 with vector, index.one_mem
  have q2 := vector.component_real at q, 2 with vector, index.two_mem
  obtain ⟨inSpace, unit⟩ := S².membership.forward with sphere
  rw normSq.value at unit
  obtain ⟨e00, e01, e02, e10, e11, e12, e20, e21, e22⟩ :=
    poleRotation.entries at q(0), q(1), q(2), w
  have matrix := poleRotation.matrix at q(0), q(1), q(2), w
  have image := matrixVector.vector at poleRotation(q(0), q(1), q(2), w), q
    with *
  have pole := basis.vector at 2 with index.two_mem
  apply FunctionBetween.ext with image, pole
  intros i hi
  have entry := matrixVector.entry at poleRotation(q(0), q(1), q(2), w), q, i
    with *
  rw entry
  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, e00, e01, e02, other]
    linear_combination ((q(0):R) * w) * unit +
      ((q(0):R) * (1 - (q(2):R))) * wz
  . have other := basis.other at 1, 2
      with index.one_mem, index.two_mem, index.one_ne_two
    rw [at1, e10, e11, e12, other]
    linear_combination ((q(1):R) * w) * unit +
      ((q(1):R) * (1 - (q(2):R))) * wz
  . have same := basis.same at 2 with index.two_mem
    rw [at2, e20, e21, e22, same]
    linear_combination unit

-- The half turn about the x-axis.
instance flipRotation.matrix
  shows flipRotation : realRing ^ 3 × 3
proof
  unfold flipRotation
  exact mat3.matrix with *

theorem flipRotation.entries
  shows flipRotation[0, 0] = (1:R) ∧ flipRotation[0, 1] = (0:R) ∧
    flipRotation[0, 2] = (0:R) ∧ flipRotation[1, 0] = (0:R) ∧
    flipRotation[1, 1] = -(1:R) ∧ flipRotation[1, 2] = (0:R) ∧
    flipRotation[2, 0] = (0:R) ∧ flipRotation[2, 1] = (0:R) ∧
    flipRotation[2, 2] = -(1:R)
proof
  have negOne : -(1:R) ∈ ℝ by typecheck
  have e00 := mat3.e00 at (1:R), (0:R), (0:R), (0:R),
    -(1:R), (0:R), (0:R), (0:R), -(1:R) with *
  have e01 := mat3.e01 at (1:R), (0:R), (0:R), (0:R),
    -(1:R), (0:R), (0:R), (0:R), -(1:R) with *
  have e02 := mat3.e02 at (1:R), (0:R), (0:R), (0:R),
    -(1:R), (0:R), (0:R), (0:R), -(1:R) with *
  have e10 := mat3.e10 at (1:R), (0:R), (0:R), (0:R),
    -(1:R), (0:R), (0:R), (0:R), -(1:R) with *
  have e11 := mat3.e11 at (1:R), (0:R), (0:R), (0:R),
    -(1:R), (0:R), (0:R), (0:R), -(1:R) with *
  have e12 := mat3.e12 at (1:R), (0:R), (0:R), (0:R),
    -(1:R), (0:R), (0:R), (0:R), -(1:R) with *
  have e20 := mat3.e20 at (1:R), (0:R), (0:R), (0:R),
    -(1:R), (0:R), (0:R), (0:R), -(1:R) with *
  have e21 := mat3.e21 at (1:R), (0:R), (0:R), (0:R),
    -(1:R), (0:R), (0:R), (0:R), -(1:R) with *
  have e22 := mat3.e22 at (1:R), (0:R), (0:R), (0:R),
    -(1:R), (0:R), (0:R), (0:R), -(1:R) with *
  have unfolded : flipRotation =
      mat3((1:R), (0:R), (0:R), (0:R), -(1:R), (0:R), (0:R), (0:R), -(1:R))
    simp [flipRotation]
  rw unfolded
  exact ⟨e00, e01, e02, e10, e11, e12, e20, e21, e22⟩

theorem flipRotation.so3
  shows flipRotation ∈ so3Matrices
proof
  obtain ⟨e00, e01, e02, e10, e11, e12, e20, e21, e22⟩ :=
    flipRotation.entries
  have orthogonal : Orthogonal(flipRotation)
    apply Orthogonal.of_entries with flipRotation.matrix
    . rw [e00, e10, e20]
      ring
    . rw [e01, e11, e21]
      ring
    . rw [e02, e12, e22]
      ring
    . rw [e00, e01, e10, e11, e20, e21]
      ring
    . rw [e00, e02, e10, e12, e20, e22]
      ring
    . rw [e01, e02, e11, e12, e21, e22]
      ring
    . rw [e00, e01, e02]
      ring
    . rw [e10, e11, e12]
      ring
    . rw [e20, e21, e22]
      ring
    . rw [e00, e01, e02, e10, e11, e12]
      ring
    . rw [e00, e01, e02, e20, e21, e22]
      ring
    . rw [e10, e11, e12, e20, e21, e22]
      ring
  have unit : determinant3(flipRotation) = (1:R)
    rw [determinant3.value, e00, e01, e02, e10, e11, e12, e20, e21, e22]
    ring
  exact so3Matrices.mk with orthogonal, unit

theorem flipRotation.pole (q : Set)
  requires sphere : q ∈ S²
  requires first : q(0) = (0:R)
  requires second : q(1) = (0:R)
  requires third : q(2) = -(1:R)
  shows flipRotation · q = basis(2)
proof
  have vector := S².function with sphere
  have matrix := flipRotation.matrix
  obtain ⟨e00, e01, e02, e10, e11, e12, e20, e21, e22⟩ :=
    flipRotation.entries
  have image := matrixVector.vector at flipRotation, q with *
  have pole := basis.vector at 2 with index.two_mem
  apply FunctionBetween.ext with image, pole
  intros i hi
  have entry := matrixVector.entry at flipRotation, q, i with *
  rw [entry, first, second, third]
  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, e00, e01, e02, other]
    ring
  . have other := basis.other at 1, 2
      with index.one_mem, index.two_mem, index.one_ne_two
    rw [at1, e10, e11, e12, other]
    ring
  . have same := basis.same at 2 with index.two_mem
    rw [at2, e20, e21, e22, same]
    ring

-- Every point of the sphere is carried to the north pole by a rotation.
theorem so3.transitive_pole (q : Set)
  requires sphere : q ∈ S²
  shows ∃ σ ∈ so3Matrices, σ · q = basis(2)
proof
  have vector := S².function with sphere
  have q0 := vector.component_real at q, 0 with vector, index.zero_mem
  have q1 := vector.component_real at q, 1 with vector, index.one_mem
  have q2 := vector.component_real at q, 2 with vector, index.two_mem
  obtain ⟨inSpace, unit⟩ := S².membership.forward with sphere
  rw normSq.value at unit
  have alternatives : (q(2):R) = -(1:R) ∨ (q(2):R) ≠ -(1:R) by tauto
  cases alternatives with south | elsewhere
  . have squares : (q(0):R) * (q(0):R) + (q(1):R) * (q(1):R) = 0
      rw south at unit
      linear_combination unit
    have nonnegativeX := realMulValue.square_nonnegative at q(0)
    have nonnegativeY := realMulValue.square_nonnegative at q(1)
    have squareX := realAddValue.nonnegative_eq_zero at (q(0):R) * (q(0):R),
      (q(1):R) * (q(1):R) with nonnegativeX, nonnegativeY, squares
    have squareY : (q(1):R) * (q(1):R) = 0 by linarith
    have first := realMulValue.square_eq_zero_iff.forward with squareX
    have second := realMulValue.square_eq_zero_iff.forward with squareY
    have rotation := flipRotation.so3
    have pole := flipRotation.pole with sphere, first, second, south
    exact ⟨flipRotation, rotation, pole⟩
  . have nonzero : (1:R) + (q(2):R) ≠ 0
      intro zero
      apply elsewhere
      linarith
    have wReal := realInvValue.real at (1:R) + (q(2):R)
    have wz := realInvValue.mul_left at (1:R) + (q(2):R) with nonzero
    have rotation := poleRotation.so3 at q(0), q(1), q(2),
      realInvValue((1:R) + (q(2):R)) with unit, wz
    have pole := poleRotation.pole at q, realInvValue((1:R) + (q(2):R))
      with sphere, wz
    witness poleRotation(q(0), q(1), q(2), realInvValue((1:R) + (q(2):R)))
    exact ⟨rotation, pole⟩