Library mathcomp.analysis.altreals.distr

(* -------------------------------------------------------------------- 
 Copyright (c) - 2015--2016 - IMDEA Software Institute                
 Copyright (c) - 2015--2018 - Inria                                   
 Copyright (c) - 2016--2018 - Polytechnique                           *)

From mathcomp Require Import all_ssreflect all_algebra.
Require Import xfinmap boolp classical_sets ereal reals discrete.
Require Import mathcomp_extra topology realseq realsum.

Set Implicit Arguments.

Import Order.TTheory GRing.Theory Num.Theory.

Local Open Scope ring_scope.

Reserved Notation "\dlet_ ( i <- d ) E"
  (at level 36, E at level 36, i, d at level 50,
     format "'[' \dlet_ ( i <- d ) '/ ' E ']'").

Reserved Notation "\dlim_ ( n ) E"
  (at level 36, E at level 36, n at level 50,
     format "'[' \dlim_ ( n ) '/ ' E ']'").

Reserved Notation "\P_[ mu ] E" (at level 2, format "\P_[ mu ] E").
Reserved Notation "\P_[ mu , A ] E" (at level 2, format "\P_[ mu , A ] E").
Reserved Notation "\E?_[ mu ] f" (at level 2, format "\E?_[ mu ] f").
Reserved Notation "\E_[ mu ] f" (at level 2, format "\E_[ mu ] f").
Reserved Notation "\E_[ mu , A ] f" (at level 2, format "\E_[ mu , A ] f").

Section Distribution.
Variables (R : realType) (T : choiceType).

Structure distr := Distr {
  mu :> T → R;
  _ : ∀ x, 0 ≤ mu x;
  _ : summable mu;
  _ : psum mu ≤ 1
}.

Definition distr_of of phant R & phant T := distr.
End Distribution.

Notation "{ 'distr' T / R }" := (distr_of (Phant R) (Phant T))
  (at level 0, T at level 2, format "{ 'distr' T / R }")
  : type_scope.

Section DistrCoreTh.
Context {R : realType} (T : choiceType) (mu : {distr T / R}).

Lemma ge0_mu : ∀ x, 0 ≤ mu x.

Lemma le1_mu : psum mu ≤ 1.

Lemma summable_mu : summable mu.
End DistrCoreTh.

#[global] Hint Resolve ge0_mu le1_mu summable_mu : core.

Section Clamp.
Context {R : realType}.

Definition clamp (x : R) :=
  Num.max (Num.min x 1) 0.

Lemma ge0_clamp x : 0 ≤ clamp x.

Lemma le1_clamp x : clamp x ≤ 1.

Definition cp01_clamp := (ge0_clamp, le1_clamp).

Lemma clamp_in01 x : 0 ≤ x ≤ 1 → clamp x = x.

Lemma clamp_id x : clamp (clamp x) = clamp x.

Lemma clamp0 : clamp 0 = 0.

Lemma clamp1 : clamp 1 = 1.
End Clamp.

Section StdDefs.
Context {R : realType} (T : choiceType).

Implicit Types (mu : {distr T / R}) (A B E : pred T) (f : T → R).

Definition dinsupp mu := fun x ⇒ mu x != 0 :> R.

Lemma in_dinsupp x (mu : {distr T / R}) :
  (x \in dinsupp mu) = (mu x != 0).

Lemma dinsuppP mu x : reflect (mu x ≠ 0) (x \in dinsupp mu).

Lemma dinsuppPn mu x : reflect (mu x = 0) (x \notin dinsupp mu).

Definition pr mu E := psum (fun x ⇒ (E x)%:R × mu x).
Definition prc mu E A := pr mu [predI E & A] / pr mu A.
Definition esp mu f := sum (fun x ⇒ f x × mu x).
Definition espc mu f A := sum (fun x ⇒ f x × prc mu (pred1 x) A).

Definition has_esp mu f := summable (fun x ⇒ f x × mu x).
End StdDefs.

Notation "\P_[ mu ] E" := (pr mu E).
Notation "\P_[ mu , A ] E" := (prc mu E A).
Notation "\E_[ mu ] f" := (esp mu f).
Notation "\E_[ mu , A ] f" := (espc mu f A).
Notation "\E?_[ mu ] f" := (has_esp mu f).
Notation dweight mu := (\P_[mu] predT).

