-- Cones over pieces of the sphere lift a decomposition of the sphere to
-- one of the punctured ball
-- 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.Euclid.Cones
import Ball
import Real.Sqrt
import Algebra.Group.Paradox

-- Scaling twice is scaling by the product.
theorem vectorScale.mul (c d v : Set)
  requires scalarC : c ∈ ℝ
  requires scalarD : d ∈ ℝ
  requires vector : v : 3 → ℝ
  shows vectorScale(c, vectorScale(d, v)) = vectorScale((c:R) * (d:R), v)
proof
  have inner := vectorScale.vector with scalarD, vector
  have left := vectorScale.vector with scalarC, inner
  have productReal : (c:R) * (d:R) ∈ ℝ by typecheck
  have right := vectorScale.vector with productReal, vector
  apply FunctionBetween.ext with left, right
  intros i hi
  have vi := vector.component_real at v, i with vector, hi
  have e1 :=
      vectorScale.entry at c, vectorScale(d, v), i with scalarC, inner, hi
  have e2 := vectorScale.entry at d, v, i with scalarD, vector, hi
  have e3 :=
      vectorScale.entry at (c:R) * (d:R), v, i with productReal, vector, hi
  rw [e1, e2, e3]
  ring

-- A vector of squared norm zero is the zero vector.
theorem normSq.eq_zero (v : Set)
  requires vector : v : 3 → ℝ
  requires zero : normSq(v) = (0:R)
  shows v = zeroVector
proof
  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
  have sq0 := realMulValue.square_nonnegative at (v(0):R)
  have sq1 := realMulValue.square_nonnegative at (v(1):R)
  have sq2 := realMulValue.square_nonnegative at (v(2):R)
  rw normSq.value at zero
  have z0 : (v(0):R) * (v(0):R) = 0 by linarith
  have z1 : (v(1):R) * (v(1):R) = 0 by linarith
  have z2 : (v(2):R) * (v(2):R) = 0 by linarith
  have c0 := realMulValue.square_eq_zero_iff.forward with z0
  have c1 := realMulValue.square_eq_zero_iff.forward with z1
  have c2 := realMulValue.square_eq_zero_iff.forward with z2
  apply FunctionBetween.ext with vector, zeroVector.vector
  intros i hi
  have zi := zeroVector.entry at i with hi
  rw zi
  have alternatives := index.cases with hi
  rcases alternatives with at0 | at1 | at2
  . rw at0
    exact c0
  . rw at1
    exact c1
  . rw at2
    exact c2

-- Positive multiples of unit vectors determine the scalar and the vector.
theorem vectorScale.sphere_injective (r s u v : Set)
  requires scalarR : r ∈ ℝ
  requires scalarS : s ∈ ℝ
  requires positiveR : (0:R) < r
  requires positiveS : (0:R) < s
  requires hu : u ∈ S²
  requires hv : v ∈ S²
  requires same : vectorScale(r, u) = vectorScale(s, v)
  shows r = s ∧ u = v
proof
  obtain ⟨uSpace, uUnit⟩ := S².membership.forward with hu
  obtain ⟨vSpace, vUnit⟩ := S².membership.forward with hv
  have vectorU := ℝ³.function with uSpace
  have vectorV := ℝ³.function with vSpace
  have normU := vectorScale.normSq with scalarR, vectorU
  have normV := vectorScale.normSq with scalarS, vectorV
  rw [same, normV, vUnit, uUnit] at normU
  have unitS : (s:R) * (s:R) * 1 = (s:R) * (s:R)
    ring
  have unitR : (r:R) * (r:R) * 1 = (r:R) * (r:R)
    ring
  rw [unitS, unitR] at normU
  have squares : (r:R) * (r:R) = (s:R) * (s:R)
    symm
    exact normU
  have nonnegR : (0:R) ≤ r by linarith
  have nonnegS : (0:R) ≤ s by linarith
  have equalScalar :=
      realMulValue.eq_of_square_eq with nonnegR, nonnegS, squares
  refine ⟨equalScalar, ?_⟩
  rw equalScalar at same
  have nonzero : (s:R) ≠ 0
    intro equal
    rw equal at positiveS
    linarith
  apply FunctionBetween.ext with vectorU, vectorV
  intros i hi
  have ui := vector.component_real at u, i with vectorU, hi
  have vi := vector.component_real at v, i with vectorV, hi
  have e1 := vectorScale.entry at s, u, i with scalarS, vectorU, hi
  have e2 := vectorScale.entry at s, v, i with scalarS, vectorV, hi
  have entries : (s:R) * (u(i):R) = (s:R) * (v(i):R)
    rw [← e1, ← e2, same]
    rfl
  have cancel := realInvValue.mul_left with nonzero
  have invReal : realInvValue(s) ∈ ℝ by typecheck
  have expandU : realInvValue(s) * ((s:R) * (u(i):R)) = (u(i):R)
    have regroup : realInvValue(s) * ((s:R) * (u(i):R)) =
      (realInvValue(s) * (s:R)) * (u(i):R)
      ring
    rw [regroup, cancel]
    ring
  have expandV : realInvValue(s) * ((s:R) * (v(i):R)) = (v(i):R)
    have regroup : realInvValue(s) * ((s:R) * (v(i):R)) =
      (realInvValue(s) * (s:R)) * (v(i):R)
      ring
    rw [regroup, cancel]
    ring
  rw [← expandU, ← expandV, entries]
  rfl

