OptionalSampling.lean

/-
Copyright (c) 2023 Rémy Degenne. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Rémy Degenne
-/
module

public import Mathlib.Order.SuccPred.LinearLocallyFinite
public import Mathlib.Probability.Martingale.Basic

/-!
# Optional sampling theorem

If `τ` is a bounded stopping time and `σ` is another stopping time, then the value of a martingale
`f` at the stopping time `min τ σ` is almost everywhere equal to
`μ[stoppedValue f τ | hσ.measurableSpace]`.

## Main results

* `stoppedValue_ae_eq_condExp_of_le_const`: the value of a martingale `f` at a stopping time `τ`
  bounded by `n` is the conditional expectation of `f n` with respect to the σ-algebra generated by
  `τ`.
* `stoppedValue_ae_eq_condExp_of_le`: if `τ` and `σ` are two stopping times with `σ ≤ τ` and `τ` is
  bounded, then the value of a martingale `f` at `σ` is the conditional expectation of its value at
  `τ` with respect to the σ-algebra generated by `σ`.
* `stoppedValue_min_ae_eq_condExp`: the optional sampling theorem. If `τ` is a bounded stopping
  time and `σ` is another stopping time, then the value of a martingale `f` at the stopping time
  `min τ σ` is almost everywhere equal to the conditional expectation of `f` stopped at `τ`
  with respect to the σ-algebra generated by `σ`.

-/

public section


open scoped MeasureTheory ENNReal

open TopologicalSpace

namespace MeasureTheory

namespace Martingale

variable {Ω E : Type*} {m : MeasurableSpace Ω} {μ : Measure Ω} [NormedAddCommGroup E]
  [NormedSpace ℝ E] [CompleteSpace E]

section FirstCountableTopology

variable {ι : Type*} [LinearOrder ι] [TopologicalSpace ι] [OrderTopology ι]
  [FirstCountableTopology ι] {ℱ : Filtration ι m} [SigmaFiniteFiltration μ ℱ] {τ σ : Ω → WithTop ι}
  {f : ι → Ω → E} {i n : ι}

