-- Open balls, closed balls, and spheres in a metric space
-- 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.Metric.Ball
import Space

def ball (X d x : Set, r : Real) : Set :=
  { y ∈ X | relates(realLess, d(x, y), r) }

theorem ball.subset (X d x : Set, r : Real)
  shows ball(X, d, x, r) ⊆ X
proof
  intros y hy
  have parts := ball.membership.forward with hy
  exact parts.left

theorem ball.powerset (X d x : Set, r : Real)
  shows ball(X, d, x, r) ∈ 𝒫(X)
proof
  exact powerset.membership_iff.backward with (ball.subset at X, d, x, r)

theorem ball.point (X d x y : Set, r : Real)
  requires member : y ∈ ball(X, d, x, r)
  shows y ∈ X
proof
  have parts := ball.membership.forward with member
  exact parts.left

theorem ball.distance (X d x y : Set, r : Real)
  requires member : y ∈ ball(X, d, x, r)
  shows d(x, y) < r
proof
  have parts := ball.membership.forward with member
  exact parts.right

theorem ball.of_distance (X d x y : Set, r : Real)
  requires hy : y ∈ X
  requires near : d(x, y) < r
  shows y ∈ ball(X, d, x, r)
proof
  exact ball.membership.backward with ⟨hy, near⟩

theorem ball.center (X d x : Set, r : Real)
  requires metric : IsMetric(X, d)
  requires hx : x ∈ X
  requires positive : 0 < r
  shows x ∈ ball(X, d, x, r)
proof
  have self := IsMetric.self with metric, hx
  apply ball.of_distance with hx
  rw self
  exact positive

theorem ball.nonempty (X d x : Set, r : Real)
  requires metric : IsMetric(X, d)
  requires hx : x ∈ X
  requires positive : 0 < r
  shows Nonempty(ball(X, d, x, r))
proof
  witness x
  exact ball.center with metric, hx, positive

theorem ball.positive_of_member (X d x y : Set, r : Real)
  requires metric : IsMetric(X, d)
  requires hx : x ∈ X
  requires member : y ∈ ball(X, d, x, r)
  shows 0 < r
proof
  have hy := ball.point with member
  have near := ball.distance with member
  have nonnegative := IsMetric.nonnegative with metric, hx, hy
  have real := IsMetric.real with metric, hx, hy
  linarith

theorem ball.empty_of_nonpositive (X d x : Set, r : Real)
  requires metric : IsMetric(X, d)
  requires hx : x ∈ X
  requires nonpositive : r ≤ 0
  shows ball(X, d, x, r) = ∅
proof
  ext y
  constructor
  . intro hy
    have positive := ball.positive_of_member with metric, hx, hy
    exfalso
    linarith
  . intro hy
    exfalso
    exact empty.membership with hy

theorem ball.symmetric_membership (X d x y : Set, r : Real)
  requires metric : IsMetric(X, d)
  requires hx : x ∈ X
  requires hy : y ∈ X
  shows y ∈ ball(X, d, x, r) ↔ x ∈ ball(X, d, y, r)
proof
  have symmetric := IsMetric.symmetric with metric, hx, hy
  constructor
  . intro member
    have near := ball.distance with member
    apply ball.of_distance with hx
    rwa symmetric at near
  . intro member
    have near := ball.distance with member
    apply ball.of_distance with hy
    rwa symmetric

theorem ball.mono (X d x : Set, r s : Real)
  requires metric : IsMetric(X, d)
  requires hx : x ∈ X
  requires radii : r ≤ s
  shows ball(X, d, x, r) ⊆ ball(X, d, x, s)
proof
  intros y hy
  have point := ball.point with hy
  have near := ball.distance with hy
  have real := IsMetric.real with metric, hx, point
  apply ball.of_distance with point
  linarith

theorem ball.mono_member (X d x y : Set, r s : Real)
  requires metric : IsMetric(X, d)
  requires hx : x ∈ X
  requires member : y ∈ ball(X, d, x, r)
  requires radii : r ≤ s
  shows y ∈ ball(X, d, x, s)
proof
  have inside := ball.mono with metric, hx, radii
  exact inside with y, member

-- Stepping from the center of one ball into another adds the radii.
theorem ball.step (X d x y z : Set, r s : Real)
  requires metric : IsMetric(X, d)
  requires hx : x ∈ X
  requires first : y ∈ ball(X, d, x, r)
  requires second : z ∈ ball(X, d, y, s)
  shows z ∈ ball(X, d, x, r + s)
proof
  have hy := ball.point with first
  have hz := ball.point with second
  have nearY := ball.distance with first
  have nearZ := ball.distance with second
  have bound := IsMetric.lt_add_of_lt with metric, hx, hy, hz, nearY, nearZ
  exact ball.of_distance with hz, bound