-- The cone over a set of directions: the multiples r u with 0 < r ≤ 1.
def cone (A : Set) : Set :=
  { v ∈ ℝ³ | ∃ u ∈ A, ∃ r ∈ ℝ,
    (0:R) < r ∧ r ≤ (1:R) ∧ v = vectorScale(r, u) }

theorem cone.subset (A : Set)
  shows cone(A) ⊆ ℝ³
proof
  intros v hv
  have parts := cone.membership.forward with hv
  exact parts.left

theorem cone.contains (A u r : Set)
  requires hu : u ∈ A
  requires uSpace : u ∈ ℝ³
  requires scalar : r ∈ ℝ
  requires positive : (0:R) < r
  requires bounded : r ≤ (1:R)
  shows vectorScale(r, u) ∈ cone(A)
proof
  have inSpace := vectorScale.space with scalar, uSpace
  exact cone.membership.backward
      with ⟨inSpace, u, hu, r, scalar, positive, bounded, rfl⟩

theorem cone.monotone (A B : Set)
  requires subset : A ⊆ B
  shows cone(A) ⊆ cone(B)
proof
  intros v hv
  obtain ⟨inSpace, u, hu, r, scalar, positive, bounded, shape⟩ :=
    cone.membership.forward with hv
  have uIn := subset with u, hu
  exact cone.membership.backward
      with ⟨inSpace, u, uIn, r, scalar, positive, bounded, shape⟩

-- The cone over directions on the sphere lies in the punctured ball.
theorem cone.subset_punctured (A : Set)
  requires subset : A ⊆ S²
  shows cone(A) ⊆ puncturedBall
proof
  intros v hv
  obtain ⟨inSpace, u, hu, r, scalar, positive, bounded, shape⟩ :=
    cone.membership.forward with hv
  obtain ⟨uSpace, uUnit⟩ := S².membership.forward with (subset with u, hu)
  have vectorU := ℝ³.function with uSpace
  have norm := vectorScale.normSq with scalar, vectorU
  rw uUnit at norm
  have unitR : (r:R) * (r:R) * 1 = (r:R) * (r:R)
    ring
  rw unitR at norm
  have nonneg : (0:R) ≤ r by linarith
  have squareBound := realMulValue.square_le_of_le with nonneg, bounded
  have oneOne : (1:R) * (1:R) = (1:R)
    ring
  rw oneOne at squareBound
  have squarePositive :=
      realLess.mul_left_of_positive at (0:R), r, r with positive, positive
  have zeroMul : (r:R) * 0 = 0
    ring
  rw zeroMul at squarePositive
  have scaledSpace := vectorScale.space with scalar, uSpace
  have vBound : normSq(vectorScale(r, u)) ≤ (1:R)
    rw norm
    exact squareBound
  have inBall : vectorScale(r, u) ∈ unitBall :=
    unitBall.membership.backward with ⟨scaledSpace, vBound⟩
  have notZero : vectorScale(r, u) ≠ zeroVector
    intro equal
    rw [equal, zeroVector.normSq] at norm
    linarith
  rw shape
  exact puncturedBall.membership.backward with ⟨inBall, notZero⟩