theorem condExp_stopping_time_ae_eq_restrict_eq_const (h : Martingale f ℱ μ)
    (hτ : IsStoppingTime ℱ τ) [SigmaFinite (μ.trim hτ.measurableSpace_le)] (hin : i ≤ n) :
    μ[f n | hτ.measurableSpace] =ᵐ[μ.restrict {x | τ x = i}] f i := by
  refine Filter.EventuallyEq.trans ?_ (ae_restrict_of_ae (h.condExp_ae_eq hin))
  refine condExp_ae_eq_restrict_of_measurableSpace_eq_on hτ.measurableSpace_le (ℱ.le i)
    (hτ.measurableSet_eq' i) fun t => ?_
  rw [Set.inter_comm _ t, IsStoppingTime.measurableSet_inter_eq_iff]

theorem condExp_stopping_time_ae_eq_restrict_eq_const_of_le_const (h : Martingale f ℱ μ)
    (hτ : IsStoppingTime ℱ τ) (hτ_le : ∀ x, τ x ≤ n)
    [SigmaFinite (μ.trim (hτ.measurableSpace_le_of_le hτ_le))] (i : ι) :
    μ[f n | hτ.measurableSpace] =ᵐ[μ.restrict {x | τ x = i}] f i := by
  by_cases hin : i ≤ n
  · refine Filter.EventuallyEq.trans ?_ (ae_restrict_of_ae (h.condExp_ae_eq hin))
    refine condExp_ae_eq_restrict_of_measurableSpace_eq_on (hτ.measurableSpace_le_of_le hτ_le)
      (ℱ.le i) (hτ.measurableSet_eq' i) fun t => ?_
    rw [Set.inter_comm _ t, IsStoppingTime.measurableSet_inter_eq_iff]
  · suffices {x : Ω | τ x = i} = ∅ by simp [this]; norm_cast
    ext1 x
    simp only [Set.mem_setOf_eq, Set.mem_empty_iff_false, iff_false]
    contrapose hin
    exact_mod_cast hin ▸ hτ_le x

variable [Nonempty ι]

theorem stoppedValue_ae_eq_restrict_eq (h : Martingale f ℱ μ) (hτ : IsStoppingTime ℱ τ)
    (hτ_le : ∀ x, τ x ≤ n) [SigmaFinite (μ.trim (hτ.measurableSpace_le_of_le hτ_le))] (i : ι) :
    stoppedValue f τ =ᵐ[μ.restrict {x | τ x = i}] μ[f n | hτ.measurableSpace] := by
  refine Filter.EventuallyEq.trans ?_
    (condExp_stopping_time_ae_eq_restrict_eq_const_of_le_const h hτ hτ_le i).symm
  rw [Filter.EventuallyEq, ae_restrict_iff' (ℱ.le _ _ (hτ.measurableSet_eq i))]
  refine Filter.Eventually.of_forall fun x hx => ?_
  rw [Set.mem_setOf_eq] at hx
  simp [stoppedValue, hx]

/-- The value of a martingale `f` at a stopping time `τ` bounded by `n` is the conditional
expectation of `f n` with respect to the σ-algebra generated by `τ`. -/
theorem stoppedValue_ae_eq_condExp_of_le_const_of_countable_range (h : Martingale f ℱ μ)
    (hτ : IsStoppingTime ℱ τ) (hτ_le : ∀ x, τ x ≤ n) (h_countable_range : (Set.range τ).Countable)
    [SigmaFinite (μ.trim (hτ.measurableSpace_le_of_le hτ_le))] :
    stoppedValue f τ =ᵐ[μ] μ[f n | hτ.measurableSpace] := by
  have : Set.univ = ⋃ i ∈ Set.range τ, {x | τ x = i} := by
    ext1 x
    simp only [Set.mem_univ, Set.mem_range, Set.iUnion_exists, Set.iUnion_iUnion_eq',
      Set.mem_iUnion, Set.mem_setOf_eq, exists_apply_eq_apply']
  nth_rw 1 [← @Measure.restrict_univ Ω _ μ]
  rw [this, ae_eq_restrict_biUnion_iff _ h_countable_range]
  intro i hi
  have h_top : i ≠ ⊤ := fun h ↦ by
    simp only [h, Set.mem_range] at hi
    obtain ⟨ω, hω⟩ := hi
    specialize hτ_le ω
    simp [hω] at hτ_le
  lift i to ι using h_top with i
  exact stoppedValue_ae_eq_restrict_eq h _ hτ_le i

omit [FirstCountableTopology ι] in
/-- The value of a martingale `f` at a stopping time `τ` bounded by `n` is the conditional
expectation of `f n` with respect to the σ-algebra generated by `τ`. -/
theorem stoppedValue_ae_eq_condExp_of_le_const [Countable ι] (h : Martingale f ℱ μ)
    (hτ : IsStoppingTime ℱ τ) (hτ_le : ∀ x, τ x ≤ n)
    [SigmaFinite (μ.trim (hτ.measurableSpace_le_of_le hτ_le))] :
    stoppedValue f τ =ᵐ[μ] μ[f n | hτ.measurableSpace] :=
  h.stoppedValue_ae_eq_condExp_of_le_const_of_countable_range hτ hτ_le (Set.to_countable _)

/-- If `τ` and `σ` are two stopping times with `σ ≤ τ` and `τ` is bounded, then the value of a
martingale `f` at `σ` is the conditional expectation of its value at `τ` with respect to the
σ-algebra generated by `σ`. -/
theorem stoppedValue_ae_eq_condExp_of_le_of_countable_range (h : Martingale f ℱ μ)
    (hτ : IsStoppingTime ℱ τ) (hσ : IsStoppingTime ℱ σ) (hσ_le_τ : σ ≤ τ) (hτ_le : ∀ x, τ x ≤ n)
    (hτ_countable_range : (Set.range τ).Countable) (hσ_countable_range : (Set.range σ).Countable)
    [SigmaFinite (μ.trim (hσ.measurableSpace_le_of_le fun x => (hσ_le_τ x).trans (hτ_le x)))] :
    stoppedValue f σ =ᵐ[μ] μ[stoppedValue f τ | hσ.measurableSpace] := by
  have : SigmaFinite (μ.trim (hτ.measurableSpace_le_of_le hτ_le)) :=
    sigmaFiniteTrim_mono _ (IsStoppingTime.measurableSpace_mono hσ hτ hσ_le_τ)
  have : μ[stoppedValue f τ | hσ.measurableSpace] =ᵐ[μ]
      μ[μ[f n | hτ.measurableSpace] | hσ.measurableSpace] := condExp_congr_ae
    (h.stoppedValue_ae_eq_condExp_of_le_const_of_countable_range hτ hτ_le hτ_countable_range)
  refine (Filter.EventuallyEq.trans ?_
    (condExp_condExp_of_le ?_ (hτ.measurableSpace_le_of_le hτ_le)).symm).trans this.symm
  · exact h.stoppedValue_ae_eq_condExp_of_le_const_of_countable_range hσ
      (fun x => (hσ_le_τ x).trans (hτ_le x)) hσ_countable_range
  · exact hσ.measurableSpace_mono hτ hσ_le_τ

omit [FirstCountableTopology ι] in
/-- If `τ` and `σ` are two stopping times with `σ ≤ τ` and `τ` is bounded, then the value of a
martingale `f` at `σ` is the conditional expectation of its value at `τ` with respect to the
σ-algebra generated by `σ`. -/
theorem stoppedValue_ae_eq_condExp_of_le [Countable ι] (h : Martingale f ℱ μ)
    (hτ : IsStoppingTime ℱ τ) (hσ : IsStoppingTime ℱ σ) (hσ_le_τ : σ ≤ τ) (hτ_le : ∀ x, τ x ≤ n)
    [SigmaFinite (μ.trim hσ.measurableSpace_le)] :
    stoppedValue f σ =ᵐ[μ] μ[stoppedValue f τ | hσ.measurableSpace] :=
  h.stoppedValue_ae_eq_condExp_of_le_of_countable_range hτ hσ hσ_le_τ hτ_le (Set.to_countable _)
    (Set.to_countable _)

end FirstCountableTopology

section SubsetOfNat

/-! In the following results the index set verifies
`[LinearOrder ι] [LocallyFiniteOrder ι] [OrderBot ι]`, which means that it is order-isomorphic to
a subset of `ℕ`. `ι` is equipped with the discrete topology, which is also the order topology,
and is a measurable space with the Borel σ-algebra. -/


variable {ι : Type*} [LinearOrder ι] [LocallyFiniteOrder ι] [OrderBot ι] [TopologicalSpace ι]
  [DiscreteTopology ι] [MeasurableSpace ι] [BorelSpace ι] [MeasurableSpace E] [BorelSpace E]
  [SecondCountableTopology E] {ℱ : Filtration ι m} {τ σ : Ω → WithTop ι} {f : ι → Ω → E} {i : ι}

theorem condExp_stoppedValue_stopping_time_ae_eq_restrict_le (h : Martingale f ℱ μ)
    (hτ : IsStoppingTime ℱ τ) (hσ : IsStoppingTime ℱ σ) [SigmaFinite (μ.trim hσ.measurableSpace_le)]
    (hτ_le : ∀ x, τ x ≤ i) :
    μ[stoppedValue f τ | hσ.measurableSpace] =ᵐ[μ.restrict {x : Ω | τ x ≤ σ x}]
      stoppedValue f τ := by
  rw [ae_eq_restrict_iff_indicator_ae_eq
    (hτ.measurableSpace_le _ (hτ.measurableSet_le_stopping_time hσ))]
  refine (condExp_indicator (integrable_stoppedValue ι hτ h.integrable hτ_le)
    (hτ.measurableSet_stopping_time_le hσ)).symm.trans ?_
  have h_int :
      Integrable ({ω : Ω | τ ω ≤ σ ω}.indicator (stoppedValue (fun n : ι => f n) τ)) μ := by
    refine (integrable_stoppedValue ι hτ h.integrable hτ_le).indicator ?_
    exact hτ.measurableSpace_le _ (hτ.measurableSet_le_stopping_time hσ)
  have h_meas : AEStronglyMeasurable[hσ.measurableSpace]
      ({ω : Ω | τ ω ≤ σ ω}.indicator (stoppedValue (fun n : ι => f n) τ)) μ := by
    refine StronglyMeasurable.aestronglyMeasurable ?_
    refine StronglyMeasurable.stronglyMeasurable_of_measurableSpace_le_on
      (hτ.measurableSet_le_stopping_time hσ) ?_ ?_ ?_
    · intro t ht
      rw [Set.inter_comm _ t] at ht ⊢
      rw [hτ.measurableSet_inter_le_iff hσ, IsStoppingTime.measurableSet_min_iff hτ hσ] at ht
      exact ht.2
    · refine StronglyMeasurable.indicator ?_ (hτ.measurableSet_le_stopping_time hσ)
      refine Measurable.stronglyMeasurable ?_
      exact measurable_stoppedValue h.stronglyAdapted.isStronglyProgressive_of_discrete hτ
    · intro x hx
      simp only [hx, Set.indicator_of_notMem, not_false_iff]
  exact condExp_of_aestronglyMeasurable' hσ.measurableSpace_le h_meas h_int

/-- **Optional Sampling theorem**. If `τ` is a bounded stopping time and `σ` is another stopping
time, then the value of a martingale `f` at the stopping time `min τ σ` is almost everywhere equal
to the conditional expectation of `f` stopped at `τ` with respect to the σ-algebra generated
by `σ`. -/
theorem stoppedValue_min_ae_eq_condExp [SigmaFiniteFiltration μ ℱ] (h : Martingale f ℱ μ)
    (hτ : IsStoppingTime ℱ τ) (hσ : IsStoppingTime ℱ σ) {n : ι} (hτ_le : ∀ x, τ x ≤ n)
    [h_sf_min : SigmaFinite (μ.trim (hτ.min hσ).measurableSpace_le)] :
    (stoppedValue f fun x => min (σ x) (τ x)) =ᵐ[μ] μ[stoppedValue f τ | hσ.measurableSpace] := by
  refine
    (h.stoppedValue_ae_eq_condExp_of_le hτ (hσ.min hτ) (fun x => min_le_right _ _) hτ_le).trans ?_
  refine ae_of_ae_restrict_of_ae_restrict_compl {x | σ x ≤ τ x} ?_ ?_
  · exact condExp_min_stopping_time_ae_eq_restrict_le hσ hτ
  · suffices μ[stoppedValue f τ | (hσ.min hτ).measurableSpace] =ᵐ[μ.restrict {x | τ x ≤ σ x}]
        μ[stoppedValue f τ | hσ.measurableSpace] by
      rw [ae_restrict_iff' (hσ.measurableSpace_le _ (hσ.measurableSet_le_stopping_time hτ).compl)]
      rw [Filter.EventuallyEq, ae_restrict_iff'] at this
      swap; · exact hτ.measurableSpace_le _ (hτ.measurableSet_le_stopping_time hσ)
      filter_upwards [this] with x hx hx_mem
      simp only [Set.mem_compl_iff, Set.mem_setOf_eq, not_le] at hx_mem
      exact hx hx_mem.le
    apply Filter.EventuallyEq.trans _ ((condExp_min_stopping_time_ae_eq_restrict_le hτ hσ).trans _)
    · exact stoppedValue f τ
    · rw [IsStoppingTime.measurableSpace_min hσ hτ,
        IsStoppingTime.measurableSpace_min hτ hσ, inf_comm]
    · have h1 : μ[stoppedValue f τ | hτ.measurableSpace] = stoppedValue f τ := by
        apply condExp_of_stronglyMeasurable hτ.measurableSpace_le
        · exact Measurable.stronglyMeasurable <|
            measurable_stoppedValue h.stronglyAdapted.isStronglyProgressive_of_discrete hτ
        · exact integrable_stoppedValue ι hτ h.integrable hτ_le
      rw [h1]
      exact (condExp_stoppedValue_stopping_time_ae_eq_restrict_le h hτ hσ hτ_le).symm

end SubsetOfNat

end Martingale

end MeasureTheory