-- The product of two metric spaces with the sum metric: its distance
-- function, the metric laws, the projections as Lipschitz maps, balls
-- against products of balls, and bounded products
-- 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.Product
import Subspace

def productDistance (X d Y e : Set) : Set :=
  { q ∈ ((X × Y) × (X × Y)) × ℝ | ∃ x ∈ X, ∃ y ∈ Y, ∃ a ∈ X, ∃ b ∈ Y,
    q = (((x, y), (a, b)) ↦ realAdd(d(x, a), e(y, b))) }

theorem productDistance.sum_real (X d Y e x y a b : Set)
  requires metricD : IsMetric(X, d)
  requires metricE : IsMetric(Y, e)
  requires hx : x ∈ X
  requires hy : y ∈ Y
  requires ha : a ∈ X
  requires hb : b ∈ Y
  shows realAdd(d(x, a), e(y, b)) ∈ ℝ
proof
  have dxa := IsMetric.real with metricD, hx, ha
  have eyb := IsMetric.real with metricE, hy, hb
  have sum : d(x, a) + e(y, b) ∈ ℝ by typecheck
  exact sum

theorem productDistance.relates (X d Y e x y a b : Set)
  requires metricD : IsMetric(X, d)
  requires metricE : IsMetric(Y, e)
  requires hx : x ∈ X
  requires hy : y ∈ Y
  requires ha : a ∈ X
  requires hb : b ∈ Y
  shows relates(productDistance(X, d, Y, e), ((x, y), (a, b)),
    realAdd(d(x, a), e(y, b)))
proof
  have sum := productDistance.sum_real with metricD, metricE, hx, hy, ha, hb
  have first := product.pair at x, y, X, Y with hx, hy
  have second := product.pair at a, b, X, Y with ha, hb
  have pair := product.pair at (x, y), (a, b), X × Y, X × Y with first, second
  have entry := product.pair at ((x, y), (a, b)), realAdd(d(x, a), e(y, b)),
    (X × Y) × (X × Y), ℝ with pair, sum
  unfold relates
  apply productDistance.membership.backward
  refine ⟨entry, x, hx, y, hy, a, ha, b, hb, ?_⟩
  rfl

theorem productDistance.value_shape (X d Y e q v : Set)
  requires edge : relates(productDistance(X, d, Y, e), q, v)
  shows ∃ x ∈ X, ∃ y ∈ Y, ∃ a ∈ X, ∃ b ∈ Y,
    q = ((x, y), (a, b)) ∧ v = realAdd(d(x, a), e(y, b))
proof
  unfold relates at edge
  obtain ⟨bound, x, hx, y, hy, a, ha, b, hb, shape⟩ :=
    productDistance.membership.forward with edge
  obtain ⟨first, second⟩ := orderedPair.injective with shape
  exact ⟨x, hx, y, hy, a, ha, b, hb, first, second⟩

theorem productDistance.relation (X d Y e : Set)
  shows RelationBetween(productDistance(X, d, Y, e), (X × Y) × (X × Y), ℝ)
proof
  intros q hq
  have parts := productDistance.membership.forward with hq
  exact parts.left

theorem productDistance.functional (X d Y e : Set)
  shows Functional(productDistance(X, d, Y, e))
proof
  intros q v first w second
  obtain ⟨x, hx, y, hy, a, ha, b, hb, shapeV, valueV⟩ :=
    productDistance.value_shape with first
  obtain ⟨x2, hx2, y2, hy2, a2, ha2, b2, hb2, shapeW, valueW⟩ :=
    productDistance.value_shape with second
  rw shapeV at shapeW
  obtain ⟨firstPair, secondPair⟩ := orderedPair.injective with shapeW
  obtain ⟨xx, yy⟩ := orderedPair.injective with firstPair
  obtain ⟨aa, bb⟩ := orderedPair.injective with secondPair
  rw [valueV, valueW, xx, yy, aa, bb]
  rfl

theorem productDistance.total (X d Y e : Set)
  requires metricD : IsMetric(X, d)
  requires metricE : IsMetric(Y, e)
  shows
    ∀ q ∈ (X × Y) × (X × Y), ∃ v ∈ ℝ, relates(productDistance(X, d, Y, e), q, v)
