Library discprob.prob.prob

Require Import Reals Psatz.
From mathcomp Require Import ssreflect ssrbool ssrfun eqtype choice fintype bigop.
From discprob.basic Require Export base Series_Ext order bigop_ext sval Reals_ext.
From discprob.prob Require Import countable double rearrange.
From Coquelicot Require Export Rcomplements Rbar Series Lim_seq Hierarchy Markov.

Ltac to_lim_seq :=
  match goal with
    | [ |- filterlim ?f eventually (locally ?n) ] ⇒ change (is_lim_seq f n)
    | [ |- real (Lim_seq ?u) = ?l ] ⇒ rewrite (is_lim_seq_unique u l); first done
  end.

Record distrib (A: countType) := mkDistrib {
  pmf :> A R;
  pmf_pos : a, pmf a 0;
  pmf_sum1 : is_series (countable_sum pmf) 1
}.


Lemma pmf_sum1_Series {A} (X: distrib A) : Series (countable_sum X) = 1.

Lemma pmf_sum_sum_n {A} (X: distrib A) n :
  sum_n (countable_sum X) n 1.

Definition pr {A: countType} (d: distrib A) (P: A bool) :=
  Series (countable_sum (λ a, if P a then d a else 0)).

Hint Resolve pmf_pos pmf_sum1 : prob.

Lemma pr_ex_series {A: countType} (d: distrib A) (P: A bool):
  ex_series (countable_sum (λ a : A, if P a then d a else 0)).

Hint Resolve pr_ex_series : prob.

Lemma ge_pr_0 {A: countType} (d: distrib A) P:
  pr d P 0.

Lemma pr_le_1 {A: countType} (d: distrib A) P:
  pr d P 1.

Lemma pr_sum_n {A} (d: distrib A) (P: A bool) n :
  sum_n (countable_sum (λ a, if P a then d a else 0)) n 1.

Lemma pr_xpredT {A: countType} (d: distrib A):
  pr d xpredT = 1.

Lemma pr_xpred0 {A: countType} (d: distrib A):
  pr d xpred0 = 0.

Lemma pr_eq_pred {A: countType} (d: distrib A) P Q:
  P =i Q
  pr d P = pr d Q.

Lemma pr_eq_pred' {A: countType} (d: distrib A) (P Q: pred A):
  ( i, P i Q i)
  pr d P = pr d Q.

Lemma pr_mono_pred {A: countType} (d: distrib A) (P Q: pred A):
  ( i, P i Q i)
  pr d P pr d Q.

Lemma pr_union {A: countType} (d: distrib A) P Q:
  pr d (predU P Q) =
  pr d P + pr d Q - pr d (predI P Q).

Lemma pr_union_le {A: countType} (d: distrib A) P Q:
  pr d (predU P Q)
  pr d P + pr d Q.

Lemma distrib_exists_support {A: countType} (d: distrib A) :
   a, d a > 0.

Record rvar {A} (Ω: distrib A) (B: eqType) := mkRvar { rvar_fun :> A B; }.


Definition rvar_const {A} {B : eqType} (Ω: distrib A) (b: B) :=
  mkRvar Ω (λ x, b).


Section product.
Variable (A B: countType).
Variable (dist1: distrib A) (dist2: distrib B).
Definition prod_pmf (ab: A × B) := (dist1 (fst ab)) × (dist2 (snd ab)).

Lemma prod_pmf_pos ab: prod_pmf ab 0.

Lemma is_series_chain s1 s2 v: is_series s2 (Series s1) is_series s1 v is_series s2 v.

Section double.
Variable P: pred (A × B).

Definition σprod :=
  λ n, match @pickle_inv [countType of (A × B)] n with
       | Some (a, b)(S (pickle a), S (pickle b))
       | None(O, O)
       end.

Definition aprod :=
  λ mn, match mn with
        | (S m, S n)
          match (@pickle_inv A m), (@pickle_inv B n) with
            | Some a, Some b
              if P (a, b) then
                prod_pmf (a, b)
              else
                0
            | _, _ ⇒ 0
          end
        | _ ⇒ 0
        end.

Lemma aprod_double_summable: double_summable aprod.

Lemma σprod_inj: n n', aprod (σprod n) 0 σprod n = σprod n' n = n'.

Lemma σprod_cov: n, aprod n 0 m, σprod m = n.

Lemma is_series_prod_row:
  is_series (countable_sum (λ ab, if P ab then prod_pmf ab else 0))
            (Series (λ j, Series (λ k, aprod (S j, S k)))).