-- A ball around a point of a ball, with the remaining radius, stays inside.
theorem ball.inside_of_le (X d x y : Set, r s : Real)
  requires metric : IsMetric(X, d)
  requires hx : x ∈ X
  requires hy : y ∈ X
  requires room : realAdd(d(x, y), s) ≤ r
  shows ball(X, d, y, s) ⊆ ball(X, d, x, r)
proof
  intros z hz
  have point := ball.point with hz
  have near := ball.distance with hz
  have triangle := IsMetric.triangle with metric, hx, hy, point
  have dxy := IsMetric.real with metric, hx, hy
  have dyz := IsMetric.real with metric, hy, point
  have dxz := IsMetric.real with metric, hx, point
  apply ball.of_distance with point
  linarith

theorem ball.inside (X d x y : Set, r : Real)
  requires metric : IsMetric(X, d)
  requires hx : x ∈ X
  requires member : y ∈ ball(X, d, x, r)
  shows ball(X, d, y, realSubValue(r, d(x, y))) ⊆ ball(X, d, x, r)
proof
  have hy := ball.point with member
  have dxy := IsMetric.real with metric, hx, hy
  have room : d(x, y) + (r - d(x, y)) ≤ r by linarith
  exact ball.inside_of_le with metric, hx, hy, room

theorem ball.inside_radius_positive (X d x y : Set, r : Real)
  requires metric : IsMetric(X, d)
  requires hx : x ∈ X
  requires member : y ∈ ball(X, d, x, r)
  shows 0 < realSubValue(r, d(x, y))
proof
  have hy := ball.point with member
  have near := ball.distance with member
  have dxy := IsMetric.real with metric, hx, hy
  linarith

-- Balls whose radii together do not reach the distance between the centers
-- are disjoint.
theorem ball.disjoint_of_separated (X d x y : Set, r s : Real)
  requires metric : IsMetric(X, d)
  requires hx : x ∈ X
  requires hy : y ∈ X
  requires separated : r + s ≤ d(x, y)
  shows Disjoint(ball(X, d, x, r), ball(X, d, y, s))
proof
  apply Disjoint.iff.backward
  intros z hzx hzy
  have point := ball.point with hzx
  have nearX := ball.distance with hzx
  have nearY := ball.distance with hzy
  have triangle := IsMetric.triangle_to with metric, hx, hy, point
  have dxy := IsMetric.real with metric, hx, hy
  have dxz := IsMetric.real with metric, hx, point
  have dyz := IsMetric.real with metric, hy, point
  linarith

-- Distinct points have disjoint balls of half their distance.
theorem ball.half_disjoint (X d x y : Set)
  requires metric : IsMetric(X, d)
  requires hx : x ∈ X
  requires hy : y ∈ X
  requires distinct : x ≠ y
  shows Disjoint(ball(X, d, x, half(d(x, y))), ball(X, d, y, half(d(x, y))))
proof
  have dxy := IsMetric.real with metric, hx, hy
  have doubled := half.add_self at d(x, y)
  have separated : half(d(x, y)) + half(d(x, y)) ≤ d(x, y) by linarith
  exact ball.disjoint_of_separated with metric, hx, hy, separated

theorem ball.half_positive (X d x y : Set)
  requires metric : IsMetric(X, d)
  requires hx : x ∈ X
  requires hy : y ∈ X
  requires distinct : x ≠ y
  shows 0 < half(d(x, y))
proof
  have positive := IsMetric.positive_of_ne with metric, hx, hy, distinct
  have dxy := IsMetric.real with metric, hx, hy
  exact half.positive with positive

-- A point of two balls has a ball inside both.
theorem ball.common_radius (X d x y z : Set, r s : Real)
  requires metric : IsMetric(X, d)
  requires hx : x ∈ X
  requires hy : y ∈ X
  requires first : z ∈ ball(X, d, x, r)
  requires second : z ∈ ball(X, d, y, s)
  shows ∃ t ∈ ℝ, 0 < t ∧
    ball(X, d, z, t) ⊆ ball(X, d, x, r) ∧ ball(X, d, z, t) ⊆ ball(X, d, y, s)
