Module FormalML.ProbTheory.SigmaAlgebras

Require Import Classical.
Require Import ClassicalChoice.
Require Import FinFun.

Require Import List Permutation.
Require Import Morphisms EquivDec Program.

Require Import Utils DVector.
Require Export Event.

Set Bullet Behavior "Strict Subproofs".

Import ListNotations.

Local Open Scope prob.

Lemma pre_event_complement_swap_classic {T} (A B:pre_event T) : pre_event_complement A === B <-> A === pre_event_complement B.

Lemma event_complement_swap_classic {T} (σ:SigmaAlgebra T) (A B:event σ) : ¬ A === B <-> A === ¬ B.

Program Instance trivial_sa (T:Type) : SigmaAlgebra T
  := {
      sa_sigma (f:pre_event T) := (f === pre_event_none \/ f === pre_Ω)
    }.

Lemma trivial_sa_sub {T} (sa:SigmaAlgebra T) :
  sa_sub (trivial_sa T) sa.

Program Instance discrete_sa (T:Type) : SigmaAlgebra T
  := {
      sa_sigma := fun _ => True
    }.
                 
Program Instance subset_sa (T:Type) (A:pre_event T) : SigmaAlgebra T
  := {
      sa_sigma f := (f === pre_event_none \/ f === A \/ f === ¬ A \/ f === pre_Ω)
    }.

Definition is_pre_partition {T}(part:nat->pre_event T)
  := pre_collection_is_pairwise_disjoint part
     /\ pre_union_of_collection part === pre_Ω.

Definition pre_sub_partition {T} (part1 part2 : nat -> pre_event T)
  := forall n, part1 n === part2 n \/ part1 n === pre_event_none.

Instance pre_sub_partition_pre {T} : PreOrder (@pre_sub_partition T).

Definition pre_in_partition_union {T} (part:nat->pre_event T) (f:pre_event T) : Prop
  := exists subpart, pre_sub_partition subpart part /\ (pre_union_of_collection subpart) === f.

Instance pre_in_partition_union_proper {T} (part:nat->pre_event T) :
  Proper (pre_event_equiv ==> iff) (pre_in_partition_union part).

Lemma pre_sub_partition_diff {T} (sub part:nat->pre_event T) :
  pre_sub_partition sub part ->
  pre_sub_partition (fun x : nat => pre_event_diff (part x) (sub x)) part.

Program Instance countable_partition_sa {T} (part:nat->pre_event T) (is_part:is_pre_partition part) : SigmaAlgebra T
  := {
      sa_sigma := pre_in_partition_union part
    }.

Program Instance sigma_algebra_intersection {T} (coll:SigmaAlgebra T->Prop): SigmaAlgebra T
  := { sa_sigma := fun e => forall sa, coll sa -> @sa_sigma _ sa e
     }.

Lemma sigma_algebra_intersection_sub {T} (coll:SigmaAlgebra T->Prop)
  : forall s, coll s -> sa_sub (sigma_algebra_intersection coll) s.

Definition all_included {T} (F:pre_event T -> Prop) : SigmaAlgebra T -> Prop
  := fun sa => forall (e:pre_event T), F e -> @sa_sigma _ sa e.

Global Instance all_included_proper {T} : Proper (equiv ==> equiv) (@all_included T).

Instance generated_sa {T} (F:pre_event T -> Prop) : SigmaAlgebra T
  := sigma_algebra_intersection (all_included F).

Lemma generated_sa_sub {T} (F:pre_event T -> Prop) :
  forall x, F x -> @sa_sigma _ (generated_sa F) x.

Lemma generated_sa_sub' {T} (F:pre_event T -> Prop) :
  pre_event_sub F (@sa_sigma _ (generated_sa F)).