Lemma Series_prod_row:
  Series (countable_sum (λ ab, if P ab then prod_pmf ab else 0)) =
  (Series (λ j, Series (λ k, aprod (S j, S k)))).

End double.

Lemma prod_pmf_pos_sum1:
  is_series (countable_sum (prod_pmf)) 1.

Definition distrib_prod := mkDistrib _ prod_pmf prod_pmf_pos prod_pmf_pos_sum1.

End product.


Definition rrvar {A} (Ω : distrib A) := @rvar A Ω [eqType of R].

Definition pr_eq {A} {B: eqType} {Ω : distrib A} (X: rvar Ω B) (b: B) :=
  pr Ω (λ a, X a == b).
Definition pr_gt {A} {Ω: distrib A} (X: rrvar Ω) (b: R) :=
  pr Ω (λ a, Rgt_dec (X a) b).
Definition pr_ge {A} {Ω: distrib A} (X: rrvar Ω) (b: R) :=
  pr Ω (λ a, Rge_dec (X a) b).

Definition rvar_dist {A} {B} {Ω : distrib A} (X: rvar Ω B) := Ω.

Lemma pr_gt_contra {A} (Ω: distrib A) (X: rrvar Ω) (r1 r2: R):
  r1 r2
  pr_gt X r1 pr_gt X r2.

