import Mathlib.MeasureTheory.Constructions.Pi
import Mathlib.MeasureTheory.Function.SpecialFunctions.Basic
import Mathlib.MeasureTheory.Measure.WithDensity

/-! # Change of variables (`docs/roleflow.tex`, `sec:cov`)

`Cov μ ν f J` says that `f` pushes `J · μ` forward to `ν`. The development needs four rules:
`Cov.pres` (measure-preserving maps), `Cov.comp`, `Cov.tri` (triangular maps, by Tonelli) and
`Cov.inv` (the inverse, with the negated log-weight). Weights are usually `expw ℓ = e^ℓ` for a
real log-weight `ℓ`. Also generic: `Cov.congr_ae` (the weight matters only almost everywhere),
`Cov.pres_comp` and `Cov.comp_pres` (composing with a measure-preserving map) and
`MeasurePreserving.of_inverse`.

The second half is about products indexed by finite types with one measure per sort: the
measure-preserving *reorders* `measurePreserving_reorder`, which rename coordinates along a
bijection of index types and transport each coordinate along an equality of sorts. Contexts
(`Val` in `RoleFlow.Sig`) and the live states of the density proof are such products. -/

open MeasureTheory ENNReal

namespace RoleFlow

/-- The weight `e^t` as an extended nonnegative real. -/
noncomputable def expw (t : ℝ) : ℝ≥0∞ := ENNReal.ofReal (Real.exp t)

@[fun_prop] theorem measurable_expw : Measurable expw := by unfold expw; fun_prop

@[simp] theorem expw_zero : expw 0 = 1 := by simp [expw]

theorem expw_add (a b : ℝ) : expw (a + b) = expw a * expw b := by
  rw [expw, Real.exp_add, ENNReal.ofReal_mul (Real.exp_pos a).le]; rfl

theorem expw_neg_mul (a : ℝ) : expw (-a) * expw a = 1 := by
  rw [← expw_add, neg_add_cancel, expw_zero]

section Cov
variable {X Y Z : Type*} [MeasurableSpace X] [MeasurableSpace Y] [MeasurableSpace Z]

/-- `f` pushes `J · μ` forward to `ν`: integrating `φ ∘ f` against `J` is integrating `φ`. -/
def Cov (μ : Measure X) (ν : Measure Y) (f : X → Y) (J : X → ℝ≥0∞) : Prop :=
  ∀ φ : Y → ℝ≥0∞, Measurable φ → ∫⁻ x, φ (f x) * J x ∂μ = ∫⁻ y, φ y ∂ν

variable {μ : Measure X} {ν : Measure Y} {τ : Measure Z}

