import PFR.TauFunctional
import PFR.HundredPercent
import PFR.Endgame

/-!
# Entropic version of polynomial Freiman-Ruzsa conjecture

Here we prove the entropic version of the polynomial Freiman-Ruzsa conjecture.

## Main results

* `entropic_PFR_conjecture`: For two $G$-valued random variables $X^0_1, X^0_2$, there is some
  subgroup $H \leq G$ such that $d[X^0_1;U_H] + d[X^0_2;U_H] \le 11 d[X^0_1;X^0_2]$.

-/

open MeasureTheory ProbabilityTheory
universe uG

variable {Ω₀₁ Ω₀₂ : Type*} [MeasureSpace Ω₀₁] [MeasureSpace Ω₀₂]

variable {Ω Ω' : Type*} [mΩ : MeasureSpace Ω] [IsProbabilityMeasure (ℙ : Measure Ω)]
  [IsProbabilityMeasure (ℙ : Measure Ω₀₁)] [IsProbabilityMeasure (ℙ : Measure Ω₀₂)]

variable {G : Type uG} [AddCommGroup G] [ElementaryAddCommGroup G 2] [Fintype G] [MeasurableSpace G]
  [MeasurableSingletonClass G]
variable (p : refPackage Ω₀₁ Ω₀₂ G) {X₁ : Ω → G} {X₂ : Ω → G} (hX₁ : Measurable X₁) (hX₂ : Measurable X₂)

/-- If $d[X_1;X_2] > 0$ then there are $G$-valued random variables $X'_1, X'_2$ such that $\tau[X'_1;X'_2] < \tau[X_1;X_2]$.
Phrased in the contrapositive form for convenience of proof. -/
theorem tau_strictly_decreases (h_min : tau_minimizes p X₁ X₂) (hpη: p.η = 1/9): d[X₁ # X₂] = 0 := by
  let ⟨A, mA, μ, Y₁, Y₂, Y₁', Y₂', hμ, h_indep, hY₁, hY₂, hY₁', hY₂', h_id1, h_id2, h_id1', h_id2'⟩
    := independent_copies4_nondep hX₁ hX₂ hX₁ hX₂ ℙ ℙ ℙ ℙ
  rw [← h_id1.rdist_eq h_id2]
  letI : MeasureSpace A := ⟨μ⟩
  have : IsProbabilityMeasure (ℙ : Measure A) := hμ
  rw [← h_id1.tau_minimizes p h_id2] at h_min
  apply tau_strictly_decreases_aux p Y₁ Y₂ Y₁' Y₂' hY₁ hY₂ hY₁' hY₂' (h_id1.trans h_id1'.symm)
    (h_id2.trans h_id2'.symm) h_indep h_min hpη

/-- `entropic_PFR_conjecture`: For two $G$-valued random variables $X^0_1, X^0_2$, there is some
    subgroup $H \leq G$ such that $d[X^0_1;U_H] + d[X^0_2;U_H] \le 11 d[X^0_1;X^0_2]$. -/
theorem entropic_PFR_conjecture (hpη : p.η = 1/9):
    ∃ H : AddSubgroup G, ∃ Ω : Type uG, ∃ mΩ : MeasureSpace Ω, ∃ U : Ω → G,
    IsProbabilityMeasure (ℙ : Measure Ω) ∧ Measurable U ∧
    IsUniform H U ∧ d[p.X₀₁ # U] + d[p.X₀₂ # U] ≤ 11 * d[p.X₀₁ # p.X₀₂] := by
  obtain ⟨Ω', mΩ', X₁, X₂, hX₁, hX₂, _, htau_min⟩ := tau_minimizer_exists p
  have hdist : d[X₁ # X₂] = 0 := tau_strictly_decreases p hX₁ hX₂ htau_min hpη
  obtain ⟨H, U, hU, hH_unif, hdistX₁, hdistX₂⟩ := exists_isUniform_of_rdist_eq_zero hX₁ hX₂ hdist
  refine ⟨H, Ω', inferInstance, U, inferInstance, hU, hH_unif , ?_⟩
  have h : τ[X₁ # X₂ | p] ≤ τ[p.X₀₂ # p.X₀₁ | p] := is_tau_min p htau_min p.hmeas2 p.hmeas1
  rw [tau, tau, hpη] at h
  norm_num at h
  have : d[p.X₀₁ # p.X₀₂] = d[p.X₀₂ # p.X₀₁] := rdist_symm
  have : d[p.X₀₁ # U] ≤ d[p.X₀₁ # X₁] + d[X₁ # U] := rdist_triangle p.hmeas1 hX₁ hU
  have : d[p.X₀₂ # U] ≤ d[p.X₀₂ # X₂] + d[X₂ # U] := rdist_triangle p.hmeas2 hX₂ hU
  linarith

theorem entropic_PFR_conjecture' (hpη : p.η = 1/9):
    ∃ H : AddSubgroup G, ∃ Ω : Type uG, ∃ mΩ : MeasureSpace Ω, ∃ U : Ω → G,
    IsUniform H U ∧ d[p.X₀₁ # U] ≤ 6 * d[p.X₀₁ # p.X₀₂] ∧
      d[p.X₀₂ # U] ≤ 6 * d[p.X₀₁ # p.X₀₂]

:= by
  have : d[p.X₀₁ # p.X₀₂] = d[p.X₀₂ # p.X₀₁] := rdist_symm
  obtain ⟨H, Ω, mΩ, U, H', hU, hUnif, h'⟩ := entropic_PFR_conjecture p hpη
  refine ⟨H, Ω, mΩ, U, hUnif, ?_⟩
  have : d[p.X₀₁ # U] ≤ d[p.X₀₁ # p.X₀₂] + d[p.X₀₂ # U] := rdist_triangle p.hmeas1 p.hmeas2 hU
  have : d[p.X₀₂ # U] ≤ d[p.X₀₂ # p.X₀₁] + d[p.X₀₁ # U] := rdist_triangle p.hmeas2 p.hmeas1 hU
  constructor
  · linarith
  · linarith


import PFR.FirstEstimate
import PFR.SecondEstimate
import PFR.ForMathlib.Elementary

/-!
# Endgame

The endgame on tau-minimizers.

Assumptions:

* $X_1, X_2$ are tau-minimizers
* $X_1, X_2, \tilde X_1, \tilde X_2$ be independent random variables, with $X_1,\tilde X_1$ copies
    of $X_1$ and $X_2,\tilde X_2$ copies of $X_2$.
* $d[X_1;X_2] = k$
* $U := X_1 + X_2$
* $V := \tilde X_1 + X_2$
* $W := X_1 + \tilde X_1$
* $S := X_1 + X_2 + \tilde X_1 + \tilde X_2$.
* $I_1 := I[U : V | S]$
* $I_2 := I[U : W | S]$
* $I_3 := I[V : W | S]$ (not explicitly defined in Lean)

# Main results:

* `sum_condMutual_le` : An upper bound on the total conditional mutual information $I_1+I_2+I_3$.
* `sum_dist_diff_le`: A sum of the "costs" of $U$, $V$, $W$.
* `construct_good`: A construction of two random variables with small Ruzsa distance between them
  given some random variables with control on total cost, as well as total mutual information.
-/

open MeasureTheory ProbabilityTheory
open scoped BigOperators

variable {G : Type*} [AddCommGroup G] [Fintype G] [hG : MeasurableSpace G]
  [MeasurableSingletonClass G] [elem: ElementaryAddCommGroup G 2] [MeasurableAdd₂ G]

variable {Ω₀₁ Ω₀₂ : Type*} [MeasureSpace Ω₀₁] [MeasureSpace Ω₀₂]
  [IsProbabilityMeasure (ℙ : Measure Ω₀₁)] [IsProbabilityMeasure (ℙ : Measure Ω₀₂)]

variable (p : refPackage Ω₀₁ Ω₀₂ G)

variable {Ω : Type*} [mΩ : MeasureSpace Ω] [IsProbabilityMeasure (ℙ : Measure Ω)]

variable (X₁ X₂ X₁' X₂' : Ω → G)
  (hX₁ : Measurable X₁) (hX₂ : Measurable X₂) (hX₁' : Measurable X₁') (hX₂' : Measurable X₂')

variable (h₁ : IdentDistrib X₁ X₁') (h₂ : IdentDistrib X₂ X₂')

variable (h_indep : iIndepFun (fun _i => hG) ![X₁, X₂, X₁', X₂'])

variable (h_min: tau_minimizes p X₁ X₂)

/-- `k := d[X₁ # X₂]`, the Ruzsa distance `rdist` between X₁ and X₂. -/
local notation3 "k" => d[X₁ # X₂]

/-- `U := X₁ + X₂` -/
local notation3 "U" => X₁ + X₂

/-- `V := X₁' + X₂` -/
local notation3 "V" => X₁' + X₂

/-- `W := X₁' + X₁` -/
local notation3 "W" => X₁' + X₁

/-- `S := X₁ + X₂ + X₁' + X₂'` -/
local notation3 "S" => X₁ + X₂ + X₁' + X₂'

/-- `I₁ := I[U : V | S]`, the conditional mutual information of `U = X₁ + X₂` and `V = X₁' + X₂`
given the quadruple sum `S = X₁ + X₂ + X₁' + X₂'`. -/
local notation3 "I₁" => I[U : V | S]

/-- `I₂ := I[U : W | S]`, the conditional mutual information of `U = X₁ + X₂` and `W = X₁' + X₁`
given the quadruple sum `S = X₁ + X₂ + X₁' + X₂'`. -/
local notation3 "I₂" => I[U : W | S]

--(Mantas) this times out in the proof below
private lemma hmeas2 :
    Measurable fun p : Fin 4 → G => ((p 0 + p 1, p 0 + p 2), p 0 + p 1 + p 2 + p 3) := by
  apply Measurable.prod
  · apply Measurable.prod
    · exact (measurable_pi_apply _).add (measurable_pi_apply _)
    · exact (measurable_pi_apply _).add (measurable_pi_apply _)
  · apply Measurable.add
    · apply Measurable.add
      · exact (measurable_pi_apply _).add (measurable_pi_apply _)
      · apply measurable_pi_apply
    · apply measurable_pi_apply

/-- The quantity $I_3 = I[V:W|S]$ is equal to $I_2$. -/
lemma I₃_eq : I[V : W | S] = I₂ := by
  have h_indep2 : iIndepFun (fun _ ↦ hG) ![X₁', X₂, X₁, X₂'] := by
    exact h_indep.reindex_four_cbad
  have hident : IdentDistrib (fun a (i : Fin 4) => ![X₁, X₂, X₁', X₂'] i a)
    (fun a (j : Fin 4) => ![X₁', X₂, X₁, X₂'] j a) := by
    exact { aemeasurable_fst := by
              apply Measurable.aemeasurable
              rw [measurable_pi_iff]
              intro x
              fin_cases x;
              all_goals aesop
            aemeasurable_snd := by
              apply Measurable.aemeasurable
              rw [measurable_pi_iff]
              intro x
              fin_cases x;
              all_goals aesop
            map_eq := by
              rw [← (ProbabilityTheory.iIndepFun_iff_pi_map_eq_map (![X₁, X₂, X₁', X₂'])
                (Fin.cases hX₁ <| Fin.cases hX₂ <| Fin.cases hX₁' <|
              Fin.cases hX₂' Fin.rec0)).mp h_indep,
              ← (ProbabilityTheory.iIndepFun_iff_pi_map_eq_map (![X₁', X₂, X₁, X₂'])
                (Fin.cases hX₁' <| Fin.cases hX₂ <| Fin.cases hX₁ <|
              Fin.cases hX₂' Fin.rec0)).mp h_indep2]
              congr
              ext i
              fin_cases i
              all_goals simp [h₁.map_eq] }
  have hmeas1 : Measurable (fun p : Fin 4 → G => (p 0 + p 1, p 0 + p 1 + p 2 + p 3)) := by
    simp_all only [Matrix.cons_val', Matrix.empty_val', Matrix.cons_val_fin_one]
    apply Measurable.prod
    · exact (measurable_pi_apply _).add (measurable_pi_apply _)
    · apply Measurable.add
      · apply Measurable.add
        · exact (measurable_pi_apply _).add (measurable_pi_apply _)
        · apply measurable_pi_apply
      · apply measurable_pi_apply
  have hUVS : IdentDistrib (prod U S) (prod V S) := by
    convert (IdentDistrib.comp hident hmeas1); simp; abel
  have hUVWS : IdentDistrib (prod (prod U W) S) (prod (prod V W) S) := by
    convert (IdentDistrib.comp hident hmeas2) <;> simp <;> abel
  have hU : Measurable U := Measurable.add hX₁ hX₂
  have hV : Measurable V := Measurable.add hX₁' hX₂
  have hW : Measurable W := Measurable.add hX₁' hX₁
  have hS : Measurable S := by measurability
  rw [condMutualInfo_eq hV hW hS, condMutualInfo_eq hU hW hS, chain_rule'' ℙ hU hS,
    chain_rule'' ℙ hV hS, chain_rule'' ℙ hW hS, chain_rule'' ℙ _ hS, chain_rule'' ℙ _ hS,
    IdentDistrib.entropy_eq hUVS, IdentDistrib.entropy_eq hUVWS]
  · exact Measurable.prod (by exact hU) (by exact hW)
  · exact Measurable.prod (by exact hV) (by exact hW)

/--
$$ I(U : V | S) + I(V : W | S) + I(W : U | S) $$
is less than or equal to
$$ 6 \eta k - \frac{1 - 5 \eta}{1-\eta} (2 \eta k - I_1).$$
-/
lemma sum_condMutual_le :
    I[U : V | S] + I[V : W | S] + I[W : U | S]
      ≤ 6 * p.η * k - (1 - 5 * p.η) / (1 - p.η) * (2 * p.η * k - I₁) := by
  have : I[W:U|S] = I₂ := by
    rw [condMutualInfo_comm]
    · exact Measurable.add' hX₁' hX₁
    · exact Measurable.add' hX₁ hX₂
  rw [I₃_eq, this]
  have h₂ := second_estimate p X₁ X₂ X₁' X₂' hX₁ hX₂ hX₁' hX₂' h₁ h₂ h_indep h_min
  have h := add_le_add (add_le_add_left h₂ I₁) h₂
  convert h using 1
  have : 1 - p.η > 0 := by linarith [p.hη']
  field_simp [this]
  ring
  all_goals { simpa }

local notation3:max "c[" A "; " μ " # " B " ; " μ' "]" =>
  d[p.X₀₁; ℙ # A; μ] - d[p.X₀₁ # X₁] + (d[p.X₀₂; ℙ # B; μ'] - d[p.X₀₂ # X₂])

local notation3:max "c[" A " # " B "]" =>
  d[p.X₀₁ # A] - d[p.X₀₁ # X₁] + (d[p.X₀₂ # B] - d[p.X₀₂ # X₂])

local notation3:max "c[" A " | " B " # " C " | " D "]" => d[p.X₀₁ # A|B] - d[p.X₀₁ # X₁] + (d[p.X₀₂ # C|D] - d[p.X₀₂ # X₂])

lemma hU : H[U] = H[X₁' + X₂'] :=
  IdentDistrib.entropy_eq (ProbabilityTheory.IdentDistrib.add h₁ h₂
    (iIndepFun.indepFun h_indep (show (0 : Fin 4) ≠ 1 by norm_cast))
     (iIndepFun.indepFun h_indep (show (2 : Fin 4) ≠ 3 by norm_cast)))

variable {X₁ X₂ X₁' X₂'} in
lemma independenceCondition1 : iIndepFun (fun _ ↦ hG) ![X₁, X₂, X₁' + X₂'] :=
  h_indep.apply_two_last hX₁ hX₂ hX₁' hX₂' measurable_add

lemma hV : H[V] = H[X₁ + X₂'] :=
IdentDistrib.entropy_eq (ProbabilityTheory.IdentDistrib.add h₁.symm h₂
  (iIndepFun.indepFun h_indep (show (2 : Fin 4) ≠ 1 by norm_cast))
  (iIndepFun.indepFun h_indep (show (0 : Fin 4) ≠ 3 by norm_cast)))

variable {X₁ X₂ X₁' X₂'} in
lemma independenceCondition2 : iIndepFun (fun _ ↦ hG) ![X₂, X₁, X₁' + X₂'] :=
  independenceCondition1 hX₂ hX₁ hX₁' hX₂' h_indep.reindex_four_bacd

variable {X₁ X₂ X₁' X₂'} in
lemma independenceCondition3 : iIndepFun (fun _ ↦ hG) ![X₁', X₂, X₁ + X₂'] :=
  independenceCondition1 hX₁' hX₂ hX₁ hX₂' h_indep.reindex_four_cbad

variable {X₁ X₂ X₁' X₂'} in
lemma independenceCondition4 : iIndepFun (fun _ ↦ hG) ![X₂, X₁', X₁ + X₂'] :=
  independenceCondition1 hX₂ hX₁' hX₁ hX₂' h_indep.reindex_four_bcad

variable {X₁ X₂ X₁' X₂'} in
lemma independenceCondition5 : iIndepFun (fun _ ↦ hG) ![X₁, X₁', X₂ + X₂'] :=
  independenceCondition1 hX₁ hX₁' hX₂ hX₂' h_indep.reindex_four_acbd

variable {X₁ X₂ X₁' X₂'} in
lemma independenceCondition6 : iIndepFun (fun _ ↦ hG) ![X₂, X₂', X₁' + X₁] :=
  independenceCondition1 hX₂ hX₂' hX₁' hX₁ h_indep.reindex_four_bdca

set_option maxHeartbeats 400000 in
/--
$$ \sum_{i=1}^2 \sum_{A\in\{U,V,W\}} \big(d[X^0_i;A|S] - d[X^0_i;X_i]\big)$$
is less than or equal to
$$ \leq (6 - 3\eta) k + 3(2 \eta k - I_1).$$
-/
lemma sum_dist_diff_le :
  c[U|S # U|S] + c[V|S # V|S] + c[W|S # W|S] ≤ (6 - 3 * p.η)*k + 3 * (2*p.η*k - I₁) := by
  let X₀₁ := p.X₀₁
  let X₀₂ := p.X₀₂
  have ineq1 : d[X₀₁ # U | S] - d[X₀₁ # X₁] ≤ (H[S ; ℙ] - H[X₁ ; ℙ])/2 := by
    have aux1 : H[S] + H[U] - H[X₁] - H[X₁' + X₂'] = H[S] - H[X₁] := by
      rw [hU X₁ X₂ X₁' X₂' h₁ h₂ h_indep]; ring
    have aux2 : d[X₀₁ # U | U + (X₁' + X₂')] - d[X₀₁ # X₁]
            ≤ (H[U + (X₁' + X₂')] + H[U] - H[X₁] - H[X₁' + X₂']) / 2 :=
      condRuzsaDist_diff_ofsum_le ℙ (hX := p.hmeas1) (hY := hX₁) (hZ := hX₂)
      (Measurable.add hX₁' hX₂') (independenceCondition1 hX₁ hX₂ hX₁' hX₂' h_indep)
    rw [← add_assoc, aux1] at aux2
    linarith [aux2]

  have ineq2 : d[X₀₂ # U | S] - d[X₀₂ # X₂] ≤ (H[S ; ℙ] - H[X₂ ; ℙ])/2 := by
    have aux1 : H[S] + H[U] - H[X₂] - H[X₁' + X₂'] = H[S] - H[X₂] := by
      rw [hU X₁ X₂ X₁' X₂' h₁ h₂ h_indep] ; ring
    have aux2 : d[X₀₂ # U | U + (X₁' + X₂')] - d[X₀₂ # X₂]
            ≤ (H[U + (X₁' + X₂')] + H[U] - H[X₂] - H[X₁' + X₂']) / 2 := by
      rw [(show U = X₂ + X₁ from add_comm _ _)]
      apply condRuzsaDist_diff_ofsum_le ℙ (p.hmeas2) (hX₂) (hX₁)
        (Measurable.add hX₁' hX₂') (independenceCondition2 hX₁ hX₂ hX₁' hX₂' h_indep)
    rw [← add_assoc, aux1] at aux2
    linarith [aux2]

  have V_add_eq : V + (X₁ + X₂') = S := by abel

  have ineq3 : d[X₀₁ # V | S] - d[X₀₁ # X₁] ≤ (H[S ; ℙ] - H[X₁ ; ℙ])/2 := by
    have aux2 : d[p.X₀₁ # V | V + (X₁ + X₂')] - d[p.X₀₁ # X₁']
            ≤ (H[V + (X₁ + X₂')] + H[V] - H[X₁'] - H[X₁ + X₂']) / 2 :=
      condRuzsaDist_diff_ofsum_le ℙ (p.hmeas1) (hX₁') (hX₂) (Measurable.add hX₁ hX₂')
      (independenceCondition3 hX₁ hX₂ hX₁' hX₂' h_indep)

    have aux1 : H[S] + H[V] - H[X₁'] - H[X₁ + X₂'] = H[S ; ℙ] - H[X₁ ; ℙ] := by
      rw [hV X₁ X₂ X₁' X₂' h₁ h₂ h_indep, h₁.entropy_eq]; ring
    rw [← ProbabilityTheory.IdentDistrib.rdist_eq (IdentDistrib.refl p.hmeas1.aemeasurable) h₁,
      V_add_eq, aux1] at aux2
    linarith [aux2]

  have ineq4 : d[X₀₂ # V | S] - d[X₀₂ # X₂] ≤ (H[S ; ℙ] - H[X₂ ; ℙ])/2 := by
    have aux2 : d[p.X₀₂ # V | V + (X₁ + X₂')] - d[p.X₀₂ # X₂]
            ≤ (H[V + (X₁ + X₂')] + H[V] - H[X₂] - H[X₁ + X₂']) / 2 := by
      rw [(show V = X₂ + X₁' from add_comm _ _)]
      apply condRuzsaDist_diff_ofsum_le ℙ (p.hmeas2) (hX₂) (hX₁') (Measurable.add hX₁ hX₂')
        (independenceCondition4 hX₁ hX₂ hX₁' hX₂' h_indep)
    have aux1 : H[S] + H[V] - H[X₂] - H[X₁ + X₂'] = H[S ; ℙ] - H[X₂ ; ℙ] := by
      rw [hV X₁ X₂ X₁' X₂' h₁ h₂ h_indep]; ring
    rw [V_add_eq, aux1] at aux2
    linarith [aux2]

  let W' := X₂ + X₂'
  have ineq5 : d[X₀₁ # W | S] - d[X₀₁ # X₁] ≤ (H[S ; ℙ] + H[W ; ℙ] - H[X₁ ; ℙ] - H[W' ; ℙ])/2 := by
    have := condRuzsaDist_diff_ofsum_le ℙ p.hmeas1 hX₁ hX₁' (Measurable.add hX₂ hX₂')
      (independenceCondition5 hX₁ hX₂ hX₁' hX₂' h_indep)
    have S_eq : X₁ + X₁' + (fun a ↦ X₂ a + X₂' a) = S := by
      rw [(show (fun a ↦ X₂ a + X₂' a) = X₂ + X₂' by rfl), ← add_assoc, add_assoc X₁, add_comm X₁',
        ← add_assoc]
    rwa [S_eq, add_comm X₁ X₁'] at this

  have ineq6 : d[X₀₂ # W' | S] - d[X₀₂ # X₂] ≤ (H[S ; ℙ] + H[W' ; ℙ] - H[X₂ ; ℙ] - H[W ; ℙ])/2 := by
    have := condRuzsaDist_diff_ofsum_le ℙ p.hmeas2 hX₂ hX₂' (Measurable.add hX₁' hX₁)
      (independenceCondition6 hX₁ hX₂ hX₁' hX₂' h_indep)
    have S_eq : X₂ + X₂' + (fun a ↦ X₁' a + X₁ a) = S := by
      rw [(show (fun a ↦ X₁' a + X₁ a) = X₁' + X₁ by rfl), add_comm, ← add_assoc, add_comm X₁',
      add_assoc X₁, add_comm X₁', ← add_assoc]
    rwa [S_eq] at this

  have dist_eq : d[X₀₂ # W' | S] = d[X₀₂ # W | S] := by
    have S_eq : S = (X₂ + X₂') + (X₁' + X₁) := by
      rw [add_comm X₁' X₁, add_assoc _ X₂', add_comm X₂', ← add_assoc X₂, ← add_assoc X₂,
        add_comm X₂]
    rw [S_eq]
    apply condRuzsaDist'_of_inj_map' p.hmeas2 (hX₂.add hX₂') (hX₁'.add hX₁)

  -- Put everything together to bound the sum of the `c` terms
  have ineq7 : c[U|S # U|S] + c[V|S # V|S] + c[W|S # W|S] ≤
    3 * H[S ; ℙ] - 3/2 * H[X₁ ; ℙ] -3/2 * H[X₂ ; ℙ] := by
    have step₁ : c[U|S # U|S] ≤ H[S ; ℙ] - (H[X₁ ; ℙ] + H[X₂ ; ℙ])/2 :=
      calc
        _ = (d[p.X₀₁ # U|S] - d[p.X₀₁ # X₁]) + (d[p.X₀₂ # U|S] - d[p.X₀₂ # X₂]) := by ring
        _ ≤ (H[S ; ℙ] - H[X₁ ; ℙ])/2 + (H[S ; ℙ] - H[X₂ ; ℙ])/2 := add_le_add ineq1 ineq2
        _ = H[S ; ℙ] - (H[X₁ ; ℙ] + H[X₂ ; ℙ])/2 := by ring
    have step₂ : c[V|S # V|S] ≤ H[S ; ℙ] - (H[X₁ ; ℙ] + H[X₂ ; ℙ])/2 :=
      calc c[V|S # V|S] =(d[p.X₀₁ # V|S] - d[p.X₀₁ # X₁]) + (d[p.X₀₂ # V|S] - d[p.X₀₂ # X₂]) := by ring
        _ ≤ (H[S ; ℙ] - H[X₁ ; ℙ])/2 + (H[S ; ℙ] - H[X₂ ; ℙ])/2 := add_le_add ineq3 ineq4
        _ = H[S ; ℙ] - (H[X₁ ; ℙ] + H[X₂ ; ℙ])/2 := by ring
    have step₃ : c[W|S # W|S] ≤ H[S ; ℙ] - (H[X₁ ; ℙ] + H[X₂ ; ℙ])/2 :=
      calc c[W|S # W|S] = (d[X₀₁ # W | S] - d[X₀₁ # X₁]) + (d[X₀₂ # W' | S] - d[X₀₂ # X₂]) :=
          by rw [dist_eq]
        _ ≤ (H[S ; ℙ] + H[W ; ℙ] - H[X₁ ; ℙ] - H[W' ; ℙ])/2 + (H[S ; ℙ] + H[W' ; ℙ] - H[X₂ ; ℙ] - H[W ; ℙ])/2
          := add_le_add ineq5 ineq6
        _ = H[S ; ℙ] - (H[X₁ ; ℙ] + H[X₂ ; ℙ])/2 := by ring
    calc c[U|S # U|S] + c[V|S # V|S] + c[W|S # W|S] ≤ (H[S ; ℙ] - (H[X₁ ; ℙ] + H[X₂ ; ℙ])/2) +
      (H[S ; ℙ] - (H[X₁ ; ℙ] + H[X₂ ; ℙ])/2) + (H[S ; ℙ] - (H[X₁ ; ℙ] + H[X₂ ; ℙ])/2) :=
        add_le_add (add_le_add step₁ step₂) step₃
    _ = 3 * H[S ; ℙ] - 3/2 * H[X₁ ; ℙ] -3/2 * H[X₂ ; ℙ] := by ring

  have h_indep' : iIndepFun (fun _i => hG) ![X₁, X₂, X₂', X₁'] := by
    apply ProbabilityTheory.iIndepFun.reindex (Equiv.swap (2 : Fin 4) 3)
    convert h_indep using 1
    ext x
    fin_cases x ; all_goals { aesop }

  have ineq8 : 3 * H[S ; ℙ] ≤ 3/2 * (H[X₁ ; ℙ] + H[X₂ ; ℙ]) + 3*(2+p.η)*k - 3*I₁ :=
    calc 3 * H[S ; ℙ] ≤ 3 * (H[X₁ ; ℙ] / 2 + H[X₂ ; ℙ] / 2 + (2+p.η)*k - I₁) := by
          apply (mul_le_mul_left (zero_lt_three' ℝ)).mpr
            (ent_ofsum_le p X₁ X₂ X₁' X₂' hX₁ hX₂ hX₁' hX₂' h₁ h₂ h_indep' h_min)
      _ =  3/2 * ( H[X₁ ; ℙ] + H[X₂ ; ℙ]) + 3*(2+p.η)*k - 3*I₁ := by ring

  -- Final computation
  calc c[U|S # U|S] + c[V|S # V|S] + c[W|S # W|S] ≤ 3 * H[S ; ℙ] - 3/2 * H[X₁ ; ℙ] -3/2 * H[X₂ ; ℙ] := ineq7
     _ = 3 * H[S ; ℙ] - (3/2 *(H[X₁ ; ℙ] + H[X₂ ; ℙ])) := by ring
     _ ≤ (3/2 * ( H[X₁ ; ℙ] + H[X₂ ; ℙ]) + 3*(2+p.η)*k - 3*I₁) - (3/2 *(H[X₁ ; ℙ] + H[X₂ ; ℙ])) :=
        sub_le_sub_right ineq8 _
     _ = (6 - 3 * p.η)*k + 3 * (2*p.η*k - I₁) := by ring

/-- $U+V+W=0$. -/
lemma sum_uvw_eq_zero : U+V+W = 0 := by
  rw [add_comm X₁' X₂, ElementaryAddCommGroup.sum_add_sum_add_sum_eq_zero]

section construct_good
variable {Ω' : Type*} [MeasureSpace Ω'] [IsProbabilityMeasure (ℙ : Measure Ω')]
variable {T₁ T₂ T₃ : Ω' → G} (hT : T₁+T₂+T₃ = 0)
variable (hT₁ : Measurable T₁) (hT₂ : Measurable T₂) (hT₃ : Measurable T₃)

local notation3:max "δ[" μ "]" => I[T₁ : T₂ ; μ] + I[T₂ : T₃ ; μ] + I[T₃ : T₁ ; μ]
local notation3:max "δ" => I[T₁ : T₂] + I[T₂ : T₃] + I[T₃ : T₁]

local notation3:max "ψ[" A " # " B "]" => d[A # B] + p.η * (c[A # B])
local notation3:max "ψ[" A "; " μ " # " B " ; " μ' "]" =>
  d[A ; μ # B ; μ'] + p.η * c[A ; μ # B ; μ']

/-- If $T_1, T_2, T_3$ are $G$-valued random variables with $T_1+T_2+T_3=0$ holds identically and
$$ \delta := \sum_{1 \leq i < j \leq 3} I[T_i;T_j]$$
Then there exist random variables $T'_1, T'_2$ such that
$$ d[T'_1;T'_2] + \eta (d[X_1^0;T'_1] - d[X_1^0;X_1]) + \eta(d[X_2^0;T'_2] - d[X_2^0;X_2]) $$
is at most
$$ \delta + \eta ( d[X^0_1;T_1]-d[X^0_1;X_1]) + \eta (d[X^0_2;T_2]-d[X^0_2;X_2]) $$
$$ + \tfrac12 \eta I[T_1: T_3] + \tfrac12 \eta I[T_2: T_3].$$
-/
lemma construct_good_prelim :
    k ≤ δ + p.η * c[T₁ # T₂] + p.η * (I[T₁: T₃] + I[T₂ : T₃])/2 := by
  let sum1 : ℝ := (Measure.map T₃ ℙ)[fun t ↦ d[T₁; ℙ[|T₃ ⁻¹' {t}] # T₂; ℙ[|T₃ ⁻¹' {t}]]]
  let sum2 : ℝ := (Measure.map T₃ ℙ)[fun t ↦ d[p.X₀₁; ℙ # T₁; ℙ[|T₃ ⁻¹' {t}]] - d[p.X₀₁ # X₁]]
  let sum3 : ℝ := (Measure.map T₃ ℙ)[fun t ↦ d[p.X₀₂; ℙ # T₂; ℙ[|T₃ ⁻¹' {t}]] - d[p.X₀₂ # X₂]]
  let sum4 : ℝ := (Measure.map T₃ ℙ)[fun t ↦ ψ[T₁; ℙ[|T₃ ⁻¹' {t}] # T₂; ℙ[|T₃ ⁻¹' {t}]]]

  have hp.η : 0 ≤ p.η := by linarith [p.hη]
  have hP : IsProbabilityMeasure (Measure.map T₃ ℙ) := isProbabilityMeasure_map hT₃.aemeasurable
  have h2T₃ : T₃ = T₁ + T₂ :=
    calc T₃ = T₁ + T₂ + T₃ - T₃ := by rw [hT, zero_sub]; simp
      _ = T₁ + T₂ := by rw [add_sub_cancel_right]
  have h2T₁ : T₁ = T₂ + T₃ := by simp [h2T₃, add_left_comm]
  have h2T₂ : T₂ = T₃ + T₁ := by simp [h2T₁, add_left_comm]

  have h1 : sum1 ≤ δ := by
    have h1 : sum1 ≤ 3 * I[T₁ : T₂] + 2 * H[T₃] - H[T₁] - H[T₂] := by
      subst h2T₃; exact ent_bsg hT₁ hT₂
    have h2 : H[⟨T₂, T₃⟩] = H[⟨T₁, T₂⟩] := by
      rw [h2T₃, entropy_add_right', entropy_comm] <;> assumption
    have h3 : H[⟨T₁, T₂⟩] = H[⟨T₃, T₁⟩] := by
      rw [h2T₃, entropy_add_left, entropy_comm] <;> assumption
    simp_rw [mutualInfo_def] at h1 ⊢; linarith

  have h2 : p.η * sum2 ≤ p.η * (d[p.X₀₁ # T₁] - d[p.X₀₁ # X₁] + I[T₁ : T₃] / 2) := by
    have : sum2 = d[p.X₀₁ # T₁ | T₃] - d[p.X₀₁ # X₁] := by
      simp [sum2, integral_sub (.of_finite _ _) (.of_finite _ _), sum2]
      simp_rw [condRuzsaDist'_eq_sum hT₁ hT₃, integral_eq_sum' _ (FiniteRange.null_of_compl _ T₃),
        Measure.map_apply hT₃ (measurableSet_singleton _), smul_eq_mul]

    gcongr
    linarith [condRuzsaDist_le' ℙ ℙ p.hmeas1 hT₁ hT₃]

  have h3 : p.η * sum3 ≤ p.η * (d[p.X₀₂ # T₂] - d[p.X₀₂ # X₂] + I[T₂ : T₃] / 2) := by
    have : sum3 = d[p.X₀₂ # T₂ | T₃] - d[p.X₀₂ # X₂] := by
      simp [sum3, integral_sub (.of_finite _ _) (.of_finite _ _)]
      simp_rw [condRuzsaDist'_eq_sum hT₂ hT₃, integral_eq_sum' _ (FiniteRange.null_of_compl _ T₃),
        Measure.map_apply hT₃ (measurableSet_singleton _), smul_eq_mul]
    gcongr
    linarith [condRuzsaDist_le' ℙ ℙ p.hmeas2 hT₂ hT₃]

  have h4 : sum4 ≤ δ + p.η * c[T₁ # T₂] + p.η * (I[T₁ : T₃] + I[T₂ : T₃]) / 2 := by
    suffices sum4 = sum1 + p.η * (sum2 + sum3) by linarith
    simp only [sum4, integral_add (.of_finite _ _) (.of_finite _ _), integral_mul_left]

  have hk : k ≤ sum4 := by
    suffices (Measure.map T₃ ℙ)[fun _ ↦ k] ≤ sum4 by simpa using this
    refine integral_mono_ae (.of_finite _ _) (.of_finite _ _) $ ae_iff_of_countable.2 fun t ht ↦ ?_
    have : IsProbabilityMeasure (ℙ[|T₃ ⁻¹' {t}]) :=
      cond_isProbabilityMeasure ℙ (by simpa [hT₃] using ht)
    dsimp only
    linarith only [distance_ge_of_min' (μ := ℙ[|T₃ ⁻¹' {t}]) (μ' := ℙ[|T₃ ⁻¹' {t}]) p h_min hT₁ hT₂]
  exact hk.trans h4

/-- If $T_1, T_2, T_3$ are $G$-valued random variables with $T_1+T_2+T_3=0$ holds identically and
-
$$ \delta := \sum_{1 \leq i < j \leq 3} I[T_i;T_j]$$

Then there exist random variables $T'_1, T'_2$ such that

$$ d[T'_1;T'_2] + \eta (d[X_1^0;T'_1] - d[X_1^0;X _1]) + \eta(d[X_2^0;T'_2] - d[X_2^0;X_2])$$

is at most

$$\delta + \frac{\eta}{3} \biggl( \delta + \sum_{i=1}^2 \sum_{j = 1}^3
    (d[X^0_i;T_j] - d[X^0_i; X_i]) \biggr).$$
-/
lemma construct_good : k ≤ δ + (p.η/3) * (δ + c[T₁ # T₁] + c[T₂ # T₂] + c[T₃ # T₃])

:= by
  have v2 := construct_good_prelim p X₁ X₂ h_min (by rw [← hT]; abel) hT₁ hT₃ hT₂
  have v3 := construct_good_prelim p X₁ X₂ h_min (by rw [← hT]; abel) hT₂ hT₁ hT₃
  have v6 := construct_good_prelim p X₁ X₂ h_min (by rw [← hT]; abel) hT₃ hT₂ hT₁
  simp only [mutualInfo, entropy_comm hT₂ hT₁, entropy_comm hT₃ hT₁, entropy_comm hT₃ hT₂]
    at *
  linarith


import Mathlib.Algebra.BigOperators.Basic
import Mathlib.MeasureTheory.Constructions.Prod.Integral
import LeanAPAP.Mathlib.Algebra.BigOperators.Basic
import PFR.ForMathlib.Elementary
import PFR.ForMathlib.Entropy.Group
import PFR.ForMathlib.Entropy.Kernel.RuzsaDist
import PFR.ForMathlib.ProbabilityMeasureProdCont
import PFR.Mathlib.Data.Fin.VecNotation
import PFR.Mathlib.Probability.IdentDistrib

/-!
# Ruzsa distance

Here we define Ruzsa distance and establish its basic properties.

## Main definitions

* `rdist`: The Ruzsa distance $d[X ; Y]$ between two random variables
* `condRuzsaDist`: The conditional Ruzsa distance $d[X | Z ; Y | W]$ (or $d[X ; Y | W]$) between two
  random variables, conditioned on additional random variables

## Main results

* `rdist_triangle`: The Ruzsa triangle inequality for three random variables.
* `kaimanovich_vershik`. $$H[X + Y + Z] - H[X + Y] \leq H[Y+ Z] - H[Y]$$
* `ent_bsg`: If $Z=A+B$, then
  $$\sum_{z} P[Z=z] d[(A | Z = z) ; (B | Z = z)] \leq 3 I[A :B] + 2 H[Z] - H[A] - H[B]$$
-/

open Filter Function MeasureTheory Measure ProbabilityTheory
open scoped BigOperators

attribute [symm] ProbabilityTheory.IdentDistrib.symm

variable {Ω Ω' Ω'' Ω''' G T : Type*}
  [mΩ : MeasurableSpace Ω] {μ : Measure Ω}
  [mΩ' : MeasurableSpace Ω'] {μ' : Measure Ω'}
  [mΩ'' : MeasurableSpace Ω''] {μ'' : Measure Ω''}
  [mΩ''' : MeasurableSpace Ω'''] {μ''' : Measure Ω'''}
  [hG : MeasurableSpace G] [MeasurableSingletonClass G] [AddCommGroup G]
  [MeasurableSub₂ G] [MeasurableAdd₂ G] [Countable G]
  [Countable S] [Nonempty S] [MeasurableSpace S]
  [Countable T] [Nonempty T] [MeasurableSpace T]

variable {X : Ω → G} {Y : Ω' → G} {Z : Ω'' → G} [FiniteRange X] [FiniteRange Y] [FiniteRange Z]

/-- The Ruzsa distance `rdist X Y` or $d[X ;Y]$ between two random variables is defined as
$H[X'- Y'] - H[X']/2 - H[Y']/2$, where $X', Y'$ are independent copies of $X, Y$. -/
noncomputable
def rdist (X : Ω → G) (Y : Ω' → G) (μ : Measure Ω := by volume_tac)
    (μ' : Measure Ω' := by volume_tac) : ℝ :=
  H[fun x ↦ x.1 - x.2 ; (μ.map X).prod (μ'.map Y)] - H[X ; μ]/2 - H[Y ; μ']/2

/- Needed a new separator here, chose `#` arbitrarily, but am open to other suggestions -/
@[inherit_doc rdist] notation3:max "d[" X " ; " μ " # " Y " ; " μ' "]" => rdist X Y μ μ'

@[inherit_doc rdist] notation3:max "d[" X " # " Y "]" => rdist X Y volume volume

/-- Explicit formula for the Ruzsa distance. -/
lemma rdist_def (X : Ω → G) (Y : Ω' → G) (μ : Measure Ω) (μ' : Measure Ω') :
    d[X ; μ # Y ; μ']
      = H[fun x ↦ x.1 - x.2 ; (μ.map X).prod (μ'.map Y)] - H[X ; μ]/2 - H[Y ; μ']/2 := rfl

/-- Entropy depends continuously on the measure. -/
-- TODO: Use notation `Hm[μ]` here? (figure out how)
lemma continuous_measureEntropy_probabilityMeasure {Ω : Type*} [Fintype Ω]
    [TopologicalSpace Ω] [DiscreteTopology Ω] [MeasurableSpace Ω] [OpensMeasurableSpace Ω] :
    Continuous (fun (μ : ProbabilityMeasure Ω) ↦ measureEntropy (S := Ω) μ) := by
  unfold measureEntropy
  simp_rw [tsum_fintype]
  apply continuous_finset_sum
  intro ω _
  apply Real.continuous_negMulLog.comp
  simp only [measure_univ, inv_one, one_smul]
  exact continuous_probabilityMeasure_apply_of_isClopen (s := {ω}) $ isClopen_discrete _

lemma continuous_entropy_restrict_probabilityMeasure [Fintype G]
    [TopologicalSpace G] [DiscreteTopology G] [BorelSpace G] :
    Continuous (fun (μ : ProbabilityMeasure G) ↦ H[id ; μ.toMeasure]) := by
  simp only [entropy_def, Measure.map_id]
  exact continuous_measureEntropy_probabilityMeasure

/-- Ruzsa distance depends continuously on the measure. -/
lemma continuous_rdist_restrict_probabilityMeasure [Fintype G]
    [TopologicalSpace G] [DiscreteTopology G] [BorelSpace G] :
    Continuous
      (fun (μ : ProbabilityMeasure G × ProbabilityMeasure G) ↦
        d[id ; μ.1.toMeasure # id ; μ.2.toMeasure]) := by
  simp [rdist_def]
  have obs₀ : Continuous (fun (μ : ProbabilityMeasure G × ProbabilityMeasure G) ↦
      H[fun x ↦ x.1 - x.2 ; μ.1.toMeasure.prod μ.2.toMeasure]) := by
    simp_rw [entropy_def]
    have diff_cts : Continuous (fun (x : G × G) ↦ x.1 - x.2) := by continuity
    have key₁ := ProbabilityMeasure.continuous_prod_of_finite (α := G) (β := G)
    have key₂ := ProbabilityMeasure.continuous_map diff_cts
    convert continuous_measureEntropy_probabilityMeasure.comp (key₂.comp key₁)
  have obs₁ : Continuous
      (fun (μ : ProbabilityMeasure G × ProbabilityMeasure G) ↦ H[id ; μ.1.toMeasure]) := by
    convert (continuous_measureEntropy_probabilityMeasure (Ω := G)).comp continuous_fst
    simp [entropy_def]
  have obs₂ : Continuous
      (fun (μ : ProbabilityMeasure G × ProbabilityMeasure G) ↦ H[id ; μ.2.toMeasure]) := by
    convert (continuous_measureEntropy_probabilityMeasure (Ω := G)).comp continuous_snd
    simp [entropy_def]
  continuity

lemma continuous_rdist_restrict_probabilityMeasure₁ [Fintype G]
    [TopologicalSpace G] [DiscreteTopology G] [BorelSpace G]
    (X : Ω → G) (P : Measure Ω := by volume_tac) [IsProbabilityMeasure P] (X_mble : Measurable X) :
    Continuous
      (fun (μ : ProbabilityMeasure G) ↦ d[id ; P.map X # id ; μ.toMeasure]) := by
  have obs : IsProbabilityMeasure (P.map X) := by
    refine ⟨by simp [Measure.map_apply X_mble MeasurableSet.univ]⟩
  let ι : ProbabilityMeasure G → ProbabilityMeasure G × ProbabilityMeasure G :=
      fun ν ↦ ⟨⟨P.map X, obs⟩, ν⟩
  have ι_cont : Continuous ι := Continuous.Prod.mk _
  convert continuous_rdist_restrict_probabilityMeasure.comp ι_cont

/-- Ruzsa distance between random variables equals Ruzsa distance between their distributions.-/
lemma rdist_eq_rdist_id_map : d[X ; μ # Y ; μ'] = d[id ; μ.map X # id ; μ'.map Y] := by
  simp only [rdist_def, entropy_def, Measure.map_id]

/-- Ruzsa distance depends continuously on the second measure. -/
lemma continuous_rdist_restrict_probabilityMeasure₁' [Fintype G]
    [TopologicalSpace G] [DiscreteTopology G] [BorelSpace G]
    (X : Ω → G) (P : Measure Ω := by volume_tac) [IsProbabilityMeasure P] (X_mble : Measurable X) :
    Continuous
      (fun (μ : ProbabilityMeasure G) ↦ d[X ; P # id ; μ.toMeasure]) := by
  simp only [@rdist_eq_rdist_id_map Ω G G mΩ P hG, Measure.map_id]
  exact continuous_rdist_restrict_probabilityMeasure₁ _ _ X_mble

/-- If $X', Y'$ are copies of $X, Y$ respectively then $d[X' ; Y']=d[X ; Y]$. -/
lemma ProbabilityTheory.IdentDistrib.rdist_eq {X' : Ω'' → G} {Y' : Ω''' →G}
    (hX : IdentDistrib X X' μ μ'') (hY : IdentDistrib Y Y' μ' μ''') :
    d[X ; μ # Y ; μ'] = d[X' ; μ'' # Y' ; μ'''] := by
  simp [rdist, hX.map_eq, hY.map_eq, hX.entropy_eq, hY.entropy_eq]

/-- If $X, Y$ are independent $G$-random variables then
$$ d[X ; Y] := H[X - Y] - H[X]/2 - H[Y]/2$$-/
lemma ProbabilityTheory.IndepFun.rdist_eq [IsFiniteMeasure μ]
    {Y : Ω → G} (h : IndepFun X Y μ) (hX : Measurable X) (hY : Measurable Y) :
    d[X ; μ # Y ; μ] = H[X - Y ; μ] - H[X ; μ]/2 - H[Y ; μ]/2 := by
  rw [rdist_def]
  congr 2
  have h_prod : (μ.map X).prod (μ.map Y) = μ.map (⟨X, Y⟩) :=
    ((indepFun_iff_map_prod_eq_prod_map_map hX.aemeasurable hY.aemeasurable).mp h).symm
  rw [h_prod, entropy_def, Measure.map_map (measurable_fst.sub measurable_snd) (hX.prod_mk hY)]
  rfl

/-- $d[X;Y] ≤ H[X]/2 + H[Y]/2$. -/
lemma rdist_le_avg_ent {X : Ω → G} {Y : Ω' → G} [FiniteRange X] [FiniteRange Y] (hX: Measurable X) (hY: Measurable Y) (μ : Measure Ω := by volume_tac) (μ' : Measure Ω' := by volume_tac) [IsProbabilityMeasure μ] [IsProbabilityMeasure μ'] :
    d[X ; μ # Y ; μ'] ≤ (H[X ; μ] + H[Y ; μ'])/2 := by
  rcases ProbabilityTheory.independent_copies_finiteRange hX hY μ μ' with ⟨ν, X', Y', hprob, hX', hY', hindep, hidentX, hidentY, hfinX, hfinY⟩
  rw [← ProbabilityTheory.IdentDistrib.rdist_eq hidentX hidentY, ← IdentDistrib.entropy_eq hidentX, ← IdentDistrib.entropy_eq hidentY, ProbabilityTheory.IndepFun.rdist_eq hindep hX' hY']
  suffices H[X' - Y'; ν] ≤ H[X'; ν] + H[Y'; ν] by linarith
  change H[(fun x ↦ x.1 - x.2) ∘ ⟨X',Y' ⟩; ν] ≤ H[X'; ν] + H[Y'; ν]
  apply (entropy_comp_le ν (by measurability) _).trans
  exact entropy_pair_le_add hX' hY' ν

/-- Applying an injective homomorphism does not affect Ruzsa distance. -/
lemma rdist_of_inj  {H:Type*} [hH : MeasurableSpace H] [MeasurableSingletonClass H] [AddCommGroup H]
  [MeasurableSub₂ H] [MeasurableAdd₂ H] [Countable H] (hX: Measurable X) (hY: Measurable Y) (φ: G →+ H) (hφ: Function.Injective φ) [IsProbabilityMeasure μ] [IsProbabilityMeasure μ']: d[φ ∘ X ; μ # φ ∘ Y ; μ'] = d[X ; μ # Y ; μ'] := by
    rw [rdist_def, rdist_def]
    congr 2
    . rw [← entropy_comp_of_injective _ (by measurability) _ hφ]
      apply IdentDistrib.entropy_eq
      constructor
      . exact Measurable.aemeasurable (measurable_of_countable _)
      . exact Measurable.aemeasurable (measurable_of_countable _)
      set g := fun x : H × H ↦ x.1 - x.2
      set f := fun x : G × G ↦ (φ x.1, φ x.2)
      have : φ ∘ (fun x ↦ x.1 - x.2) = g ∘ f := by
        ext x
        simp
      rw [this, ← MeasureTheory.Measure.map_map (g := g) (f := f) (measurable_of_countable _) (measurable_of_countable _), ← MeasureTheory.Measure.map_map (measurable_of_countable _) hX, ← MeasureTheory.Measure.map_map (measurable_of_countable _) hY]
      congr
      convert MeasureTheory.Measure.map_prod_map _ _ (measurable_of_countable _) (measurable_of_countable _)
      . exact instSFiniteMap μ X
      . exact instSFiniteMap μ' Y
      all_goals infer_instance
    . congr 1
      exact entropy_comp_of_injective _ hX _ hφ
    exact entropy_comp_of_injective _ hY _ hφ

/-- $$ d[X ; 0] = H[X] / 2 $$ -/
lemma rdist_zero_eq_half_ent [IsFiniteMeasure μ] [IsProbabilityMeasure μ'] :
    d[X ; μ # fun _ ↦ 0 ; μ'] = H[X ; μ]/2 := by
  have aux : H[fun x => x.1 - x.2 ; Measure.prod (Measure.map X μ) (Measure.map (fun x => 0) μ')]
            = H[X ; μ] := by
    have h: Measure.map (fun x => x.1 - x.2)
                        (Measure.prod (Measure.map X μ) (Measure.map (fun x => 0) μ'))
            = Measure.map X μ := by
              simp [MeasureTheory.Measure.map_const, MeasureTheory.Measure.prod_dirac]
              rw [Measure.map_map]
              have helper : ((fun (x : G × G) => x.1 - x.2) ∘ fun x => (x, (0 : G))) = id := by
                funext; simp
              rw [helper, Measure.map_id]
              measurability
              measurability
    simp [entropy_def, h]
  simp [rdist_def, entropy_const (0 : G), aux]
  ring

/-- $$ d[X ; Y] = d[Y ; X]$$ -/
lemma rdist_symm [IsFiniteMeasure μ] [IsFiniteMeasure μ'] :
    d[X ; μ # Y ; μ'] = d[Y ; μ' # X ; μ] := by
  rw [rdist_def, rdist_def, sub_sub, sub_sub, add_comm]
  congr 1
  rw [← entropy_neg (measurable_fst.sub measurable_snd)]
  have : (-fun x : G × G ↦ x.1 - x.2) = (fun x ↦ x.1 - x.2) ∘ Prod.swap := by ext ; simp
  rw [this, entropy_def, ← Measure.map_map (measurable_fst.sub measurable_snd) measurable_swap,
    Measure.prod_swap]
  rfl

/-- $$|H[X] - H[Y]| \leq 2 d[X ; Y]$$ -/
lemma diff_ent_le_rdist [IsProbabilityMeasure μ] [IsProbabilityMeasure μ']
    (hX : Measurable X) (hY : Measurable Y) :
    |H[X ; μ] - H[Y ; μ']| ≤ 2 * d[X ; μ # Y ; μ'] := by
  obtain ⟨ν, X', Y', _, hX', hY', hind, hIdX, hIdY, _, _⟩ := independent_copies_finiteRange hX hY μ μ'
  rw [← hIdX.rdist_eq hIdY, hind.rdist_eq hX' hY', ← hIdX.entropy_eq, ← hIdY.entropy_eq, abs_le]
  have := max_entropy_le_entropy_sub hX' hY' hind
  constructor
  · linarith[le_max_right H[X'; ν] H[Y'; ν]]
  · linarith[le_max_left H[X'; ν] H[Y'; ν]]

/-- $$H[X - Y] - H[X] \leq 2d[X ; Y]$$ -/
lemma diff_ent_le_rdist' [IsProbabilityMeasure μ] {Y : Ω → G}
    (hX : Measurable X) (hY : Measurable Y) (h : IndepFun X Y μ) [FiniteRange Y]:
    H[X - Y ; μ] - H[X ; μ] ≤ 2 * d[X ; μ # Y ; μ] := by
  rw [h.rdist_eq hX hY]
  linarith[max_entropy_le_entropy_sub hX hY h, le_max_right H[X ; μ] H[Y; μ]]

/-- $$H[X - Y] - H[Y] \leq 2d[X ; Y]$$ -/
lemma diff_ent_le_rdist'' [IsProbabilityMeasure μ] {Y : Ω → G}
    (hX : Measurable X) (hY : Measurable Y) (h : IndepFun X Y μ) [FiniteRange Y]:
    H[X - Y ; μ] - H[Y ; μ] ≤ 2 * d[X ; μ # Y ; μ] := by
  rw [h.rdist_eq hX hY]
  linarith[max_entropy_le_entropy_sub hX hY h, le_max_left H[X ; μ] H[Y; μ]]

/-- $$ d[X ; Y] \geq 0$$ -/
lemma rdist_nonneg [IsProbabilityMeasure μ] [IsProbabilityMeasure μ']
    (hX : Measurable X) (hY : Measurable Y) : 0 ≤ d[X ; μ # Y ; μ'] := by
  linarith [ge_trans (diff_ent_le_rdist hX hY) (abs_nonneg (H[X ; μ] - H[Y ; μ']))]

/-- If $G$ is an additive group and $X$ is a $G$-valued random variable and
$H\leq G$ is a finite subgroup then, with $\pi:G\to G/H$ the natural homomorphism we have
(where $U_H$ is uniform on $H$) $\mathbb{H}(\pi(X))\leq 2d[X;U_H].$ -/
lemma ent_of_proj_le {UH: Ω' → G} [FiniteRange X] [FiniteRange UH]
    [IsProbabilityMeasure μ] [IsProbabilityMeasure μ']
    (hX : Measurable X) (hU: Measurable UH) {H: AddSubgroup G} [Finite H] -- TODO: infer from [FiniteRange UH]?
    (hunif: IsUniform H UH μ')
    : H[(QuotientAddGroup.mk' H) ∘ X; μ] ≤ 2 * d[X; μ # UH ; μ'] := by
  obtain ⟨ν, X', UH', hν, hX', hUH', h_ind, h_id_X', h_id_UH', _, _⟩ :=
    independent_copies_finiteRange hX hU μ μ'
  replace hunif : IsUniform H UH' ν :=
    IsUniform.of_identDistrib hunif h_id_UH'.symm (measurableSet_discrete _)
  rewrite [← (h_id_X'.comp (measurable_discrete _)).entropy_eq, ← h_id_X'.rdist_eq h_id_UH']
  let π := ⇑(QuotientAddGroup.mk' H)
  let νq := Measure.map (π ∘ X') ν
  haveI : Countable (HasQuotient.Quotient G H) := Quotient.countable
  haveI : MeasurableSingletonClass (HasQuotient.Quotient G H) :=
    { measurableSet_singleton := fun _ ↦ measurableSet_quotient.mpr (measurableSet_discrete _) }
  have : H[X' - UH' | π ∘ X' ; ν] = H[UH' ; ν] := by
    have h_meas_le : ∀ y ∈ FiniteRange.toFinset (π ∘ X'),
        (νq {y}).toReal * H[X' - UH' | (π ∘ X') ← y ; ν] ≤ (νq {y}).toReal * H[UH' ; ν] := by
      intro x _
      refine mul_le_mul_of_nonneg_left ?_ ENNReal.toReal_nonneg
      let ν' := ν[|π ∘ X' ← x]
      let π' := QuotientAddGroup.mk (s := H)
      have h_card : Nat.card (π' ⁻¹' {x}) = Nat.card H := Nat.card_congr <|
        (QuotientAddGroup.preimageMkEquivAddSubgroupProdSet H _).trans <| Equiv.prodUnique H _
      haveI : Finite (π' ⁻¹' {x}) :=
        Nat.finite_of_card_ne_zero <| h_card.trans_ne <| Nat.pos_iff_ne_zero.mp Nat.card_pos
      let H_x := (π' ⁻¹' {x}).toFinite.toFinset
      have h : ∀ᵐ ω ∂ν', (X' - UH') ω ∈ H_x := by
        let T : Set (G × G) := ((π' ∘ X') ⁻¹' {x})ᶜ
        let U : Set (G × G) := UH' ⁻¹' Hᶜ
        have h_subset : (X' - UH') ⁻¹' H_xᶜ ⊆ T ∪ U :=
          fun ω hω ↦ Classical.byContradiction fun h ↦ by simp_all [not_or, T, U, H_x, π']
        refine MeasureTheory.mem_ae_iff.mpr (le_zero_iff.mp ?_)
        calc
          _ ≤ (ν' T) + (ν' U) := (measure_mono h_subset).trans (measure_union_le T U)
          _ = (ν' T) + 0 := congrArg _ <| by
            show _ * _ = 0
            rw [le_zero_iff.mp <| (restrict_apply_le _ U).trans_eq hunif.measure_preimage_compl,
              mul_zero]
          _ = 0 := (add_zero _).trans <| by
            have : restrict ν (π ∘ X' ⁻¹' {x}) T = 0 := by
              simp [restrict_apply (measurableSet_discrete T), T, π', π]
            show _ * _ = 0
            rw [this, mul_zero]
      convert entropy_le_log_card_of_mem (Measurable.sub hX' hUH') h
      simp_rw [hunif.entropy_eq' hUH', H_x, Set.Finite.mem_toFinset, h_card, SetLike.coe_sort_coe]
    have h_one : (∑ x in FiniteRange.toFinset (π ∘ X'), (νq {x}).toReal) = 1 := by
      rewrite [Finset.sum_toReal_measure_singleton]
      apply (ENNReal.toReal_eq_one_iff _).mpr
      haveI := isProbabilityMeasure_map <| (measurable_discrete (π ∘ X')).aemeasurable (μ := ν)
      rewrite [← measure_univ (μ := νq), ← FiniteRange.range]
      let rng := Set.range (π ∘ X')
      have h_compl : νq rngᶜ = 0 := ae_map_mem_range (π ∘ X') (measurableSet_discrete _) ν
      rw [← MeasureTheory.measure_add_measure_compl (measurableSet_discrete rng), h_compl, add_zero]
    haveI := FiniteRange.sub X' UH'
    have h_ge : H[X' - UH' | π ∘ X' ; ν] ≥ H[UH' ; ν] := calc
      _ ≥ H[X' - UH' | X' ; ν] := condEntropy_comp_ge ν hX' (hX'.sub hUH') π
      _ = H[UH' | X' ; ν] := condEntropy_sub_left hUH' hX'
      _ = H[UH' ; ν] := h_ind.symm.condEntropy_eq_entropy hUH' hX'
    have h_le : H[X' - UH' | π ∘ X' ; ν] ≤ H[UH' ; ν] := by
      rewrite [condEntropy_eq_sum _ _ _ (measurable_discrete _)]
      apply (Finset.sum_le_sum h_meas_le).trans
      rewrite [← Finset.sum_mul, h_one, one_mul]
      rfl
    exact h_le.ge_iff_eq.mp h_ge
  have : H[X' - UH' ; ν] = H[π ∘ X' ; ν] + H[UH' ; ν] := by calc
    _ = H[⟨X' - UH', π ∘ (X' - UH')⟩ ; ν] := (entropy_prod_comp (hX'.sub hUH') ν π).symm
    _ = H[⟨X' - UH', π ∘ X'⟩ ; ν] := by
      apply IdentDistrib.entropy_eq
      apply IdentDistrib.of_ae_eq (Measurable.aemeasurable (measurable_discrete _))
      apply MeasureTheory.mem_ae_iff.mpr
      convert hunif.measure_preimage_compl
      ext; simp [π]
    _ = H[π ∘ X' ; ν] + H[UH' ; ν] := by
      rewrite [chain_rule ν (by exact hX'.sub hUH') (measurable_discrete _)]
      congr
  have : d[X' ; ν # UH' ; ν] = H[π ∘ X' ; ν] + (H[UH' ; ν] - H[X' ; ν]) / 2 := by
    rewrite [h_ind.rdist_eq hX' hUH']
    linarith only [this]
  linarith only [this, (abs_le.mp (diff_ent_le_rdist hX' hUH' (μ := ν) (μ' := ν))).2]

/-- Adding a constant to a random variable does not change the Rusza distance. -/
lemma rdist_add_const [IsProbabilityMeasure μ] [IsProbabilityMeasure μ']
    (hX : Measurable X) (hY : Measurable Y) :
    d[X ; μ # Y + fun _ ↦ c; μ'] = d[X ; μ # Y ; μ'] := by
  obtain ⟨ν, X', Y', _, hX', hY', hind, hIdX, hIdY, _, _⟩ := independent_copies_finiteRange hX hY μ μ'
  have A : IdentDistrib (Y' + fun _ ↦ c) (Y + fun _ ↦ c) ν μ' := by
    change IdentDistrib (fun ω ↦ Y' ω + c) (fun ω ↦ Y ω + c) ν μ'
    apply hIdY.comp (measurable_add_const c)
  have B : IndepFun X' (Y' + fun _ ↦ c) ν := by
    change IndepFun X' (fun ω ↦ Y' ω + c) ν
    apply hind.comp measurable_id (measurable_add_const c)
  have C : X' - (Y' + fun _ ↦ c) = (X' - Y') + (fun _ ↦ -c) := by
    ext ω; simp; abel
  rw [← hIdX.rdist_eq hIdY, ← hIdX.rdist_eq A, hind.rdist_eq hX' hY',
    B.rdist_eq hX' (hY'.add_const _), entropy_add_const hY' c, C, entropy_add_const]
  exact hX'.sub hY'

/-- A variant of `rdist_add_const` where one adds constants to both variables. -/
lemma rdist_add_const' [IsProbabilityMeasure μ] [IsProbabilityMeasure μ'] (c: G) (c': G)
    (hX : Measurable X) (hY : Measurable Y) :
    d[X + fun _ ↦ c; μ # Y + fun _ ↦ c'; μ'] = d[X ; μ # Y ; μ'] := by
    rw [rdist_add_const _ hY, rdist_symm, rdist_add_const hY hX, rdist_symm]
    measurability

/-- The **improved entropic Ruzsa triangle inequality**. -/
lemma ent_of_diff_le (X : Ω → G) (Y : Ω → G) (Z : Ω → G)
    (hX : Measurable X) (hY : Measurable Y) (hZ : Measurable Z)
    (h : IndepFun (⟨X, Y⟩) Z μ) [IsProbabilityMeasure μ] [FiniteRange X] [FiniteRange Y] [FiniteRange Z]:
    H[X - Y; μ] ≤ H[X - Z; μ] + H[Z - Y; μ] - H[Z; μ] := by
  have h1 : H[⟨X - Z, ⟨Y, X - Y⟩⟩; μ] + H[X - Y; μ] ≤ H[⟨X - Z, X - Y⟩; μ] + H[⟨Y, X - Y⟩; μ] :=
    entropy_triple_add_entropy_le μ (hX.sub hZ) hY (hX.sub hY)
  have h2 : H[⟨X - Z, X - Y⟩ ; μ] ≤ H[X - Z ; μ] + H[Y - Z ; μ] := by
    calc H[⟨X - Z, X - Y⟩ ; μ] ≤ H[⟨X - Z, Y - Z⟩ ; μ] := by
          have : ⟨X - Z, X - Y⟩ = (fun p ↦ (p.1, p.1 - p.2)) ∘ ⟨X - Z, Y - Z⟩ := by ext1; simp
          rw [this]
          apply entropy_comp_le μ (by measurability)
    _ ≤ H[X - Z ; μ] + H[Y - Z ; μ] := by
          have h : 0 ≤ H[X - Z ; μ] + H[Y - Z ; μ] - H[⟨X - Z, Y - Z⟩ ; μ] := by
            apply mutualInfo_nonneg (by measurability) (by measurability) μ
          linarith
  have h3 : H[⟨Y, X - Y⟩ ; μ] ≤ H[⟨X, Y⟩ ; μ] := by
    have : ⟨Y, X - Y⟩ = (fun p ↦ (p.2, p.1 - p.2)) ∘ ⟨X, Y⟩ := by ext1; simp
    rw [this]
    exact entropy_comp_le μ (hX.prod_mk hY) _
  have h4 : H[⟨X - Z, ⟨Y, X - Y⟩⟩; μ] = H[⟨X, ⟨Y, Z⟩⟩ ; μ] := by
    refine entropy_of_comp_eq_of_comp μ ((hX.sub hZ).prod_mk (hY.prod_mk (hX.sub hY)))
      (hX.prod_mk (hY.prod_mk hZ))
      (fun p : G × (G × G) ↦ (p.2.2 + p.2.1, p.2.1, -p.1 + p.2.2 + p.2.1))
      (fun p : G × G × G ↦ (p.1 - p.2.2, p.2.1, p.1 - p.2.1)) ?_ ?_
    · ext1; simp
    · ext1; simp
  have h5 : H[⟨X, ⟨Y, Z⟩⟩ ; μ] = H[⟨X, Y⟩ ; μ] + H[Z ; μ] := by
    rw [entropy_assoc hX hY hZ, entropy_pair_eq_add (hX.prod_mk hY) hZ]
    exact h
  rw [h4, h5] at h1
  calc H[X - Y; μ] ≤ H[X - Z; μ] + H[Y - Z; μ] - H[Z; μ] := by linarith
  _ = H[X - Z; μ] + H[Z - Y; μ] - H[Z; μ] := by
    congr 2
    rw [entropy_sub_comm hY hZ]

/-- The **entropic Ruzsa triangle inequality** -/
lemma rdist_triangle {X : Ω → G} {Y : Ω' → G} {Z : Ω'' → G}
  (hX : Measurable X) (hY : Measurable Y) (hZ : Measurable Z)
  [hμ : IsProbabilityMeasure μ] [hμ' : IsProbabilityMeasure μ'] [hμ'' : IsProbabilityMeasure μ''] [FiniteRange X] [FiniteRange Y] [FiniteRange Z]:
    d[X ; μ # Z ; μ''] ≤ d[X ; μ # Y ; μ'] + d[Y ; μ' # Z ; μ''] := by
  obtain ⟨A, mA, μA, X', Y', Z', hμA, hInd, hX', hY', hZ', HX, HY, HZ, _, _, _⟩ :=
    independent_copies3_nondep_finiteRange hX hY hZ μ μ' μ''
  suffices d[X' ; μA # Z' ; μA] ≤ d[X' ; μA # Y' ; μA] + d[Y' ; μA # Z' ; μA] by
    rwa [HX.rdist_eq HY, HY.rdist_eq HZ, HX.rdist_eq HZ] at this
  have IndepLem : IndepFun (⟨X', Z'⟩) Y' μA :=
    iIndepFun.indepFun_prod_mk hInd (fun i => by fin_cases i ; all_goals { simpa }) 0 2 1
      (by norm_cast) (by norm_cast)
  calc d[X' ; μA # Z' ; μA] = H[X' - Z'; μA] - (H[X'; μA] / 2 + H[Z'; μA] / 2) := by
        rw [ProbabilityTheory.IndepFun.rdist_eq
          (by simpa using hInd.indepFun (show 0 ≠ 2 by norm_cast)) hX' hZ'] ; ring
    _ ≤ (H[X' - Y' ; μA] + H[Y' - Z' ; μA] - H[Y' ; μA]) - (H[X'; μA] / 2 + H[Z'; μA] / 2) :=
          sub_le_sub_right (ent_of_diff_le _ _ _ hX' hZ' hY' IndepLem) _
    _ = (H[X' - Y' ; μA] - H[X'; μA] / 2 - H[Y' ; μA] / 2) +
          (H[Y' - Z' ; μA] - H[Y' ; μA] / 2 - H[Z'; μA] / 2) := by ring
    _ = d[X' ; μA # Y' ; μA] + d[Y' ; μA # Z' ; μA] := by
        rw [ProbabilityTheory.IndepFun.rdist_eq (by simpa using (iIndepFun.indepFun hInd
          (show 0 ≠ 1 by norm_cast))) hX' hY', ProbabilityTheory.IndepFun.rdist_eq
          (by simpa using (iIndepFun.indepFun hInd (show 1 ≠ 2 by norm_cast))) hY' hZ']

variable [MeasurableSingletonClass S] [MeasurableSingletonClass T]

/-- The conditional Ruzsa distance `d[X|Z ; Y|W]`. -/
noncomputable
def condRuzsaDist (X : Ω → G) (Z : Ω → S) (Y : Ω' → G) (W : Ω' → T)
    (μ : Measure Ω := by volume_tac) [IsFiniteMeasure μ]
    (μ' : Measure Ω' := by volume_tac) [IsFiniteMeasure μ'] : ℝ :=
  dk[condDistrib X Z μ ; μ.map Z # condDistrib Y W μ' ; μ'.map W]

@[inherit_doc condRuzsaDist]
notation3:max "d[" X " | " Z " ; " μ " # " Y " | " W " ; " μ'"]" => condRuzsaDist X Z Y W μ μ'

@[inherit_doc condRuzsaDist]
notation3:max "d[" X " | " Z " # " Y " | " W "]" => condRuzsaDist X Z Y W volume volume

lemma condRuzsaDist_def (X : Ω → G) (Z : Ω → S) (Y : Ω' → G) (W : Ω' → T)
    (μ : Measure Ω) [IsFiniteMeasure μ] (μ' : Measure Ω') [IsFiniteMeasure μ'] :
    d[X | Z ; μ # Y | W ; μ']
      = dk[condDistrib X Z μ ; μ.map Z # condDistrib Y W μ' ; μ'.map W] := rfl

/-- $$ d[X|Z; Y|W] = d[Y|W; X|Z]$$-/
lemma condRuzsaDist_symm {X : Ω → G} {Z : Ω → S} {Y : Ω' → G} {W : Ω' → T} (hZ : Measurable Z)
    (hW : Measurable W) [IsProbabilityMeasure μ] [IsProbabilityMeasure μ'] [FiniteRange Z]
    [FiniteRange W] :
    d[X | Z ; μ # Y | W ; μ'] = d[Y | W ; μ' # X | Z ; μ] := by
  have : IsProbabilityMeasure (μ.map Z) := isProbabilityMeasure_map hZ.aemeasurable
  have : IsProbabilityMeasure (μ'.map W) := isProbabilityMeasure_map hW.aemeasurable
  rw [condRuzsaDist_def, condRuzsaDist_def, kernel.rdist_symm]

@[simp] lemma condRuszaDist_zero_right (X : Ω → G) (Z : Ω → S) (Y : Ω' → G) (W : Ω' → T)
    (μ : Measure Ω) [IsFiniteMeasure μ] :
    d[X | Z ; μ # Y | W ; 0] = 0 := by
  simp only [condRuzsaDist, aemeasurable_zero_measure, not_true_eq_false, Measure.map_zero,
    kernel.rdist_zero_right]

@[simp] lemma condRuszaDist_zero_left (X : Ω → G) (Z : Ω → S) (Y : Ω' → G) (W : Ω' → T)
    (μ' : Measure Ω') [IsFiniteMeasure μ'] :
    d[X | Z ; 0 # Y | W ; μ'] = 0 := by
  simp [condRuzsaDist]

lemma condRuzsaDist_nonneg {X : Ω → G} (hX : Measurable X) [FiniteRange X]
    {Z : Ω → S} (hZ : Measurable Z) [FiniteRange Z]
    {Y : Ω' → G} (hY : Measurable Y) [FiniteRange Y]
    {W : Ω' → T} (hW : Measurable W) [FiniteRange W]
    [IsProbabilityMeasure μ] [IsProbabilityMeasure μ'] :
    0 ≤ d[X | Z ; μ # Y | W ; μ'] := by
  rw [condRuzsaDist_def]
  have : IsProbabilityMeasure (μ.map Z) := isProbabilityMeasure_map hZ.aemeasurable
  have : IsProbabilityMeasure (μ'.map W) := isProbabilityMeasure_map hW.aemeasurable
  refine kernel.rdist_nonneg ?_ ?_
  · exact kernel.aefiniteKernelSupport_condDistrib _ _ _ hX hZ
  · exact kernel.aefiniteKernelSupport_condDistrib _ _ _ hY hW

/-- Ruzsa distance of random variables equals Ruzsa distance of the kernels. -/
lemma rdist_eq_rdistm : d[X ; μ # Y ; μ'] = kernel.rdistm (μ.map X) (μ'.map Y) := rfl

/-- Explicit formula for conditional Ruzsa distance $d[X|Z; Y|W]$. -/
lemma condRuzsaDist_eq_sum {X : Ω → G} {Z : Ω → S} {Y : Ω' → G} {W : Ω' → T}
    (hX : Measurable X) (hZ : Measurable Z) (hY : Measurable Y) (hW : Measurable W)
    (μ : Measure Ω) [IsFiniteMeasure μ] (μ' : Measure Ω') [IsFiniteMeasure μ']
    [FiniteRange Z] [FiniteRange W] :
    d[X | Z ; μ # Y | W ; μ']
      = ∑ z in FiniteRange.toFinset Z, ∑ w in FiniteRange.toFinset W,
        (μ (Z ⁻¹' {z})).toReal * (μ' (W ⁻¹' {w})).toReal
          * d[X ; (μ[|Z ← z]) # Y ; (μ'[|W ← w])] := by
  have : Measure.prod (μ.map Z) (μ'.map W) ((((FiniteRange.toFinset Z)
      ×ˢ (FiniteRange.toFinset W)) : Finset (S × T)): Set (S × T))ᶜ = 0 := by
    apply Measure.prod_of_full_measure_finset
    all_goals {
      rw [Measure.map_apply ‹_›]
      convert measure_empty
      simp [← FiniteRange.range]
      measurability
    }
  rw [condRuzsaDist_def, kernel.rdist, integral_eq_sum' _ this]
  simp_rw [Measure.prod_apply_singleton, ENNReal.toReal_mul, smul_eq_mul, Finset.sum_product,
    Measure.map_apply hZ (measurableSet_singleton _),
    Measure.map_apply hW (measurableSet_singleton _)]
  congr with z
  congr with w
  by_cases hz : μ (Z ⁻¹' {z}) = 0
  · simp only [mul_eq_mul_left_iff, mul_eq_zero]
    refine Or.inr (Or.inl ?_)
    simp [ENNReal.toReal_eq_zero_iff, measure_ne_top μ, hz]
  by_cases hw : μ' (W ⁻¹' {w}) = 0
  · simp only [mul_eq_mul_left_iff, mul_eq_zero]
    refine Or.inr (Or.inr ?_)
    simp [ENNReal.toReal_eq_zero_iff, measure_ne_top μ', hw]
  congr 1
  rw [rdist_eq_rdistm]
  rw [condDistrib_apply hX hZ _ _ hz, condDistrib_apply hY hW _ _ hw]

/-- Explicit formula for conditional Ruzsa distance $d[X|Z; Y|W]$ in a fintype. -/
lemma condRuzsaDist_eq_sum' {X : Ω → G} {Z : Ω → S} {Y : Ω' → G} {W : Ω' → T}
    (hX : Measurable X) (hZ : Measurable Z) (hY : Measurable Y) (hW : Measurable W)
    (μ : Measure Ω) [IsFiniteMeasure μ] (μ' : Measure Ω') [IsFiniteMeasure μ']
    [Fintype S] [Fintype T] :
    d[X | Z ; μ # Y | W ; μ']
      = ∑ z, ∑ w, (μ (Z ⁻¹' {z})).toReal * (μ' (W ⁻¹' {w})).toReal
          * d[X ; (μ[|Z ← z]) # Y ; (μ'[|W ← w])] := by
  rw [condRuzsaDist_def, kernel.rdist, integral_eq_sum]
  simp_rw [Measure.prod_apply_singleton, ENNReal.toReal_mul, smul_eq_mul, Fintype.sum_prod_type,
    Measure.map_apply hZ (measurableSet_singleton _),
    Measure.map_apply hW (measurableSet_singleton _)]
  congr with z
  congr with w
  by_cases hz : μ (Z ⁻¹' {z}) = 0
  · simp only [mul_eq_mul_left_iff, mul_eq_zero]
    refine Or.inr (Or.inl ?_)
    simp [ENNReal.toReal_eq_zero_iff, measure_ne_top μ, hz]
  by_cases hw : μ' (W ⁻¹' {w}) = 0
  · simp only [mul_eq_mul_left_iff, mul_eq_zero]
    refine Or.inr (Or.inr ?_)
    simp [ENNReal.toReal_eq_zero_iff, measure_ne_top μ', hw]
  congr 1
  rw [rdist_eq_rdistm]
  rw [condDistrib_apply hX hZ _ _ hz, condDistrib_apply hY hW _ _ hw]

/-- The conditional Ruzsa distance `d[X ; Y|W]`. -/
noncomputable
def condRuzsaDist' (X : Ω → G) (Y : Ω' → G) (W : Ω' → T)
    (μ : Measure Ω := by volume_tac) (μ' : Measure Ω' := by volume_tac) [IsFiniteMeasure μ'] : ℝ :=
  dk[kernel.const Unit (μ.map X) ; Measure.dirac () # condDistrib Y W μ' ; μ'.map W]

@[inherit_doc condRuzsaDist']
notation3:max "d[" X " ; " μ " # " Y " | " W " ; " μ' "]" => condRuzsaDist' X Y W μ μ'
@[inherit_doc condRuzsaDist']
notation3:max "d[" X " # " Y " | " W "]" => condRuzsaDist' X Y W volume volume

/-- Conditional Ruzsa distance equals Ruzsa distance of associated kernels. -/
lemma condRuzsaDist'_def (X : Ω → G) (Y : Ω' → G) (W : Ω' → T) (μ : Measure Ω) (μ' : Measure Ω')
    [IsFiniteMeasure μ'] :
    d[X ; μ # Y | W ; μ'] =
      dk[kernel.const Unit (μ.map X) ; Measure.dirac () # condDistrib Y W μ' ; μ'.map W] :=
  rfl

@[simp] lemma condRuzsaDist'_zero_right (X : Ω → G) (Y : Ω' → G) (W : Ω' → T) (μ : Measure Ω) :
    d[X ; μ # Y | W ; 0] = 0 := by
  simp only [condRuzsaDist'_def, aemeasurable_zero_measure, not_true_eq_false, Measure.map_zero,  kernel.rdist_zero_right]

/-- Explicit formula for conditional Ruzsa distance $d[X ; Y|W]$. -/
lemma condRuzsaDist'_eq_sum {X : Ω → G} {Y : Ω' → G} {W : Ω' → T} (hY : Measurable Y)
    (hW : Measurable W) (μ : Measure Ω) (μ' : Measure Ω') [IsFiniteMeasure μ'] [FiniteRange W] :
    d[X ; μ # Y | W ; μ']
      = ∑ w in FiniteRange.toFinset W, (μ' (W ⁻¹' {w})).toReal * d[X ; μ # Y ; (μ'[|W ← w])] := by
  have : Measure.prod (dirac ()) (μ'.map W) ((Finset.univ (α := Unit) ×ˢ FiniteRange.toFinset W : Finset (Unit × T)) : Set (Unit × T))ᶜ = 0 := by
    apply Measure.prod_of_full_measure_finset
    . simp
    rw [Measure.map_apply ‹_›]
    convert measure_empty
    simp [← FiniteRange.range]
    measurability
  rw [condRuzsaDist'_def, kernel.rdist, integral_eq_sum' _ this]
  simp_rw [Measure.prod_apply_singleton, smul_eq_mul, Finset.sum_product]
  simp only [Finset.univ_unique, PUnit.default_eq_unit, MeasurableSpace.measurableSet_top,
    Measure.dirac_apply', Set.mem_singleton_iff, Set.indicator_of_mem, Pi.one_apply, one_mul,
    Finset.sum_singleton]
  simp_rw [Measure.map_apply hW (measurableSet_singleton _)]
  congr with w
  by_cases hw : μ' (W ⁻¹' {w}) = 0
  · simp only [mul_eq_mul_left_iff]
    refine Or.inr ?_
    simp [ENNReal.toReal_eq_zero_iff, measure_ne_top μ', hw]
  congr 1
  rw [rdist_eq_rdistm, condDistrib_apply hY hW _ _ hw]
  congr

/-- Alternate formula for conditional Ruzsa distance $d[X ; Y|W]$ when T is a Fintype. -/
lemma condRuzsaDist'_eq_sum' {X : Ω → G} {Y : Ω' → G} {W : Ω' → T} (hY : Measurable Y)
    (hW : Measurable W) (μ : Measure Ω) (μ' : Measure Ω') [IsFiniteMeasure μ'] [Fintype T] :
    d[X ; μ # Y | W ; μ']
      = ∑ w, (μ' (W ⁻¹' {w})).toReal * d[X ; μ # Y ; (μ'[|W ← w])] := by
  rw [condRuzsaDist'_eq_sum hY hW μ μ']
  apply Finset.sum_subset
  . simp
  intro w _ hw
  simp at hw
  have : W⁻¹' {w} = ∅ := by
    ext ω; simp [hw ω]
  simp [this]

open scoped ENNReal

lemma condRuzsaDist'_prod_eq_sum {X : Ω → G} {Y : Ω' → G} {W W' : Ω' → T}
    (μ : Measure Ω) (μ' : Measure Ω') (hY : Measurable Y) (hW' : Measurable W') (hW : Measurable W)
    [IsFiniteMeasure μ'] [FiniteRange W] [FiniteRange W']:
    d[X ; μ # Y | ⟨W', W⟩; μ']
      = ∑ w in FiniteRange.toFinset W, (μ' (W ⁻¹' {w})).toReal * d[X ; μ # Y | W' ; (μ'[|W ← w])] := by
  have : d[X ; μ # Y | ⟨W', W⟩; μ'] = ∑ w in ((FiniteRange.toFinset W') ×ˢ FiniteRange.toFinset W), (μ' ((fun a => (W' a, W a)) ⁻¹' {w})).toReal * d[X ; μ # Y ; μ'[|(fun a => (W' a, W a)) ⁻¹' {w}]] := by
    rw [condRuzsaDist'_eq_sum hY (hW'.prod_mk hW)]
    apply Finset.sum_subset
    . intro (t, t')
      simp
      intro ω h1 h2
      exact ⟨⟨ω, h1⟩, ⟨ω, h2⟩⟩
    intro (t, t') _ ht
    simp at ht
    have : (fun a => (W' a, W a)) ⁻¹' {(t, t')} = ∅ := by
      ext ω
      simp
      exact ht ω
    simp [this]
  rw [this, Finset.sum_product_right]
  congr 1 with w
  rw [condRuzsaDist'_eq_sum hY hW', Finset.mul_sum]
  congr 1 with w'
  rw [← mul_assoc]
  have A : (fun a ↦ (W' a, W a)) ⁻¹' {(w', w)} = W' ⁻¹' {w'} ∩ W⁻¹' {w} := by ext; simp
  congr 1
  · simp only [ProbabilityTheory.cond, smul_toOuterMeasure, OuterMeasure.coe_smul, Pi.smul_apply,
      smul_eq_mul, ENNReal.toReal_mul, A, restrict_apply (hW' (MeasurableSet.singleton w'))]
    rcases eq_bot_or_bot_lt (μ' (W ⁻¹' {w})) with hw|hw
    · have : μ' (W' ⁻¹' {w'} ∩ W ⁻¹' {w}) = 0 :=
        le_antisymm (le_trans (measure_mono (Set.inter_subset_right _ _)) hw.le) bot_le
      simp [hw, this]
    · rw [← mul_assoc, ← ENNReal.toReal_mul, ENNReal.mul_inv_cancel, ENNReal.one_toReal, one_mul]
      exacts [hw.ne', by finiteness]
  · congr 1
    rw [A, cond_cond_eq_cond_inter' _ (hW (.singleton w)) (hW' (.singleton w')), Set.inter_comm]
    finiteness

/-- Version of `condRuzsaDist'_prod_eq_sum` when `W` has finite codomain. -/
lemma condRuzsaDist'_prod_eq_sum' {X : Ω → G} {Y : Ω' → G} {W W' : Ω' → T}
    (μ : Measure Ω) (μ' : Measure Ω') (hY : Measurable Y) (hW' : Measurable W') (hW : Measurable W)
    [IsFiniteMeasure μ'] [Fintype T]:
    d[X ; μ # Y | ⟨W', W⟩; μ']
      = ∑ w, (μ' (W ⁻¹' {w})).toReal * d[X ; μ # Y | W' ; (μ'[|W ← w])] := by
  rw [condRuzsaDist'_prod_eq_sum μ μ' hY hW' hW]
  apply Finset.sum_subset
  . simp
  intro w _ hw
  simp at hw
  have : W⁻¹' {w} = ∅ := by
    ext ω; simp [hw ω]
  simp [this]

/-- Explicit formula for conditional Ruzsa distance $d[X ; Y|W]$, in integral form. -/
lemma condRuzsaDist'_eq_integral (X : Ω → G) {Y : Ω' → G} {W : Ω' → T}
    (hY : Measurable Y) (hW : Measurable W)
    (μ : Measure Ω) (μ' : Measure Ω') [IsFiniteMeasure μ'] [FiniteRange W] :
    d[X ; μ # Y | W ; μ']
      = (μ'.map W)[fun w ↦ d[X ; μ # Y ; (μ'[|W ← w])]] := by
  rw [condRuzsaDist'_eq_sum hY hW]
  simp_rw [← smul_eq_mul]
  have : (μ'.map W) (FiniteRange.toFinset W : Set T)ᶜ = 0 := by
    rw [Measure.map_apply ‹_›]
    convert measure_empty
    simp [← FiniteRange.range]
    measurability
  convert symm $ integral_eq_sum' (μ'.map W) this (fun w ↦ d[X ; μ # Y ; (μ'[|W ← w])])
  rw [Measure.map_apply hW (MeasurableSet.singleton _)]

/-- Conditioning by a constant does not affect Ruzsa distance. -/
lemma condRuzsaDist_of_const {X : Ω → G} (hX : Measurable X) (Y : Ω' → G) (W : Ω' → T) (c : S)
    [IsProbabilityMeasure μ] [IsProbabilityMeasure μ'] [FiniteRange W] :
    d[X|(fun _ ↦ c) ; μ # Y | W ; μ'] = d[X ; μ # Y | W ; μ'] := by
  rw [condRuzsaDist_def, condRuzsaDist'_def, Measure.map_const,measure_univ,one_smul, kernel.rdist,
    kernel.rdist, integral_prod, integral_dirac, integral_prod,integral_dirac]
  dsimp; congr; ext x; congr
  rw [condDistrib_apply hX measurable_const]
  · simp
  · simp
  · exact integrable_of_finiteSupport _
  · exact integrable_of_finiteSupport _

/-- If $(X,Z)$ and $(Y,W)$ are independent, then
$$ d[X | Z ; Y | W] = H[X'- Y'|Z', W'] - H[X'|Z']/2 - H[Y'|W']/2$$
-/
lemma condRuzsaDist_of_indep
    {X : Ω → G} {Z : Ω → S} {Y : Ω → G} {W : Ω → T}
    (hX : Measurable X) (hZ : Measurable Z) (hY : Measurable Y) (hW : Measurable W)
    (μ : Measure Ω) [IsProbabilityMeasure μ]
    (h : IndepFun (⟨X, Z⟩) (⟨Y, W⟩) μ) [FiniteRange Z] [FiniteRange W] :
    d[X | Z ; μ # Y | W ; μ] = H[X - Y | ⟨Z, W⟩ ; μ] - H[X | Z ; μ]/2 - H[Y | W ; μ]/2 := by
  have : IsProbabilityMeasure (μ.map Z) := isProbabilityMeasure_map hZ.aemeasurable
  have : IsProbabilityMeasure (μ.map W) := isProbabilityMeasure_map hW.aemeasurable
  rw [condRuzsaDist_def, kernel.rdist_eq', condEntropy_eq_kernel_entropy _ (hZ.prod_mk hW),
    condEntropy_eq_kernel_entropy hX hZ, condEntropy_eq_kernel_entropy hY hW]
  swap; · exact hX.sub hY
  congr 2
  have hZW : IndepFun Z W μ := by
    have h' := IndepFun.comp h measurable_snd measurable_snd
    exact h'
  have hZW_map : μ.map (⟨Z, W⟩) = (μ.map Z).prod (μ.map W) :=
    (indepFun_iff_map_prod_eq_prod_map_map hZ.aemeasurable hW.aemeasurable).mp hZW
  rw [← hZW_map]
  refine kernel.entropy_congr ?_
  have : kernel.map (condDistrib (⟨X, Y⟩) (⟨Z, W⟩) μ) (fun x ↦ x.1 - x.2) measurable_sub
      =ᵐ[μ.map (⟨Z, W⟩)] condDistrib (X - Y) (⟨Z, W⟩) μ :=
    (condDistrib_comp (hX.prod_mk hY) (hZ.prod_mk hW) _ _).symm
  refine (this.symm.trans ?_).symm
  suffices kernel.prodMkRight T (condDistrib X Z μ)
        ×ₖ kernel.prodMkLeft S (condDistrib Y W μ)
      =ᵐ[μ.map (⟨Z, W⟩)] condDistrib (⟨X, Y⟩) (⟨Z, W⟩) μ by
    filter_upwards [this] with x hx
    rw [kernel.map_apply, kernel.map_apply, hx]
  . exact (condDistrib_eq_prod_of_indepFun hX hZ hY hW μ h).symm

/-- Formula for conditional Ruzsa distance for independent sets of variables. -/
lemma condRuzsaDist'_of_indep {X : Ω → G} {Y : Ω → G} {W : Ω → T}
    (hX : Measurable X) (hY : Measurable Y) (hW : Measurable W)
    (μ : Measure Ω) [IsProbabilityMeasure μ]
    (h : IndepFun X (⟨Y, W⟩) μ) [FiniteRange W] :
    d[X ; μ # Y | W ; μ] = H[X - Y | W ; μ] - H[X ; μ]/2 - H[Y | W ; μ]/2 := by
  have : IsProbabilityMeasure (μ.map W) := isProbabilityMeasure_map hW.aemeasurable
  rw [condRuzsaDist'_def, kernel.rdist_eq', condEntropy_eq_kernel_entropy _ hW,
    condEntropy_eq_kernel_entropy hY hW, entropy_eq_kernel_entropy]
  rotate_left
  · exact hX.sub hY
  congr 2
  let Z : Ω → Unit := fun _ ↦ ()
  rw [← condDistrib_unit_right hX μ]
  have h' : IndepFun (⟨X,Z⟩) (⟨Y, W⟩) μ := by
    rw [indepFun_iff_measure_inter_preimage_eq_mul]
    intro s t hs ht
    have : ⟨X, Z⟩ ⁻¹' s = X ⁻¹' ((fun c ↦ (c, ())) ⁻¹' s) := by ext1 y; simp
    rw [this]
    rw [indepFun_iff_measure_inter_preimage_eq_mul] at h
    exact h _ _ (measurable_prod_mk_right hs) ht
  have h_indep := condDistrib_eq_prod_of_indepFun hX measurable_const hY hW _ h'
  have h_meas_eq : μ.map (⟨Z, W⟩) = (Measure.dirac ()).prod (μ.map W) := by
    ext s hs
    rw [Measure.map_apply (measurable_const.prod_mk hW) hs, Measure.prod_apply hs, lintegral_dirac,
      Measure.map_apply hW (measurable_prod_mk_left hs)]
    congr
  rw [← h_meas_eq]
  have : kernel.map (kernel.prodMkRight T (condDistrib X Z μ)
        ×ₖ kernel.prodMkLeft Unit (condDistrib Y W μ)) (fun x ↦ x.1 - x.2) measurable_sub
      =ᵐ[μ.map (⟨Z, W⟩)] kernel.map (condDistrib (⟨X, Y⟩) (⟨Z, W⟩) μ)
        (fun x ↦ x.1 - x.2) measurable_sub := by
    filter_upwards [h_indep] with y hy
    conv_rhs => rw [kernel.map_apply, hy]
  rw [kernel.entropy_congr this]
  have : kernel.map (condDistrib (⟨X, Y⟩) (⟨Z, W⟩) μ) (fun x ↦ x.1 - x.2) measurable_sub
      =ᵐ[μ.map (⟨Z, W⟩)] condDistrib (X - Y) (⟨Z, W⟩) μ :=
    (condDistrib_comp (hX.prod_mk hY) (measurable_const.prod_mk hW) _ _).symm
  rw [kernel.entropy_congr this]
  have h_meas : μ.map (⟨Z, W⟩) = (μ.map W).map (Prod.mk ()) := by
    ext s hs
    rw [Measure.map_apply measurable_prod_mk_left hs, h_meas_eq, Measure.prod_apply hs,
      lintegral_dirac]
  have h_ker : condDistrib (X - Y) (⟨Z, W⟩) μ
      =ᵐ[μ.map (⟨Z, W⟩)] kernel.prodMkLeft Unit (condDistrib (X - Y) W μ) := by
    rw [Filter.EventuallyEq, ae_iff_of_countable]
    intro x hx
    rw [Measure.map_apply (measurable_const.prod_mk hW) (measurableSet_singleton _)] at hx
    ext s hs
    have h_preimage_eq : (fun a ↦ (PUnit.unit, W a)) ⁻¹' {x} = W ⁻¹' {x.2} := by
      conv_lhs => rw [← Prod.eta x, ← Set.singleton_prod_singleton, Set.mk_preimage_prod]
      ext1 y
      simp
    rw [kernel.prodMkLeft_apply, condDistrib_apply' _ (measurable_const.prod_mk hW) _ _ hx hs,
      condDistrib_apply' _ hW _ _ _ hs]
    rotate_left
    · exact hX.sub hY
    · convert hx
      exact h_preimage_eq.symm
    · exact hX.sub hY
    congr
  rw [kernel.entropy_congr h_ker, h_meas, kernel.entropy_prodMkLeft_unit]

/-- The conditional Ruzsa distance is unchanged if the sets of random variables are replaced with
copies. -/
lemma condRuzsaDist_of_copy {X : Ω → G} (hX : Measurable X) {Z : Ω → S} (hZ : Measurable Z)
    {Y : Ω' → G} (hY : Measurable Y) {W : Ω' → T} (hW : Measurable W)
    {X' : Ω'' → G} (hX' : Measurable X') {Z' : Ω'' → S} (hZ' : Measurable Z')
    {Y' : Ω''' → G} (hY' : Measurable Y') {W' : Ω''' → T} (hW' : Measurable W')
    [IsFiniteMeasure μ] [IsFiniteMeasure μ'] [IsFiniteMeasure μ''] [IsFiniteMeasure μ''']
    (h1 : IdentDistrib (⟨X, Z⟩) (⟨X', Z'⟩) μ μ'') (h2 : IdentDistrib (⟨Y, W⟩) (⟨Y', W'⟩) μ' μ''')
    [FiniteRange Z] [FiniteRange W] [FiniteRange Z'] [FiniteRange W'] :
    d[X | Z ; μ # Y | W ; μ'] = d[X' | Z' ; μ'' # Y' | W' ; μ''']

:= by
  classical
  set A := (FiniteRange.toFinset Z) ∪ (FiniteRange.toFinset Z')
  set B := (FiniteRange.toFinset W) ∪ (FiniteRange.toFinset W')
  have hfull : Measure.prod (μ.map Z) (μ'.map W) ((A ×ˢ B : Finset (S × T)): Set (S × T))ᶜ = 0 := by
    simp only [A, B]
    apply Measure.prod_of_full_measure_finset
    all_goals {
      rw [Measure.map_apply ‹_›]
      convert measure_empty
      simp [← FiniteRange.range]
      measurability
    }
  have hfull' : Measure.prod (μ''.map Z') (μ'''.map W') ((A ×ˢ B : Finset (S × T)): Set (S × T))ᶜ = 0 := by
    simp only [A, B]
    apply Measure.prod_of_full_measure_finset
    all_goals {
      rw [Measure.map_apply ‹_›]
      convert measure_empty
      simp [← FiniteRange.range]
      measurability
    }
  rw [condRuzsaDist_def, condRuzsaDist_def, kernel.rdist, kernel.rdist,
    integral_eq_sum' _ hfull, integral_eq_sum' _ hfull']
  have hZZ' : μ.map Z = μ''.map Z' := (h1.comp measurable_snd).map_eq
  have hWW' : μ'.map W = μ'''.map W' := (h2.comp measurable_snd).map_eq
  simp_rw [Measure.prod_apply_singleton, ENNReal.toReal_mul, ← hZZ', ← hWW',
    Measure.map_apply hZ (measurableSet_singleton _),
    Measure.map_apply hW (measurableSet_singleton _)]
  congr with x
  by_cases hz : μ (Z ⁻¹' {x.1}) = 0
  · simp only [smul_eq_mul, mul_eq_mul_left_iff, mul_eq_zero]
    refine Or.inr (Or.inl ?_)
    simp [ENNReal.toReal_eq_zero_iff, measure_ne_top, hz]
  by_cases hw : μ' (W ⁻¹' {x.2}) = 0
  · simp only [smul_eq_mul, mul_eq_mul_left_iff, mul_eq_zero]
    refine Or.inr (Or.inr ?_)
    simp [ENNReal.toReal_eq_zero_iff, measure_ne_top, hw]
  congr 2
  · have hZZ'x : μ (Z ⁻¹' {x.1}) = μ'' (Z' ⁻¹' {x.1}) := by
      have : μ.map Z {x.1} = μ''.map Z' {x.1} := by rw [hZZ']
      rwa [Measure.map_apply hZ (measurableSet_singleton _),
        Measure.map_apply hZ' (measurableSet_singleton _)] at this
    ext s hs
    rw [condDistrib_apply' hX hZ _ _ hz hs, condDistrib_apply' hX' hZ' _ _ _ hs]
    swap; · rwa [hZZ'x] at hz
    congr
    have : μ.map (⟨X, Z⟩) (s ×ˢ {x.1}) = μ''.map (⟨X', Z'⟩) (s ×ˢ {x.1}) := by rw [h1.map_eq]
    rwa [Measure.map_apply (hX.prod_mk hZ) (hs.prod (measurableSet_singleton _)),
      Measure.map_apply (hX'.prod_mk hZ') (hs.prod (measurableSet_singleton _)),
      Set.mk_preimage_prod, Set.mk_preimage_prod, Set.inter_comm,
      Set.inter_comm ((fun a ↦ X' a) ⁻¹' s)] at this
  · have hWW'x : μ' (W ⁻¹' {x.2}) = μ''' (W' ⁻¹' {x.2}) := by
      have : μ'.map W {x.2} = μ'''.map W' {x.2} := by rw [hWW']
      rwa [Measure.map_apply hW (measurableSet_singleton _),
        Measure.map_apply hW' (measurableSet_singleton _)] at this
    ext s hs
    rw [condDistrib_apply' hY hW _ _ hw hs, condDistrib_apply' hY' hW' _ _ _ hs]
    swap; · rwa [hWW'x] at hw
    congr
    have : μ'.map (⟨Y, W⟩) (s ×ˢ {x.2}) = μ'''.map (⟨Y', W'⟩) (s ×ˢ {x.2}) := by rw [h2.map_eq]
    rwa [Measure.map_apply (hY.prod_mk hW) (hs.prod (measurableSet_singleton _)),
      Measure.map_apply (hY'.prod_mk hW') (hs.prod (measurableSet_singleton _)),
      Set.mk_preimage_prod, Set.mk_preimage_prod, Set.inter_comm,
      Set.inter_comm ((fun a ↦ Y' a) ⁻¹' s)] at this