proof
  intros q hq
  obtain ⟨bound, p, hp, p2, hp2, rfl⟩ := product.membership.forward with hq
  obtain ⟨boundP, x, hx, y, hy, rfl⟩ := product.membership.forward with hp
  obtain ⟨boundP2, a, ha, b, hb, rfl⟩ := product.membership.forward with hp2
  have sum := productDistance.sum_real with metricD, metricE, hx, hy, ha, hb
  refine ⟨realAdd(d(x, a), e(y, b)), sum, ?_⟩
  exact productDistance.relates with metricD, metricE, hx, hy, ha, hb

theorem productDistance.function (X d Y e : Set)
  requires metricD : IsMetric(X, d)
  requires metricE : IsMetric(Y, e)
  shows productDistance(X, d, Y, e) : (X × Y) × (X × Y) → ℝ
proof
  constructor
  . exact productDistance.relation at X, d, Y, e
  . constructor
    . exact productDistance.functional at X, d, Y, e
    . exact productDistance.total with metricD, metricE

theorem productDistance.apply (X d Y e x y a b : Set)
  requires metricD : IsMetric(X, d)
  requires metricE : IsMetric(Y, e)
  requires hx : x ∈ X
  requires hy : y ∈ Y
  requires ha : a ∈ X
  requires hb : b ∈ Y
  shows productDistance(X, d, Y, e)((x, y), (a, b)) = realAdd(d(x, a), e(y, b))
proof
  have function := productDistance.function with metricD, metricE
  have first := product.pair at x, y, X, Y with hx, hy
  have second := product.pair at a, b, X, Y with ha, hb
  have pair := product.pair at (x, y), (a, b), X × Y, X × Y with first, second
  have edge := productDistance.relates with metricD, metricE, hx, hy, ha, hb
  have left := Metric.apply_pair at productDistance(X, d, Y, e), (x, y), (a, b)
  rw ← left
  exact FunctionBetween.apply_value with function, pair, edge

theorem productDistance.metric (X d Y e : Set)
  requires metricD : IsMetric(X, d)
  requires metricE : IsMetric(Y, e)
  shows IsMetric(X × Y, productDistance(X, d, Y, e))
proof
  constructor
  . exact productDistance.function with metricD, metricE
  . constructor
    . intros p hp q hq
      obtain ⟨boundP, x, hx, y, hy, rfl⟩ := product.membership.forward with hp
      obtain ⟨boundQ, a, ha, b, hb, rfl⟩ := product.membership.forward with hq
      have value := productDistance.apply with metricD, metricE, hx, hy, ha, hb
      rw value
      have dxa := IsMetric.real with metricD, hx, ha
      have eyb := IsMetric.real with metricE, hy, hb
      have nonnegativeD := IsMetric.nonnegative with metricD, hx, ha
      have nonnegativeE := IsMetric.nonnegative with metricE, hy, hb
      linarith
    . constructor
      . intros p hp q hq
        obtain ⟨boundP, x, hx, y, hy, rfl⟩ := product.membership.forward with hp
        obtain ⟨boundQ, a, ha, b, hb, rfl⟩ := product.membership.forward with hq
        have value := productDistance.apply with
          metricD, metricE, hx, hy, ha, hb
        rw value
        have dxa := IsMetric.real with metricD, hx, ha
        have eyb := IsMetric.real with metricE, hy, hb
        have nonnegativeD := IsMetric.nonnegative with metricD, hx, ha
        have nonnegativeE := IsMetric.nonnegative with metricE, hy, hb
        constructor
        . intro zero
          have zeroD : d(x, a) = (0:R) by linarith
          have zeroE : e(y, b) = (0:R) by linarith
          have xa := IsMetric.eq_of_zero with metricD, hx, ha, zeroD
          have yb := IsMetric.eq_of_zero with metricE, hy, hb, zeroE
          rw [xa, yb]
          rfl
        . intro same
          obtain ⟨xa, yb⟩ := orderedPair.injective with same
          rw [xa, yb]
          have selfA := IsMetric.self with metricD, ha
          have selfB := IsMetric.self with metricE, hb
          rw [selfA, selfB]
          linarith
      . constructor
        . intros p hp q hq
          obtain ⟨boundP, x, hx, y, hy, rfl⟩ :=
            product.membership.forward with hp
          obtain ⟨boundQ, a, ha, b, hb, rfl⟩ :=
            product.membership.forward with hq
          have value := productDistance.apply with
            metricD, metricE, hx, hy, ha, hb
          have valueBack := productDistance.apply with
            metricD, metricE, ha, hb, hx, hy
          rw [value, valueBack]
          have symmetricD := IsMetric.symmetric with metricD, hx, ha
          have symmetricE := IsMetric.symmetric with metricE, hy, hb
          rw [symmetricD, symmetricE]
          rfl
        . intros p hp q hq s hs
          obtain ⟨boundP, x, hx, y, hy, rfl⟩ :=
            product.membership.forward with hp
          obtain ⟨boundQ, a, ha, b, hb, rfl⟩ :=
            product.membership.forward with hq
          obtain ⟨boundS, c, hc, g, hg, rfl⟩ :=
            product.membership.forward with hs
          have valuePS := productDistance.apply with
            metricD, metricE, hx, hy, hc, hg
          have valuePQ := productDistance.apply with
            metricD, metricE, hx, hy, ha, hb
          have valueQS := productDistance.apply with
            metricD, metricE, ha, hb, hc, hg
          rw [valuePS, valuePQ, valueQS]
          have triangleD := IsMetric.triangle with metricD, hx, ha, hc
          have triangleE := IsMetric.triangle with metricE, hy, hb, hg
          have dxc := IsMetric.real with metricD, hx, hc
          have dxa := IsMetric.real with metricD, hx, ha
          have dac := IsMetric.real with metricD, ha, hc
          have eyg := IsMetric.real with metricE, hy, hg
          have eyb := IsMetric.real with metricE, hy, hb
          have ebg := IsMetric.real with metricE, hb, hg
          linarith

