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]


import PFR.ForMathlib.Entropy.Kernel.Group
import PFR.Mathlib.MeasureTheory.Measure.MeasureSpace

/-!
# Ruzsa distance between kernels

## Main definitions

*

## Notations

* `dk[κ ; μ # η ; ν] = `

-/

open Real MeasureTheory

open scoped ENNReal NNReal Topology ProbabilityTheory BigOperators

namespace ProbabilityTheory.kernel

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

/-- The Rusza distance between two measures, defined as `H[X - Y] - H[X]/2 - H[Y]/2` where `X`
and `Y` are independent variables distributed according to the two measures. -/
noncomputable
def rdistm (μ : Measure G) (ν : Measure G) : ℝ :=
    Hm[(μ.prod ν).map (fun x ↦ x.1 - x.2)] - Hm[μ]/2 - Hm[ν]/2

/-- The Rusza distance between two kernels taking values in the same space, defined as the average
Rusza distance between the image measures. -/
noncomputable
def rdist (κ : kernel T G) (η : kernel T' G) (μ : Measure T) (ν : Measure T') : ℝ :=
  (μ.prod ν)[fun p ↦ rdistm (κ p.1) (η p.2)]

/-- The Rusza distance between two kernels taking values in the same space, defined as the average
Rusza distance between the image measures. -/
notation3:max "dk[" κ " ; " μ " # " η " ; " μ' "]" => rdist κ η μ μ'