-- Disjoint sets of directions have disjoint cones.
theorem cone.disjoint (A B : Set)
  requires boundedA : A ⊆ S²
  requires boundedB : B ⊆ S²
  requires disjoint : Disjoint(A, B)
  shows Disjoint(cone(A), cone(B))
proof
  have apart := Disjoint.iff.forward with disjoint
  apply Disjoint.iff.backward
  intros v hvA hvB
  obtain ⟨inSpace, u, hu, r, scalarR, positiveR, boundedR, shapeU⟩ :=
    cone.membership.forward with hvA
  obtain ⟨inSpace2, w, hw, s, scalarS, positiveS, boundedS, shapeW⟩ :=
    cone.membership.forward with hvB
  have same : vectorScale(r, u) = vectorScale(s, w) by cc
  have uSphere := boundedA with u, hu
  have wSphere := boundedB with w, hw
  have both := vectorScale.sphere_injective with
    scalarR, scalarS, positiveR, positiveS, uSphere, wSphere, same
  rw both.right at hu
  exact apart with w, hu, hw

-- A rotation carries the cone over S to the cone over the rotated S.
theorem cone.rotate (M : realRing ^ 3 × 3) (S : Set)
  requires hM : M ∈ so3Matrices
  requires subset : S ⊆ S²
  shows translate(spaceAction, ℝ³, M, cone(S)) = cone(translate(sphereAction,
      S², M, S))
proof
  obtain ⟨sphereSubset, closed⟩ := Invariant.sphere
  ext y
  constructor
  . intro hy
    obtain ⟨yIn, x, hx, equal⟩ := translate.membership.forward with hy
    obtain ⟨xIn, u, hu, r, scalar, positive, bounded, shape⟩ :=
      cone.membership.forward with hx
    have uSphere := subset with u, hu
    obtain ⟨uSpace, uUnit⟩ := S².membership.forward with uSphere
    have vectorU := ℝ³.function with uSpace
    have space := spaceAction.apply with hM, xIn
    have sphere := sphereAction.apply with hM, uSphere
    have scaled := matrixVector.scale at M, r, u with scalar, vectorU
    have imageIn := closed with M, hM, u, uSphere
    have shapeBack : M · u = sphereAction(M, u)
      symm
      exact sphere
    have movedIn : M · u ∈ translate(sphereAction, S², M, S) :=
      translate.membership.backward with ⟨imageIn, u, hu, shapeBack⟩
    have imageSpace := S².subset with M · u, imageIn
    rw [equal, space, shape, scaled]
    exact cone.contains with movedIn, imageSpace, scalar, positive, bounded
  . intro hy
    obtain ⟨yIn, w, hw, r, scalar, positive, bounded, shape⟩ :=
      cone.membership.forward with hy
    obtain ⟨wIn, u, hu, wShape⟩ := translate.membership.forward with hw
    have uSphere := subset with u, hu
    obtain ⟨uSpace, uUnit⟩ := S².membership.forward with uSphere
    have vectorU := ℝ³.function with uSpace
    have sphere := sphereAction.apply with hM, uSphere
    rw sphere at wShape
    have scaledIn := cone.contains with hu, uSpace, scalar, positive, bounded
    have scaledSpace := vectorScale.space with scalar, uSpace
    have space := spaceAction.apply with hM, scaledSpace
    have scaled := matrixVector.scale at M, r, u with scalar, vectorU
    have yShape : y = spaceAction(M, vectorScale(r, u))
      rw [space, scaled, ← wShape]
      exact shape
    exact translate.membership.backward
        with ⟨yIn, vectorScale(r, u), scaledIn, yShape⟩

-- The cone over the sphere is the punctured ball: v = ‖v‖ · (v / ‖v‖).
theorem cone.sphere
  shows cone(S²) = puncturedBall
