/-
Copyright (c) 2026 Shing Hin Ho. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Shing Hin Ho, Edwin Fernando
-/
/-
This file is obtained from another project which uses an Apache 2.0 license. The exact link is below:
https://github.com/Verified-zkEVM/iris-lean/blob/822420f0e630774f577643e30e0f6be2deb7e372/src/Bluebell/ProbabilityTheory/MeasureOnSpace.lean
The first half of above file has been copied. It has since been golfed.
-/

import Mathlib.Probability.Independence.Conditional
import Mathlib.Data.Set.Basic

/-! ## Independent product of probability measures

The common foundation shared by Lilac and Bluebell (the project from which this file was adapted)
-/

open ProbabilityTheory MeasureTheory

alias MeasureTheory.MeasureSpace.σAlg := MeasureSpace.toMeasurableSpace
alias MeasureTheory.MeasureSpace.μ := MeasureSpace.volume

variable {Ω Ω' : Type*}
         {α β : Type*}

noncomputable section

namespace MeasureTheory
universe u

abbrev Measure.map! (ms₁ : MeasurableSpace α) (ms₂ : MeasurableSpace β)
                    (f : α → β) (μ : Measure α) : Measure β := @μ.map _ _ ms₁ ms₂ f

abbrev Measure.cast {β : Type u} {ms₁ : MeasurableSpace β}
  (μ : Measure β) (ms₂ : MeasurableSpace β) := μ.map! ms₁ ms₂ id

@[simp]
lemma Measure.cast_eq_self {β : Type u} {ms₁ : MeasurableSpace β}
  (μ : Measure β) (ms₂ : MeasurableSpace β) : μ.cast ms₂ = @μ.map _ _ ms₁ ms₂ id := rfl

/--
An induction principle with respect to the trivial σ-algebra
to show that a property P u holds where u is a measurable set
with respect to the trivial σ-algebra, it suffices to show that
P ∅ and P Ω
-/
lemma MeasurableSet.measurableSet_bot_induction
  {P : Set Ω → Prop}
  (h_base : P ∅) (h_ind : P Set.univ)
  : (u : Set Ω) → MeasurableSet[⊥] u → P u := by
  intro u h_u
  rcases MeasurableSpace.measurableSet_bot_iff.1 h_u with rfl | rfl <;> assumption

end MeasureTheory

@[simp, reducible]
def MeasurableSpace.sum (m₁ : MeasurableSpace Ω) (m₂ : MeasurableSpace Ω) : MeasurableSpace Ω :=
  MeasurableSpace.generateFrom (MeasurableSet[m₁] ∪ MeasurableSet[m₂])

lemma sum_eq_sup {Ω : Type*} (m₁ m₂ : MeasurableSpace Ω) :
    MeasurableSpace.sum m₁ m₂ = m₁ ⊔ m₂ := by
  apply le_antisymm
  · apply MeasurableSpace.generateFrom_le
    rintro s (h | h)
    · exact @le_sup_left _ _ m₁ m₂ s h
    · exact @le_sup_right _ _ m₁ m₂ s h
  · apply sup_le <;> intro s hs <;>
    exact MeasurableSpace.measurableSet_generateFrom
      (by first | left; exact hs | right; exact hs)

def MeasurableSpace.cast
  (m₁ m₂ : MeasurableSpace Ω) (p : m₁ = m₂) (E : Set Ω) (h : MeasurableSet[m₁] E) :
  MeasurableSet[m₂] E := by
  subst p
  simp_all only

end

noncomputable section

section GeneratorMembership

variable {Ω : Type*} (m₁ m₂ : MeasurableSpace Ω)

def MeasurableSpace.sumGenerator (m₁ m₂ : MeasurableSpace Ω) : Set (Set Ω) :=
  {S : Set Ω | ∃ F G, S = F ∩ G ∧ MeasurableSet[m₁] F ∧ MeasurableSet[m₂] G}

variable
  {u v w : Set Ω}

/-- Every set in the generating set `w` is of the form `u ∩ v` -/
lemma exists_inter_measurableSet_of_mem_sumGenerator
  (hw : w ∈ MeasurableSpace.sumGenerator m₁ m₂)
  : ∃ u v, w = u ∩ v ∧ MeasurableSet[m₁] u ∧ MeasurableSet[m₂] v := by
  rcases hw with ⟨u, v, rfl, hu, hv⟩
  exact ⟨u, v, rfl, hu, hv⟩

@[aesop 50% apply]
lemma mem_sumGenerator_l (hu : MeasurableSet[m₁] u) :
  u ∈ MeasurableSpace.sumGenerator m₁ m₂ :=
  ⟨u, ⊤, (Set.inter_univ u).symm, hu, MeasurableSet.univ⟩

@[aesop 50% apply]
lemma mem_sumGenerator_r (hu : MeasurableSet[m₂] u) :
  u ∈ MeasurableSpace.sumGenerator m₁ m₂ :=
  ⟨⊤, u, (Set.univ_inter u).symm, MeasurableSet.univ, hu⟩

lemma inter_mem_sumGenerator
  (hu : MeasurableSet[m₁] u) (hv : MeasurableSet[m₂] v) :
  u ∩ v ∈ MeasurableSpace.sumGenerator m₁ m₂ :=
  ⟨u, v, rfl, hu, hv⟩

theorem MeasurableSpace.generateFrom_sumGenerator_eq_sum :
  MeasurableSpace.generateFrom (MeasurableSpace.sumGenerator m₁ m₂)
    = MeasurableSpace.sum m₁ m₂ := by
  ext s
  refine ⟨?p, by apply MeasurableSpace.generateFrom_mono (fun _ _ ↦ by aesop)⟩
  let sumSp := MeasurableSet[m₁.sum m₂]
  apply MeasurableSpace.generateFrom_le
  rintro t ⟨u, ⟨v, _, _, _⟩⟩
  have h₂ : MeasurableSpace.sumGenerator m₁ m₂ ⊆ sumSp := fun w hw ↦ by
    obtain ⟨u, v, rfl, hu, hv⟩ :=
      @exists_inter_measurableSet_of_mem_sumGenerator Ω m₁ m₂ (w := w) hw
    unfold sumSp
    have h₁ : MeasurableSet[m₁.sum m₂] u := by
      have : u ∈ MeasurableSpace.sumGenerator m₁ m₂ := by aesop
      apply MeasurableSpace.measurableSet_generateFrom
      aesop
    have h₂ : MeasurableSet[m₁.sum m₂] v := by
      have : v ∈ MeasurableSpace.sumGenerator m₁ m₂ := by aesop
      apply MeasurableSpace.measurableSet_generateFrom
      aesop
    apply MeasurableSet.inter h₁ h₂
  aesop (add simp (show u ∩ v ∈ sumGenerator m₁ m₂ by use u, v))

end GeneratorMembership

section Sum

variable {Ω : Type*} (m m₁ m₂ : MeasurableSpace Ω)

@[simp, reducible]
def MeasurableSpace.sumUnit : MeasurableSpace Ω := ⊥

lemma MeasurableSpace.sum_identity
  : MeasurableSpace.sum sumUnit m = m := by
  let f : Set (Set Ω) := MeasurableSet[sumUnit]
  let g : Set (Set Ω) := MeasurableSet[m]
  have : f ∪ g = g := by
    ext u
    simp only [Set.mem_union, or_iff_right_iff_imp]
    unfold f
    intro hu
    apply MeasurableSet.measurableSet_bot_induction
    · exact MeasurableSet.empty
    · exact MeasurableSet.univ
    · exact hu
  simp only [sum]
  have h : MeasurableSpace.generateFrom (f ∪ g) = MeasurableSpace.generateFrom g := by
    grind
  have : MeasurableSpace.generateFrom MeasurableSet[m] = m := by
    apply MeasurableSpace.generateFrom_measurableSet
  aesop

lemma MeasurableSpace.sum_comm
  : m₁.sum m₂ = m₂.sum m₁ := by
  unfold MeasurableSpace.sum; congr 1; exact Set.union_comm _ _

@[aesop 50% apply]
lemma mem_sum_l {u : Set Ω} (hu : MeasurableSet[m₁] u) :
    MeasurableSet[m₁.sum m₂] u := by
  unfold MeasurableSpace.sum
  exact MeasurableSpace.measurableSet_generateFrom (Set.mem_union_left _ hu)

@[aesop 50% apply]
lemma mem_sum_r {u : Set Ω} (hu : MeasurableSet[m₂] u) :
    MeasurableSet[m₁.sum m₂] u := by
  rw [MeasurableSpace.sum_comm]; apply mem_sum_l; exact hu

lemma mem_sum_inter {u v : Set Ω}
  (hu : MeasurableSet[m₁] u) (hv : MeasurableSet[m₂] v)
  : MeasurableSet[m₁.sum m₂] (u ∩ v) := by
  have : MeasurableSet[m₁.sum m₂] u := by aesop
  have : MeasurableSet[m₁.sum m₂] v := by aesop
  aesop

lemma subset_sum_l : m₁ ≤ m₁.sum m₂ := by intro _ ht; apply mem_sum_l; exact ht

lemma subset_sum_r : m₂ ≤ m₁.sum m₂ := by intro _ ht; apply mem_sum_r; exact ht

lemma MeasurableSpace.sum_assoc_left {m₁ m₂ m₃ : MeasurableSpace Ω}
  : (m₁.sum m₂).sum m₃ ≤ m₁.sum (m₂.sum m₃) := by
  have : m₁.sum m₂ ≤ m₁.sum (m₂.sum m₃) := by
    apply MeasurableSpace.generateFrom_le
    intro t ht
    rcases ht with h₁ | h₂
    · aesop
    · aesop
  have : m₃ ≤ m₁.sum (m₂.sum m₃) := by
    have : m₃ ≤ m₂.sum m₃ := by apply subset_sum_r
    have : m₂.sum m₃ ≤ m₁.sum (m₂.sum m₃) := by apply subset_sum_r
    grind
  apply MeasurableSpace.generateFrom_le
  intro t ht
  rcases ht with h₁ | h₂
  · aesop
  · aesop

lemma MeasurableSpace.sum_assoc_right {m₁ m₂ m₃ : MeasurableSpace Ω}
  : (m₁.sum m₂).sum m₃ ≥ m₁.sum (m₂.sum m₃) := by
  have : m₁ ≤ (m₁.sum m₂).sum m₃ := by
    have : m₁ ≤ m₁.sum m₂ := by apply subset_sum_l
    have : m₁.sum m₂ ≤ (m₁.sum m₂).sum m₃ := by apply subset_sum_l
    grind
  have : m₂.sum m₃ ≤ (m₁.sum m₂).sum m₃ := by
    have : m₂ ≤ (m₁.sum m₂).sum m₃ := by
      have : m₂ ≤ m₁.sum m₂ := by apply subset_sum_r
      have : m₁.sum m₂ ≤ (m₁.sum m₂).sum m₃ := by apply subset_sum_l
      grind
    have : m₃ ≤ (m₁.sum m₂).sum m₃ := by apply subset_sum_r
    apply MeasurableSpace.generateFrom_le
    intro t ht
    rcases ht with h₁ | h₂
    · aesop
    · aesop
  apply MeasurableSpace.generateFrom_le
  intro t ht
  rcases ht with h₁ | h₂
  · aesop
  · aesop

lemma MeasurableSpace.sum_assoc {m₁ m₂ m₃ : MeasurableSpace Ω}
  : (m₁.sum m₂).sum m₃ = m₁.sum (m₂.sum m₃) := by
  have : (m₁.sum m₂).sum m₃ ≤ m₁.sum (m₂.sum m₃) := by
    apply MeasurableSpace.sum_assoc_left
  have : (m₁.sum m₂).sum m₃ ≥ m₁.sum (m₂.sum m₃) := by
    apply MeasurableSpace.sum_assoc_right
  aesop

lemma MeasurableSpace.sum_mono {m₁ m₂ : MeasurableSpace Ω}
  : m₁ ≤ m₂.sum m₁ := by
  intro _ hu; apply mem_sum_r; exact hu

lemma MeasurableSpace.sum_functorial {m m₁ m₂ : MeasurableSpace Ω}
  (h : m₁ ≤ m₂)
  : m₁.sum m ≤ m₂.sum m := by
  apply MeasurableSpace.generateFrom_le
  intro u hu
  rcases hu with h₁ | h₂
  · aesop
  · aesop

lemma MeasurableSpace.sum_functorial_right {m m₁ m₂ : MeasurableSpace Ω}
  (h : m₁ ≤ m₂)
  : m.sum m₁ ≤ m.sum m₂ := by
  rw [sum_comm m m₁, sum_comm m m₂]
  exact sum_functorial h

lemma MeasurableSpace.sum_mono_both {m₁ m₂ n₁ n₂ : MeasurableSpace Ω}
  (h₁ : m₁ ≤ m₂) (h₂ : n₁ ≤ n₂) : m₁.sum n₁ ≤ m₂.sum n₂ :=
  le_trans (sum_functorial h₁) (sum_functorial_right h₂)

end Sum

/- We define `(𝓕, μ) ≤ (𝓖, ν)` if `𝓕 ⊆ 𝓖` and `μ` is the restriction of `ν` to `𝓕` -/
@[ext]
structure MeasureOnSpace (Ω : Type*) where
  ms : MeasurableSpace Ω
  μ : Measure[ms] Ω

instance (Ω : Type*) : Preorder (MeasureOnSpace Ω) where
  le (ps₁ ps₂) := ps₁.ms ≤ ps₂.ms ∧ ps₁.μ = ps₂.μ.cast _
  le_refl := by simp
  le_trans (h₁ h₂) := by
    aesop (add safe forward le_trans) (add safe (by rw [MeasureTheory.Measure.map_map]))

structure PSpace (Ω : Type*) extends MeasureOnSpace Ω where
  is_prob : IsProbabilityMeasure μ

instance (Ω : Type*) : Preorder (PSpace Ω) where
  le (ps₁ ps₂) := (ps₁.1.ms ≤ ps₂.1.ms) ∧ ps₁.1.μ = ps₂.1.μ.cast _
  le_refl := by simp
  le_trans {a b c} (h₁ h₂) := by
    aesop (add safe forward le_trans) (add safe (by rw [MeasureTheory.Measure.map_map]))

/-- Holds if `r` is the independent product of `p` and `q` -/
def PSpace.isIndependentProduct (r p q : PSpace Ω) :=
  r.ms = p.ms.sum q.ms ∧
  ∀ E (_ : MeasurableSet[p.ms] E) F (_ : MeasurableSet[q.ms] F),
    r.μ (E ∩ F) = p.μ E * q.μ F

lemma PSpace.isIndependentProduct_def {r p q : PSpace Ω} :
  isIndependentProduct r p q ↔
  r.1.ms = p.1.ms.sum q.1.ms ∧
  ∀ E (_ : MeasurableSet[p.1.ms] E)
    F (_ : MeasurableSet[q.1.ms] F),
  r.1.μ (E ∩ F) = p.1.μ E * q.1.μ F := Iff.rfl

open PSpace

notation:60 (priority := high) r " =ᵢ " p " ⊕ᵢ " q:arg => isIndependentProduct r p q

lemma PSpace.ms_eq_of_isIndependentProduct {r r' p q : PSpace Ω}
  (h₁ : r =ᵢ p ⊕ᵢ q) (h₂ : r' =ᵢ p ⊕ᵢ q) :
  r.1.ms = r'.1.ms := by
  rcases h₁ with ⟨a, _⟩
  rcases h₂ with ⟨c, _⟩
  grind

section GeneratingPiSystem

variable {Ω : Type*} (p q : MeasureOnSpace Ω)

/--
Given `p q : MeasureOnSpace Ω`, `generator p q` is a set of subsets of Ω that
generates the smallest σ-algebra containing both σ-algebras
-/
def generator (p q : MeasureOnSpace Ω) : Set (Set Ω) :=
  {S : Set Ω | ∃ F G, S = F ∩ G ∧ MeasurableSet[p.ms] F ∧ MeasurableSet[q.ms] G}

section

variable {p q : MeasureOnSpace Ω}
         {u v w : Set Ω}

/-- Every set in the generating set `w` is of the form `u ∩ v` -/
lemma exists_inter_measurableSet_of_mem_generator (hw : w ∈ generator p q) :
  ∃ u v, w = u ∩ v ∧ MeasurableSet[p.ms] u ∧ MeasurableSet[q.ms] v := by
  rcases hw with ⟨u, v, rfl, hu, hv⟩
  exact ⟨u, v, rfl, hu, hv⟩

@[aesop 50% apply]
lemma mem_generator_l (hu : MeasurableSet[p.ms] u) :
  u ∈ generator p q :=
  ⟨u, ⊤, (Set.inter_univ u).symm, hu, MeasurableSet.univ⟩

@[aesop 50% apply]
lemma mem_generator_r (hu : MeasurableSet[q.ms] u) :
  u ∈ generator p q :=
  ⟨⊤, u, (Set.univ_inter u).symm, MeasurableSet.univ, hu⟩

lemma inter_mem_generator
  (hu : MeasurableSet[p.ms] u) (hv : MeasurableSet[q.ms] v) :
  u ∩ v ∈ generator p q :=
  ⟨u, v, rfl, hu, hv⟩

attribute [local aesop safe apply] MeasurableSpace.measurableSet_generateFrom

theorem MeasureOnSpace.generateFrom_generator_eq_sum :
  MeasurableSpace.generateFrom (generator p q) = p.ms.sum q.ms := by
  ext s
  refine ⟨?p, by apply MeasurableSpace.generateFrom_mono (fun _ _ ↦ by aesop)⟩
  let sumSp := MeasurableSet[p.ms.sum q.ms]
  apply MeasurableSpace.generateFrom_le
  rintro t ⟨u, ⟨v, _, _, _⟩⟩
  have h₂ : generator p q ⊆ sumSp := fun u hu ↦ by
    obtain ⟨u₁, v₁, rfl, _, _⟩ := exists_inter_measurableSet_of_mem_generator hu
    exact (show sumSp u₁ by aesop).inter (by aesop)
  apply MeasurableSpace.generateFrom_le (by convert h₂)
  aesop (add simp (show u ∩ v ∈ generator p q by use u, v))

lemma MeasureOnSpace.isPiSystem_generator : IsPiSystem (generator p q) := fun _ hu _ hv _ ↦ by
  obtain ⟨_, _, rfl, _, _⟩ := exists_inter_measurableSet_of_mem_generator hu
  obtain ⟨_, _, rfl, _, _⟩ := exists_inter_measurableSet_of_mem_generator hv
  rw [Set.inter_inter_inter_comm]
  aesop (add safe apply inter_mem_generator)

lemma mem_generator_imp_mem_sum (h : u ∈ generator p q)
  : MeasurableSet[p.ms.sum q.ms] u := by
  have h₁ := @MeasureOnSpace.generateFrom_generator_eq_sum Ω p q
  rw [← h₁]
  apply MeasurableSpace.measurableSet_generateFrom h
end

end GeneratingPiSystem

section Uniqueness

lemma MeasurableSpace.eq_cast_ext_ext
  {m₁ m₂ : MeasurableSpace Ω}
  (μ : Measure[m₁] Ω) (ν : Measure[m₂] Ω)
  (h_eq : m₁ = m₂) (h₂ : ∀ w, MeasurableSet[m₁] w → μ w = ν w) :
  μ = ν.cast _ := by aesop

@[ext]
lemma MeasureOnSpace.ext_ms {p q : MeasureOnSpace Ω}
  (h_eq_alg : p.ms = q.ms)
  (h_eq_mea : ∀ E, MeasurableSet[p.ms] E → p.μ E = q.μ E) :
  p = q := by
  aesop (add safe cases MeasureOnSpace)

open MeasurableSpace in
lemma Measure.le_preserves_measure {p q : MeasurableSpace Ω}
  {h : p ≤ q} {μ : Measure[q] Ω}
  {u : Set Ω}
  {hu : MeasurableSet[p] u}
  : μ.cast p u = μ u := by
  have hu' : MeasurableSet[q] u := by apply h; exact hu
  unfold Measure.cast Measure.map!
  rw [Measure.map_apply] <;> aesop

lemma MeasureOnSpace.le_preserves_measure
  {p q : MeasureOnSpace Ω} (h : p ≤ q) {u} (hp : MeasurableSet[p.ms] u) : p.μ u = q.μ u := by
  rcases h with ⟨h₁, h₂⟩
  aesop (add simp Measure.le_preserves_measure)

instance [Inhabited Ω] : One (MeasureOnSpace Ω) where
  one := {ms := ⊥, μ := @Measure.dirac Ω ⊥ default}

instance [Inhabited Ω] : PartialOrder (MeasureOnSpace Ω) where
  le_antisymm m₁ m₂ h₁ h₂ := by
    have h : m₁.ms = m₂.ms := le_antisymm h₁.1 h₂.1
    apply MeasureOnSpace.ext_ms h
    intro u hu
    exact MeasureOnSpace.le_preserves_measure h₁ hu

instance (Ω : Type*) [Inhabited Ω] : PartialOrder (PSpace Ω) where
  le_antisymm := by
    rintro ⟨m₁, p₁⟩ ⟨m₂, p₂⟩ h₁ h₂
    have : m₁ ≤ m₂ := by aesop
    have : m₂ ≤ m₁ := by aesop
    have : m₁ = m₂ := @le_antisymm (MeasureOnSpace Ω) _ m₁ m₂ (by aesop) (by aesop)
    aesop

@[ext]
lemma PSpace.ext_ms {p q : PSpace Ω}
  (h_eq_alg : p.1.ms = q.1.ms)
  (h_eq_mea : ∀ E, MeasurableSet[p.1.ms] E → p.1.μ E = q.1.μ E) : p = q := by
  rcases p with ⟨a, _⟩
  rcases q with ⟨b, _⟩
  have : a = b := MeasureOnSpace.ext_ms h_eq_alg h_eq_mea
  aesop

@[simp, grind .]
lemma PSpace.measure_ne_top {m : PSpace Ω} {u : Set Ω} : m.1.μ u ≠ ⊤ := by
  apply ne_of_lt
  have h₁ : m.1.μ Set.univ = 1 := m.2.measure_univ
  have h₂ : u ⊆ Set.univ := by aesop
  have h₃ : m.1.μ u ≤ m.1.μ Set.univ := measure_mono h₂
  exact lt_of_le_of_lt (b := 1) (by aesop) (by aesop)

theorem PSpace.uniqueness (p q r r' : PSpace Ω)
  (h₁ : r =ᵢ p ⊕ᵢ q) (h₂ : r' =ᵢ p ⊕ᵢ q) : r = r' := by
  apply PSpace.ext_ms (h₁.1 ▸ h₂.1 ▸ rfl)
  -- have : IsPiSystem (generator p.1 q.1) := MeasureOnSpace.isPiSystem_generator p.1 q.1
  -- Applying the π-λ theorem: the σ-algebra is by definition a λ-system,
  -- so we just need to show that the measures agree on a generating π-system
  rw [PSpace.isIndependentProduct_def] at h₁ h₂
  apply MeasurableSpace.induction_on_inter
  · exact @MeasureOnSpace.isPiSystem_generator Ω p.1 q.1
  · simp
  · intro t ht
    obtain ⟨u, v, rfl, hu, hv⟩ := exists_inter_measurableSet_of_mem_generator ht
    grind
  · intro t ht ih
    have hms : r.1.ms = r'.1.ms := h₁.1.trans h₂.1.symm
    rw [measure_compl ht r.measure_ne_top, r.2.measure_univ,
        measure_compl (hms ▸ ht) r'.measure_ne_top, r'.2.measure_univ, ih]
  · intro us disjoint hus prf
    have hms : r.1.ms = r'.1.ms := h₁.1.trans h₂.1.symm
    have h_sum1 : r.1.μ (⋃ i, us i) = ∑' i, r.1.μ (us i) :=
      @Measure.m_iUnion (α := Ω) (f := us) r.1.ms r.1.μ hus disjoint
    have h_sum2 : r'.1.μ (⋃ i, us i) = ∑' i, r'.1.μ (us i) :=
      @Measure.m_iUnion (α := Ω) (f := us) r'.1.ms r'.1.μ (hms ▸ hus) disjoint
    simp only [h_sum1, h_sum2]
    congr 1; ext i; exact prf i
  · aesop (add simp MeasureOnSpace.generateFrom_generator_eq_sum)

end Uniqueness

section Trim

@[simp]
def MeasureOnSpace.trim
  {p : MeasureOnSpace Ω} {f : MeasurableSpace Ω} (h : f ≤ p.ms)
  : MeasureOnSpace Ω := {
  ms := f
  μ := p.μ.trim h
}

lemma MeasureOnSpace.trim_eq
  {p : MeasureOnSpace Ω} {f : MeasurableSpace Ω} (h : f ≤ p.ms)
  {u : Set Ω} (hu : MeasurableSet[f] u)
  : (p.trim h).μ u = p.μ u := by
  simp only [trim]
  exact trim_measurableSet_eq h hu

lemma Measure.trim_preserves_prob
  (f g : MeasurableSpace Ω)
  {μ : Measure[g] Ω}
  (hf : f ≤ g) (hp : IsProbabilityMeasure μ)
  : IsProbabilityMeasure (μ.trim hf) := by
  constructor
  rw [trim_measurableSet_eq hf MeasurableSet.univ]
  exact hp.measure_univ

@[simp]
def PSpace.trim
  {p : PSpace Ω} {f : MeasurableSpace Ω} (h : f ≤ p.1.ms)
  : PSpace Ω := ⟨p.1.trim h, Measure.trim_preserves_prob f p.1.ms h p.2⟩

/-
Trimming a PSpace to a coarser σ-algebra yields a smaller PSpace
-/
lemma PSpace.trim_le {p : PSpace Ω} {f : MeasurableSpace Ω} (h : f ≤ p.1.ms) :
    p.trim h ≤ p := ⟨h, trim_eq_map h⟩

end Trim

section Identity

lemma dirac_is_prob [Inhabited Ω] : IsProbabilityMeasure (@Measure.dirac Ω ⊥ default) := by
  apply isProbabilityMeasure_iff.2
  simp

def PSpace.unit [Inhabited Ω] : PSpace Ω := ⟨{
  ms := ⊥
  μ := @Measure.dirac Ω ⊥ default
}, dirac_is_prob⟩

instance [Inhabited Ω] : One (PSpace Ω) where
  one := PSpace.unit

lemma empty_sigma_algebra_is_identity [Inhabited Ω] (p : MeasureOnSpace Ω)
  : p.ms = MeasurableSpace.generateFrom (unit.1.ms.MeasurableSet' ∪ p.ms.MeasurableSet') := by
  let a : Set (Set Ω) := p.ms.MeasurableSet'
  let b : Set (Set Ω) := unit.1.ms.MeasurableSet'
  have h : a = b ∪ a := by
    ext u
    constructor
    · grind
    · rintro (hu | hu)
      · apply MeasurableSet.measurableSet_bot_induction
        · exact p.ms.measurableSet_empty
        · letI := p.ms; exact MeasurableSet.univ
        · exact hu
      · exact hu
  rw [← h]
  exact (@MeasurableSpace.generateFrom_measurableSet Ω p.ms).symm

theorem PSpace.indepenendentProduct_identity [Inhabited Ω] {p : PSpace Ω}
  : p =ᵢ unit ⊕ᵢ p := by
  refine ⟨(MeasurableSpace.sum_identity p.1.ms).symm, ?_⟩
  intro E hE F hF
  rcases MeasurableSpace.measurableSet_bot_iff.1 hE with rfl | rfl
  · simp
  · simp [unit.2.measure_univ]

end Identity

section Commutativity

theorem PSpace.independentProduct_comm [Inhabited Ω] {r p q : PSpace Ω}
  (h : r =ᵢ p ⊕ᵢ q)
  : r =ᵢ q ⊕ᵢ p := by
  constructor
  · have h₁ : MeasurableSpace.sum p.1.ms q.1.ms
      = MeasurableSpace.sum q.1.ms p.1.ms := by
      apply MeasurableSpace.sum_comm
    have : r.1.ms = MeasurableSpace.sum p.1.ms q.1.ms := h.1
    grind
  · intro u hu v hv
    let μ := r.1.μ
    let μ₁ := q.1.μ
    let μ₂ := p.1.μ
    have : μ (v ∩ u) = μ₂ v * μ₁ u := h.2 v hv u hu
    grind

end Commutativity

section Associativity

-- Recall the definiton of partial associativity (Kleene equality):
--  (a * b) * c ≃ a * (b * c) means:
-- If (a * b) and (a * b) * c are defined then
--   1. (b * c) and a * (b * c) are defined
--   2. (a * b) * c = a * (b * c)
-- The above definition suffices because we proved commutativity
theorem PSpace.independentProduct_assoc [Inhabited Ω] {pq p q s r : PSpace Ω}
  (h_pq : pq =ᵢ p ⊕ᵢ q)
  (h_pq_r : s =ᵢ pq ⊕ᵢ r)
  : ∃ qr, qr =ᵢ q ⊕ᵢ r ∧ s =ᵢ p ⊕ᵢ qr
  := by
  let qr_ms : MeasurableSpace Ω := MeasurableSpace.sum q.1.ms r.1.ms
  have h : qr_ms <= s.1.ms := by
    unfold qr_ms
    rw [h_pq_r.1, h_pq.1]
    have : q.1.ms.sum r.1.ms ≤ p.1.ms.sum (q.1.ms.sum r.1.ms) := by
      apply MeasurableSpace.sum_mono
    rw [MeasurableSpace.sum_assoc]
    assumption
  let qr : PSpace Ω := @s.trim Ω qr_ms h
  have h_qr : qr =ᵢ q ⊕ᵢ r := by
    constructor
    · simp
      aesop
    · intro u hu v hv
      have hou : MeasurableSet[pq.1.ms] (Set.univ ∩ u) := by
        simp only [Set.univ_inter]
        have h : pq.1.ms = p.1.ms.sum q.1.ms := h_pq.1
        rw [h]
        have h₂ : u ∈ generator p.1 q.1 := mem_generator_r hu
        apply @mem_generator_imp_mem_sum Ω p.1 q.1 u h₂
      have h := h_pq_r.2 (Set.univ ∩ u) hou v hv
      have h₁ : pq.1.μ (Set.univ ∩ u) = q.1.μ u := by
        have h := h_pq.2 Set.univ MeasurableSet.univ u hu
        simp only [p.2.measure_univ, one_mul] at h; exact h
      have h₂ : s.1.μ (Set.univ ∩ u ∩ v) = qr.1.μ (u ∩ v) := by
        have := h_pq_r.2 (Set.univ ∩ u) hou v hv
        have h₃ : s.1.μ (Set.univ ∩ u ∩ v) = pq.1.μ (Set.univ ∩ u) * r.1.μ v := by
          grind
        have h₄ : pq.1.μ (Set.univ ∩ u) = p.1.μ Set.univ * q.1.μ u :=
          h_pq.2 Set.univ MeasurableSet.univ u hu
        rw [h₄] at h₃
        have h₅ : p.1.μ Set.univ = 1 := p.2.measure_univ
        rw [h₅] at h₃
        unfold qr
        apply Eq.symm
        have h₇ : MeasurableSet[q.1.ms.sum r.1.ms] (u ∩ v) := by
          apply mem_generator_imp_mem_sum
          apply inter_mem_generator hu hv
        have h₈ : q.1.ms.sum r.1.ms ≤ s.1.ms := by
          rw [h_pq_r.1, h_pq.1, MeasurableSpace.sum_assoc]
          apply MeasurableSpace.sum_mono
        have h_trim := @s.1.trim_eq Ω qr_ms h₈ (u ∩ v) h₇
        have h_set : s.1.μ (Set.univ ∩ u ∩ v) = s.1.μ (u ∩ v) := by
          rw [show Set.univ ∩ u ∩ v = u ∩ v by grind]
        exact h_trim.trans h_set.symm
      rw [← h₂, h, h₁]
  refine ⟨qr, h_qr, ?_⟩
  have h_p_qr : s =ᵢ p ⊕ᵢ qr := by
    constructor
    · rw [h_pq_r.1, h_pq.1, h_qr.1]
      apply @MeasurableSpace.sum_assoc Ω p.1.ms q.1.ms r.1.ms
    · have h :
        ∀ u (hu : MeasurableSet[p.1.ms] u)
          vw (hvw : MeasurableSet[q.1.ms.sum r.1.ms] vw),
        s.1.μ (u ∩ vw) = p.1.μ u * qr.1.μ vw
      := by
        intro u hu
        apply MeasurableSpace.induction_on_inter
        · apply MeasureOnSpace.isPiSystem_generator (p := q.1) (q := r.1)
        · simp
        · intro t ht
          obtain ⟨v, w, rfl, hv, hw⟩ := exists_inter_measurableSet_of_mem_generator ht
          have hmsuvpq : MeasurableSet[pq.1.ms] (u ∩ v) := by
            rw [h_pq.1]; exact @mem_sum_inter Ω p.1.ms q.1.ms u v hu hv
          have h1 := h_pq_r.2 (u ∩ v) hmsuvpq w hw
          have h2 := h_pq.2 u hu v hv
          have h3 := h_qr.2 v hv w hw
          calc s.1.μ (u ∩ (v ∩ w))
              = s.1.μ ((u ∩ v) ∩ w) := by rw [← Set.inter_assoc]
            _ = pq.1.μ (u ∩ v) * r.1.μ w := h1
            _ = p.1.μ u * q.1.μ v * r.1.μ w := by rw [h2]
            _ = p.1.μ u * (q.1.μ v * r.1.μ w) := mul_assoc _ _ _
            _ = p.1.μ u * qr.1.μ (v ∩ w) := by rw [← h3]
        · intro t ht ih
          have hmst : MeasurableSet[s.1.ms] (u ∩ t) := by
            rw [h_pq_r.1, h_pq.1, MeasurableSpace.sum_assoc]
            exact @mem_sum_inter Ω p.1.ms (q.1.ms.sum r.1.ms) u t hu ht
          have hsdiff : s.1.μ (u \ (u ∩ t)) = s.1.μ u - s.1.μ (u ∩ t) :=
            @measure_diff Ω s.1.ms s.1.μ u (u ∩ t) (by grind)
              hmst.nullMeasurableSet s.measure_ne_top
          have hmsuv : MeasurableSet[pq.1.ms] (u ∩ Set.univ) := by
            rw [h_pq.1]; exact @mem_sum_inter Ω p.1.ms q.1.ms u Set.univ hu MeasurableSet.univ
          have hsu_eq_pu : s.1.μ u = p.1.μ u := by
            have h1 := h_pq_r.2 (u ∩ Set.univ) hmsuv Set.univ MeasurableSet.univ
            have h2 := h_pq.2 u hu Set.univ MeasurableSet.univ
            simp only [Set.inter_univ, q.2.measure_univ, r.2.measure_univ, mul_one] at h1 h2
            grind
          have h₃ : qr.1.μ tᶜ = 1 - qr.1.μ t := by
            have h := @measure_compl Ω qr.1.ms (μ := qr.1.μ) (s := t) ht qr.measure_ne_top
            simp only [qr.2.measure_univ] at h; exact h
          have hfact : p.1.μ u - p.1.μ u * qr.1.μ t = p.1.μ u * (1 - qr.1.μ t) := by
            have h := (@ENNReal.mul_sub (p.1.μ u) 1 (qr.1.μ t) (fun _ _ ↦ p.measure_ne_top)).symm
            simp only [mul_one] at h; exact h
          calc s.1.μ (u ∩ tᶜ)
              = s.1.μ (u \ (u ∩ t)) := by rw [show u ∩ tᶜ = u \ (u ∩ t) by grind]
            _ = s.1.μ u - s.1.μ (u ∩ t) := hsdiff
            _ = s.1.μ u - p.1.μ u * qr.1.μ t := by rw [ih]
            _ = p.1.μ u - p.1.μ u * qr.1.μ t := by rw [hsu_eq_pu]
            _ = p.1.μ u * (1 - qr.1.μ t) := hfact
            _ = p.1.μ u * qr.1.μ tᶜ := by rw [h₃]
        · intro us hdis hus ih
          have hus' : ∀ i, MeasurableSet[s.1.ms] (u ∩ us i) := fun i ↦ by
            rw [h_pq_r.1, h_pq.1, MeasurableSpace.sum_assoc]
            exact @mem_sum_inter Ω p.1.ms (q.1.ms.sum r.1.ms) u (us i) hu (hus i)
          have hdis' : Pairwise (Function.onFun Disjoint (fun k ↦ u ∩ us k)) := by
            intro k l p
            have h₂ : Disjoint (us k) (us l) := hdis p
            have h₃ := @Disjoint.inter_left' Ω (us k) (us l) u h₂
            have := @Disjoint.inter_right' Ω (u ∩ us k) (us l) u h₃
            aesop
          have hsiUnion : s.1.μ (⋃ i, u ∩ us i) = ∑' i, s.1.μ (u ∩ us i) :=
            @Measure.m_iUnion Ω s.1.ms s.1.μ (fun i ↦ u ∩ us i) hus' hdis'
          calc s.1.μ (u ∩ ⋃ i, us i)
              = s.1.μ (⋃ i, u ∩ us i) := by rw [Set.inter_iUnion]
            _ = ∑' i, s.1.μ (u ∩ us i) := hsiUnion
            _ = ∑' i, p.1.μ u * qr.1.μ (us i) := by
                congr 1; ext i; exact ih i
            _ = p.1.μ u * ∑' i, qr.1.μ (us i) := ENNReal.tsum_mul_left
            _ = p.1.μ u * qr.1.μ (⋃ i, us i) := by
                congr 1; exact (@qr.1.μ.m_iUnion us hus hdis).symm
        · have := @MeasurableSpace.generateFrom_sumGenerator_eq_sum Ω q.1.ms r.1.ms
          exact Eq.symm MeasureOnSpace.generateFrom_generator_eq_sum
      aesop
  exact h_p_qr

theorem independentProduct_assoc_right [Inhabited Ω] {p q r qr s : PSpace Ω}
  (h_qr : qr =ᵢ q ⊕ᵢ r)
  (h_p_qr : s =ᵢ p ⊕ᵢ qr)
  : ∃ pq, pq =ᵢ p ⊕ᵢ q ∧ s =ᵢ pq ⊕ᵢ r := by
  have h₁ : qr =ᵢ r ⊕ᵢ q := by
    apply @independentProduct_comm Ω _ qr q r h_qr
  have h₂ : s =ᵢ qr ⊕ᵢ p := by
    apply @independentProduct_comm Ω _ s p qr h_p_qr
  have h₃ := @independentProduct_assoc Ω _ qr r q s p h₁ h₂
  obtain ⟨qp, h₃⟩ := h₃
  have h₄ : qp =ᵢ p ⊕ᵢ q := by
    have : qp =ᵢ p ⊕ᵢ q := by
      have := @independentProduct_comm Ω _ qp q p
      aesop
    have : s =ᵢ r ⊕ᵢ qp := by
      have := h₃.2
      aesop
    aesop
  have h₅ : s =ᵢ qp ⊕ᵢ r := by
    have := @independentProduct_comm Ω _
    aesop
  aesop

end Associativity

section PSpace.Auxiliary

/-- Two `PSpace`s are dependent if the independent combination does
    not exist -/
@[simp, grind]
def PSpace.dependent (p q : PSpace Ω) :=
  ¬∃r : PSpace Ω, r =ᵢ p ⊕ᵢ q

theorem PSpace.dependent_symm [Inhabited Ω] {p q : PSpace Ω}
  (h : p.dependent q) : q.dependent p := by
  simp_all only [dependent, not_exists]
  intro x hx
  have : x =ᵢ p ⊕ᵢ q := by apply independentProduct_comm hx
  have := h x (by aesop)
  contradiction

theorem PSpace.dependent_mono_left {p q r qr : PSpace Ω}
  (hinc : p.dependent q) (hqr : qr =ᵢ q ⊕ᵢ r)
  : p.dependent qr := by
  simp_all only [dependent, not_exists]
  intro x hx
  have : ∃ y : PSpace Ω, y =ᵢ p ⊕ᵢ q := by
    let pqms := p.1.ms.sum q.1.ms
    have hmsle : pqms ≤ x.1.ms := by
      unfold pqms
      have := @MeasurableSpace.sum_comm Ω
      have h₁ := @MeasurableSpace.sum_mono Ω
      rw [hx.1, hqr.1]
      have h₂ : p.1.ms.sum (q.1.ms.sum r.1.ms) = (p.1.ms.sum q.1.ms).sum (r.1.ms) := by
        have := @MeasurableSpace.sum_assoc Ω
        aesop
      rw [h₂]
      grind
    let x' := @x.trim Ω pqms (by aesop)
    have : x' =ᵢ p ⊕ᵢ q := by
      unfold PSpace.isIndependentProduct
      constructor
      · aesop
      · intro u hu v hv
        have : x'.1.μ ((u ∩ v) ∩ Set.univ) = p.1.μ u * qr.1.μ (v ∩ Set.univ) := by
          have h₁ : v ∩ Set.univ = v := by grind
          have h₂ : (u ∩ v) ∩ Set.univ = u ∩ (v ∩ Set.univ) := by grind
          rw [h₂, h₁]
          have : MeasurableSet[pqms] (u ∩ v) := by
            unfold pqms
            have : MeasurableSet[pqms] u := by apply @mem_sum_l Ω p.1.ms q.1.ms u hu
            have : MeasurableSet[pqms] v := by apply @mem_sum_r Ω p.1.ms q.1.ms v hv
            have := @MeasurableSet.inter Ω pqms u v (by aesop) (by aesop)
            assumption
          have h₃ := @x.1.trim_eq Ω pqms (by aesop) (u ∩ v) (by aesop)
          have : x.1.μ (u ∩ v) = p.1.μ u * qr.1.μ v := by
            have hv' : MeasurableSet[qr.1.ms] v := by
              rw [hqr.1]; exact @mem_sum_l Ω q.1.ms r.1.ms v hv
            exact hx.2 u hu v hv'
          have : x'.1.μ (u ∩ v) = x.1.μ (u ∩ v) := by aesop
          grind
        have : p.1.μ u * qr.1.μ (v ∩ Set.univ) = p.1.μ u * (q.1.μ v * r.1.μ Set.univ) := by
          congr 1; exact hqr.2 v hv Set.univ MeasurableSet.univ
        calc
            x'.1.μ (u ∩ v)
        _ = x'.1.μ ((u ∩ v) ∩ Set.univ) := by aesop
        _ = p.1.μ u * qr.1.μ (v ∩ Set.univ) := by aesop
        _ = p.1.μ u * (q.1.μ v * r.1.μ Set.univ) := by aesop
        _ = p.1.μ u * (q.1.μ v * 1) := by have := r.2.measure_univ; aesop
        _ = p.1.μ u * q.1.μ v := by aesop
    aesop
  aesop

lemma PSpace.measure_compl {p : PSpace Ω} {u : Set Ω}
  (h : MeasurableSet[p.1.ms] u)
  : p.1.μ uᶜ = 1 - p.1.μ u := by
  have : p.1.μ Set.univ = 1 := p.2.measure_univ
  rw [← this]
  exact MeasureTheory.measure_compl h p.measure_ne_top

theorem PSpace.functoriality [Inhabited Ω] {p q a pa qa : PSpace Ω}
  (hpa : pa =ᵢ p ⊕ᵢ a)
  (hqa : qa =ᵢ q ⊕ᵢ a)
  (hlt : p ≤ q)
  : pa ≤ qa := by
  have hms : pa.1.ms ≤ qa.1.ms := by
    have := @MeasurableSpace.sum_functorial Ω a.1.ms p.1.ms q.1.ms hlt.1
    have := hpa.1
    have := hqa.1
    aesop
  constructor
  · exact hms
  have := hlt.2
  have : ∀ {u : Set Ω}, MeasurableSet[pa.1.ms] u → pa.1.μ u = qa.1.μ u := by
    apply MeasurableSpace.induction_on_inter
    · exact @MeasureOnSpace.isPiSystem_generator Ω p.1 a.1
    · aesop
    · intro t ht
      obtain ⟨u, v, h, hu, hv⟩ :=
        @exists_inter_measurableSet_of_mem_sumGenerator Ω p.1.ms a.1.ms t ht
      have hhu : MeasurableSet[q.1.ms] u := hlt.1 u hu
      rw [h, hpa.2 u hu v hv, hqa.2 u hhu v hv]
      have hpu_eq : p.1.μ u = q.1.μ u := by
        have hcast : p.1.μ = q.1.μ.cast p.1.ms := hlt.2
        rw [hcast]
        exact @Measure.le_preserves_measure Ω p.1.ms q.1.ms hlt.1 q.1.μ u hu
      rw [hpu_eq]
    · intro t ht h
      calc
            pa.1.μ tᶜ
        _ = 1 - pa.1.μ t := by
          have := @PSpace.measure_compl Ω pa t (by
            have := @MeasurableSet.compl Ω t pa.1.ms ht; aesop)
          aesop
        _ = 1 - qa.1.μ t := by aesop
        _ = qa.1.μ tᶜ := by
          have := @PSpace.measure_compl Ω qa t
          have := hms t ht
          aesop
    · intro us hdis hus heq
      calc
            pa.1.μ (⋃ i, us i)
        _ = ∑' i, pa.1.μ (us i) := measure_iUnion hdis hus
        _ = ∑' i, qa.1.μ (us i) := by aesop
        _ = qa.1.μ (⋃ i, us i) := Eq.symm (measure_iUnion hdis fun i => hms (us i) (hus i))
    · have := @MeasurableSpace.generateFrom_sumGenerator_eq_sum Ω p.1.ms a.1.ms
      have : p.1.ms.sumGenerator a.1.ms = generator p.1 a.1 := by rfl
      rw [hpa.1]
      aesop
  apply @Measure.ext Ω pa.1.ms
  intro u hu
  have := @Measure.le_preserves_measure Ω pa.1.ms qa.1.ms hms qa.1.μ u hu
  aesop

lemma PSpace.dependent_mono_right [Inhabited Ω] {p q r pq : PSpace Ω}
  (hinc : q.dependent r) (hpq : pq =ᵢ p ⊕ᵢ q)
  : pq.dependent r := by
  have : r.dependent q := by apply PSpace.dependent_symm (by aesop)
  have : pq =ᵢ q ⊕ᵢ p := by apply independentProduct_comm (by aesop)
  have := @PSpace.dependent_mono_left Ω r q (by aesop) pq (by aesop) (by aesop)
  have : pq.dependent r := by
    apply PSpace.dependent_symm (by aesop)
  assumption

lemma PSpace.le_of_isIndependentProduct_left
  [Inhabited Ω] {x y xy : PSpace Ω} (h : xy =ᵢ x ⊕ᵢ y)
  : x ≤ xy := by
  refine ⟨?_, ?_⟩
  · rw [h.1]; apply subset_sum_l
  · have hms : x.1.ms ≤ xy.1.ms := by rw [h.1]; apply subset_sum_l
    apply @Measure.ext Ω x.1.ms
    intro u hu
    have hcast : xy.1.μ.cast x.1.ms u = xy.1.μ u :=
      @Measure.le_preserves_measure Ω x.1.ms xy.1.ms hms xy.1.μ u hu
    rw [hcast]
    have h2 : xy.1.μ (u ∩ Set.univ) = x.1.μ u * y.1.μ Set.univ :=
      h.2 u hu Set.univ MeasurableSet.univ
    have h3 : u ∩ Set.univ = u := by simp
    have h4 : y.1.μ Set.univ = 1 := y.2.measure_univ
    rw [h3, h4, mul_one] at h2
    exact h2.symm

lemma PSpace.le_of_isIndependentProduct_right
  [Inhabited Ω] {x y xy : PSpace Ω} (h : xy =ᵢ y ⊕ᵢ x)
  : x ≤ xy := by
  refine ⟨?_, ?_⟩
  · rw [h.1]; apply subset_sum_r
  · have hms : x.1.ms ≤ xy.1.ms := by rw [h.1]; apply subset_sum_r
    apply @Measure.ext Ω x.1.ms
    intro u hu
    have hcast : xy.1.μ.cast x.1.ms u = xy.1.μ u :=
      @Measure.le_preserves_measure Ω x.1.ms xy.1.ms hms xy.1.μ u hu
    rw [hcast]
    have h2 : xy.1.μ (Set.univ ∩ u) = y.1.μ Set.univ * x.1.μ u :=
      h.2 Set.univ MeasurableSet.univ u hu
    have h3 : Set.univ ∩ u = u := by simp
    have h4 : y.1.μ Set.univ = 1 := y.2.measure_univ
    rw [h3, h4, one_mul] at h2
    exact h2.symm

end PSpace.Auxiliary
