Joint law for the finite-copy suspension

import Definitions.Def_KLSModular_FunctionalInequalities_KLS_BKL_Suspension_Basic
import Definitions.Def_KLSModular_FunctionalInequalities_Probability_Line_Laplace
import Mathlib.MeasureTheory.Constructions.Pi
import Mathlib.Analysis.InnerProductSpace.PiL2
import Definitions.Def_KLSModular_FunctionalInequalities_KLS_Isotropic
import Definitions.Def_KLSModular_FunctionalInequalities_Probability_LogConcave
import Definitions.Def_KLSModular_FunctionalInequalities_KLS_LogConcave_Density_Tilted
import Mathlib.MeasureTheory.Measure.Tilted
import Definitions.Def_KLSModular_FunctionalInequalities_Analysis_BakryEmery_Cumulant_Basic
import Definitions.Def_KLSModular_FunctionalInequalities_KLS_LogConcave_Transport

noncomputable section
open MeasureTheory
open FunctionalInequalities.Probability.Line
namespace FunctionalInequalities
namespace KLS
namespace LogConcave
namespace Suspension
/-- The joint space `ℝ^{Nn+1}`, indexed by `none` (the coordinate `S`) and pairs `(i, j)`. -/
abbrev JointSpace (N n : ℕ) : Type := EuclideanSpace ℝ (Option (Fin N × Fin n))
/-- The source's `σ_β = √(1 + 2/β²)`. -/
def sigma (β : ℝ) : ℝ := Real.sqrt (1 + 2 / β ^ 2)
/-- `F_N = N^{-1/2} Σ_i f(x_i)`. -/
def copyStat {n : ℕ} (f : EuclideanSpace ℝ (Fin n) → ℝ) (N : ℕ)
    (x : Fin N → EuclideanSpace ℝ (Fin n)) : ℝ :=
  (∑ i, f (x i)) / Real.sqrt N
/-- The map `(x₁, …, x_N, s) ↦ ((F_N(x) + s)/σ, x₁, …, x_N)` into the joint space. -/
def embed {n : ℕ} (f : EuclideanSpace ℝ (Fin n) → ℝ) (σ : ℝ) (N : ℕ)
    (q : (Fin N → EuclideanSpace ℝ (Fin n)) × ℝ) : JointSpace N n :=
  WithLp.toLp 2 fun o => Option.elim o ((copyStat f N q.1 + q.2) / σ) fun ij => q.1 ij.1 ij.2
/-- The joint law of `(S, X₁, …, X_N)`: the image of `μ^{⊗N} ⊗ η_β` under `embed`. -/
def jointLaw {n : ℕ} (μ : Measure (EuclideanSpace ℝ (Fin n))) (f : EuclideanSpace ℝ (Fin n) → ℝ)
    (β σ : ℝ) (N : ℕ) : Measure (JointSpace N n) :=
  ((Measure.pi fun _ : Fin N => μ).prod (laplaceMeasure β)).map (embed f σ N)
/-- `jointLaw μ f β σ N` is a probability measure when `μ` and `laplaceMeasure β` are and `f` is
measurable. -/
theorem isProbabilityMeasure_jointLaw {n : ℕ} (μ : Measure (EuclideanSpace ℝ (Fin n)))
    [IsProbabilityMeasure μ] {f : EuclideanSpace ℝ (Fin n) → ℝ} (hf : Measurable f) {β : ℝ}
    [IsProbabilityMeasure (laplaceMeasure β)] (σ : ℝ) (N : ℕ) :
    IsProbabilityMeasure (jointLaw μ f β σ N) := by
    have hg :
        Measurable
          (fun q : (Fin N → EuclideanSpace ℝ (Fin n)) × ℝ =>
            fun o : Option (Fin N × Fin n) =>
              o.elim
                ((copyStat f N q.1 + q.2) / σ)
                (fun ij => (q.1 ij.1).ofLp ij.2)) := by
      rw [measurable_pi_iff]
      intro o
      cases o with
      | none =>
          unfold copyStat
          measurability
      | some ij =>
          measurability
    have h_embed : Measurable (embed f σ N) := by
      unfold embed
      exact (WithLp.measurable_toLp 2 _).comp hg
    refine ⟨?_⟩
    rw [jointLaw,
      Measure.map_apply_of_aemeasurable h_embed.aemeasurable MeasurableSet.univ]
    simp
/-- The joint law is a probability measure for every rate `β > 0` (the Laplace law's mass is
`isProbabilityMeasure_laplaceMeasure`). -/
theorem isProbabilityMeasure_jointLaw_of_pos {n : ℕ} (μ : Measure (EuclideanSpace ℝ (Fin n)))
    [IsProbabilityMeasure μ] {f : EuclideanSpace ℝ (Fin n) → ℝ} (hf : Measurable f) {β : ℝ}
    (hβ : 0 < β) (σ : ℝ) (N : ℕ) : IsProbabilityMeasure (jointLaw μ f β σ N) :=
  haveI := isProbabilityMeasure_laplaceMeasure hβ
  isProbabilityMeasure_jointLaw μ hf σ N
open scoped RealInnerProductSpace in
/-- The hypothesis of BKL Proposition 7.1: `R^ν_{d+1} ≤ b Id` for every isotropic log-concave
probability measure `ν` in every dimension, i.e. `Σ_J κ^ν_{d+1}(u, e_{J 1}, …, e_{J d})² ≤ b |u|²`
in the standard basis of `ℝ^m`. This names the hypothesis; it asserts nothing. -/
def CumulantBound (d : ℕ) (b : ℝ) : Prop :=
  ∀ (m : ℕ) (ν : ProbabilityMeasure (EuclideanSpace ℝ (Fin m))),
    KLS.Isotropic ν → _root_.FunctionalInequalities.KLS.LogConcave ν →
      ∀ u : EuclideanSpace ℝ (Fin m),
      ∑ J : Fin d → Fin m, Analysis.BakryEmery.cumulant (ν : Measure (EuclideanSpace ℝ (Fin m))) (d + 1)
        (Fin.cons u fun k => EuclideanSpace.basisFun (Fin m) ℝ (J k)) ^ 2 ≤ b * ‖u‖ ^ 2
/-- The index set of the joint space has `N n + 1` elements. -/
theorem card_option_prod (N n : ℕ) : Fintype.card (Option (Fin N × Fin n)) = N * n + 1 := by
  simp
/-- The coordinate reindexing of the joint space onto `ℝ^{Nn+1}`. -/
def coordinateEquiv (N n : ℕ) : JointSpace N n ≃ₗᵢ[ℝ] EuclideanSpace ℝ (Fin (N * n + 1)) :=
  LinearIsometryEquiv.piLpCongrLeft 2 ℝ ℝ (Fintype.equivFinOfCardEq (card_option_prod N n))
end Suspension
end LogConcave
end KLS
end FunctionalInequalities