Uniform Distribution Proofs

import Mathlib.Probability.Distributions.Uniform
import Mathlib.Data.Fintype.Card

-- leandown
-- [meta]
-- title = "Uniform distribution proofs"
-- group = "Basic proofs"
-- [content]

-- # {{theorem}}: Bijections preserve the uniform distribution
--
-- **given:**
-- - a finite nonempty type $\alpha$
-- - a finite nonempty type $\beta$
-- - $e$ a bijection between $\alpha$ and $\beta$
-- - $A$ a uniform distribution on $\alpha$
-- - $B$ a uniform distribution on $\beta$
--
-- **then:**
-- - applying $e$ to all samples from $A$ gives $B$
-- - for all $a \sim A$ and $b : \beta$, $\Pr[e(a) = b] = \Pr[b \sim B]$
--
/-- The uniform distribution is preserved by any bijection:
    if `e : α ≃ β` then pushing forward
    `uniformOfFintype α` along `e` gives `uniformOfFintype β`.
-/
theorem PMF.uniformOfFintype_map_equiv {α β : Type*} [Fintype α] [Fintype β]
    [Nonempty α] [Nonempty β] (e : α ≃ β) :
    (PMF.uniformOfFintype α).map e = PMF.uniformOfFintype β := by
  -- fix an arbitrary output $b : \beta$; it suffices to show both sides
  -- assign $b$ the same probability
  ext b
  -- unfold the push-forward:
  -- $\Pr[e \text{ maps to } b] = \sum_{a} \Pr[a] \cdot \mathbf{1}[e(a)=b]$
  simp only [PMF.map_apply]
  -- replace each $\Pr[a]$ with $1/|\alpha|$ (uniform distribution)
  simp only [PMF.uniformOfFintype_apply]
  -- convert the infinite sum (tsum) to a finite sum over Finset.univ,
  -- since $\alpha$ is finite
  simp only [tsum_fintype]
  -- our goal is to show this sum:
  --   $\sum_{a} \frac{1}{|\alpha|} \cdot \mathbf{1}[e(a)=b]$
  -- is equal to $\Pr(b) = \frac{1}{|\beta|}$

  -- the sum has exactly one nonzero term (when $e(a) = b$); we show:
  -- * exactly one term is $\frac{1}{|\alpha|}$ (when $e(a)=b$)
  -- * all other terms are $0$ (when $e(a) \neq b$)
  -- and therefore the sum is $\frac{1}{|\alpha|}$
  -- and by $e$ being a bijection it must be true that
  -- $\frac{1}{|\alpha|} = \frac{1}{|\beta|}$
  rw [Finset.sum_eq_single_of_mem (e.symm b) (Finset.mem_univ _)]
  -- in the case where $e(a) = b$:
  -- $a = e^{-1}(b)$ so $e(e^{-1}(b)) = b$ and the indicator is $1$,
  -- and $|\alpha| = |\beta|$ since $e$ is a bijection,
  -- so $\frac{1}{|\alpha|} \cdot 1 = \frac{1}{|\beta|}$
  · simp [e.apply_symm_apply, Fintype.card_congr e]
  -- in the case where $e(a) \neq b$ (i.e. $a \neq e^{-1}(b)$):
  -- since $e$ is injective, $a \neq e^{-1}(b)$ implies $e(a) \neq b$,
  -- so the indicator $\mathbf{1}[e(a)=b] = 0$,
  -- making the whole term $\frac{1}{|\alpha|} \cdot 0 = 0$
  · intro a _ ha
    rw [if_neg (fun heq => ha (by
      have h := congr_arg e.symm heq
      rw [e.symm_apply_apply] at h
      exact h.symm))]