\begin{lemma}[Conditional multidistance chain rule]\label{multidist-chain-rule-cond}\lean{cond_multiDist_chainRule}\leanok
    Let $\pi \colon G \to H$ be a homomorphism of abelian groups.
    Let $I$ be a finite index set and let $X_{[m]}$ be a tuple of $G$-valued random variables.
    Let $Y_{[m]}$ be another tuple of random variables (not necessarily $G$-valued).
    Suppose that the pairs $(X_i, Y_i)$ are jointly independent of one another (but $X_i$ need not be independent of $Y_i$).
    Then
    \begin{align}\nonumber
        D[ X_{[m]} | Y_{[m]} ] &=  D[ X_{[m]} \,|\, \pi(X_{[m]}), Y_{[m]}] + D[ \pi(X_{[m]}) \,|\, Y_{[m]}] \\
         &\quad\qquad + \bbI[ \sum_{i=1}^m X_i : \pi(X_{[m]}) \; \big| \;  \pi\bigl(\sum_{i=1}^m X_i \bigr), Y_{[m]} ].\label{chain-eq-cond}
    \end{align}
  \end{lemma}
\begin{proof}\uses{multidist-chain-rule}\leanok
  For each $y_i$ in the support of $p_{Y_i}$, apply \Cref{multidist-chain-rule} with $X_i$ replaced by the conditioned random variable $(X_i|Y_i=y_i)$, and the claim~\eqref{chain-eq-cond} follows by averaging~\eqref{chain-eq} in the $y_i$ using the weights $p_{Y_i}$.
  \end{proof}
