AutoSamplingTheory.TechnicalLemmas.Measure.ChewiTheorem1_3_23

import AutoSamplingTheory.TechnicalLemmas.Measure.DisplacementInterpolationConstantSpeed
import AutoSamplingTheory.TechnicalLemmas.Measure.WassersteinSymmetry

/-!
# Chewi Theorem 1.3.23: constant-speed displacement interpolation

This is the source-facing closure of the Chapter 1.3 transport spine.

For an optimal quadratic coupling `γ` of two `P₂,ac` endpoint laws, the
McCann/displacement interpolation travels at constant `W₂` speed:

`W₂(μ_s, μ_t) = |s-t| W₂(μ₀, μ₁)` for `s,t ∈ [0,1]`.

The hard ingredients are already proved in separate reusable modules:

* finite-second-moment endpoints give finite `W₂`;
* the exact `W₂` triangle inequality;
* exact cost scaling of the canonical two-time interpolation coupling;
* ordered-time constant speed;
* symmetry of `W₂`.

This file only closes the order split and packages the result through the
source-side `IsWassersteinGeodesic` predicate.
-/

namespace AutoSamplingTheory
namespace TechnicalLemmas
namespace Measure
namespace ChewiTheorem1_3_23

open MeasureTheory Set
open scoped ENNReal

noncomputable section

variable {E : Type*} [NormedAddCommGroup E] [InnerProductSpace ℝ E]
  [FiniteDimensional ℝ E] [MeasurableSpace E] [BorelSpace E]
  [SecondCountableTopology E] [StandardBorelSpace E] [Nonempty E]

open DisplacementInterpolation

/-- Arbitrary-time form of the constant-speed identity for the displacement
interpolation generated by one quadratic-optimal coupling. -/
theorem wassersteinDistance_displacementInterpolation_eq_abs
    {γ : Measure (E × E)} {μ₀ μ₁ : Measure E}
    (hμ₀ : WassersteinSpace.IsAbsolutelyContinuousFiniteSecondMoment μ₀)
    (hμ₁ : WassersteinSpace.IsAbsolutelyContinuousFiniteSecondMoment μ₁)
    (hγ : IsQuadraticOptimalCoupling γ μ₀ μ₁)
    {s t : ℝ} (hs : s ∈ Icc (0 : ℝ) 1) (ht : t ∈ Icc (0 : ℝ) 1) :
    WassersteinSpace.wassersteinDistance
        (displacementInterpolation γ s)
        (displacementInterpolation γ t) =
      ENNReal.ofReal |s - t| *
        WassersteinSpace.wassersteinDistance μ₀ μ₁ := by
  rcases le_total s t with hst | hts
  · have h :=
      DisplacementInterpolationConstantSpeed.wassersteinDistance_interpolation_eq_of_le
        (E := E) hμ₀ hμ₁ hγ hs.1 hst ht.2
    have habs : |s - t| = t - s := by
      rw [abs_sub_comm, abs_of_nonneg (sub_nonneg.mpr hst)]
    simpa [habs] using h
  · have h :=
      DisplacementInterpolationConstantSpeed.wassersteinDistance_interpolation_eq_of_le
        (E := E) hμ₀ hμ₁ hγ ht.1 hts hs.2
    calc
      WassersteinSpace.wassersteinDistance
          (displacementInterpolation γ s)
          (displacementInterpolation γ t) =
        WassersteinSpace.wassersteinDistance
          (displacementInterpolation γ t)
          (displacementInterpolation γ s) :=
        WassersteinSymmetry.wassersteinDistance_comm _ _
      _ = ENNReal.ofReal (s - t) *
          WassersteinSpace.wassersteinDistance μ₀ μ₁ := h
      _ = ENNReal.ofReal |s - t| *
          WassersteinSpace.wassersteinDistance μ₀ μ₁ := by
        rw [abs_of_nonneg (sub_nonneg.mpr hts)]

/-- Chewi Theorem 1.3.23 in the repository's source-side geodesic interface:
any curve represented by an optimal displacement interpolation has constant
`W₂` speed on `[0,1]`. -/
theorem chewi_theorem_1_3_23_constant_speed
    {μ₀ μ₁ : Measure E} {curve : ℝ → Measure E}
    (hcurve : IsWassersteinGeodesic μ₀ μ₁ curve)
    {s t : ℝ} (hs : s ∈ Icc (0 : ℝ) 1) (ht : t ∈ Icc (0 : ℝ) 1) :
    WassersteinSpace.wassersteinDistance (curve s) (curve t) =
      ENNReal.ofReal |s - t| *
        WassersteinSpace.wassersteinDistance μ₀ μ₁ := by
  rcases hcurve with ⟨hμ₀, hμ₁, γ, hγ, hrepr⟩
  rw [hrepr s hs, hrepr t ht]
  exact wassersteinDistance_displacementInterpolation_eq_abs
    (E := E) hμ₀ hμ₁ hγ hs ht

end

end ChewiTheorem1_3_23
end Measure
end TechnicalLemmas
end AutoSamplingTheory