theorem Cov.congr {f f' : X → Y} {J J' : X → ℝ≥0∞} (hf : ∀ x, f x = f' x)
    (hJ : ∀ x, J x = J' x) (h : Cov μ ν f J) : Cov μ ν f' J' := by
  rwa [show f = f' from funext hf, show J = J' from funext hJ] at h

/-- The weight matters only almost everywhere. -/
theorem Cov.congr_ae {f : X → Y} {J J' : X → ℝ≥0∞} (hJ : J =ᵐ[μ] J') (h : Cov μ ν f J) :
    Cov μ ν f J' := fun φ hφ => by
  rw [← h φ hφ]; exact lintegral_congr_ae (hJ.mono fun x hx => by simp only [hx])

/-- **Pres**: a measure-preserving map has weight `1`. -/
theorem Cov.pres {f : X → Y} (hf : MeasurePreserving f μ ν) : Cov μ ν f fun _ => 1 :=
  fun φ hφ => by simpa using hf.lintegral_comp hφ

/-- **Comp**: weights multiply along composition. -/
theorem Cov.comp {f : X → Y} {g : Y → Z} {J : X → ℝ≥0∞} {K : Y → ℝ≥0∞} (hg : Measurable g)
    (hK : Measurable K) (h₁ : Cov μ ν f J) (h₂ : Cov ν τ g K) :
    Cov μ τ (fun x => g (f x)) (fun x => J x * K (f x)) := fun φ hφ => by
  rw [← h₂ φ hφ, ← h₁ (fun y => φ (g y) * K y) ((hφ.comp hg).mul hK)]
  congr 1; funext x; ring

/-- **Tri**: the triangular map `(p, a) ↦ (p, g p a)`, by Tonelli and the fibre at each `p`.
Only joint measurability is needed, no derivative in the parameter. -/
theorem Cov.tri [SFinite ν] [SFinite τ] {g : X → Y → Z} {J : X → Y → ℝ≥0∞}
    (hg : Measurable fun q : X × Y => g q.1 q.2) (hJ : Measurable fun q : X × Y => J q.1 q.2)
    (h : ∀ p, Cov ν τ (g p) (J p)) :
    Cov (μ.prod ν) (μ.prod τ) (fun q => (q.1, g q.1 q.2)) (fun q => J q.1 q.2) := fun φ hφ => by
  rw [lintegral_prod (fun q => φ (q.1, g q.1 q.2) * J q.1 q.2)
      (((hφ.comp (measurable_fst.prodMk hg)).mul hJ).aemeasurable),
    lintegral_prod _ hφ.aemeasurable]
  exact lintegral_congr fun p => h p _ (hφ.comp measurable_prodMk_left)

/-- **Inv**: a left inverse `g` of `f` pushes `ν` back with the negated log-weight. Only the
left law `g ∘ f = id` is needed: `ν` lives on the range of `f`, where `g` is determined. Since
the log-weight is real, no reciprocity or nonvanishing hypothesis appears. -/
theorem Cov.inv {f : X → Y} {g : Y → X} {ℓ : X → ℝ} (hg : Measurable g) (hℓ : Measurable ℓ)
    (hgf : ∀ x, g (f x) = x) (h : Cov μ ν f fun x => expw (ℓ x)) :
    Cov ν μ g fun y => expw (-ℓ (g y)) := fun φ hφ => by
  rw [← h (fun y => φ (g y) * expw (-ℓ (g y))) (by fun_prop)]
  refine lintegral_congr fun x => ?_
  rw [hgf, mul_assoc, expw_neg_mul, mul_one]

theorem Cov.pres_comp {f : X → Y} {g : Y → Z} {K : Y → ℝ≥0∞} (hf : MeasurePreserving f μ ν)
    (hg : Measurable g) (hK : Measurable K) (h : Cov ν τ g K) :
    Cov μ τ (fun x => g (f x)) (fun x => K (f x)) :=
  (Cov.comp hg hK (Cov.pres hf) h).congr (fun _ => rfl) (fun _ => one_mul _)

theorem Cov.comp_pres {f : X → Y} {J : X → ℝ≥0∞} {g : Y → Z} (h : Cov μ ν f J)
    (hg : MeasurePreserving g ν τ) : Cov μ τ (fun x => g (f x)) J :=
  (Cov.comp hg.measurable measurable_const h (Cov.pres hg)).congr (fun _ => rfl)
    (fun _ => mul_one _)

/-- With a measurable left inverse `g`, `J` is the density of `g_* ν` with respect to `μ`. -/
theorem Cov.map_eq {f : X → Y} {g : Y → X} {J : X → ℝ≥0∞} (hg : Measurable g)
    (hgf : ∀ x, g (f x) = x) (h : Cov μ ν f J) : ν.map g = μ.withDensity J := by
  ext s hs
  rw [Measure.map_apply hg hs, withDensity_apply _ hs, ← lintegral_indicator hs,
    ← lintegral_indicator_one (hg hs)]
  have hφ : Measurable fun y => s.indicator (1 : X → ℝ≥0∞) (g y) :=
    (measurable_one.indicator hs).comp hg
  calc ∫⁻ y, (g ⁻¹' s).indicator 1 y ∂ν = ∫⁻ y, s.indicator 1 (g y) ∂ν := by
        rfl
    _ = ∫⁻ x, s.indicator 1 (g (f x)) * J x ∂μ := (h _ hφ).symm
    _ = _ := by congr 1; funext x; by_cases hx : x ∈ s <;> simp [hx, hgf]

end Cov

/-- If `g` preserves measure and `f` is a measurable two-sided inverse, `f` preserves measure. -/
theorem _root_.MeasureTheory.MeasurePreserving.of_inverse {Y Z : Type*} [MeasurableSpace Y]
    [MeasurableSpace Z] {μ' : Measure Y} {ν' : Measure Z} {f : Z → Y} {g : Y → Z}
    (hg : MeasurePreserving g μ' ν') (hf : Measurable f) (hfg : ∀ y, f (g y) = y) :
    MeasurePreserving f ν' μ' :=
  ⟨hf, by rw [← hg.map_eq, Measure.map_map hf hg.measurable]; simp [Function.comp_def, hfg]⟩

/-! ## Products with one measure per sort, and reorders -/

section Reorder
variable {S : Type*} (X : S → Type*)

/-- Transport along an equality of sorts. -/
def tr {a b : S} (h : a = b) (x : X a) : X b := h ▸ x

variable {X}

@[simp] theorem tr_rfl {a : S} (x : X a) : tr X rfl x = x := rfl

@[simp] theorem tr_tr {a b c : S} (h : a = b) (h' : b = c) (x : X a) :
    tr X h' (tr X h x) = tr X (h.trans h') x := by subst h h'; rfl

/-- Transport along an equality of values of a family. -/
theorem tr_apply_eq {ι : Type*} {σ : ι → S} (s : ∀ i, X (σ i)) {i j : ι} (hij : i = j)
    (h : σ i = σ j) : tr X h (s i) = s j := by subst hij; rfl

variable [∀ s, MeasurableSpace (X s)]

@[fun_prop] theorem measurable_tr {a b : S} (h : a = b) : Measurable (tr X h) := by
  subst h; exact measurable_id

theorem measurePreserving_tr (μ : ∀ s, Measure (X s)) {a b : S} (h : a = b) :
    MeasurePreserving (tr X h) (μ a) (μ b) := by
  subst h; exact MeasurePreserving.id _

/-- **Reorder**: rename the coordinates of a product along a bijection `e`, transporting each
along the equality of its sorts. -/
def reorderEquiv {ι κ : Type*} (α : ι → S) (β : κ → S) (e : ι ≃ κ) (h : ∀ i, β (e i) = α i) :
    (∀ k, X (β k)) ≃ᵐ (∀ i, X (α i)) where
  toFun f i := tr X (h i) (f (e i))
  invFun g k := tr X ((h _).symm.trans (congrArg β (e.apply_symm_apply k))) (g (e.symm k))
  left_inv f := funext fun k => (tr_tr _ _ _).trans (tr_apply_eq f (e.apply_symm_apply k) _)
  right_inv g := funext fun i => (tr_tr _ _ _).trans (tr_apply_eq g (e.symm_apply_apply i) _)
  measurable_toFun := measurable_pi_lambda _ fun _ => (measurable_tr _).comp (measurable_pi_apply _)
  measurable_invFun :=
    measurable_pi_lambda _ fun _ => (measurable_tr _).comp (measurable_pi_apply _)

/-- **Reorder** preserves the product measure. -/
theorem measurePreserving_reorder (μ : ∀ s, Measure (X s)) [∀ s, SigmaFinite (μ s)]
    {ι κ : Type*} [Fintype ι] [Fintype κ] (α : ι → S) (β : κ → S) (e : ι ≃ κ)
    (h : ∀ i, β (e i) = α i) :
    MeasurePreserving (reorderEquiv α β e h)
      (Measure.pi fun k => μ (β k)) (Measure.pi fun i => μ (α i)) := by
  have h₁ := (measurePreserving_piCongrLeft (fun k => μ (β k)) e).symm
  have h₂ := measurePreserving_pi (fun i => μ (β (e i))) (fun i => μ (α i))
    (fun i => measurePreserving_tr μ (h i))
  convert h₂.comp h₁ using 1
  funext f i
  simp [MeasurableEquiv.piCongrLeft, reorderEquiv]

end Reorder

end RoleFlow