proof
  apply Subset.antisymmetric
  . exact cone.subset_punctured with (Subset.reflexive at S²)
  . intros v hv
    obtain ⟨inBall, notZero⟩ := puncturedBall.membership.forward with hv
    obtain ⟨inSpace, bound⟩ := unitBall.membership.forward with inBall
    have vector := ℝ³.function with inSpace
    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
    have sq0 := realMulValue.square_nonnegative at (v(0):R)
    have sq1 := realMulValue.square_nonnegative at (v(1):R)
    have sq2 := realMulValue.square_nonnegative at (v(2):R)
    have normReal : normSq(v) ∈ ℝ by typecheck
    have expand := normSq.value at v
    have normNonneg : (0:R) ≤ normSq(v)
      rw expand
      linarith
    have normNonzero : normSq(v) ≠ (0:R)
      intro zero
      exact notZero with (normSq.eq_zero with vector, zero)
    have nonnegative := NonnegativeReal.of_le with normNonneg
    have rReal := sqrt.real at normSq(v)
    have rIn : sqrt(normSq(v)) ∈ ℝ by typecheck
    have rPositive := sqrt.positive with normNonzero
    have rSquare := sqrt.square at normSq(v)
    have zeroLeOne : (0:R) ≤ (1:R) by linarith
    have oneNonneg := NonnegativeReal.of_le with zeroLeOne
    have rLe := sqrt.le_of_le with nonnegative, oneNonneg, bound
    have oneSquare : (1:R) * (1:R) = (1:R)
      ring
    have sqrtOne := sqrt.of_square at (1:R) with zeroLeOne
    rw oneSquare at sqrtOne
    rw sqrtOne at rLe
    have rNonzero : sqrt(normSq(v)) ≠ 0
      intro equal
      rw equal at rPositive
      linarith
    have invReal : realInvValue(sqrt(normSq(v))) ∈ ℝ by typecheck
    have cancel := realInvValue.mul_left with rNonzero
    have cancel2 := realInvValue.mul_right with rNonzero
    have uSpace := vectorScale.space with invReal, inSpace
    have uNorm := vectorScale.normSq with invReal, vector
    have uUnit : normSq(vectorScale(realInvValue(sqrt(normSq(v))), v)) = (1:R)
      rw uNorm
      calc
        realInvValue(sqrt(normSq(v))) * realInvValue(sqrt(normSq(v))) *
            normSq(v) =
            realInvValue(sqrt(normSq(v))) * realInvValue(sqrt(normSq(v))) *
            (sqrt(normSq(v)) * sqrt(normSq(v))) := by rw rSquare
        _ = (realInvValue(sqrt(normSq(v))) * sqrt(normSq(v))) *
            (realInvValue(sqrt(normSq(v))) * sqrt(normSq(v))) := by ring
        _ = (1:R) * (1:R) := by rw cancel
        _ = (1:R) := by ring
    have uSphere : vectorScale(realInvValue(sqrt(normSq(v))), v) ∈ S² :=
      S².membership.backward with ⟨uSpace, uUnit⟩
    have combine := vectorScale.mul at sqrt(normSq(v)),
        realInvValue(sqrt(normSq(v))), v with rIn, invReal, vector
    have unit := vectorScale.one at v with vector
    have vShape : v = vectorScale(sqrt(normSq(v)),
        vectorScale(realInvValue(sqrt(normSq(v))), v))
      rw [combine, cancel2, unit]
      rfl
    rw vShape
    exact cone.contains with uSphere, uSpace, rIn, rPositive, rLe

-- Lifting a decomposition of the sphere: each piece and its element are
-- sent to the cone over the piece with the same element.
defgraph conePair (p ∈ 𝒫(S²) × SO3.carrier) : 𝒫(ℝ³) × SO3.carrier :=
  (cone(fst(p)), snd(p))
proof
  obtain ⟨bound, S, hS, g, hg, shape⟩ :=
    product.membership.forward with member
  have coneIn := powerset.membership_iff.backward with (cone.subset at S)
  have firstP := fst.pair at S, g
  have secondP := snd.pair at S, g
  rw [shape, firstP, secondP]
  exact product.pair with coneIn, hg

theorem conePair.apply (S g : Set)
  requires hS : S ∈ 𝒫(S²)
  requires hg : g ∈ Group.carrier(SO3)
  shows conePair((S, g)) = (cone(S), g)
proof
  have pair := product.pair at S, g, 𝒫(S²), Group.carrier(SO3) with hS, hg
  have value := conePair.value at (S, g) with pair
  have firstP := fst.pair at S, g
  have secondP := snd.pair at S, g
  rw [value, firstP, secondP]
  rfl

theorem conePair.first (p : Set)
  requires hp : p ∈ 𝒫(S²) × Group.carrier(SO3)
  shows fst(conePair(p)) = cone(fst(p))