Lemma generated_sa_minimal {T} (F:pre_event T -> Prop) (sa':SigmaAlgebra T) :
  (forall x, F x -> @sa_sigma _ sa' x) ->
  (forall x, @sa_sigma _ (generated_sa F) x -> @sa_sigma _ sa' x).

Inductive prob_space_closure {T} : (pre_event T->Prop) -> pre_event T -> Prop
  :=
  | psc_all p : prob_space_closure p pre_Ω
  | psc_refl (p:pre_event T->Prop) q : p q -> prob_space_closure p q
  | psc_countable p (collection: nat -> pre_event T) :
      (forall n, prob_space_closure p (collection n)) ->
      prob_space_closure p (pre_union_of_collection collection)
  | psc_complement p q : prob_space_closure p q -> prob_space_closure p (pre_event_complement q)
.

Program Instance closure_sigma_algebra {T} (F:pre_event T -> Prop) : SigmaAlgebra T
  := { sa_sigma := prob_space_closure F }.

Theorem generated_sa_closure {T} (F:pre_event T -> Prop) : generated_sa F === closure_sigma_algebra F.

Global Instance closure_sigma_algebra_proper {T} :
  Proper (Equivalence.equiv ==> Equivalence.equiv) (@closure_sigma_algebra T).

Global Instance prob_space_closure_proper {T} :
  Proper (Equivalence.equiv ==> Equivalence.equiv ==> iff) (@prob_space_closure T).

Lemma generated_sa_sub_sub {X:Type} (E1 E2:pre_event X -> Prop) :
  sa_sub (generated_sa E1) (generated_sa E2) <->
  (pre_event_sub E1 (sa_sigma (SigmaAlgebra:=(generated_sa E2)))).
    
Lemma generated_sa_equiv_subs {X:Type} (E1 E2:pre_event X -> Prop) :
  generated_sa E1 === generated_sa E2 <->
  (pre_event_sub E1 (sa_sigma (SigmaAlgebra:=(generated_sa E2))) /\
   pre_event_sub E2 (sa_sigma (SigmaAlgebra:=(generated_sa E1)))).

Global Instance generated_sa_sub_proper {X:Type} :
  Proper (pre_event_sub ==> sa_sub) (@generated_sa X).

Global Instance generated_sa_proper {X:Type} :
  Proper (pre_event_equiv ==> sa_equiv) (@generated_sa X).

Definition pre_event_set_product {T₁ T₂} (s₁ : pre_event T₁ -> Prop) (s₂ : pre_event T₂ -> Prop) : pre_event (T₁ * T₂) -> Prop
  := fun (e:pre_event (T₁ * T₂)) =>
       exists e₁ e₂,
         s₁ e₁ /\ s₂ e₂ /\
         e === (fun '(x₁, x₂) => e₁ x₁ /\ e₂ x₂).

Instance pre_event_set_product_proper {T1 T2} : Proper (equiv ==> equiv ==> equiv) (@pre_event_set_product T1 T2).

Instance product_sa {T₁ T₂} (sa₁:SigmaAlgebra T₁) (sa₂:SigmaAlgebra T₂) : SigmaAlgebra (T₁ * T₂)
  := generated_sa (pre_event_set_product (@sa_sigma _ sa₁) (@sa_sigma _ sa₂)).

Global Instance product_sa_proper {T1 T2} : Proper (equiv ==> equiv ==> equiv) (@product_sa T1 T2).

Theorem product_sa_sa {T₁ T₂} {sa1:SigmaAlgebra T₁} {sa2:SigmaAlgebra T₂} (a:event sa1) (b:event sa2) :
  sa_sigma (SigmaAlgebra:=product_sa sa1 sa2) (fun '(x,y) => a x /\ b y).

Definition product_sa_event {T₁ T₂} {sa1:SigmaAlgebra T₁} {sa2:SigmaAlgebra T₂} (a:event sa1) (b:event sa2)
  := exist _ _ (product_sa_sa a b).

Definition pre_event_set_dep_product {T₁:Type} {T₂:T₁->Type} (s₁ : pre_event T₁ -> Prop) (s₂ : forall x, pre_event (T₂ x) -> Prop) : pre_event (sigT T₂) -> Prop
  := fun (e:pre_event (sigT T₂)) =>
       exists (e₁:pre_event T₁) (e₂:forall x, pre_event (T₂ x)),
         s₁ e₁ /\ (forall x, e₁ x -> s₂ x (e₂ x)) /\
         e === (fun x12 => e₁ (projT1 x12) /\ e₂ _ (projT2 x12)).


Lemma pre_event_set_dep_product_proper {T1:Type} {T2:T1->Type}
      (s1x s1y : pre_event T1 -> Prop)
      (s1equiv : equiv s1x s1y)
      (s2x s2y : forall x, pre_event (T2 x) -> Prop)
      (s2equiv : (forall x y, s2x x y <-> s2y x y)) :
    pre_event_equiv (pre_event_set_dep_product s1x s2x)
                    (pre_event_set_dep_product s1y s2y).

Instance dep_product_sa {T₁:Type} {T₂:T₁->Type}
         (sa₁:SigmaAlgebra T₁) (sa₂:forall x, SigmaAlgebra (T₂ x)) : SigmaAlgebra (sigT T₂)
  := generated_sa (pre_event_set_dep_product (@sa_sigma _ sa₁) (fun x => @sa_sigma _ (sa₂ x))).

Lemma dep_product_sa_proper {T1 T2}
      (s1x s1y : SigmaAlgebra T1)
      (s1equiv : equiv s1x s1y)
      (s2x s2y : forall x, SigmaAlgebra (T2 x))
      (s2equiv : (forall x, equiv (s2x x) (s2y x))) :
  sa_equiv (dep_product_sa s1x s2x) (dep_product_sa s1y s2y).

Lemma dep_product_sa_sa {T₁:Type} {T₂:T₁->Type} {sa₁:SigmaAlgebra T₁} {sa₂:forall x, SigmaAlgebra (T₂ x)}
       (a:event sa₁) (b:forall a, event (sa₂ a)) :
  sa_sigma (SigmaAlgebra:=dep_product_sa sa₁ sa₂) (fun '(existT x y) => a x /\ b x y).

Definition dep_product_sa_event {T₁:Type} {T₂:T₁->Type} {sa₁:SigmaAlgebra T₁} {sa₂:forall x, SigmaAlgebra (T₂ x)}
  (a:event sa₁) (b:forall a, event (sa₂ a)) :
  event (dep_product_sa sa₁ sa₂)
  := exist _ _ (dep_product_sa_sa a b).


Definition pre_event_push_forward {X Y:Type} (f:X->Y) (ex:pre_event X) : pre_event Y
  := fun y => exists x, f x = y /\ ex x.

Instance pre_event_push_forward_proper {X Y:Type} (f:X->Y) : Proper (equiv ==> equiv) (pre_event_push_forward f).

Definition pullback_sa_sigma {X Y:Type} (sa:SigmaAlgebra Y) (f:X->Y) : pre_event X -> Prop
  := fun (xe:pre_event X) =>
       exists ye:pre_event Y,
         sa_sigma (SigmaAlgebra:=sa) ye /\
         forall a, xe a <-> ye (f a).

Program Instance pullback_sa {X Y:Type} (sa:SigmaAlgebra Y) (f:X->Y) : SigmaAlgebra X
  := {|
  sa_sigma := pullback_sa_sigma sa f |}.

Instance pullback_sa_sigma_proper {X Y:Type} :
  Proper (equiv ==> (pointwise_relation X equiv) ==> equiv)
         (@pullback_sa_sigma X Y).
  
Instance pullback_sa_proper {X Y:Type} :
  Proper (equiv ==> (pointwise_relation X equiv) ==> equiv)
         (@pullback_sa X Y).

Lemma pullback_sa_pullback {X Y:Type} (sa:SigmaAlgebra Y) (f:X->Y) (y:pre_event Y) :
  sa_sigma y -> sa_sigma (SigmaAlgebra:=pullback_sa sa f) (pre_event_preimage f y).

Lemma nested_pullback_sa_equiv {X Y Z : Type} (sa : SigmaAlgebra Z) (f : X -> Y) (g : Y -> Z) :
  sa_equiv (pullback_sa sa (fun x => g (f x))) (pullback_sa (pullback_sa sa g) f).

Lemma pullback_sa_compose_equiv {X Y Z : Type} (sa : SigmaAlgebra Z) (f : X -> Y) (g : Y -> Z) :
  sa_equiv (pullback_sa sa (compose g f)) (pullback_sa (pullback_sa sa g) f).

Section isos.
  Context {Ts:Type} {Td:Type} {cod: SigmaAlgebra Td}.

  Lemma pullback_sa_iso_l_sub (rv1 : Ts -> Td) (f g : Td -> Td)
    (inv:forall x, g (f x) = x)
    (g_sigma:forall s, sa_sigma s -> sa_sigma (fun x => s (g x))) :
    sa_sub (pullback_sa _ rv1) (pullback_sa _ (fun x => f (rv1 x))).

  Lemma pullback_sa_f_sub (rv1 : Ts -> Td) (f : Td -> Td)
    (f_sigma:forall s, sa_sigma s -> sa_sigma (fun x => s (f x))) :
    sa_sub (pullback_sa _ (fun x => f (rv1 x))) (pullback_sa _ rv1).

  Lemma pullback_sa_isos (rv1 : Ts -> Td) (f g : Td -> Td)
    (inv:forall x, g (f x) = x)
    (f_sigma:forall s, sa_sigma s -> sa_sigma (fun x => s (g x)))
    (g_sigma:forall s, sa_sigma s -> sa_sigma (fun x => s (f x))) :
    sa_equiv (pullback_sa _ rv1) (pullback_sa _ (fun x => f (rv1 x))).

End isos.


Definition union_sa {T : Type} (sa1 sa2:SigmaAlgebra T) :=
  generated_sa (pre_event_union (@sa_sigma _ sa1)
                                (@sa_sigma _ sa2)).

Global Instance union_sa_proper {T:Type} :
  Proper (sa_equiv ==> sa_equiv ==> sa_equiv) (@union_sa T).

Global Instance union_sa_sub_proper {T:Type} :
  Proper (sa_sub ==> sa_sub ==> sa_sub) (@union_sa T).

Lemma union_sa_sub_l {T : Type} (sa1 sa2:SigmaAlgebra T) :
  sa_sub sa1 (union_sa sa1 sa2).

Lemma union_sa_sub_r {T : Type} (sa1 sa2:SigmaAlgebra T) :
  sa_sub sa2 (union_sa sa1 sa2).

Lemma union_sa_comm {T : Type} (sa1 sa2:SigmaAlgebra T) :
  sa_equiv (union_sa sa1 sa2) (union_sa sa2 sa1).

Lemma union_sa_generated_l_simpl {T : Type} (a b:pre_event T -> Prop) :
  generated_sa
    (pre_event_union (sa_sigma (SigmaAlgebra:=(generated_sa a))) b)
    ===
    generated_sa
    (pre_event_union a b).

Lemma union_sa_generated_r_simpl {T : Type} (a b:pre_event T -> Prop) :
  generated_sa
    (pre_event_union a (sa_sigma (SigmaAlgebra:=(generated_sa b))))
    ===
    generated_sa
    (pre_event_union a b).

Lemma union_sa_assoc {T : Type} (sa1 sa2 sa3:SigmaAlgebra T) :
  sa_equiv (union_sa sa1 (union_sa sa2 sa3)) (union_sa (union_sa sa1 sa2) sa3).

Lemma union_sa_sub_both {T : Type} {sa1 sa2 dom:SigmaAlgebra T} :
  sa_sub sa1 dom ->
  sa_sub sa2 dom ->
  sa_sub (union_sa sa1 sa2) dom.

Definition countable_union_sa {T : Type} (sas:nat->SigmaAlgebra T) :=
  generated_sa (pre_union_of_collection (fun n => (@sa_sigma _ (sas n)))).

Global Instance countable_union_sa_sub_proper {T:Type} :
  Proper (pointwise_relation _ sa_sub ==> sa_sub) (@countable_union_sa T).

Lemma countable_union_sa_proper {T:Type} :
  Proper (pointwise_relation _ sa_equiv ==> sa_equiv) (@countable_union_sa T).

Lemma countable_union_sa_sub {T : Type} (sas:nat -> SigmaAlgebra T) :
  forall n, sa_sub (sas n) (countable_union_sa sas).

Lemma countable_union_sa_sub_all {T : Type} {sas:nat -> SigmaAlgebra T} {dom: SigmaAlgebra T} :
  (forall n, sa_sub (sas n) dom) ->
  sa_sub (countable_union_sa sas) dom.

Definition is_countable {T} (e:pre_event T)
  := exists (coll:nat -> T -> Prop),
    (forall n t1 t2, coll n t1 -> coll n t2 -> t1 = t2) /\
    e === (fun a => exists n, coll n a).

Lemma is_countable_empty {T} : @is_countable T pre_event_none.

Instance is_countable_proper_sub {T} : Proper (pre_event_sub --> Basics.impl) (@is_countable T).

Global Instance is_countable_proper {T} : Proper (equiv ==> iff) (@is_countable T).

Lemma is_countable_singleton {A:Type} {σ:SigmaAlgebra A} x pf :
  is_countable (event_singleton x pf).

Definition Finj_event {T} (coll:nat->pre_event T) (n:nat) : ({x:T | coll n x} -> nat) -> Prop
  := fun f => Injective f.

Lemma union_of_collection_is_countable {T} (coll:nat->pre_event T) :
  (forall n : nat, is_countable (coll n)) -> is_countable (pre_union_of_collection coll).

Program Instance countable_sa (T:Type) : SigmaAlgebra T
  := {
      sa_sigma (f:pre_event T) := is_countable f \/ is_countable (¬ f)
    }.

Definition pre_event_set_vector_product {T} {n} (v:vector ((pre_event T)->Prop) n) : pre_event (vector T n) -> Prop
  := fun (e:pre_event (vector T n)) =>
       exists (sub_e:vector (pre_event T) n),
         (forall i pf, (vector_nth i pf v) (vector_nth i pf sub_e))
         /\
         e === (fun (x:vector T n) => forall i pf, (vector_nth i pf sub_e) (vector_nth i pf x)).

Instance pre_event_set_vector_product_proper {n} {T} : Proper (equiv ==> equiv) (@pre_event_set_vector_product T n).

Instance vector_sa {n} {T} (sav:vector (SigmaAlgebra T) n) : SigmaAlgebra (vector T n)
  := generated_sa (pre_event_set_vector_product (vector_map (@sa_sigma _) sav)).

Global Instance vector_sa_proper {T} {n} : Proper (equiv ==> equiv) (@vector_sa T n).

Definition is_pre_partition_list {T} (l:list (pre_event T)) :=
  ForallOrdPairs pre_event_disjoint l /\ pre_list_union l === pre_Ω.

Lemma pre_is_partition_list_partition {T} {l:list (pre_event T)} :
  is_pre_partition_list l ->
  is_pre_partition (pre_list_collection l pre_event_none).

Instance list_partition_sa {T} (l:list (pre_event T)) (is_part:is_pre_partition_list l) :
  SigmaAlgebra T := countable_partition_sa
                      (pre_list_collection l pre_event_none)
                      (pre_is_partition_list_partition is_part).

Definition is_partition_list {T} {σ:SigmaAlgebra T} (l:list (event σ)) :=
  ForallOrdPairs event_disjoint l /\ list_union l === Ω.

Lemma is_partition_list_pre {T} {σ:SigmaAlgebra T}
      (l : list (event σ))
      (isp:is_pre_partition_list (map event_pre l)) :
  is_partition_list l <->
  is_pre_partition_list (map event_pre l).

Global Instance is_partition_list_perm {T} {σ:SigmaAlgebra T} :
  Proper (@Permutation _ ==> iff) (@is_partition_list T σ).

Global Instance is_partition_list_event_equiv {T} {σ:SigmaAlgebra T}:
  Proper (Forall2 event_equiv ==> iff) (@is_partition_list T σ).

Lemma is_partition_list_trivial {T} {σ:SigmaAlgebra T} :
  is_partition_list (Ω :: nil) (σ:=σ).

Section dec.

  Context {Ts:Type} {dom:SigmaAlgebra Ts}.

  Lemma is_partition_refine (l1 l2 : list dec_sa_event) :
    is_partition_list (map dsa_event l1) ->
    is_partition_list (map dsa_event l2) ->
    is_partition_list (map dsa_event (refine_dec_sa_partitions l1 l2)).

End dec.

Section filtration.
  Context {Ts:Type}.

  Class IsSubAlgebras (dom:SigmaAlgebra Ts) (sas:nat->SigmaAlgebra Ts)
    := is_sub_algebras : forall n, sa_sub (sas n) dom.
    
  Class IsFiltration (sas:nat->SigmaAlgebra Ts)
    := is_filtration : forall n, sa_sub (sas n) (sas (S n)).

  Lemma is_filtration_le (sas:nat->SigmaAlgebra Ts) {isf:IsFiltration sas}
    : forall k n, k <= n -> sa_sub (sas k) (sas n).

  Global Instance is_filtration_const (c:SigmaAlgebra Ts) : IsFiltration (fun _ : nat => c).

  Global Instance IsSubAlgebras_proper :
    Proper (sa_sub ==> pointwise_relation _ sa_sub --> impl) IsSubAlgebras.

  Global Instance IsSubAlgebras_eq_proper' :
    Proper (sa_equiv ==> pointwise_relation _ sa_equiv ==> impl) IsSubAlgebras.

  Global Instance IsSubAlgebras_eq_proper :
    Proper (sa_equiv ==> pointwise_relation _ sa_equiv ==> iff) IsSubAlgebras.


  Global Instance IsFiltration_proper' :
    Proper (pointwise_relation _ sa_equiv ==> impl) IsFiltration.
  
  Global Instance IsFiltration_proper :
    Proper (pointwise_relation _ sa_equiv ==> iff) IsFiltration.

  Section fs.
    Context (sas:nat->SigmaAlgebra Ts).
    Fixpoint filtrate_sa (n : nat) : SigmaAlgebra Ts
      := match n with
         | 0%nat => sas 0%nat
         | S k => union_sa (sas (S k)) (filtrate_sa k)
         end.
  End fs.

  Global Instance filtrate_sa_is_filtration (sas:nat->SigmaAlgebra Ts) :
    IsFiltration (filtrate_sa sas).

  Lemma filtrate_sa_sub (sas:nat->SigmaAlgebra Ts) (n : nat) : sa_sub (sas n) (filtrate_sa sas n).

  Lemma filtrate_sa_sub_all (sas:nat->SigmaAlgebra Ts) (dom : SigmaAlgebra Ts) (n : nat) :
    (forall k, k <= n -> sa_sub (sas k) dom) -> sa_sub (filtrate_sa sas n) dom.

  Global Instance filtrate_sa_is_sub_algebra
         (sas:nat->SigmaAlgebra Ts) (dom : SigmaAlgebra Ts) {subs : IsSubAlgebras dom sas} :
    IsSubAlgebras dom (filtrate_sa sas).
  
  Lemma filtrate_sa_countable_union (sas:nat->SigmaAlgebra Ts) :
    sa_equiv (countable_union_sa (filtrate_sa sas)) (countable_union_sa sas).

  Lemma filtrate_sa_countable_union_sub (sas:nat->SigmaAlgebra Ts) n :
    sa_sub (filtrate_sa sas n) (countable_union_sa sas).

  Global Instance filtrate_sa_sub_proper : Proper (pointwise_relation _ sa_sub ==> pointwise_relation _ sa_sub) filtrate_sa.

  Global Instance filtrate_sa_proper : Proper (pointwise_relation _ sa_equiv ==> pointwise_relation _ sa_equiv) filtrate_sa.

End filtration.