-- The projections are 1-Lipschitz, hence continuous.
theorem firstProjection.lipschitz (X d Y e : Set)
  requires metricD : IsMetric(X, d)
  requires metricE : IsMetric(Y, e)
  shows IsLipschitz(firstProjection(X, Y), X × Y, productDistance(X, d, Y, e),
    X, d, (1:R))
proof
  have one : (0:R) ≤ 1 by linarith
  unfold IsLipschitz
  constructor
  . exact firstProjection.function at X, Y
  . constructor
    . exact one
    . intros p hp q hq
      obtain ⟨boundP, x, hx, y, hy, rfl⟩ := product.membership.forward with hp
      obtain ⟨boundQ, a, ha, b, hb, rfl⟩ := product.membership.forward with hq
      have valueP := firstProjection.pair at X, Y, x, y with hx, hy
      have valueQ := firstProjection.pair at X, Y, a, b with ha, hb
      have value := productDistance.apply with metricD, metricE, hx, hy, ha, hb
      rw [valueP, valueQ, value]
      have dxa := IsMetric.real with metricD, hx, ha
      have eyb := IsMetric.real with metricE, hy, hb
      have nonnegativeE := IsMetric.nonnegative with metricE, hy, hb
      have unit := realMulValue.one_left at d(x, a) + e(y, b)
      rw unit
      linarith

theorem secondProjection.lipschitz (X d Y e : Set)
  requires metricD : IsMetric(X, d)
  requires metricE : IsMetric(Y, e)
  shows IsLipschitz(secondProjection(X, Y), X × Y, productDistance(X, d, Y, e),
    Y, e, (1:R))
proof
  have one : (0:R) ≤ 1 by linarith
  unfold IsLipschitz
  constructor
  . exact secondProjection.function at X, Y
  . constructor
    . exact one
    . intros p hp q hq
      obtain ⟨boundP, x, hx, y, hy, rfl⟩ := product.membership.forward with hp
      obtain ⟨boundQ, a, ha, b, hb, rfl⟩ := product.membership.forward with hq
      have valueP := secondProjection.pair at X, Y, x, y with hx, hy
      have valueQ := secondProjection.pair at X, Y, a, b with ha, hb
      have value := productDistance.apply with metricD, metricE, hx, hy, ha, hb
      rw [valueP, valueQ, value]
      have dxa := IsMetric.real with metricD, hx, ha
      have eyb := IsMetric.real with metricE, hy, hb
      have nonnegativeD := IsMetric.nonnegative with metricD, hx, ha
      have unit := realMulValue.one_left at d(x, a) + e(y, b)
      rw unit
      linarith

theorem firstProjection.metric_continuous (X d Y e : Set)
  requires metricD : IsMetric(X, d)
  requires metricE : IsMetric(Y, e)
  shows IsMetricContinuous(firstProjection(X, Y), X × Y,
    productDistance(X, d, Y, e), X, d)