lemma rdist_eq {κ : kernel T G} {η : kernel T' G} {μ : Measure T} {ν : Measure T'}
    [IsProbabilityMeasure μ] [IsProbabilityMeasure ν]
    [FiniteSupport μ] [FiniteSupport ν] :
    dk[κ ; μ # η ; ν] = (μ.prod ν)[fun p ↦ Hm[((κ p.1).prod (η p.2)).map (fun x ↦ x.1 - x.2)]]
      - Hk[κ, μ]/2 - Hk[η, ν]/2 := by
  rw [rdist]
  simp_rw [rdistm]
  rw [integral_sub, integral_sub]
  · simp_rw [div_eq_mul_inv, integral_mul_right, integral_prod _ (integrable_of_finiteSupport _)]
    simp [entropy]
  all_goals { exact integrable_of_finiteSupport _ }

lemma rdist_eq' {κ : kernel T G} {η : kernel T' G} [IsFiniteKernel κ] [IsFiniteKernel η]
    {μ : Measure T} {ν : Measure T'} [IsProbabilityMeasure μ] [IsProbabilityMeasure ν]
    [FiniteSupport μ] [FiniteSupport ν] :
    dk[κ ; μ # η ; ν] =
      Hk[map ((prodMkRight T' κ) ×ₖ (prodMkLeft T η)) (fun x ↦ x.1 - x.2) measurable_sub, μ.prod ν]
      - Hk[κ, μ]/2 - Hk[η, ν]/2 := by
  rw [rdist_eq]
  congr with p
  simp only
  rw [map_apply, prod_apply, prodMkLeft_apply, prodMkRight_apply]

lemma rdist_symm {κ : kernel T G} {η : kernel T' G} [IsFiniteKernel κ] [IsFiniteKernel η]
    {μ : Measure T} {ν : Measure T'} [IsProbabilityMeasure μ] [IsProbabilityMeasure ν]
    [FiniteSupport μ] [FiniteSupport ν] :
    dk[κ ; μ # η ; ν] = dk[η ; ν # κ ; μ] := by
  rw [rdist_eq', rdist_eq', sub_sub, sub_sub, add_comm]
  congr 1
  rw [← entropy_comap_swap, comap_map_comm, entropy_sub_comm, Measure.comap_swap, Measure.prod_swap,
    comap_prod_swap, map_map]
  congr

@[simp] lemma rdist_zero_right (κ : kernel T G) (η : kernel T' G) (μ : Measure T) :
    dk[κ ; μ # η ; 0] = 0 := by
  simp [rdist]

@[simp] lemma rdist_zero_left (κ : kernel T G) (η : kernel T' G) (ν' : Measure T') :
    dk[κ ; 0 # η ; ν'] = 0 := by
  simp [rdist]

@[simp] lemma rdist_zero_kernel_right {κ : kernel T G} [IsFiniteKernel κ]
    {μ : Measure T} {ν : Measure T'} [IsProbabilityMeasure μ] [IsProbabilityMeasure ν]
    [FiniteSupport μ] [FiniteSupport ν] :
    dk[κ ; μ # 0 ; ν] = - Hk[κ, μ] / 2 := by
  rw [rdist_eq']
  simp only [prodMkLeft_zero, entropy_zero_kernel, zero_div, sub_zero]
  rw [sub_eq_iff_eq_add]
  ring_nf
  have : map (prodMkRight T' κ ×ₖ 0) (fun x ↦ x.1 - x.2) measurable_sub
      = 0 := by
    ext1 x
    rw [map_apply, prod_apply]
    simp
  rw [this, entropy_zero_kernel]

@[simp] lemma rdist_zero_kernel_left {η : kernel T' G} [IsFiniteKernel η]
    {μ : Measure T} {ν : Measure T'} [IsProbabilityMeasure μ] [IsProbabilityMeasure ν]
    [FiniteSupport μ] [FiniteSupport ν] :
    dk[0 ; μ # η ; ν] = - Hk[η, ν] / 2 := by
  rw [rdist_symm, rdist_zero_kernel_right]

@[simp] lemma rdist_dirac_zero_right {κ : kernel T G} [IsFiniteKernel κ]
    {μ : Measure T} {ν : Measure T'} [IsProbabilityMeasure μ] [IsProbabilityMeasure ν]
    [FiniteSupport μ] [FiniteSupport ν] :
    dk[κ ; μ # const T' (Measure.dirac 0) ; ν] = Hk[κ, μ] / 2 := by
  rw [rdist_eq']
  simp only [entropy_const, measure_univ, ENNReal.one_toReal, measureEntropy_dirac, mul_zero,
    zero_div, sub_zero]
  rw [sub_eq_iff_eq_add]
  ring_nf
  have : map (prodMkRight T' κ ×ₖ prodMkLeft T (const T' (Measure.dirac 0)))
        (fun x ↦ x.1 - x.2) measurable_sub
      = prodMkRight T' κ := by
    ext x s hs
    rw [prodMkRight_apply, map_apply, prod_apply, prodMkLeft_apply, const_apply,
      prodMkRight_apply, Measure.map_apply measurable_sub hs,
      Measure.prod_apply (measurable_sub hs)]
    simp only [Set.mem_preimage, sub_zero, Measure.dirac_apply]
    have : ∀ x : G, Prod.mk x ⁻¹' ((fun p : G × G ↦ p.1 - p.2) ⁻¹' s) = {y | x - y ∈ s} := by
      intro x
      ext y
      simp
    simp_rw [this]
    have : ∀ x : G, Set.indicator {y | x - y ∈ s} (1 : G → ℝ≥0∞) 0 = s.indicator (fun _ ↦ 1) x := by
      intro x
      by_cases hx : x ∈ s <;> simp [hx]
    simp_rw [this]
    rw [lintegral_indicator_const hs, one_mul]
  rw [this, entropy_prodMkRight']

@[simp] lemma rdist_dirac_zero_left {η : kernel T' G} [IsFiniteKernel η]
    {μ : Measure T} {ν : Measure T'} [IsProbabilityMeasure μ] [IsProbabilityMeasure ν]
    [FiniteSupport μ] [FiniteSupport ν] :
    dk[const T (Measure.dirac 0) ; μ # η ; ν] = Hk[η, ν] / 2 := by
  rw [rdist_symm, rdist_dirac_zero_right]

lemma ruzsa_triangle_aux (κ : kernel T (G × G)) (η : kernel T G)
    [IsMarkovKernel κ] [IsMarkovKernel η] :
    map (κ ×ₖ η) (fun p ↦ p.2 - p.1.2) (measurable_of_countable _)
      = map (η ×ₖ snd κ) (fun p ↦ p.1 - p.2) (measurable_of_countable _) := by
  have : (fun p : G × G ↦ p.1 - p.2) = (fun p ↦ p.2 - p.1) ∘ Prod.swap := by ext1 p; simp
  rw [this, ← map_map _ measurable_swap (measurable_of_countable _), map_prod_swap]
  ext x s hs
  rw [map_apply' _ _ _ hs, map_apply' _ _ _ hs, prod_apply' _ _ _ (measurable_of_countable _ hs),
    prod_apply' _ _ _ (measurable_of_countable _ hs), lintegral_snd _ _ (measurable_of_countable _)]
  congr

lemma abs_sub_entropy_le_rdist {κ : kernel T G} {η : kernel T' G}
    [IsMarkovKernel κ] [IsMarkovKernel η]
    {μ : Measure T} {ν : Measure T'} [IsProbabilityMeasure μ] [IsProbabilityMeasure ν]
    [FiniteSupport μ] [FiniteSupport ν]
    (hκ : AEFiniteKernelSupport κ μ) (hη : AEFiniteKernelSupport η ν) :
    |Hk[κ, μ] - Hk[η, ν]| ≤ 2 * dk[κ ; μ # η ; ν] := by
  have h := max_entropy_le_entropy_sub_prod (prodMkRight T' κ) (prodMkLeft T η) (μ.prod ν)
    (hκ.prodMkRight ν) (hη.prodMkLeft μ)
  rw [entropy_prodMkRight', entropy_prodMkLeft] at h
  rw [rdist_eq', abs_le]
  constructor
  · linarith [le_max_right (Hk[κ, μ]) (Hk[η, ν])]
  · linarith [le_max_left (Hk[κ, μ]) (Hk[η, ν])]

lemma rdist_nonneg {κ : kernel T G} {η : kernel T' G}
    [IsMarkovKernel κ] [IsMarkovKernel η]
    {μ : Measure T} {ν : Measure T'} [IsProbabilityMeasure μ] [IsProbabilityMeasure ν]
    [FiniteSupport μ] [FiniteSupport ν]
    (hκ : AEFiniteKernelSupport κ μ) (hη : AEFiniteKernelSupport η ν) :
    0 ≤ dk[κ ; μ # η ; ν] := by
  suffices 0 ≤ 2 * dk[κ ; μ # η ; ν] by linarith
  calc 0 ≤ |Hk[κ, μ] - Hk[η, ν]| := abs_nonneg _
  _ ≤ 2 * dk[κ ; μ # η ; ν] := abs_sub_entropy_le_rdist hκ hη

-- Kernel equivalent of `H[X - Y; μ] ≤ H[X - Z; μ] + H[Z - Y; μ] - H[Z; μ]`
-- `κ` is `⟨X,Y⟩`, `η` is `Z`. Independence is expressed through the product `×ₖ`.

/-- The **improved entropic Ruzsa triangle inequality**. -/
lemma ent_of_diff_le (κ : kernel T (G × G)) (η : kernel T G) [IsMarkovKernel κ] [IsMarkovKernel η]
    (μ : Measure T) [IsProbabilityMeasure μ] [FiniteSupport μ]
    (hκ : FiniteKernelSupport κ) (hη : FiniteKernelSupport η) :
    Hk[map κ (fun p : G × G ↦ p.1 - p.2) measurable_sub, μ]
      ≤ Hk[map ((fst κ) ×ₖ η) (fun p : G × G ↦ p.1 - p.2) measurable_sub, μ]
        + Hk[map (η ×ₖ (snd κ)) (fun p : G × G ↦ p.1 - p.2) measurable_sub, μ]
        - Hk[η, μ]

:= by
  have hκη := kernel.finiteKernelSupport_of_prod hκ hη
  have h1 : Hk[map (κ ×ₖ η) (fun p ↦ (p.1.1 - p.2, (p.1.2, p.1.1 - p.1.2)))
        (measurable_of_countable _), μ] + Hk[map κ (fun p ↦ p.1 - p.2) measurable_sub, μ]
      ≤ Hk[map (κ ×ₖ η) (fun p ↦ (p.1.1 - p.2, p.1.1 - p.1.2)) (measurable_of_countable _), μ]
        + Hk[map κ (fun p ↦ (p.2, p.1 - p.2)) (measurable_of_countable _), μ] := by
    have h := entropy_triple_add_entropy_le
      (map (κ ×ₖ η) (fun p ↦ (p.1.1 - p.2, (p.1.2, p.1.1 - p.1.2))) (measurable_of_countable _)) μ
    simp only [snd_map_prod _ (measurable_of_countable _) (measurable_of_countable _)] at h
    rw [deleteMiddle_map_prod _ (measurable_of_countable _) (measurable_of_countable _)
        (measurable_of_countable _)] at h
    have : map (κ ×ₖ η) (fun x ↦ x.1.1 - x.1.2) (measurable_of_countable _)
        = map κ (fun p ↦ p.1 - p.2) measurable_sub := by
      have : (fun x : (G × G) × G ↦ x.1.1 - x.1.2)
        = (fun x ↦ x.1 - x.2) ∘ Prod.fst := by ext1 y; simp
      rw [this, ← map_map, ← kernel.fst, fst_prod]
    rw [this] at h
    refine (h ?_).trans_eq ?_
    . apply FiniteKernelSupport.aefiniteKernelSupport
      apply kernel.finiteKernelSupport_of_map
      exact kernel.finiteKernelSupport_of_prod hκ hη
    congr 2
    have : (fun x : (G × G) × G ↦ (x.1.2, x.1.1 - x.1.2))
      = (fun x ↦ (x.2, x.1 - x.2)) ∘ Prod.fst := by ext1 y; simp
    rw [this, ← map_map, ← kernel.fst, fst_prod]
  have h2 : Hk[map (κ ×ₖ η) (fun p ↦ (p.1.1 - p.2, p.1.1 - p.1.2)) (measurable_of_countable _), μ]
      ≤ Hk[map (κ ×ₖ η) (fun p ↦ p.1.1 - p.2) (measurable_of_countable _), μ]
        + Hk[map (κ ×ₖ η) (fun p ↦ p.1.2 - p.2) (measurable_of_countable _), μ] := by
    calc Hk[map (κ ×ₖ η) (fun p ↦ (p.1.1 - p.2, p.1.1 - p.1.2)) (measurable_of_countable _), μ]
      ≤ Hk[map (κ ×ₖ η) (fun p ↦ (p.1.1 - p.2, p.1.2 - p.2)) (measurable_of_countable _), μ] := by
          have : (fun p : (G × G) × G ↦ (p.1.1 - p.2, p.1.1 - p.1.2))
            = (fun p ↦ (p.1, p.1 - p.2)) ∘ (fun p ↦ (p.1.1 - p.2, p.1.2 - p.2)) := by ext1; simp
          rw [this, ← map_map]
          apply entropy_map_le _ _
          apply FiniteKernelSupport.aefiniteKernelSupport
          apply kernel.finiteKernelSupport_of_map hκη
    _ ≤ Hk[map (κ ×ₖ η) (fun p ↦ p.1.1 - p.2) (measurable_of_countable _), μ]
        + Hk[map (κ ×ₖ η) (fun p ↦ p.1.2 - p.2) (measurable_of_countable _), μ] := by
          have h : 0 ≤ Hk[map (κ ×ₖ η) (fun p ↦ p.1.1 - p.2) (measurable_of_countable _), μ]
              + Hk[map (κ ×ₖ η) (fun p ↦ p.1.2 - p.2) (measurable_of_countable _), μ]
              - Hk[map (κ ×ₖ η) (fun p ↦ (p.1.1 - p.2, p.1.2 - p.2))
                (measurable_of_countable _), μ] := by
            have h' := mutualInfo_nonneg (κ := map (κ ×ₖ η) (fun p ↦ (p.1.1 - p.2, p.1.2 - p.2))
                (measurable_of_countable _)) (μ := μ) ?_
            rwa [mutualInfo, fst_map_prod _ (measurable_of_countable _) (measurable_of_countable _),
              snd_map_prod _ (measurable_of_countable _) (measurable_of_countable _)] at h'
            apply FiniteKernelSupport.aefiniteKernelSupport
            apply kernel.finiteKernelSupport_of_map hκη
          linarith
  have h3 : Hk[map κ (fun p : G × G ↦ (p.2, p.1 - p.2)) (measurable_of_countable _), μ]
      ≤ Hk[κ, μ] := by
    exact entropy_map_le _ (hκ.aefiniteKernelSupport _)
  have h4 : Hk[map (κ ×ₖ η) (fun p ↦ (p.1.1 - p.2, (p.1.2, p.1.1 - p.1.2)))
      (measurable_of_countable _), μ] = Hk[κ ×ₖ η, μ] := by
    refine entropy_of_map_eq_of_map
      (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.1 - p.2, (p.1.2, p.1.1 - p.1.2))) ?_ ?_ ?_
        (hκη.aefiniteKernelSupport _)
    · rw [map_map]
      suffices ((fun p : G × G × G ↦ ((p.2.2 + p.2.1, p.2.1), -p.1 + p.2.2 + p.2.1))
          ∘ fun p ↦ (p.1.1 - p.2, p.1.2, p.1.1 - p.1.2)) = id by
        simp_rw [this, map_id]
      ext1 p
      simp
    · rfl
    apply FiniteKernelSupport.aefiniteKernelSupport
    apply kernel.finiteKernelSupport_of_map hκη
  have h5 : Hk[κ ×ₖ η, μ] = Hk[κ, μ] + Hk[η, μ] := by
    rw [entropy_prod (hκ.aefiniteKernelSupport _) (hη.aefiniteKernelSupport _)]
  rw [h4, h5] at h1
  calc Hk[map κ (fun p : G × G ↦ p.1 - p.2) measurable_sub, μ]
    ≤ Hk[map (κ ×ₖ η) (fun p ↦ p.1.1 - p.2) (measurable_of_countable _), μ]
      + Hk[map (κ ×ₖ η) (fun p ↦ p.1.2 - p.2) (measurable_of_countable _), μ] - Hk[η, μ] := by
        linarith
  _ = Hk[map (κ ×ₖ η) (fun p ↦ p.1.1 - p.2) (measurable_of_countable _), μ]
      + Hk[map (κ ×ₖ η) (fun p ↦ p.2 - p.1.2) (measurable_of_countable _), μ] - Hk[η, μ] := by
        congr 2
        rw [← entropy_neg, map_map]
        congr with p
        simp
  _ = Hk[map ((fst κ) ×ₖ η) (fun p : G × G ↦ p.1 - p.2) measurable_sub, μ]
      + Hk[map (η ×ₖ (snd κ)) (fun p : G × G ↦ p.1 - p.2) measurable_sub, μ]
      - Hk[η, μ] := by
        congr 3
        · ext x s hs
          rw [map_apply' _ _ _ hs, map_apply' _ _ _ hs, prod_apply', prod_apply', lintegral_fst]
          · congr with x
          · exact measurable_of_countable _
          · exact measurable_sub hs
          · exact measurable_of_countable _ hs
        · exact ruzsa_triangle_aux κ η


import PFR.Main

/-!
# Improved PFR

An improvement to PFR that lowers the exponent from 12 to 11.

## Main results

*
-/

open MeasureTheory ProbabilityTheory

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

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

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

variable (Y : Ω₀ → G) (hY : Measurable Y)

variable (Z₁ Z₂ Z₃ Z₄ : Ω → G)
  (hZ₁ : Measurable Z₁) (hZ₂ : Measurable Z₂) (hZ₃ : Measurable Z₃) (hZ₄ : Measurable Z₄)

variable (h_indep : iIndepFun (fun _i => hG) ![Z₁, Z₂, Z₃, Z₄])

local notation3 "Sum" => Z₁ + Z₂ + Z₃ + Z₄

open scoped BigOperators

lemma gen_ineq_aux1 :
    d[Y # Z₁ + Z₂ | ⟨Z₁ + Z₃, Sum⟩] ≤ d[Y # Z₁]
      + (d[Z₁ # Z₂] + d[Z₁ # Z₃] + d[Z₂ # Z₄] - d[Z₁ | Z₁ + Z₂ # Z₃ | Z₃ + Z₄]) / 2
      + (H[Z₁ + Z₂] - H[Z₃ + Z₄] + H[Z₂] - H[Z₁]) / 4

:= by
  have hS : Measurable Sum := ((hZ₁.add' hZ₂).add' hZ₃).add' hZ₄
  have C : d[Z₁ # Z₃] + d[Z₂ # Z₄] = d[Z₁ + Z₂ # Z₃ + Z₄]
           + d[Z₁|Z₁ + Z₂ # Z₃|Z₃ + Z₄] + I[Z₁ + Z₂ : Z₁ + Z₃ | Sum] := by
    have M : d[Z₃ # Z₁] + d[Z₄ # Z₂] = d[Z₃ + Z₄ # Z₁ + Z₂]
           + d[Z₃|Z₃ + Z₄ # Z₁|Z₁ + Z₂] + I[Z₃ + Z₁ : Z₁ + Z₂ | Z₃ + Z₁ + Z₄ + Z₂] := by
      apply sum_of_rdist_eq_char_2 ![Z₃, Z₁, Z₄, Z₂] h_indep.reindex_four_cadb (fun i ↦ ?_)
      fin_cases i <;> assumption
    have J1 : Z₃ + Z₁ + Z₄ + Z₂ = Z₁ + Z₂ + Z₃ + Z₄ := by abel
    have J2 : Z₃ + Z₁ = Z₁ + Z₃ := by abel
    simp_rw [J1, J2] at M
    simpa only [rdist_symm (Y := Z₁), rdist_symm (X := Z₄), rdist_symm (X := Z₃ + Z₄),
      condRuzsaDist_symm (hZ₃.add' hZ₄) (hZ₁.add' hZ₂),
      condMutualInfo_comm (hZ₁.add' hZ₃) (hZ₁.add' hZ₂)] using M
  calc
  d[Y # Z₁ + Z₂ | ⟨Z₁ + Z₃, Sum⟩]
    ≤ d[Y # Z₁ + Z₂ | Sum] + I[Z₁ + Z₂ : Z₁ + Z₃ | Sum]/2 :=
      condRuzsaDist_le'_prod (ℙ : Measure Ω₀) (ℙ : Measure Ω) hY (hZ₁.add hZ₂) (hZ₁.add hZ₃) hS
  _ ≤ d[Y # Z₁ + Z₂] + (d[Z₁ + Z₂ # Z₃ + Z₄] + I[Z₁ + Z₂ : Z₁ + Z₃ | Sum]) / 2
        + (H[Z₁ + Z₂] - H[Z₃ + Z₄]) / 4 := by
      have I : IndepFun (Z₁ + Z₂) (Z₃ + Z₄) := by
        exact h_indep.indepFun_add_add (ι := Fin 4) (by intro i; fin_cases i <;> assumption)
          0 1 2 3 (by decide) (by decide) (by decide) (by decide)
      have J1 : (fun a ↦ Z₁ a + Z₂ a) + (fun a ↦ Z₃ a + Z₄ a) = Sum := by ext; simp; abel
      have J2 : (fun a ↦ Z₁ a + Z₂ a) = Z₁ + Z₂ := rfl
      have J3 : (fun a ↦ Z₃ a + Z₄ a) = Z₃ + Z₄ := rfl
      have A := condRuzsaDist_diff_le''' (ℙ : Measure Ω₀) (μ' := (ℙ : Measure Ω)) hY (hZ₁.add hZ₂)
        (hZ₃.add hZ₄) I
      rw [J1, J2, J3] at A
      linarith
  _ ≤ d[Y # Z₁] + (d[Z₁ # Z₂] + d[Z₁ + Z₂ # Z₃ + Z₄] + I[Z₁ + Z₂ : Z₁ + Z₃ | Sum]) / 2
        + (H[Z₁ + Z₂] - H[Z₃ + Z₄] + H[Z₂] - H[Z₁]) / 4 := by
      have I : IndepFun Z₁ Z₂ := by exact h_indep.indepFun (show 0 ≠ 1 by decide)
      have A := condRuzsaDist_diff_le' (ℙ : Measure Ω₀) (μ' := (ℙ : Measure Ω)) hY hZ₁ hZ₂ I
      linarith
  _ = _ := by linarith


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


import PFR.FirstEstimate
import PFR.Mathlib.Probability.Independence.FourVariables

/-!
# Second estimate

The second estimate 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$
* $I_1 := I_1 [X_1+X_2 : \tilde X_1 + X_2 | X_1+X_2+\tilde X_1+\tilde X_2]$
* $I_2 := I[X_1+X_2 : X_1 + \tilde X_1 | X_1+X_2+\tilde X_1+\tilde X_2]$

## Main results

* `second_estimate` : $$ I_2 \leq 2 \eta k + \frac{2 \eta (2 \eta k - I_1)}{1 - \eta}.$$
-/

open MeasureTheory ProbabilityTheory ElementaryAddCommGroup

variable {G : Type*} [addgroup: 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*} [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₂]

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

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

:= by
  refine' LE.le.ge (LE.le.trans _ (distance_ge_of_min p h_min (hX₁.add hX₁') (hX₂.add hX₂')))
  rw [sub_sub, sub_le_sub_iff_left k, ← mul_add,mul_div_assoc]
  refine' mul_le_mul_of_nonneg_left _ (by linarith [p.hη])
  have h₁' := condRuzsaDist_diff_le' ℙ p.hmeas1 hX₁ hX₁' (h_indep.indepFun (show 0 ≠ 2 by decide))
  have h₂' := condRuzsaDist_diff_le' ℙ p.hmeas2 hX₂ hX₂' (h_indep.indepFun (show 1 ≠ 3 by decide))
  rw [h₁.entropy_eq, add_sub_cancel_right, ← (IdentDistrib.refl hX₁.aemeasurable).rdist_eq h₁] at h₁'
  rw [h₂.entropy_eq, add_sub_cancel_right, ← (IdentDistrib.refl hX₂.aemeasurable).rdist_eq h₂] at h₂'
  linarith