proof
  have hz := ball.point with first
  have dxz := IsMetric.real with metric, hx, hz
  have dyz := IsMetric.real with metric, hy, hz
  have positiveX := ball.inside_radius_positive with metric, hx, first
  have positiveY := ball.inside_radius_positive with metric, hy, second
  have insideX := ball.inside with metric, hx, first
  have insideY := ball.inside with metric, hy, second
  have order := realLess.trichotomy at r - d(x, z), s - d(y, z)
  rcases order with smaller | equal | larger
  . have le := realLess.implies_le with smaller
    have realT : r - d(x, z) ∈ ℝ by typecheck
    refine ⟨r - d(x, z), realT, positiveX, insideX, ?_⟩
    have mono := ball.mono with metric, hz, le
    exact Subset.transitive with mono, insideY
  . have realT : r - d(x, z) ∈ ℝ by typecheck
    refine ⟨r - d(x, z), realT, positiveX, insideX, ?_⟩
    rw equal
    exact insideY
  . have le := realLess.implies_le with larger
    have realT : s - d(y, z) ∈ ℝ by typecheck
    refine ⟨s - d(y, z), realT, positiveY, ?_, insideY⟩
    have mono := ball.mono with metric, hz, le
    exact Subset.transitive with mono, insideX

-- Closed balls

def closedBall (X d x : Set, r : Real) : Set :=
  { y ∈ X | relates(realLe, d(x, y), r) }

theorem closedBall.subset (X d x : Set, r : Real)
  shows closedBall(X, d, x, r) ⊆ X
proof
  intros y hy
  have parts := closedBall.membership.forward with hy
  exact parts.left

theorem closedBall.point (X d x y : Set, r : Real)
  requires member : y ∈ closedBall(X, d, x, r)
  shows y ∈ X
proof
  have parts := closedBall.membership.forward with member
  exact parts.left

theorem closedBall.distance (X d x y : Set, r : Real)
  requires member : y ∈ closedBall(X, d, x, r)
  shows d(x, y) ≤ r
proof
  have parts := closedBall.membership.forward with member
  exact parts.right

theorem closedBall.of_distance (X d x y : Set, r : Real)
  requires hy : y ∈ X
  requires near : d(x, y) ≤ r
  shows y ∈ closedBall(X, d, x, r)
proof
  exact closedBall.membership.backward with ⟨hy, near⟩

theorem closedBall.center (X d x : Set, r : Real)
  requires metric : IsMetric(X, d)
  requires hx : x ∈ X
  requires nonnegative : 0 ≤ r
  shows x ∈ closedBall(X, d, x, r)
proof
  have self := IsMetric.self with metric, hx
  apply closedBall.of_distance with hx
  rw self
  exact nonnegative

theorem ball.subset_closedBall (X d x : Set, r : Real)
  requires metric : IsMetric(X, d)
  requires hx : x ∈ X
  shows ball(X, d, x, r) ⊆ closedBall(X, d, x, r)
proof
  intros y hy
  have point := ball.point with hy
  have near := ball.distance with hy
  have real := IsMetric.real with metric, hx, point
  apply closedBall.of_distance with point
  exact realLess.implies_le with near

theorem closedBall.subset_ball (X d x : Set, r s : Real)
  requires metric : IsMetric(X, d)
  requires hx : x ∈ X
  requires radii : r < s
  shows closedBall(X, d, x, r) ⊆ ball(X, d, x, s)
proof
  intros y hy
  have point := closedBall.point with hy
  have near := closedBall.distance with hy
  have real := IsMetric.real with metric, hx, point
  apply ball.of_distance with point
  linarith

theorem closedBall.mono (X d x : Set, r s : Real)
  requires metric : IsMetric(X, d)
  requires hx : x ∈ X
  requires radii : r ≤ s
  shows closedBall(X, d, x, r) ⊆ closedBall(X, d, x, s)
proof
  intros y hy
  have point := closedBall.point with hy
  have near := closedBall.distance with hy
  have real := IsMetric.real with metric, hx, point
  apply closedBall.of_distance with point
  linarith

theorem closedBall.empty_of_negative (X d x : Set, r : Real)
  requires metric : IsMetric(X, d)
  requires hx : x ∈ X
  requires negative : r < 0
  shows closedBall(X, d, x, r) = ∅
proof
  ext y
  constructor
  . intro hy
    have point := closedBall.point with hy
    have near := closedBall.distance with hy
    have nonnegative := IsMetric.nonnegative with metric, hx, point
    have real := IsMetric.real with metric, hx, point
    exfalso
    linarith
  . intro hy
    exfalso
    exact empty.membership with hy

theorem closedBall.zero (X d x : Set)
  requires metric : IsMetric(X, d)
  requires hx : x ∈ X
  shows closedBall(X, d, x, 0) = {x}
proof
  ext y
  constructor
  . intro hy
    have point := closedBall.point with hy
    have near := closedBall.distance with hy
    have nonnegative := IsMetric.nonnegative with metric, hx, point
    have real := IsMetric.real with metric, hx, point
    have zero : d(x, y) = (0:R) by linarith
    have same := IsMetric.eq_of_zero with metric, hx, point, zero
    rw same
    exact singleton.self at y
  . intro hy
    have same := singleton.membership.forward with hy
    rw same
    have reflexive : (0:R) ≤ 0 by linarith
    exact closedBall.center with metric, hx, reflexive