proof
  have product := productDistance.metric with metricD, metricE
  have lipschitz := firstProjection.lipschitz with metricD, metricE
  exact IsLipschitz.continuous with product, metricD, lipschitz

theorem secondProjection.metric_continuous (X d Y e : Set)
  requires metricD : IsMetric(X, d)
  requires metricE : IsMetric(Y, e)
  shows IsMetricContinuous(secondProjection(X, Y), X × Y,
    productDistance(X, d, Y, e), Y, e)
proof
  have product := productDistance.metric with metricD, metricE
  have lipschitz := secondProjection.lipschitz with metricD, metricE
  exact IsLipschitz.continuous with product, metricE, lipschitz

-- Balls in the product against products of balls.
theorem productDistance.ball_subset_product (X d Y e x y : Set, r : Real)
  requires metricD : IsMetric(X, d)
  requires metricE : IsMetric(Y, e)
  requires hx : x ∈ X
  requires hy : y ∈ Y
  shows ball(X × Y, productDistance(X, d, Y, e), (x, y), r) ⊆
    ball(X, d, x, r) × ball(Y, e, y, r)
proof
  intros q hq
  have hqP := ball.point with hq
  have near := ball.distance with hq
  obtain ⟨boundQ, a, ha, b, hb, rfl⟩ := product.membership.forward with hqP
  have value := productDistance.apply with metricD, metricE, hx, hy, ha, hb
  rw value at near
  have dxa := IsMetric.real with metricD, hx, ha
  have eyb := IsMetric.real with metricE, hy, hb
  have nonnegativeD := IsMetric.nonnegative with metricD, hx, ha
  have nonnegativeE := IsMetric.nonnegative with metricE, hy, hb
  have nearA : d(x, a) < r by linarith
  have nearB : e(y, b) < r by linarith
  apply product.pair
  . exact ball.of_distance with ha, nearA
  . exact ball.of_distance with hb, nearB

theorem productDistance.product_subset_ball (X d Y e x y : Set, r s : Real)
  requires metricD : IsMetric(X, d)
  requires metricE : IsMetric(Y, e)
  requires hx : x ∈ X
  requires hy : y ∈ Y
  shows ball(X, d, x, r) × ball(Y, e, y, s) ⊆
    ball(X × Y, productDistance(X, d, Y, e), (x, y), r + s)
proof
  intros q hq
  obtain ⟨boundQ, a, ha, b, hb, rfl⟩ := product.membership.forward with hq
  have haX := ball.point with ha
  have hbY := ball.point with hb
  have nearA := ball.distance with ha
  have nearB := ball.distance with hb
  have value := productDistance.apply with metricD, metricE, hx, hy, haX, hbY
  have member := product.pair at a, b, X, Y with haX, hbY
  apply ball.of_distance with member
  rw value
  have dxa := IsMetric.real with metricD, hx, haX
  have eyb := IsMetric.real with metricE, hy, hbY
  linarith

theorem IsBounded.product (X d Y e S T : Set)
  requires metricD : IsMetric(X, d)
  requires metricE : IsMetric(Y, e)
  requires boundedS : IsBounded(X, d, S)
  requires boundedT : IsBounded(Y, e, T)
  shows IsBounded(X × Y, productDistance(X, d, Y, e), S × T)
proof
  obtain ⟨insideS, r, hr, boundS⟩ := boundedS
  obtain ⟨insideT, s, hs, boundT⟩ := boundedT
  have typedR : Real(r) by exact hr
  have typedS : Real(s) by exact hs
  apply IsBounded.of_bound at
    X × Y, productDistance(X, d, Y, e), S × T, r + s with
    (product.monotone with insideS, insideT)
  intros p hp q hq
  obtain ⟨boundP, x, hx, y, hy, rfl⟩ := product.membership.forward with hp
  obtain ⟨boundQ, a, ha, b, hb, rfl⟩ := product.membership.forward with hq
  have hxX := insideS with x, hx
  have haX := insideS with a, ha
  have hyY := insideT with y, hy
  have hbY := insideT with b, hb
  have value := productDistance.apply with metricD, metricE, hxX, hyY, haX, hbY
  rw value
  have dxa := IsMetric.real with metricD, hxX, haX
  have eyb := IsMetric.real with metricE, hyY, hbY
  have estimateD := boundS with x, hx, a, ha
  have estimateE := boundT with y, hy, b, hb
  linarith