Section DistrTheory.
Context {R : realType} {T : choiceType} (mu : T → R).

Definition isdistr :=
  (∀ x, 0 ≤ mu x) ∧ (∀ J, uniq J → \sum_(j <- J) mu j ≤ 1).

Hypothesis isd : isdistr.




Definition mkdistr := @Distr R T mu isd1 isd2 isd3.

Lemma mkdistrE : mkdistr =1 mu.

Definition ispredistr {T : choiceType} (mu : T → R) :=
  [/\ ∀ x, 0 ≤ mu x & summable mu].
End DistrTheory.

Lemma isdistr_finP {R : realType} {I : finType} (mu : I → R) :
  (isdistr mu) ↔ (∀ x, 0 ≤ mu x) ∧ (\sum_j mu j ≤ 1).

Lemma le1_mu1
  {R : realType} {T : choiceType} (mu : {distr T / R}) x : mu x ≤ 1.

Section DistrD.
Context {R : realType} {T : choiceType} (mu : T → R).

Definition mnull := fun x : T ⇒ (0 : R).

Lemma isd_mnull : isdistr mnull.

Definition dnull := locked (mkdistr isd_mnull).

Lemma dnullE x : dnull x = 0.

Definition mkdistrd : {distr T / R} :=
  if @idP `[< isdistr mu >] is ReflectT Px then
    mkdistr (asboolW Px)
  else dnull.
End DistrD.

Lemma lef_dnull {R : realType} {T : choiceType} (mu : {distr T / R}) :
  dnull <=1 mu.

Section Restr.
Context {R : realType} {T : choiceType} (p : pred T).

Definition mrestr (mu : {distr T / R}) :=
  fun x ⇒ if p x then mu x else 0.

Lemma isd_mrestr (mu : {distr T / R}) : isdistr (mrestr mu).

Definition drestr (mu : {distr T / R}) := locked (mkdistr (isd_mrestr mu)).

Lemma drestrE (mu : {distr T / R}) x :
  drestr mu x = if p x then mu x else 0.
End Restr.

Section RestrTheory.
Context {R : realType} {T : choiceType}.

Lemma drestrD (mu : {distr T / R}) (p : pred T) x :
  mu x = drestr p mu x + drestr (predC p) mu x.

Lemma dinsupp_restr (mu : {distr T / R}) (p : pred T) x :
  (x \in dinsupp (drestr p mu)) = (x \in dinsupp mu) && p x.
End RestrTheory.

Section DRat.
Context {R : realType} (T : choiceType).


Implicit Types (s : seq T).

Definition mrat (s : seq T) : T → R :=
  fun x : T ⇒ (count (pred1 x) s)%:R / (size s)%:R.

Lemma ge0_mrat s : ∀ x, 0 ≤ mrat s x.




Lemma isd_mrat s : isdistr (mrat s).

Definition drat s := locked (mkdistr (isd_mrat s)).

Lemma drat1E s x :
  drat s x = (count_mem x s)%:R / (size s)%:R.

Definition dunit x := locked (drat [:: x]).
Definition duni s := locked (drat (undup s)).

Lemma dunit1E x y : (dunit x) y = (x == y)%:R.

Lemma duni1E s x : (duni s) x = (x \in s)%:R / (size (undup s))%:R.

Lemma in_dunit t t' :
  t' \in dinsupp (dunit t) → t' = t :> T.
End DRat.

Section Flip.
Context {R : realType}.

Definition mflip (xt : R) :=
  fun b ⇒ if b then clamp xt else 1 - clamp xt.

Lemma isd_mflip xt : isdistr (mflip xt).

Definition dflip (xt : R) := locked (mkdistr (isd_mflip xt)).

Lemma dflip1E xt : dflip xt =1 mflip xt.
End Flip.

Section Std.
Context {R : realType}.


Implicit Types (T U : choiceType).

Section Bind.
Context {T U : choiceType} (f : T → distr U) (mu : distr T).

Definition mlet := fun y : U ⇒ psum (fun x ⇒ mu x × f x y).

Lemma isd_mlet : isdistr mlet.

Definition dlet := locked (mkdistr isd_mlet).

Lemma dletE y : dlet y = psum (fun x ⇒ mu x × f x y).
End Bind.

Notation "\dlet_ ( i <- d ) E" := (dlet (fun i ⇒ E) d).

Definition dlift {A : choiceType} (f : A → {distr A / R}) :=
  fun d ⇒ \dlet_(x <- d) f x.

Definition diter {A : choiceType} n (f : A → {distr A / R}) :=
  fun a ⇒ (iter n (dlift f) (dunit a)).

Section BindTheory.
Variables (T U : choiceType).

Implicit Types (f g : T → distr U) (mu nu : distr T).

Lemma dlet_null f : dlet f dnull =1 dnull.

Lemma dlet_unit f v : \dlet_(y <- dunit v) f y =1 f v.

Lemma dlet_dunit_id mu : \dlet_(t <- mu) (dunit t) =1 mu.

Lemma eq_in_dlet f g mu nu : {in dinsupp mu, f =2 g} → mu =1 nu →
  dlet f mu =1 dlet g nu.

Lemma summable_mu_wgtd (f : T → R) mu :
  (∀ x, 0 ≤ f x ≤ 1) → summable (fun x ⇒ mu x × f x).

Lemma summable_mlet f mu y : summable (fun x : T ⇒ mu x × (f x) y).

Lemma le_in_dlet f g mu : {in dinsupp mu, f <=2 g} →
  dlet f mu <=1 dlet g mu.

Lemma le_mu_dlet f mu nu : mu <=1 nu → dlet f mu <=1 dlet f nu.

Lemma le_dlet f g mu nu :
    mu <=1 nu
  → {in dinsupp mu, ∀ x, f x <=1 g x}
  → \dlet_(x <- mu) f x <=1 \dlet_(x <- nu) g x.

Lemma dletC (mu : {distr T / R}) (nu : {distr U / R}) y :
  (\dlet_(_ <- mu) nu) y = (dweight mu) × (nu y).

Lemma dinsupp_dlet f mu y :
  y \in dinsupp (\dlet_(x <- mu) f x) →
    exists2 x, x \in dinsupp mu & f x y != 0.

Lemma dlet_dinsupp f mu x y :
  x \in dinsupp mu → f x y != 0 → y \in dinsupp (dlet f mu).

Lemma dlet_eq0 (f : T → U) mu y :
  {in dinsupp mu, ∀ x, f x != y} → (\dlet_(x <- mu) dunit (f x)) y = 0.

Lemma eq0_dlet (mu : {distr T / R}) (F : T → {distr U / R}) y :
  (\dlet_(x <- mu) F x) y = 0 → ∀ x, x \in dinsupp mu → F x y = 0.
End BindTheory.

Section DLetDLet.
Context {T U V : choiceType} (f1 : T → distr U) (f2 : U → distr V).

Lemma dlet_dlet (mu : {distr T / R}) :
     \dlet_(x <- \dlet_(y <- mu) f1 y) f2 x
  =1 \dlet_(y <- mu) (\dlet_(x <- f1 y) f2 x).
End DLetDLet.

Section DLetAlg.
Context {T U : choiceType} (mu mu1 mu2 : {distr T / R}).

Lemma dlet_additive (f : T → {distr U / R}) z :
  (∀ x, mu x = mu1 x + mu2 x) → (\dlet_(x <- mu) f x) z =
    (\dlet_(x <- mu1) f x) z + (\dlet_(x <- mu2) f x) z.
End DLetAlg.

Definition mlim T (f : nat → distr T) : T → R :=
  fun x ⇒ fine (nlim (fun n ⇒ f n x)).

Lemma isd_mlim T (f : nat → distr T) : isdistr (mlim f).

Definition dlim T (f : nat → distr T) :=
  locked (mkdistr (isd_mlim f)).

Notation "\dlim_ ( n ) E" := (dlim (fun n ⇒ E)).

Lemma dlimE T (f : nat → distr T) x :
  (\dlim_(n) f n) x = fine (nlim (fun n ⇒ f n x)).

Section DLimTheory.
Variables (T U : choiceType).

Implicit Types (f g : nat → distr T) (h : T → {distr U / R}).
Implicit Types (mu : {distr T / R}).

Lemma dlimC mu : \dlim_(n) mu =1 mu.

Lemma eq_dlim f g : f =2 g → dlim f =1 dlim g.

Lemma eq_from_dlim K f g :
  (∀ n, (K ≤ n)%N → f n =1 g n) → dlim f =1 dlim g.

Definition dlim_bump (mu : nat → {distr T / R}) :
  dlim (fun n ⇒ mu n.+1) =1 dlim mu.

Definition dlim_lift (mu : nat → {distr T / R}) p :
  dlim (fun n ⇒ mu (n + p)%N) =1 dlim mu.

Definition dcvg {T : choiceType} (f : nat → {distr T / R}) :=
  ∀ x, ∃ l, ncvg (fun n ⇒ f n x) l.

Definition ducvg {T : choiceType} (f : nat → {distr T / R}) :=
  ∃ l, ∀ x, ncvg (fun n ⇒ f n x) l.

CoInductive dlim_spec f (x : T) : R → Type :=
| DLimCvg : ∀ l : R, 0 ≤ l → l ≤ 1 →
    ncvg (fun n ⇒ f n x) l%:E → dlim_spec f x l

| DLimOut : ¬ (∃ l : \bar R, ncvg (fun n ⇒ f n x) l) →
    dlim_spec f x 0.

Lemma dlimP f x : dlim_spec f x (dlim f x).

Lemma dcvgP f : dcvg f → ∀ x,
  exists2 l, (0 ≤ l ≤ 1) & ncvg (f^~ x) l%:E.

Lemma dcvg_homo f :
  (∀ n m, (n ≤ m)%N → f n <=1 f m) → dcvg f.

Lemma ge0_dlim f : ∀ x, 0 ≤ dlim f x.

Lemma le1_dlim f : ∀ x, dlim f x ≤ 1.

Lemma le_dlim f g :
  (∀ n, f n <=1 g n) → dcvg g → dlim f <=1 dlim g.

Lemma leub_dlim f mu : (∀ n, f n <=1 mu) → dlim f <=1 mu.

Lemma dlim_ub f k :
  (∀ n m, (n ≤ m)%N → f n <=1 f m) → f k <=1 dlim f.

Lemma dlet_lim f h : (∀ n m, (n ≤ m)%N → f n <=1 f m) →
  \dlet_(x <- dlim f) h x =1 \dlim_(n) \dlet_(x <- f n) h x.

Lemma dlim_let (f : nat → T → {distr U / R}) (mu : {distr T / R}) :
  (∀ x n m, (n ≤ m)%N → f n x <=1 f m x) →
  \dlim_(n) \dlet_(x <- mu) (f n x) =1
    \dlet_(x <- mu) \dlim_(n) (f n x).
End DLimTheory.

Section Marginals.
Variable (T U : choiceType) (h : T → U) (mu : distr T).

Definition dmargin := \dlet_(x <- mu) (dunit (h x)).

Lemma dmarginE : dmargin = \dlet_(y <- mu) (dunit (h y)).
End Marginals.

Section MarginalsTh.
Variable (T U V : choiceType).

Lemma dmargin_psumE (mu : {distr T / R}) (f : T → U) y :
  (dmargin f mu) y = psum (fun x ⇒ (f x == y)%:R × mu x).

Lemma dlet_dmargin (mu : {distr T / R}) (f : T → U) (g : U → {distr V / R}):
  \dlet_(u <- dmargin f mu) g u =1 \dlet_(t <- mu) (g (f t)).

Lemma dmargin_dlet (mu : {distr T / R}) (f : U → V) (g : T → {distr U / R}):
  dmargin f (\dlet_(t <- mu) g t) =1 \dlet_(t <- mu) (dmargin f (g t)).

Lemma dmargin_dunit (t : T) (f : T → U):
  dmargin f (dunit t) =1 dunit (f t) :> {distr U / R}.
End MarginalsTh.
End Std.

Notation dfst mu := (dmargin fst mu).
Notation dsnd mu := (dmargin snd mu).

Notation "\dlet_ ( i <- d ) E" := (dlet (fun i ⇒ E) d).
Notation "\dlim_ ( n ) E" := (dlim (fun n ⇒ E)).

Section DSwap.
Context {R : realType} {A B : choiceType} (mu : {distr (A × B) / R}).

Definition dswap : {distr (B × A) / R} :=
  dmargin (fun xy ⇒ (xy.2, xy.1)) mu.
End DSwap.

Section DSwapCoreTheory.
Context {R : realType} {A B : choiceType} (mu : {distr (A × B) / R}).

Lemma dswapE xy : dswap mu xy = mu (xy.2, xy.1).
End DSwapCoreTheory.

Section DSwapTheory.
Context {R : realType} {A B : choiceType} (mu : {distr (A × B) / R}).

Lemma dswapK : dswap (dswap mu) =1 mu.

Lemma dinsupp_swap xy : (xy.2, xy.1) \in dinsupp mu →
  xy \in dinsupp (dswap mu).

Lemma dfst_dswap : dfst (dswap mu) =1 dsnd mu.

Lemma dsnd_dswap : dsnd (dswap mu) =1 dfst mu.
End DSwapTheory.

Section DFst.
Context {R : realType} {T U : choiceType}.

Lemma dfstE (mu : {distr (T × U) / R}) x :
  dfst mu x = psum (fun y ⇒ mu (x, y)).

Lemma summable_fst (mu : {distr (T × U) / R}) x :
  summable (fun y ⇒ mu (x, y)).
End DFst.

Section DSnd.
Context {R : realType} {T U : choiceType}.

Lemma dsndE (mu : {distr (T × U) / R}) y :
  dsnd mu y = psum (fun x ⇒ mu (x, y)).

Lemma summable_snd (mu : {distr (T × U) / R}) y :
  summable (fun x ⇒ mu (x, y)).
End DSnd.

Section PrCoreTheory.
Context {R : realType} {T : choiceType}.

Implicit Types (mu : {distr T / R}) (A B E : pred T).

Lemma summable_pr E mu : summable (fun x ⇒ (E x)%:R × mu x).

Lemma pr_pred0 mu : \P_[mu] pred0 = 0.

Lemma pr_pred1 mu x : mu x = \P_[mu] (pred1 x).

Lemma pr_exp mu (E : pred T) : \P_[mu] E = \E_[mu] (fun m ⇒ (E m)%:R).

Lemma pr_predT mu : \P_[mu] predT = psum mu.

Lemma pr_dunit E x : \P_[dunit x] E = (E x)%:R :> R.

Lemma exp_dunit (f : T → R) (x : T) : \E_[dunit x] f = f x.

Lemma exp_cst mu r : \E_[mu] (fun _ ⇒ r) = \P_[mu] predT × r.

Lemma exp0 mu : \E_[mu] (fun _ ⇒ 0) = 0.

Lemma has_expC mu c : \E?_[mu] (fun _ ⇒ c).

Lemma has_exp0 mu : \E?_[mu] (fun _ ⇒ 0).

Lemma has_exp1 mu : \E?_[mu] (fun _ ⇒ 1).

Lemma has_expZ mu c F : \E?_[mu] F → \E?_[mu] (c \*o F).

Lemma expZ mu F c : \E_[mu] (c \*o F) = c × \E_[mu] F.

Lemma ge0_pr A mu : 0 ≤ \P_[mu] A.

Lemma ge0_prc A B mu : 0 ≤ \P_[mu, B] A.

Lemma eq_in_pr A B mu :
  {in dinsupp mu, A =i B} → \P_[mu] A = \P_[mu] B.

Lemma eq_pr A B mu : A =i B → \P_[mu] A = \P_[mu] B.

Lemma eq_exp mu (f1 f2 : T → R):
   {in dinsupp mu, f1 =1 f2} → \E_[mu] f1 = \E_[mu] f2.

Lemma pr_pred0_eq (mu : {distr T / R}) (E : pred T) :
  E =1 pred0 → \P_[mu] E = 0.
End PrCoreTheory.

Section PrTheory.
Context {R : realType} {T U : choiceType} {I : eqType}.

Implicit Types (mu : {distr T / R}) (A B E : pred T).

Lemma pr_dlet E f (mu : {distr U / R}) :
  \P_[dlet f mu] E = \E_[mu] (fun x ⇒ \P_[f x] E).

Lemma pr_dmargin E f (mu : {distr U / R}) :
  \P_[dmargin f mu] E = \P_[mu] [pred x | f x \in E].

Lemma eq0_pr A mu :
  (∀ x, x \in dinsupp mu → x \notin A) → \P_[mu] A = 0.

Lemma eq0_prc A B mu :
    (∀ x, x \in dinsupp mu → x \in B → x \notin A)
  → \P_[mu, B] A = 0.

Lemma subset_pr A B mu : {subset B ≤ A} → \P_[mu] B ≤ \P_[mu] A.

Lemma le1_pr A mu : \P_[mu] A ≤ 1.

Lemma le_exp mu f1 f2: \E?_[mu] f1 → \E?_[mu] f2 →
  f1 <=1 f2 → \E_[mu] f1 ≤ \E_[mu] f2.

Lemma le_in_pr E1 E2 mu :
  (∀ x, x \in dinsupp mu → x \in E1 → x \in E2) →
    \P_[mu] E1 ≤ \P_[mu] E2.

Lemma le_mu_pr A mu nu :
    (∀ x, x \in dinsupp nu → x \in A → nu x ≤ mu x)
  → \P_[nu] A ≤ \P_[mu] A.

Lemma le1_prc A B mu : \P_[mu, B] A ≤ 1.

Lemma prc_sum A mu : 0 < \P_[mu] A →
  psum (fun x ⇒ \P_[mu, A] (pred1 x)) = 1.

Lemma pr_eq0 mu E : \P_[mu] E = 0 → ∀ x, x \in E → mu x = 0.

Lemma prID A B mu :
  \P_[mu] A = \P_[mu] [predI A & B] + \P_[mu] [predI A & predC B].

Lemma pr_or_indep (A B : pred T) (mu : {distr T / R}) :
  (∀ x, x \in A → x \notin B) →
    \P_[mu] [predU A & B] = \P_[mu] A + \P_[mu] B.

Lemma pr_mem_map f mu (r : seq U) : uniq r →
    \P_[mu] [pred x | f x \in r]
  = \sum_(y <- r) \P_[mu] [pred x | f x == y].

Lemma pr_mem mu (r : seq T) : uniq r →
  \P_[mu] [pred x | x \in r] = \sum_(x <- r) mu x.

Lemma pr_bigor_indep mu (P : I → pred T) (r : seq I) :
    uniq r
  → (∀ p1 p2 x, p1 != p2 → p1 \in r → p2 \in r → x \in P p1 → x \notin P p2)
  → \P_[mu] [pred x | has [pred p | x \in P p] r]
  = \sum_(p <- r) \P_[mu] (P p).

Lemma pr_or A B mu : \P_[mu] [predU A & B] =
  \P_[mu] A + \P_[mu] B - \P_[mu] [predI A & B].

Lemma pr_and A B mu : \P_[mu] [predI A & B] =
  \P_[mu] A + \P_[mu] B - \P_[mu] [predU A & B].

Lemma ler_pr_or A B mu :
  \P_[mu] [predU A & B] ≤ \P_[mu] A + \P_[mu] B.

Lemma ler_pr_and A B mu :
  \P_[mu] [predI A & B] ≤ \P_[mu] A + \P_[mu] B.

Lemma pr_predC E mu: \P_[mu](predC E) = \P_[mu] predT - \P_[mu] E.

Lemma pr_split B A mu : \P_[mu] A =
    \P_[mu] B × \P_[mu, B] A
  + \P_[mu] (predC B) × \P_[mu, predC B] A.

Lemma exp_split A f mu : \E?_[mu] f → \E_[mu] f =
    \P_[mu] A × \E_[mu, A] f
  + \P_[mu] (predC A) × \E_[mu, predC A] f.

Lemma has_esp_bounded f mu :
  (∃ M, ∀ x, `|f x| < M) → \E?_[mu] f.

Lemma bounded_has_exp mu F :
  (∃ M, ∀ x, `|F x| ≤ M) → \E?_[mu] F.

Lemma summable_has_exp mu F : summable F → \E?_[mu] F.

Lemma exp_le_bd mu F (M : R) :
  0 ≤ M → (∀ x, `|F x| ≤ M) → \E_[mu] F ≤ M.

Lemma exp_dlet mu (nu : T → {distr U / R}) F :
  (∀ eta, \E?_[eta] F) →
    \E_[dlet nu mu] F = \E_[mu] (fun x ⇒ \E_[nu x] F).
End PrTheory.