lemma cond_multiDist_chainRule {G H : Type*} [hG : MeasurableSpace G] [MeasurableSingletonClass G]
    [AddCommGroup G] [Fintype G]
    [hH : MeasurableSpace H] [MeasurableSingletonClass H] [AddCommGroup H]
    [Fintype H] (π : G →+ H)
    {S : Type*} [Fintype S] [hS : MeasurableSpace S] [MeasurableSingletonClass S]
    {m : ℕ} {Ω : Type*} [hΩ : MeasureSpace Ω]
    {X : Fin m → Ω → G} (hX : ∀ i, Measurable (X i))
    {Y : Fin m → Ω → S} (hY : ∀ i, Measurable (Y i))
    (h_indep : iIndepFun (fun i ↦ ⟨X i, Y i⟩)) :
    D[X | Y; fun _ ↦ hΩ] = D[X | fun i ↦ ⟨π ∘ X i, Y i⟩; fun _ ↦ hΩ]
      + D[fun i ↦ π ∘ X i | Y; fun _ ↦ hΩ]
      + I[∑ i, X i : fun ω ↦ (fun i ↦ π (X i ω)) |
            ⟨π ∘ (∑ i, X i), fun ω ↦ (fun i ↦ Y i ω)⟩] := by
  have : IsProbabilityMeasure (ℙ : Measure Ω) := h_indep.isProbabilityMeasure
  set E' := fun (y : Fin m → S) ↦ ⋂ i, Y i ⁻¹' {y i}
  set f := fun (y : Fin m → S) ↦ (ℙ (E' y)).toReal
  set hΩc : (Fin m → S) → MeasureSpace Ω := fun y ↦ ⟨cond ℙ (E' y)⟩

  calc
    _ = ∑ y, (f y) * D[X; fun _ ↦ hΩc y] := condMultiDist_eq' hX hY h_indep
    _ = ∑ y, (f y) * D[X | fun i ↦ π ∘ X i; fun _ ↦ hΩc y]
        + ∑ y, (f y) * D[fun i ↦ π ∘ X i; fun _ ↦ hΩc y]
        + ∑ y, (f y) * I[∑ i, X i : fun ω ↦ (fun i ↦ π (X i ω)) |
          π ∘ (∑ i, X i); (hΩc y).volume] := by
      simp_rw [← Finset.sum_add_distrib, ← left_distrib]
      congr with y
      by_cases hf : f y = 0
      . simp only [hf, zero_mul]
      congr 1
      convert multiDist_chainRule π (hΩc y) hX _
      refine h_indep.cond hY ?_ fun _ ↦ .singleton _
      apply prob_nonzero_of_prod_prob_nonzero
      convert hf
      rw [← ENNReal.toReal_prod]
      congr
      exact (iIndepFun.meas_iInter h_indep fun _ ↦ mes_of_comap <| .singleton _).symm
    _ = _ := by
      have hmes : Measurable (π ∘ ∑ i : Fin m, X i) := by
        apply Measurable.comp .of_discrete
        convert Finset.measurable_sum (f := X) Finset.univ _ with ω
        . exact Fintype.sum_apply ω X
        exact (fun i _ ↦ hX i)
      have hpi_indep : iIndepFun (fun i ↦ ⟨π ∘ X i, Y i⟩) ℙ := by
        set g : G × S → H × S := fun p ↦ ⟨π p.1, p.2⟩
        convert iIndepFun.comp h_indep (fun _ ↦ g) _
        intro i
        exact .of_discrete
      have hpi_indep' : iIndepFun (fun i ↦ ⟨X i, ⟨π ∘ X i, Y i⟩⟩) ℙ := by
        set g : G × S → G × (H × S) := fun p ↦ ⟨p.1, ⟨π p.1, p.2⟩⟩
        convert iIndepFun.comp h_indep (fun _ ↦ g) _
        intro i
        exact .of_discrete

      have hey_mes : ∀ y, MeasurableSet (E' y) := by
          intro y
          apply MeasurableSet.iInter
          intro i
          exact MeasurableSet.preimage (MeasurableSet.singleton (y i)) (hY i)

      congr 2
      . rw [condMultiDist_eq' hX _ hpi_indep']
        . rw [← Equiv.sum_comp (Equiv.arrowProdEquivProdArrow _ _ _).symm, Fintype.sum_prod_type, Finset.sum_comm]
          congr with y
          by_cases pey : ℙ (E' y) = 0
          . simp only [pey, ENNReal.zero_toReal, zero_mul, f]
            apply (Finset.sum_eq_zero _).symm
            intro s _
            convert zero_mul _
            convert ENNReal.zero_toReal
            apply measure_mono_null _ pey
            intro ω hω
            simp [E', Equiv.arrowProdEquivProdArrow] at hω ⊢
            intro i
            exact (hω i).2
          rw [condMultiDist_eq' (hΩ := hΩc y) hX, Finset.mul_sum]
          . congr with s
            dsimp [f, E', Equiv.arrowProdEquivProdArrow]
            rw [← mul_assoc, ← ENNReal.toReal_mul]
            congr 2
            . rw [mul_comm]
              convert ProbabilityTheory.cond_mul_eq_inter (hey_mes y) ?_ _
              . rw [← Set.iInter_inter_distrib]
                apply Set.iInter_congr
                intro i
                ext ω
                simp only [Set.mem_preimage, Set.mem_singleton_iff, Prod.mk.injEq, comp_apply, Set.mem_inter_iff]
                exact And.comm
              infer_instance
            funext _
            congr 1
            dsimp [hΩc, E']
            rw [ProbabilityTheory.cond_cond_eq_cond_inter (hey_mes y), ← Set.iInter_inter_distrib]
            . congr 1
              apply Set.iInter_congr
              intro i
              ext ω
              simp only [Set.mem_inter_iff, Set.mem_preimage, Set.mem_singleton_iff, comp_apply, Prod.mk.injEq]
              exact And.comm
            apply MeasurableSet.iInter
            intro i
            apply MeasurableSet.preimage (MeasurableSet.singleton _)
            exact Measurable.comp .of_discrete (hX i)
          . intro i
            exact Measurable.comp .of_discrete (hX i)
          set g : G → G × H := fun x ↦ ⟨x, π x⟩
          refine iIndepFun.comp ?_ (fun _ ↦ g) fun _ ↦ .of_discrete
          . refine h_indep.cond hY ?_ fun _ ↦ .singleton _
            rw [iIndepFun.meas_iInter h_indep fun _ ↦ mes_of_comap <| .singleton _] at pey
            contrapose! pey
            obtain ⟨i, hi⟩ := pey
            exact Finset.prod_eq_zero (Finset.mem_univ i) hi
        intro i
        exact Measurable.prodMk (Measurable.comp .of_discrete (hX i)) (hY i)
      . rw [condMultiDist_eq' _ hY hpi_indep]
        intro i
        apply Measurable.comp .of_discrete (hX i)
      rw [condMutualInfo_eq_sum', Fintype.sum_prod_type, Finset.sum_comm]
      . congr with y
        by_cases pey : ℙ (E' y) = 0
        . simp only [pey, ENNReal.zero_toReal, zero_mul, f]
          apply (Finset.sum_eq_zero _).symm
          intro s _
          convert zero_mul _
          convert ENNReal.zero_toReal
          apply measure_mono_null _ pey
          intro ω hω
          simp [E'] at hω ⊢
          rw [← hω.2]
          simp only [implies_true]
        have : IsProbabilityMeasure (hΩc y).volume := cond_isProbabilityMeasure pey
        rw [condMutualInfo_eq_sum' hmes, Finset.mul_sum]
        congr with x
        dsimp [f, E']
        rw [← mul_assoc, ← ENNReal.toReal_mul]
        congr 2
        . rw [mul_comm]
          convert ProbabilityTheory.cond_mul_eq_inter (hey_mes y) ?_ _
          . ext ω
            simp only [Set.mem_preimage, Set.mem_singleton_iff, Prod.mk.injEq, comp_apply,
              Finset.sum_apply, _root_.map_sum, Set.mem_inter_iff, Set.mem_iInter, E']
            rw [and_comm]
            apply and_congr_left
            intro _
            exact funext_iff
          infer_instance
        dsimp [hΩc, E']
        rw [ProbabilityTheory.cond_cond_eq_cond_inter (hey_mes y)]
        . congr
          ext ω
          simp only [Set.mem_inter_iff, Set.mem_iInter, Set.mem_preimage, Set.mem_singleton_iff,
            comp_apply, Finset.sum_apply, _root_.map_sum, Prod.mk.injEq, E']
          rw [and_comm]
          apply and_congr_right
          intro _
          exact Iff.symm funext_iff
        exact MeasurableSet.preimage (MeasurableSet.singleton x) hmes
      exact Measurable.prodMk hmes (measurable_pi_lambda (fun ω i ↦ Y i ω) hY)