Lemma pr_eq_rvar_ext {A} {B: eqType} {Ω: distrib A} (X X': rvar Ω B) (b: B):
  X =1 X'
  pr_eq X b = pr_eq X' b.

Lemma pr_eq_ge_0 {A} {B: eqType} {Ω: distrib A} (X: rvar Ω B) (b: B) :
  pr_eq X b 0.

Lemma pr_eq_le_1 {A} {B: eqType} {Ω: distrib A} (X: rvar Ω B) (b: B) :
  pr_eq X b 1.


Section Ex.
Definition Ex {A} {Ω: distrib A} (X: rrvar Ω) : R :=
  Series (countable_sum (λ r: imgT X, (pr_eq X (sval r) × (sval r)))).

Definition is_Ex {A} {Ω : distrib A} (X: rrvar Ω) (v: R) :=
  ex_series (countable_sum (λ r: imgT X, (pr_eq X (sval r) × Rabs (sval r))))
  is_series (countable_sum (λ r: imgT X, (pr_eq X (sval r) × sval r))) v.

Definition ex_Ex {A} {Ω : distrib A} (X: rrvar Ω) :=
  ex_series (countable_sum (λ r: imgT X, (pr_eq X (sval r) × Rabs (sval r)))).

Lemma Ex_correct {A} {Ω : distrib A} (X: rrvar Ω) :
  ex_Ex X is_Ex X (Ex X).

Lemma is_Ex_unique {A} {Ω: distrib A} (X: rrvar Ω) (v : R) :
  is_Ex X v Ex X = v.

Lemma ex_Ex' {A} {Ω : distrib A} (X: rrvar Ω) :
  ex_Ex X
  ex_series (Rabs \o countable_sum (λ r: imgT X, (pr_eq X (sval r) × (sval r)))).

Lemma ex_Ex_nonabs {A} {Ω : distrib A} (X: rrvar Ω) :
  ex_Ex X
  ex_series (countable_sum (λ r: imgT X, (pr_eq X (sval r) × (sval r)))).

Hint Resolve ex_Ex' ex_Ex_nonabs.

Definition rvar_comp {A: countType} {B C: eqType} {Ω : distrib A} (X: rvar Ω B) (f: B C)
  : rvar Ω C :=
  mkRvar (rvar_dist X) (f \o X).

Section Ex_pt_gen.
Variable (A: countType).
Variable (B: eqType).
Variable (Ω: distrib A).
Variable (X: rvar Ω B).
Variable (f: B R).

Definition σpt : nat nat × nat.

Definition apt :=
  λ mn, match mn with
        | (S m, S n)
          match (@pickle_inv [countType of imgT X] m), (@pickle_inv A n) with
            | Some v, Some a(if (X a) == (sval v) then rvar_dist X a else 0) × f (sval v)
            | _, _ ⇒ 0
          end
        | _ ⇒ 0
        end.

Lemma Series_correct_ext a a':
  ( n, a n = a' n) is_series a (Series a) is_series a' (Series a').

Lemma Series_correct' a v:
  Series a = v ex_series a is_series a v.

Lemma apt_double_summable_by_row
      (EX: ex_series (countable_sum (λ r: imgT X, (pr_eq X (sval r) × Rabs (f (sval r))))))
  : double_summable apt.

Lemma apt_double_summable_by_column
 (EX_pt: ex_series (countable_sum (λ a, Rabs (f (X a)) × (rvar_dist X) a))):
 double_summable apt.

Lemma σpt_inj: n n', apt (σpt n) 0 σpt n = σpt n' n = n'.

Lemma σpt_cov: n, apt n 0 m, σpt m = n.

Lemma apt_row_rewrite j:
  Series (λ k : nat, apt (S j, S k)) = countable_sum (λ r : imgT X, pr_eq X (sval r) × f (sval r)) j.

Lemma apt_row_rewrite_abs j:
  Series (λ k : nat, Rabs (apt (S j, S k)))
  = countable_sum (λ r : imgT X, pr_eq X (sval r) × Rabs (f (sval r))) j.

Lemma is_series_Ex_pt_f
  (EX: ex_series (countable_sum (λ r: imgT X, (pr_eq X (sval r) × Rabs (f (sval r)))))):
  is_series (countable_sum (λ a : A, f (X a) × (rvar_dist X) a))
    (Series (countable_sum (λ a : A, f (X a) × (rvar_dist X) a))).

Lemma is_series_Ex_pt_f_abs
  (EX: ex_series (countable_sum (λ r: imgT X, (pr_eq X (sval r) × Rabs (f (sval r)))))):
  is_series (countable_sum (λ a : A, Rabs (f (X a)) × (rvar_dist X) a))
    (Series (countable_sum (λ a : A, Rabs (f (X a)) × (rvar_dist X) a))).

Lemma is_series_Ex_pt_f'
  (EX: ex_series (countable_sum (λ r: imgT X, (pr_eq X (sval r) × Rabs (f (sval r)))))):
  is_series (countable_sum (λ r: imgT X, (pr_eq X (sval r) × f (sval r))))
            (Series (countable_sum (λ a, f (X a) × (rvar_dist X) a))).

Lemma is_series_Ex_by_pt_f':
  ex_series (countable_sum (λ a, Rabs (f (X a)) × (rvar_dist X) a))
  is_series (countable_sum (λ r: imgT X, (pr_eq X (sval r) × f (sval r))))
            (Series (countable_sum (λ a, f (X a) × (rvar_dist X) a))).

Lemma is_series_Ex_abs_by_pt_f':
  ex_series (countable_sum (λ a, Rabs (f (X a)) × (rvar_dist X) a))
  is_series (countable_sum (λ r: imgT X, pr_eq X (sval r) × Rabs (f (sval r))))
            (Series (countable_sum (λ a, Rabs (f (X a)) × (rvar_dist X) a))).

Lemma is_series_Ex_by_pt_f:
  ex_series (countable_sum (λ a, Rabs (f (X a)) × (rvar_dist X) a))
  is_series (countable_sum (λ r: imgT X, (pr_eq X (sval r) × f (sval r))))
            (Series (countable_sum (λ r: imgT X, (pr_eq X (sval r) × f (sval r))))).

Lemma ex_Ex_pt_f
  (EX: ex_series (countable_sum (λ r: imgT X, (pr_eq X (sval r) × Rabs (f (sval r)))))):
  ex_series (countable_sum (λ a, f (X a) × (rvar_dist X) a)).

Lemma ex_Ex_pt_f_abs
  (EX: ex_series (countable_sum (λ r: imgT X, (pr_eq X (sval r) × Rabs (f (sval r)))))):
  ex_series (countable_sum (λ a, Rabs (f (X a)) × (rvar_dist X) a)).

Lemma Ex_pt_f
  (EX: ex_series (countable_sum (λ r: imgT X, (pr_eq X (sval r) × Rabs (f (sval r)))))):
  Series (countable_sum (λ r: imgT X, (pr_eq X (sval r) × (f (sval r)))))
  = Series (countable_sum (λ a, f (X a) × (rvar_dist X) a)).

Lemma ex_Ex_by_pt_f:
  ex_series (countable_sum (λ a, Rabs (f (X a)) × (rvar_dist X) a))
  ex_series (countable_sum (λ r: imgT X, (pr_eq X (sval r) × Rabs (f (sval r))))).

Lemma Ex_pt_by_column_f:
  ex_series (countable_sum (λ a, Rabs (f (X a)) × (rvar_dist X) a))
  Series (countable_sum (λ r: imgT X, (pr_eq X (sval r) × (f (sval r)))))
  = Series (countable_sum (λ a, f (X a) × (rvar_dist X) a)).

End Ex_pt_gen.

Section Ex_pt.
Variable (A: countType).
Variable (Ω: distrib A).
Variable (X: rrvar Ω).
Lemma is_series_Ex_pt:
  ex_Ex X
  is_series (countable_sum (λ a : A, (X a) × (rvar_dist X) a))
    (Series (countable_sum (λ a : A, X a × (rvar_dist X) a))).

Lemma is_series_Ex_pt_abs:
  ex_Ex X
  is_series (countable_sum (λ a : A, Rabs (X a) × (rvar_dist X) a))
    (Series (countable_sum (λ a : A, Rabs (X a) × (rvar_dist X) a))).

Lemma is_series_Ex_pt':
  ex_Ex X
  is_series (countable_sum (λ r: imgT X, (pr_eq X (sval r) × sval r)))
            (Series (countable_sum (λ a, X a × (rvar_dist X) a))).

Lemma is_series_Ex_by_pt':
  ex_series (countable_sum (λ a, Rabs (X a) × (rvar_dist X) a))
  is_series (countable_sum (λ r: imgT X, (pr_eq X (sval r) × sval r)))
            (Series (countable_sum (λ a, X a × (rvar_dist X) a))).

Lemma is_series_Ex_abs_by_pt':
  ex_series (countable_sum (λ a, Rabs (X a) × (rvar_dist X) a))
  is_series (countable_sum (λ r: imgT X, pr_eq X (sval r) × Rabs (sval r)))
            (Series (countable_sum (λ a, Rabs (X a) × (rvar_dist X) a))).

Lemma is_series_Ex_by_pt:
  ex_series (countable_sum (λ a, Rabs (X a) × (rvar_dist X) a))
  is_series (countable_sum (λ r: imgT X, (pr_eq X (sval r) × sval r)))
            (Series (countable_sum (λ r: imgT X, (pr_eq X (sval r) × sval r)))).

Lemma ex_Ex_pt:
  ex_Ex X
  ex_series (countable_sum (λ a, X a × (rvar_dist X) a)).

Lemma ex_Ex_pt_abs:
  ex_Ex X
  ex_series (countable_sum (λ a, Rabs (X a) × (rvar_dist X) a)).

Lemma Ex_pt:
  ex_Ex X
  Ex X = Series (countable_sum (λ a, X a × (rvar_dist X) a)).

Lemma ex_Ex_by_pt:
  ex_series (countable_sum (λ a, Rabs (X a) × (rvar_dist X) a))
  ex_Ex X.

Lemma Ex_pt_by_column:
  ex_series (countable_sum (λ a, Rabs (X a) × (rvar_dist X) a))
  Ex X = Series (countable_sum (λ a, X a × (rvar_dist X) a)).
End Ex_pt.

Lemma ex_Ex_ext {A} {Ω: distrib A} (X Y: rrvar Ω) :
  ex_Ex X
  ( a, X a = Y a)
  ex_Ex Y.

Lemma Ex_ext {A} {Ω: distrib A} (X Y: rrvar Ω) :
  ex_Ex X
  ( a, X a = Y a)
  Ex X = Ex Y.

Lemma Ex_ext_le {A} {Ω: distrib A} (X Y: rrvar Ω) :
  ex_Ex X
  ex_Ex Y
  ( a, X a Y a)
  Ex X Ex Y.

End Ex.

Hint Resolve ex_Ex_pt.

Definition rvar_pair {A1 A2: countType} {B1 B2: eqType} {Ω1 : distrib A1} {Ω2 : distrib A2}
           (X: rvar Ω1 B1) (Y: rvar Ω2 B2) :=
  mkRvar (distrib_prod Ω1 Ω2) (λ a, (X (fst a), Y (snd a))).

Lemma pair_joint_pred {A1 A2: countType} (Ω1 : distrib A1) (Ω2 : distrib A2) P Q:
  pr (distrib_prod Ω1 Ω2) (λ x, P (fst x) && Q (snd x)) =
  pr Ω1 P × pr Ω2 Q.

Lemma pair_marginal1 {A1 A2: countType} (Ω1 : distrib A1) (Ω2 : distrib A2) P :
  pr (distrib_prod Ω1 Ω2) (λ x, P (fst x)) =
  pr Ω1 P.

Lemma pair_marginal2 {A1 A2: countType} {Ω1 : distrib A1} {Ω2 : distrib A2} P:
  pr (distrib_prod Ω1 Ω2) (λ x, P (snd x)) =
  pr Ω2 P.

Lemma rvar_pair_marginal1 {A1 A2: countType} {B1 B2: eqType} {Ω1 : distrib A1} {Ω2 : distrib A2}
           (X: rvar Ω1 B1) (Y: rvar Ω2 B2) P :
  pr (rvar_dist (rvar_pair X Y)) (λ x, P (fst x)) =
  pr (rvar_dist X) P.

Lemma rvar_pair_marginal2 {A1 A2: countType} {B1 B2: eqType} {Ω1 : distrib A1} {Ω2 : distrib A2}
           (X: rvar Ω1 B1) (Y: rvar Ω2 B2) P :
  pr (rvar_dist (rvar_pair X Y)) (λ x, P (snd x)) =
  pr (rvar_dist Y) P.

Lemma rvar_comp_assoc {A: countType} {Ω: distrib A} {B C D: eqType}
      (T: rvar Ω B) (f: B C) (g: C D):
  rvar_comp (rvar_comp T f) g = rvar_comp T (g \o f).

Lemma rvar_comp_comp {A: countType} {Ω: distrib A} {B C D: eqType}
      (T: rvar Ω B) (f: B C) (g: C D):
  rvar_comp (rvar_comp T f) g = rvar_comp T (g \o f).

Module Ex_comp_pt.
Section Ex_comp_pt.

Variable (A: countType).
Variable (B: eqType).
Variable (Ω: distrib A).
Variable (X: rvar Ω B).
Variable (f: B R).
Variable (EX: ex_Ex (rvar_comp X f)).

Definition σcomppt : nat nat × nat.

Definition acomppt :=
  λ mn, match mn with
        | (S m, S n)
          match (@pickle_inv [countType of imgT X] m), (@pickle_inv A n) with
            | Some v, Some a(if (X a) == (sval v) then rvar_dist X a else 0) × f (sval v)
            | _, _ ⇒ 0
          end
        | _ ⇒ 0
        end.

Lemma σcomppt_inj: n n', acomppt (σcomppt n) 0 σcomppt n = σcomppt n' n = n'.

Lemma σcomppt_cov: n, acomppt n 0 m, σcomppt m = n.

Lemma acomppt_row_rewrite j:
  Series (λ k : nat, acomppt (S j, S k)) = countable_sum (λ r : imgT X, pr_eq X (sval r) × f (sval r)) j.

Lemma acomppt_row_rewrite_abs j:
  Series (λ k : nat, Rabs (acomppt (S j, S k)))
  = countable_sum (λ r : imgT X, pr_eq X (sval r) × Rabs (f (sval r))) j.

Lemma is_series_Ex_comp_pt:
  is_series (countable_sum (λ r: imgT X, (pr_eq X (sval r) × f (sval r))))
            (Series (countable_sum (λ a, f (X a) × (rvar_dist X) a))).

Lemma is_series_Ex_comp_pt_abs:
  is_series (countable_sum (λ r: imgT X, (pr_eq X (sval r) × Rabs (f (sval r)))))
            (Series (countable_sum (λ a, Rabs (f (X a)) × (rvar_dist X) a))).

Lemma ex_series_Ex_comp_pt:
  ex_series (countable_sum (λ r: imgT X, (pr_eq X (sval r) × f (sval r)))).

Lemma ex_series_Ex_comp_pt_abs:
  ex_series (countable_sum (λ r: imgT X, (pr_eq X (sval r) × Rabs (f (sval r))))).

Lemma is_series_Ex_comp_pt' v:
  (Series (countable_sum (λ a, f (X a) × (rvar_dist X) a))) = v
  is_series (countable_sum (λ r: imgT X, (pr_eq X (sval r) × f (sval r)))) v.

Lemma Ex_comp:
  Ex (rvar_comp X f) = Series (countable_sum (λ r: imgT X, (pr_eq X (sval r) × f (sval r)))).
End Ex_comp_pt.
End Ex_comp_pt.

Lemma ex_Ex_comp_by_alt {A} {B} {Ω: distrib A} (X: rvar Ω B) f:
  ex_series (countable_sum (λ r : imgT X, pr_eq X (sval r) × Rabs (f (sval r))))
  ex_Ex (rvar_comp X f).

Lemma Ex_comp_by_alt {A} {B} {Ω: distrib A} (X: rvar Ω B) f:
  ex_series (countable_sum (λ r : imgT X, pr_eq X (sval r) × Rabs (f (sval r))))
  Ex (rvar_comp X f) = Series (countable_sum (λ r : imgT X, pr_eq X (sval r) × f (sval r))).

Section Ex_sum.

Variable (A: countType).
Variable (Ω: distrib A).
Variable (X Y: rrvar Ω).
Variable (EX_X: ex_Ex X).
Variable (EX_Y: ex_Ex Y).

Lemma ex_Ex_sum :
  ex_Ex (mkRvar Ω (λ x, X x + Y x)).

Lemma Ex_sum :
  Ex (mkRvar Ω (λ x, X x + Y x)) = Ex X + Ex Y.

Lemma is_Ex_sum :
  is_Ex (mkRvar Ω (λ x, X x + Y x)) (Ex X + Ex Y).

End Ex_sum.

Section Ex_prop.

Variable (A: countType).
Variable (Ω: distrib A).
Variable (X: rrvar Ω).
Variable (EX_X: ex_Ex X).

Lemma ex_Ex_const (x: R):
  ex_Ex (rvar_const Ω x).

Hint Resolve ex_Ex_const.

Lemma Ex_const (x: R):
  Ex (rvar_const Ω x) = x.

Lemma is_Ex_const (x: R):
  is_Ex (rvar_const Ω x) x.

Lemma Ex_const_ge c :
  ( a, c X a)
  c Ex X.

Lemma ex_Ex_plus_l (x: R):
  ex_Ex (rvar_comp X (Rplus x)).

Lemma Ex_plus_l (x: R):
  Ex (rvar_comp X (Rplus x)) = x + Ex X.

Lemma ex_Ex_plus_r (x: R):
  ex_Ex (rvar_comp X (λ y, y + x)).

Lemma ex_Ex_plus_r' (x: R):
  ex_Ex (mkRvar Ω (λ a, X a + x)).

Lemma Ex_plus_r (x: R):
  Ex (rvar_comp X (λ y, y + x)) = Ex X + x.

Lemma ex_Ex_scal_l (c: R):
  ex_Ex (rvar_comp X (λ x, c × x)).

Lemma Ex_scal_l (c: R):
  Ex (rvar_comp X (λ x, c × x)) = c × Ex X.

Lemma ex_Ex_scal_r (c: R):
  ex_Ex (rvar_comp X (λ x, x × c)).

Lemma Ex_scal_r (c: R):
  Ex (rvar_comp X (λ x, x × c)) = Ex X × c.

Lemma ex_Ex_sum_inv (Y: rrvar Ω):
  ex_Ex (mkRvar Ω (λ a, X a + Y a))
  ex_Ex Y.

Lemma ex_Ex_le (Y: rrvar Ω) :
  ( i, Rabs (Y i) Rabs (X i)) ex_Ex Y.

Lemma ex_Ex_comp_abs:
  ex_Ex (rvar_comp X Rabs).

End Ex_prop.


Definition id_rvar {A: countType} (Ω: distrib A) : rvar Ω A :=
  mkRvar Ω (λ a, a).
Definition indicator {A: countType} {B: eqType} {Ω : distrib A} (X: rvar Ω B) (P: B bool) :=
  rvar_comp X (λ a, if P a then 1 else 0).

Lemma ex_Ex_indicator {A} {B: eqType} {Ω : distrib A} (X: rvar Ω B) P:
  ex_Ex (indicator X P).

Lemma Ex_indicator {A} {B: eqType} {Ω : distrib A} (X: rvar Ω B) P:
  Ex (indicator X P) = pr_eq (indicator X P) 1.

Lemma Ex_indicator_le_1 {A} {B : eqType} {Ω: distrib A} (X : rvar Ω B) P:
  Ex (indicator X P) 1.

Lemma ex_Ex_abs {A} {Ω: distrib A} (X: rrvar Ω) :
  ex_Ex X ex_Ex (rvar_comp X (λ x, Rabs x)).

Lemma ex_Ex_pow_nat_le {A} {Ω: distrib A} (X: rrvar Ω) (r s: nat) :
  (0 < r)%nat (r s)%nat
  ex_Ex (rvar_comp X (λ x, (Rabs x)^s))
  ex_Ex (rvar_comp X (λ x, (Rabs x)^r)).

Definition Var {A} {Ω: distrib A} (X: rrvar Ω) :=
  Ex (rvar_comp X (λ x, (x - Ex X)^2)).

Definition ex_Var {A} {Ω: distrib A} (X: rrvar Ω) :=
  ex_Ex X ex_Ex (rvar_comp X (λ x, (x - Ex X)^2)).

Lemma ex_Ex_second_moment {A} {Ω: distrib A} (X: rrvar Ω):
  ex_Ex (rvar_comp X (pow^~ 2%nat))
  ex_Ex X.

Lemma ex_Var_second_moment {A} {Ω: distrib A} (X: rrvar Ω) :
  ex_Var X ex_Ex (rvar_comp X (λ x, x^2)).

Lemma ex_Var_first_moment {A} {Ω: distrib A} (X: rrvar Ω) :
  ex_Var X ex_Ex X.


Lemma Var_Ex_diff {A} {Ω: distrib A} (X: rrvar Ω) :
  ex_Var X
  Var X = Ex (rvar_comp X (λ x, x^2)) - (Ex X)^2.

Definition Cov {A} {Ω: distrib A} (X Y : rrvar Ω) :=
  Ex (mkRvar Ω (λ a, (X a - Ex X) × (Y a - Ex Y))).

Definition ex_Cov {A} {Ω: distrib A} (X Y: rrvar Ω) :=
  ex_Ex X ex_Ex Y ex_Ex (mkRvar Ω (λ a, (X a - Ex X) × (Y a - Ex Y))).


Lemma distrib_of_rvar_pmf_sum1 {A: countType} {Ω: distrib A} {B: eqType} (X: rvar Ω B):
  is_series (countable_sum (λ x : [countType of imgT X], pr_eq X (sval x))) 1.

Definition distrib_of_rvar {A: countType} {Ω: distrib A} {B: eqType} (X: rvar Ω B) :=
  mkDistrib [countType of imgT X] (λ x, pr_eq X (sval x)) (λ x, pr_eq_ge_0 X (sval x))
            (distrib_of_rvar_pmf_sum1 X).

Lemma pr_pred_rvar {A} {B: eqType} {Ω: distrib A} (X: rvar Ω B) P :
  pr Ω (λ a, P (X a)) =
  Series (countable_sum (λ a : imgT X, if P (sval a) then pr_eq X (sval a) else 0)).

Lemma Ex_sum_n_le {A} {Ω: distrib A} (X: rrvar Ω) n:
  ex_Ex X
  ( a, pr_eq X a > 0 a 0)
  sum_n (countable_sum (λ a : imgT X, pr_eq X (sval a) × sval a)) n Ex X.

Lemma Ex_sum_n_le_comp {A} {B} {Ω: distrib A} (X: rvar Ω B) f n:
  ex_Ex (rvar_comp X f)
  ( a, pr_eq X a > 0 f a 0)
  sum_n (countable_sum (λ a : imgT X, pr_eq X (sval a) × f (sval a))) n Ex (rvar_comp X f).

Lemma pr_eq_sum_n_le_1 {A} {B} {Ω: distrib A} (X: rvar Ω B) P n:
  sum_n (countable_sum (λ a : imgT X, if P (sval a) then pr_eq X (sval a) else 0)) n 1.

Definition sum_list_rrvar {A} {Ω: distrib A} (l: list (rrvar Ω)) : rrvar Ω :=
  mkRvar Ω (λ a, fold_right Rplus 0 (map (λ X, (rvar_fun _ _ X) a) l)).

Definition prod_list_rrvar {A} {Ω: distrib A} (l: list (rrvar Ω)) : rrvar Ω :=
  mkRvar Ω (λ a, fold_right Rmult 1 (map (λ X, (rvar_fun _ _ X) a) l)).

From mathcomp Require Import seq.

Definition rvar_list {A} {B: eqType} {Ω: distrib A} (lX: list (rvar Ω B))
  : (rvar Ω [eqType of seq B]) :=
  mkRvar Ω (λ a, map (λ f, (rvar_fun _ _ f) a) lX).

Lemma sum_list_rrvar_alt {A} {Ω: distrib A} (l: list (rrvar Ω)):
  sum_list_rrvar l = rvar_comp (rvar_list l) (fold_right Rplus 0).

Lemma prod_list_rrvar_alt {A} {Ω: distrib A} (l: list (rrvar Ω)):
  prod_list_rrvar l = rvar_comp (rvar_list l) (fold_right Rmult 1).

Lemma ex_Ex_sum_list {A} {Ω: distrib A} (l: list (rrvar Ω)) :
  ( X, In X l ex_Ex X)
  ex_Ex (sum_list_rrvar l).

Lemma Ex_sum_list {A} {Ω: distrib A} (l: list (rrvar Ω)) :
  ( X, In X l ex_Ex X)
  Ex (sum_list_rrvar l) = fold_right Rplus 0 (map Ex l).

Lemma rvar_exists_support {A: countType} {B: eqType} {Ω: distrib A} (X: rvar Ω B):
   a, pr_eq X a > 0.

Lemma is_series_pr_eq_over_range {A : countType} {B: eqType} {Ω: distrib A} (X: rvar Ω B):
  is_series (countable_sum (λ v : imgT X, pr_eq X (sval v))) 1.

Lemma Series_pr_eq_over_range {A : countType} {B: eqType} {Ω: distrib A} (X: rvar Ω B):
  Series (countable_sum (λ v : imgT X, pr_eq X (sval v))) = 1.

Module total_prob.
Section total_prob.

Variable (A: countType).
Variable (B: eqType).
Variable (Ω: distrib A).
Variable (X: rvar Ω B).
Variable (P: pred A).

Definition atotal :=
  λ mn, match mn with
        | (m, n)
          match (@pickle_inv [countType of imgT X] m), (@pickle_inv A n) with
            | Some v, Some a ⇒ (if (X a == sval v) && P a then Ω a else 0)
            | _, _ ⇒ 0
          end
        end.

Lemma ex_total_column_abs v:
  ex_series
    (λ k : nat,
     Rabs
       match @pickle_inv A k with
       | Some aif (X a == v) && P a then Ω a else 0
       | None ⇒ 0
       end).

Lemma atotal_double_summable_by_column:
 double_summable atotal.

Lemma atotal_row_rewrite j:
  Series (λ k : nat, atotal (j, k))
  = countable_sum (λ b : imgT X, pr Ω (λ a, (X a == sval b) && P a)) j.

Lemma atotal_column_rewrite k:
  Series (λ j : nat, atotal (j, k))
  = countable_sum (λ a, if P a then Ω a else 0) k.

Lemma pr_total_prob:
  pr Ω P = Series (countable_sum (λ r: imgT X, pr Ω (λ a, (X a == sval r) && P a))).
End total_prob.
End total_prob.

Lemma pr_total_prob {A : countType} {B: eqType} {Ω: distrib A} (X: rvar Ω B) P:
  pr Ω P = Series (countable_sum (λ r: imgT X, pr Ω (λ a, (X a == sval r) && P a))).

Lemma img_alt {A: countType} {B: eqType} (f: A B) i: (i \in img f) x, f x = i.

Lemma img_alt' {A: finType} {B: eqType} (f: A B) i:
  (i \in undup [seq f i | i <- Finite.enum A]) x, f x = i.

Lemma pr_img {A} {B} {Ω : distrib A} (X: rvar Ω B) i: pr_eq X i > 0 i \in img X.

Lemma pr_img0 {A} {B} {Ω : distrib A} (X: rvar Ω B) i : pr_eq X i 0 i \in img X.

Lemma pr_img_nin {A B} {Ω: distrib A} (X: rvar Ω B) i: ¬ (i \in img X) pr_eq X i = 0.

Lemma pr_img_gt {A} {Ω: distrib A} (X: rrvar Ω) i: pr_gt X i > 0 i', i' > i i' \in img X.

Lemma pr_img_ge {A} {Ω: distrib A} (X: rrvar Ω) i: pr_ge X i > 0 i', i' i i' \in img X.

Lemma SeriesC_pred_sat_le {A: countType} (f : A R) (x : A) (P: pred A) :
  ( x, f x 0)
  P x
  ex_series (countable_sum (λ a : A, if P a then f a else 0))
  f x Series (countable_sum (λ a : A, if P a then f a else 0)).

Lemma if_case_match {A: Type} (b b': bool) (x y: A):
  b == b' (if b then x else y) = (if b' then x else y).

Lemma pr_eq_sum_list {A: countType} {B: eqType} {Ω: distrib A} (X : rvar Ω B) l:
    Series (countable_sum (λ n : img_countType X,
                                          if sval n \in l then pr_eq X (sval n) else 0))
             = \big[Rplus/0]_(i <- undup l) pr_eq X i.

Lemma pr_eq_le_sum_list {A: countType} {B: eqType} {Ω: distrib A} (X: rvar Ω B)
      (l: list B) b :
  ¬ (b \in l)
  pr_eq X b 1 - \big[Rplus/0]_(i <- undup l) (pr_eq X i).