proof
  obtain ⟨bound, S, hS, g, hg, shape⟩ := product.membership.forward with hp
  have value := conePair.apply with hS, hg
  have firstP := fst.pair at S, g
  have firstC := fst.pair at cone(S), g
  rw [shape, value, firstP, firstC]
  rfl

theorem conePair.moved (p : Set)
  requires hp : p ∈ 𝒫(S²) × Group.carrier(SO3)
  shows moved(spaceAction, ℝ³, conePair(p)) =
    cone(moved(sphereAction, S², p))
proof
  obtain ⟨bound, S, hS, g, hg, shape⟩ := product.membership.forward with hp
  have subset := powerset.membership_iff.forward with hS
  have gIn := hg
  rw SO3.carrierOf at gIn
  have matrixG := so3Matrices.matrix with gIn
  have value := conePair.apply with hS, hg
  have rotate := cone.rotate with gIn, subset
  have movedC := moved.value at spaceAction, ℝ³, cone(S), g
  have movedS := moved.value at sphereAction, S², S, g
  rw [shape, value, movedC, movedS]
  exact rotate

theorem cone.equidecomposable (A B : Set)
  requires boundedA : A ⊆ S²
  requires boundedB : B ⊆ S²
  requires equidecomposable : Equidecomposable(sphereAction, SO3, S², A, B)
  shows Equidecomposable(spaceAction, SO3, ℝ³, cone(A), cone(B))
proof
  obtain ⟨F, finite, bounded, pieces, movedPieces⟩ := equidecomposable
  obtain ⟨piecesBounded, piecesCover, piecesDisjoint⟩ := pieces
  obtain ⟨movedBounded, movedCover, movedDisjoint⟩ := movedPieces
  have function := conePair.function
  have sphereSubset : ∀ p ∈ F, fst(p) ⊆ S²
    intros p hp
    obtain ⟨bound, S, hS, g, hg, shape⟩ :=
        product.membership.forward with (bounded with p, hp)
    rw [shape, fst.pair]
    exact powerset.membership_iff.forward with hS
  have movedSubset : ∀ p ∈ F, moved(sphereAction, S², p) ⊆ S²
    intros p hp
    rw moved.unfold
    exact translate.subset at sphereAction, S², snd(p), fst(p)
  apply Equidecomposable.of_reindex with function, finite, bounded
  . intros p hp
    have first := conePair.first with (bounded with p, hp)
    rw first
    exact cone.monotone with (piecesBounded with p, hp)
  . intros x hx
    obtain ⟨inSpace, u, hu, r, scalar, positive, boundedR, shape⟩ :=
      cone.membership.forward with hx
    obtain ⟨p, hp, member⟩ := piecesCover with u, hu
    have first := conePair.first with (bounded with p, hp)
    have uSpace := S².subset with u, (boundedA with u, hu)
    have scaledIn :=
        cone.contains with member, uSpace, scalar, positive, boundedR
    witness p
    rw [first, shape]
    exact ⟨hp, scaledIn⟩
  . intros p hp q hq different
    have firstP := conePair.first with (bounded with p, hp)
    have firstQ := conePair.first with (bounded with q, hq)
    have apart := piecesDisjoint with p, hp, q, hq, different
    rw [firstP, firstQ]
    exact cone.disjoint with (sphereSubset with p, hp), (sphereSubset with q,
        hq), apart
  . intros p hp
    have move := conePair.moved with (bounded with p, hp)
    rw move
    exact cone.monotone with (movedBounded with p, hp)
  . intros y hy
    obtain ⟨inSpace, w, hw, r, scalar, positive, boundedR, shape⟩ :=
      cone.membership.forward with hy
    obtain ⟨p, hp, member⟩ := movedCover with w, hw
    have move := conePair.moved with (bounded with p, hp)
    have wSpace := S².subset with w, (boundedB with w, hw)
    have scaledIn :=
        cone.contains with member, wSpace, scalar, positive, boundedR
    witness p
    rw [move, shape]
    exact ⟨hp, scaledIn⟩
  . intros p hp q hq different
    have moveP := conePair.moved with (bounded with p, hp)
    have moveQ := conePair.moved with (bounded with q, hq)
    have apart := movedDisjoint with p, hp, q, hq, different
    rw [moveP, moveQ]
    exact cone.disjoint with (movedSubset with p, hp), (movedSubset with q,
        hq), apart