theorem closedBall.symmetric_membership (X d x y : Set, r : Real)
  requires metric : IsMetric(X, d)
  requires hx : x ∈ X
  requires hy : y ∈ X
  shows y ∈ closedBall(X, d, x, r) ↔ x ∈ closedBall(X, d, y, r)
proof
  have symmetric := IsMetric.symmetric with metric, hx, hy
  constructor
  . intro member
    have near := closedBall.distance with member
    apply closedBall.of_distance with hx
    rwa symmetric at near
  . intro member
    have near := closedBall.distance with member
    apply closedBall.of_distance with hy
    rwa symmetric

theorem closedBall.step (X d x y z : Set, r s : Real)
  requires metric : IsMetric(X, d)
  requires hx : x ∈ X
  requires first : y ∈ closedBall(X, d, x, r)
  requires second : z ∈ closedBall(X, d, y, s)
  shows z ∈ closedBall(X, d, x, r + s)
proof
  have hy := closedBall.point with first
  have hz := closedBall.point with second
  have nearY := closedBall.distance with first
  have nearZ := closedBall.distance with second
  have bound := IsMetric.le_add_of_le with metric, hx, hy, hz, nearY, nearZ
  exact closedBall.of_distance with hz, bound

-- Spheres

def sphere (X d x : Set, r : Real) : Set := { y ∈ X | d(x, y) = r }

theorem sphere.subset (X d x : Set, r : Real)
  shows sphere(X, d, x, r) ⊆ X
proof
  intros y hy
  have parts := sphere.membership.forward with hy
  exact parts.left

theorem sphere.subset_closedBall (X d x : Set, r : Real)
  requires metric : IsMetric(X, d)
  requires hx : x ∈ X
  shows sphere(X, d, x, r) ⊆ closedBall(X, d, x, r)
proof
  intros y hy
  obtain ⟨point, equal⟩ := sphere.membership.forward with hy
  have real := IsMetric.real with metric, hx, point
  apply closedBall.of_distance with point
  rw equal
  exact realLe.reflexive

theorem sphere.disjoint_ball (X d x : Set, r : Real)
  shows Disjoint(sphere(X, d, x, r), ball(X, d, x, r))
proof
  apply Disjoint.iff.backward
  intros y onSphere inBall
  obtain ⟨point, equal⟩ := sphere.membership.forward with onSphere
  have near := ball.distance with inBall
  rw equal at near
  exact realLess.not_self with near

theorem closedBall.eq_ball_union_sphere (X d x : Set, r : Real)
  requires metric : IsMetric(X, d)
  requires hx : x ∈ X
  shows closedBall(X, d, x, r) = ball(X, d, x, r) ∪ sphere(X, d, x, r)
proof
  apply Subset.antisymmetric
  . intros y hy
    have point := closedBall.point with hy
    have near := closedBall.distance with hy
    have real := IsMetric.real with metric, hx, point
    have order := realLess.trichotomy at d(x, y), r
    rcases order with less | equal | greater
    . exact union2.left with (ball.of_distance with point, less)
    . have onSphere : y ∈ sphere(X, d, x, r) := sphere.membership.backward with
      ⟨point, equal⟩
      exact union2.right with onSphere
    . exfalso
      linarith
  . apply union2.subset
    . exact ball.subset_closedBall with metric, hx
    . exact sphere.subset_closedBall with metric, hx

theorem sphere.zero (X d x : Set)
  requires metric : IsMetric(X, d)
  requires hx : x ∈ X
  shows sphere(X, d, x, 0) = {x}
proof
  ext y
  constructor
  . intro hy
    obtain ⟨point, zero⟩ := sphere.membership.forward with hy
    have same := IsMetric.eq_of_zero with metric, hx, point, zero
    rw same
    exact singleton.self at y
  . intro hy
    have same := singleton.membership.forward with hy
    rw same
    have self := IsMetric.self with metric, hx
    exact sphere.membership.backward with ⟨hx, self⟩

theorem sphere.empty_of_negative (X d x : Set, r : Real)
  requires metric : IsMetric(X, d)
  requires hx : x ∈ X
  requires negative : r < 0
  shows sphere(X, d, x, r) = ∅
proof
  ext y
  constructor
  . intro hy
    obtain ⟨point, equal⟩ := sphere.membership.forward with hy
    have nonnegative := IsMetric.nonnegative with metric, hx, point
    rw equal at nonnegative
    exfalso
    linarith
  . intro hy
    exfalso
    exact empty.membership with hy