(* infotheo: information theory and error-correcting codes in Coq             *)
(* Copyright (C) 2020 infotheo authors, license: LGPL-2.1-or-later            *)
From mathcomp Require Import all_ssreflect ssralg fingroup finalg matrix.
Require Import Reals.
Require Import ssrR Reals_ext logb ssr_ext ssralg_ext.

(******************************************************************************)
(*                     Additional lemmas about bigops                         *)
(******************************************************************************)

Set Implicit Arguments.
Unset Strict Implicit.
Import Prenex Implicits.

Lemma leq_bigmax_seq (A : eqType) (F : A -> nat) x (l : seq A) :
  x \in l -> F x <= \max_(i <- l) F i.
Proof.
elim: l => // y l ih; rewrite inE big_cons => /predU1P[->|xl].
  by rewrite leq_maxl.
by rewrite (leq_trans (ih xl))// leq_maxr.
Qed.

Section bigop_no_law.
Variables (R : Type) (idx : R) (op : R -> R -> R).

Lemma big_tcast n m (n_m : n = m) (A : finType) F (P : pred {: n.-tuple A}) :
  \big[op/idx]_(p in {: n.-tuple A} | P p) F p =
  \big[op/idx]_(p in {: m.-tuple A} | P (tcast (esym n_m) p)) F (tcast (esym n_m) p).
Proof. subst m; apply eq_bigr => ta => /andP[_ H]; by rewrite tcast_id. Qed.

Lemma big_cast_rV n m (n_m : n = m) (A : finType) F (P : pred {: 'rV[A]_n}) :
  \big[op/idx]_(p in {: 'rV_n} | P p) F p =
  \big[op/idx]_(p in {: 'rV_m} | P (castmx (erefl, esym n_m) p)) F (castmx (erefl, esym n_m) p).
Proof. by subst m; apply eq_bigr => ta => /andP[_ H]. Qed.

End bigop_no_law.
Arguments big_tcast {R} {idx} {op} {n} {m} _ {A} _ _.
Arguments big_cast_rV {R} {idx} {op} {n} {m} _ {A} _ _.

Section bigop_add_law.
Variables (R : Type) (idx : R) (op : R -> R -> R) (M : Monoid.add_law idx op).

Lemma Set2sumE (A : finType) (f : A -> R) (card_A : #|A| = 2%nat) :
 \big[M/idx]_(i in A) (f i) = M (f (Set2.a card_A)) (f (Set2.b card_A)).
Proof.
by rewrite /index_enum -enumT Set2.enumE !big_cons big_nil (Monoid.addm0 M) !enum_valP.
Qed.

Lemma big_bool (f : bool -> R) : \big[M/idx]_(i in {:bool}) f i = M (f false) (f true).
Proof.
set h : 'I_2 -> bool := [eta (fun=>false) with ord0 |-> false, lift ord0 ord0 |-> true].
set h' : bool -> 'I_2 := fun x => match x with false => ord0 | true => lift ord0 ord0 end.
rewrite (reindex_onto h h') /=; last by move=> i _; rewrite /h /h'; case: ifP.
rewrite (eq_bigl xpredT); last first.
  move=> i; move: (ord2 i) => /orP[|] /eqP -> /=; first by rewrite eqxx.
  exact/eqP/val_inj.
by rewrite big_ord_recl /= big_ord_recl big_ord0 Monoid.addm0.
Qed.

Variable (A : finType).
Local Open Scope ring_scope.

Lemma big_rV0_row_of_tuple f (P : pred _) :
  \big[M/idx]_(v in 'rV[A]_0 | P v) f v =
  if P (row_of_tuple [tuple]) then f (row_of_tuple [tuple]) else idx.
Proof.
rewrite -big_map /= /index_enum -enumT /=.
set e := enum _.
rewrite (_ : e = [:: row_of_tuple [tuple]]).
  by rewrite /= big_cons big_nil Monoid.addm0.
rewrite /e.
apply (@eq_from_nth _ (row_of_tuple [tuple])).
  by rewrite -cardE card_mx muln0 expn0.
move=> i.
rewrite -cardE card_mx muln0 expn0 ltnS leqn0 => /eqP ->{i}.
rewrite -/e.
destruct e => //.
apply val_inj => /=.
by apply/ffunP => /= -[? []].
Qed.

End bigop_add_law.

(** Switching from a sum on the domain to a sum on the image of function *)
Section partition_big_finType_eqType.
Variables (A : finType) (B : eqType).
Variables (R : Type) (idx : R) (op : R -> R -> R) (M : Monoid.add_law idx op).

Lemma partition_big_undup_map (p : seq A) (f : A -> B) : forall g, uniq p ->
  \big[M/idx]_(i <- p) g i =
  \big[M/idx]_(r <- undup (map f p)) \big[M/idx]_(i <- p | f i == r) (g i).
Proof.
move Hn : (undup (map f (p))) => n.
elim: n p f Hn => [p f H F ? | b bs IH p f H F ?].
  rewrite big_nil.
  suff -> : p = [::] by rewrite big_nil.
  move/undup_nil : H => /(congr1 size) /=; rewrite size_map.
  by move/eqP; rewrite size_eq0 => /eqP.
rewrite big_cons.
have [h [t [H1 [H2 H3]]]] : exists h t,
    perm_eq p (h ++ t) /\ undup (map f h) = [:: b] /\ undup (map f t) = bs.
  exact: undup_perm.
transitivity (\big[M/idx]_(i <- h ++ t) F i); first exact: perm_big.
transitivity (M
    (\big[M/idx]_(i <- h ++ t | f i == b) F i)
    (\big[M/idx]_(j <- bs) \big[M/idx]_(i <- h ++ t | f i == j) F i)); last first.
  congr (M _ _).
    by apply: perm_big; rewrite perm_sym.
  by apply eq_bigr => b0 _ /=; apply: perm_big; rewrite perm_sym.
have -> : \big[M/idx]_(j <- bs) \big[M/idx]_(i <- h ++ t | f i == j) F i =
    \big[M/idx]_(j <- bs) \big[M/idx]_(i <- t | f i == j) F i.
  rewrite [in LHS]big_seq [in RHS]big_seq /=.
  apply/esym/eq_bigr => b0 b0bs.
  rewrite big_cat /=.
  rewrite (_ : \big[M/idx]_(i0 <- h | f i0 == b0) F i0 = idx) ?Monoid.add0m//.
  transitivity (\big[M/idx]_(i0 <- h | false) F i0); last by rewrite big_pred0.
  rewrite big_seq_cond; apply eq_bigl => /= a.
  apply/negP => /andP[ah /eqP fai]; subst b0.
  have fab : f a \in [:: b].
    have H' : f a \in map f h by apply/mapP; exists a.
    have : f a \in undup (map f h) by rewrite mem_undup.
    by rewrite H2.
  have : uniq (b :: bs) by rewrite -H undup_uniq.
  move: fab; rewrite /= in_cons in_nil orbC /= => /eqP ?; subst b.
  by rewrite b0bs.
rewrite -IH //; last first.
  have : uniq (h ++ t) by rewrite -(perm_uniq H1).
  by rewrite cat_uniq => /and3P[].
suff -> : \big[M/idx]_(i <- h ++ t | f i == b) F i = \big[M/idx]_(i <- h) F i.
  by rewrite big_cat.
rewrite big_cat /=.
have -> : \big[M/idx]_(i <- t | f i == b) F i = idx.
  transitivity (\big[M/idx]_(i0 <- t | false) F i0); last by rewrite big_pred0.
  rewrite big_seq_cond; apply eq_bigl => /= a.
  apply/negP => /andP[ta /eqP fab]; subst b.
  have fabs : f a \in bs.
    have : f a \in undup (map f t) by rewrite mem_undup; apply/mapP; exists a.
    by rewrite H3.
  have : uniq (f a :: bs) by rewrite -H undup_uniq.
  by rewrite /= fabs.
rewrite Monoid.addm0 big_seq_cond /=.
apply/esym.
rewrite big_seq /=; apply congr_big => //= a.
case/boolP : (a \in h) => ah //=; apply/esym.
have : f a \in [:: b] by rewrite -H2 mem_undup; apply/mapP; exists a.
by rewrite in_cons /= in_nil orbC.
Qed.

(* NB: compare with finset.partition_big_imset *)
Lemma partition_big_fin_img (f : A -> B) g :
  \big[M/idx]_(i in A) (g i) =
  \big[M/idx]_(r <- fin_img f) \big[M/idx]_(i in A | f i == r) (g i).
Proof.
have /= := @partition_big_undup_map (enum A) f g.
rewrite enum_uniq => /(_ isT) H.
transitivity (\big[M/idx]_(i <- enum A) g i); first by rewrite enumT.
by rewrite H; apply eq_bigr => i _; apply congr_big => //; rewrite enumT.
Qed.

End partition_big_finType_eqType.

Section BigOps.
Variables (R : Type) (idx : R).
Variables (op : Monoid.law idx) (aop : Monoid.com_law idx).
Variables I J : finType.
Implicit Type A B : {set I}.
Implicit Type h : I -> J.
Implicit Type P : pred I.
Implicit Type F : I -> R.

Lemma partition_big_preimset h (B : {set J}) F :
  \big[aop/idx]_(i in h @^-1: B) F i =
     \big[aop/idx]_(j in B) \big[aop/idx]_(i in I | h i == j) F i.
Proof.
have HA : [disjoint B :&: [set h x | x in I] & B :\: [set h x | x in I]]
    by rewrite -setI_eq0 -setIA setIDA [in _ :&: B]setIC -setIDA setDv !setI0.
have Hha : [disjoint h @^-1: (B :&: [set h x | x in I])
                             & h @^-1: (B :\: [set h x | x in I])].
  rewrite -setI_eq0 -preimsetI.
  suff : [disjoint B :&: [set h x | x in I] & B :\: [set h x | x in I]].
    by rewrite -setI_eq0; move/eqP => ->; rewrite preimset0.
  by [].
rewrite -(setID B (h @: I)) /= preimsetU.
under eq_bigl do rewrite in_setU.
rewrite bigU //.
under [in RHS]eq_bigl do rewrite in_setU.
rewrite bigU //.
have -> : h @^-1: (B :\: [set h x | x in I]) = set0.
  apply/setP/subset_eqP/andP; rewrite sub0set; split => //.
  apply/subsetP=> i; rewrite !inE; case/andP.
  move/imsetP=> H _; elimtype False; apply H.
  by exists i; rewrite ?inE.
rewrite big_set0 Monoid.mulm1.
have -> : \big[aop/idx]_(x in B :\: [set h x | x in I]) \big[aop/idx]_(i | h i == x) F i =
          \big[aop/idx]_(x in B :\: [set h x | x in I]) idx.
  apply eq_bigr => j.
  rewrite inE; case/andP => Hj Hj'.
  apply big_pred0 => i.
  apply/negP => /eqP hij.
  move: Hj; rewrite -hij.
  move/imsetP; apply.
  by exists i.
rewrite big1_eq Monoid.mulm1.
set B' := B :&: [set h x | x in I].
set A := h @^-1: B'.
have -> : B' = h @: A by rewrite imset_preimset //; apply subsetIr.
have Hright j : j \in h @: A -> \big[aop/idx]_(i in I | h i == j) F i =
                               \big[aop/idx]_(i in A | h i == j) F i.
  move=> Hj; apply eq_bigl => i; apply andb_id2r; move/eqP => hij.
  move: Hj; rewrite -hij !inE.
  case/imsetP => x; rewrite /A /B' !inE => /andP [H H0] ->.
  by rewrite H H0.
rewrite [in RHS](eq_bigr _ Hright).
exact: partition_big_imset.
Qed.

End BigOps.

Section big_pred1_inj.
Variables (R : Type) (idx : R) (op : Monoid.law idx).

Lemma big_pred1_inj (A C : finType) h i (k : A -> C) : injective k ->
  \big[op/idx]_(a | k a == k i) h a = h i.
Proof. by move=> ?; rewrite (big_pred1 i) // => ?; rewrite eqtype.inj_eq. Qed.

End big_pred1_inj.
Arguments big_pred1_inj [R] [idx] [op] [A] [C] [h] [i] [k] _.

Section bigop_com_law.
Variables (R : Type) (idx : R) (M : Monoid.com_law idx).
Variable A : finType.

Lemma big_union (X1 X2 : {set A}) f :
  [disjoint X2 & X1] ->
  \big[M/idx]_(a | a \in X1 :|: X2) f a =
  M (\big[M/idx]_(a | a \in X1) f a) (\big[M/idx]_(a | a \in X2) f a).
Proof.
move=> Hdisj.
rewrite (@big_setID _ _ _ _ _ X1) /= setUK setDUl setDv set0U.
suff : X2 :\: X1 = X2 by move=> ->.
by apply/setDidPl.
Qed.

Variable B : finType.

(** Big sums lemmas for cartesian products *)

Lemma pair_big_fst (F : {: A * B} -> R) P Q :
  P =1 Q \o fst ->
  \big[M/idx]_(i in A | Q i) \big[M/idx]_(j in B) F (i, j) =
  \big[M/idx]_(i in {: A * B} | P i) F i.
Proof.
move=> /= PQ; rewrite pair_big /=; apply eq_big; last by case.
case=> /= ? ?; by rewrite inE andbC PQ.
Qed.

Lemma pair_big_snd (F : {: A * B} -> R) P Q :
  P =1 Q \o snd ->
  \big[M/idx]_(i in A) \big[M/idx]_(j in B | Q j) F (i, j) =
  \big[M/idx]_(i in {: A * B} | P i) F i.
Proof.
move=> /= PQ; rewrite pair_big /=; apply eq_big; last by case.
case=> /= ? ?; by rewrite PQ.
Qed.

Lemma big_setX (a : {set A}) (b : {set B}) f :
  \big[M/idx]_(x in a `* b) f x = \big[M/idx]_(x in a) \big[M/idx]_(y in b) f (x, y).
Proof.
rewrite (eq_bigl (fun x => (x.1 \in a) && (x.2 \in b))); last first.
  by case=> x y; rewrite in_setX.
rewrite (eq_bigr (fun x => f (x.1, x.2))); last by case.
by rewrite -(pair_big _ _ (fun a b => f (a, b))).
Qed.

Lemma big_rV_prod n f (X : {set 'rV[A * B]_n}) :
  \big[M/idx]_(a in 'rV[A * B]_n | a \in X) f a =
  \big[M/idx]_(a in {: 'rV[A]_n * 'rV[B]_n} | (prod_rV a) \in X) f (prod_rV a).
Proof.
rewrite (reindex_onto (@rV_prod _ _ _) (@prod_rV _ _ _)) //=; last first.
  move=> ? _; by rewrite prod_rVK.
apply eq_big => [?|? _]; by rewrite rV_prodK // eqxx andbC.
Qed.

Local Open Scope vec_ext_scope.
Local Open Scope ring_scope.

Lemma big_rV_1 f g (P : pred _) (Q : pred _):
  (forall i : 'rV[A]_1, f i = g (i ``_ ord0)) ->
  (forall i : 'rV[A]_1, P i = Q (i ``_ ord0)) ->
  \big[M/idx]_(i in 'rV[A]_1 | P i) f i = \big[M/idx]_(i in A | Q i) g i.
Proof.
move=> FG PQ.
rewrite (reindex_onto (fun i => \row_(j < 1) i) (fun p => p ``_ ord0)) /=; last first.
  move=> m Pm.
  apply/rowP => a; by rewrite {a}(ord1 a) mxE.
apply eq_big => a.
  by rewrite PQ mxE eqxx andbT.
by rewrite FG !mxE.
Qed.

Lemma big_rV1_ord0 (f : A -> R) k :
  \big[M/idx]_(i in A) f i = k -> \big[M/idx]_(i in 'rV[A]_1) f (i ``_ ord0) = k.
Proof.
move=> <-.
rewrite (reindex_onto (fun j => \row_(i < 1) j) (fun p => p ``_ ord0)) /=.
- apply eq_big => a; first by rewrite mxE eqxx inE.
  move=> _; by rewrite mxE.
- move=> t _; apply/rowP => a; by rewrite (ord1 a) mxE.
Qed.

Local Open Scope vec_ext_scope.
Local Open Scope ring_scope.

Lemma big_rV_cons n (F : 'rV[A]_n.+1 -> R) (a : A) (i0 : 'I_n.+1) : i0 = ord0 ->
  \big[M/idx]_(v in 'rV[A]_n) (F (row_mx (\row_(k < 1) a) v)) =
  \big[M/idx]_(v in 'rV[A]_n.+1 | v ``_ i0 == a) (F v).
Proof.
move=> i00.
rewrite [in RHS](reindex_onto (row_mx (\row_(k < 1) a)) rbehead) /=; last first.
  move=> m /eqP <-; by rewrite i00 row_mx_rbehead.
apply eq_bigl => ?; by rewrite i00 rbehead_row_mx eqxx andbT row_mx_row_ord0 eqxx.
Qed.

Lemma big_rV_behead n (F : 'rV[A]_n.+1 -> R) (w : 'rV[A]_n) :
  \big[M/idx]_(a in A) (F (row_mx (\row_(k < 1) a) w)) =
  \big[M/idx]_(v in 'rV[A]_n.+1 | rbehead v == w) (F v).
Proof.
rewrite [in RHS](reindex_onto
  (fun p => row_mx (\row_(k < 1) p) w) (fun p => p ``_ ord0) ) /=; last first.
  move=> i /eqP <-; by rewrite row_mx_rbehead.
apply eq_bigl => ?; by rewrite rbehead_row_mx eqxx /= row_mx_row_ord0 eqxx.
Qed.

Lemma big_rV_cons_behead_support n (F : 'rV[A]_n.+1 -> R)
  (X1 : {set A}) (X2 : {set {: 'rV[A]_n}}) :
  \big[M/idx]_(a in X1) \big[M/idx]_(v in X2) (F (row_mx (\row_(k < 1) a) v)) =
  \big[M/idx]_(w in 'rV[A]_n.+1 | (w ``_ ord0 \in X1) && (rbehead w \in X2)) (F w).
Proof.
rewrite [in RHS](partition_big (fun x : 'rV_n.+1 => x ``_ ord0) (mem X1)) /=; last first.
  by move=> i /andP[].
apply eq_bigr => i Hi.
rewrite (reindex_onto (fun j => row_mx (\row_(k < 1) i) j) rbehead) /=; last first.
  move=> j /andP[] => _ /eqP => <-; by rewrite row_mx_rbehead.
apply eq_big => //= x; by rewrite row_mx_row_ord0 rbehead_row_mx !eqxx Hi !andbT.
Qed.

Lemma big_rV_cons_behead n (F : 'rV[A]_n.+1 -> R)
  (P1 : pred A) (P2 : pred 'rV[A]_n) :
  \big[M/idx]_(i in A | P1 i)
    \big[M/idx]_(j in 'rV[A]_n | P2 j) (F (row_mx (\row_(k < 1) i) j)) =
  \big[M/idx]_(p in 'rV[A]_n.+1 | (P1 (p ``_ ord0)) && (P2 (rbehead p)) ) (F p).
Proof.
rewrite [in RHS](partition_big (fun x : 'rV_n.+1 => x ``_ ord0) P1) /=; last first.
  by move=> i /andP[].
apply eq_bigr => i Hi.
rewrite (reindex_onto (fun j => row_mx (\row_(k < 1) i) j) rbehead) /=; last first.
    move=> j /andP[] Hj1 /eqP => <-; by rewrite row_mx_rbehead.
apply eq_big => //= x; by rewrite row_mx_row_ord0 rbehead_row_mx 2!eqxx Hi !andbT.
Qed.

Lemma big_rV_belast_last n (F : 'rV[A]_n.+1 -> R)
  (P1 : pred 'rV[A]_n) (P2 : pred A) :
  \big[M/idx]_(i in 'rV[A]_n | P1 i)
    \big[M/idx]_(j in A | P2 j) (F (castmx (erefl, addn1 n) (row_mx i (\row_(k < 1) j)))) =
  \big[M/idx]_(p in 'rV[A]_n.+1 | (P1 (rbelast p)) && (P2 (rlast p)) ) (F p).
Proof.
rewrite [in RHS](partition_big (fun x : 'rV_n.+1 => rlast x) P2) /=; last first.
  by move=> i /andP[].
rewrite exchange_big.
apply eq_bigr => i Hi.
rewrite (reindex_onto (fun j => (castmx (erefl 1%nat, addn1 n) (row_mx j (\row_(k < 1) i)))) rbelast) /=; last first.
    move=> j /andP[] Hj1 /eqP => <-; by rewrite row_mx_rbelast.
apply eq_big => //= x.
by rewrite row_mx_row_ord_max rbelast_row_mx 2!eqxx !andbT Hi andbT.
Qed.

End bigop_com_law.
Arguments pair_big_fst {R} {idx} {M} {A} {B} _.
Arguments pair_big_snd {R} {idx} {M} {A} {B} _.
Arguments big_rV_belast_last {R} {idx} {M} {A} {n} _ _ _.
Arguments big_rV_cons_behead {R} {idx} {M} {A} {n} _ _ _.

Section bigcup_ext.
Variables I A B : finType.

Lemma bigcup_set0 i (D : I -> {set B}) (F : B -> I -> {set A}) b :
  b \in D i -> F b i != set0 ->
  \bigcup_(b' in D i) F b' i == set0 -> D i == set0.
Proof.
move=> bDi Fbi0 /set0Pn F0; apply/set0Pn => -[b' b'i]; apply F0 => {F0}.
by case/set0Pn : Fbi0 => a tA; exists a; apply/bigcupP; exists b.
Qed.

Lemma bigcup_setX (E : I -> {set A}) (F : {set B}) :
  \bigcup_(i in I) (F `* E i) = F `* (\bigcup_(i in I) (E i)).
Proof.
apply/setP => -[b a] /=; rewrite !inE /=.
apply/bigcupP/andP => [[/= i _]|[K1] /bigcupP[/= i _ aEi]].
  rewrite !inE /= => /andP[xb yai]; rewrite xb; split => //.
  by apply/bigcupP; exists i.
by exists i => //; rewrite !inE /= K1.
Qed.

Lemma bigcup_preimset (P : pred I) (F : A -> B) (E : I -> {set B}) :
  \bigcup_(i | P i) F @^-1: E i = F @^-1: \bigcup_(i | P i) E i.
Proof.
rewrite/preimset.
apply/setP=> x; rewrite inE; apply/bigcupP/bigcupP => -[] i HPi; rewrite ?inE => HFxEi; exists i => //=; by rewrite inE.
Qed.

Lemma bigcup_set1 (E : {set A}) : E = \bigcup_(a in E) [set a].
Proof.
apply/setP => a; apply/idP/bigcupP => [aE|[a' a'E]].
by exists a => //; rewrite inE.
by rewrite inE => /eqP ->.
Qed.

End bigcup_ext.

Section MyPartitions.

Variables (R : Type) (idx : R) (op : Monoid.com_law idx).

Variables T I : finType.

Lemma big_bigcup_partition (F : I -> {set T}) E (V : {set I}):
  (forall i j, i != j -> [disjoint F i & F j]) ->
  \big[op/idx]_(x in \bigcup_(i in V) F i) E x =
    \big[op/idx]_(i in V) \big[op/idx]_(x in F i) E x.
Proof.
move=> disjF; pose Q := [set F i | i in V & F i != set0].
have trivP: trivIset Q.
  apply/trivIsetP=> _ _ /imsetP[i _ ->] /imsetP[j _ ->] neqFij.
  by apply: disjF; apply: contraNneq neqFij => ->.
have ->: \bigcup_(i in V) F i = cover Q.
  apply/esym.
  rewrite cover_imset big_mkcond /=.
  apply/esym.
  rewrite big_mkcond /=.
  apply: eq_bigr => i _.
  case: ifP.
    rewrite inE => -> /=.
    by case: ifP => // /eqP.
  by rewrite inE => ->.
rewrite big_trivIset // big_imset => [|i j _ /setIdP[_ notFj0] eqFij].
  rewrite big_mkcond [in X in _ = X]big_mkcond.
  apply: eq_bigr => i _.
  rewrite inE.
  case: ifP.
    by case/andP=> ->.
  move/negbT.
  rewrite negb_and.
  case/orP.
    by move/negbTE => ->.
  rewrite negbK => /eqP Fi.
  case: ifP => //.
  by rewrite Fi big_set0.
by apply: contraNeq (disjF _ _) _; rewrite -setI_eq0 eqFij setIid.
Qed.

End MyPartitions.

Section bigcap_ext.
Variable A : finType.

Lemma bigcap_seq_const I (B : {set A}) (r : seq.seq I) :
  (0 < size r)%nat -> \bigcap_(i <- r) B = B.
Proof.
elim: r => [//|a r IHr] _; rewrite big_cons.
case: r IHr => [|b r] IHr; first by rewrite big_nil setIT.
by rewrite IHr // setIid.
Qed.

Lemma bigcap_ord_const n' (B : {set A}) :
  \bigcap_(i < n'.+1) B = B.
Proof. by rewrite bigcap_seq_const // /index_enum -enumT size_enum_ord. Qed.

Lemma bigcap_const (I : eqType) (B : {set A}) (r : seq.seq I) (p : pred I) :
  (exists2 i : I, i \in r & p i) -> \bigcap_(i <- r | p i) B = B.
Proof.
case=> i H1 H2; rewrite -big_filter bigcap_seq_const //.
rewrite size_filter- has_count.
by apply/hasP; exists i.
Qed.

End bigcap_ext.

Section big_tuple_ffun.
Import Monoid.Theory.
Variable R : Type.
Variable times : Monoid.mul_law R0.
Local Notation "*%M" := times (at level 0).
Variable plus : Monoid.add_law R0 *%M.
Local Notation "+%M" := plus (at level 0).

Lemma big_tuple_ffun (I J : finType) (F : {ffun I -> J} -> R)
  (G : _ -> _ -> _) (jdef : J) (idef : I) :
  \big[+%M/R0]_(j : #|I|.-tuple J) G (F [ffun x => tnth j (enum_rank x)]) (nth jdef j 0)
    = \big[+%M/R0]_(f : {ffun I -> J}) G (F f) (f (nth idef (enum I) 0)).
Proof.
rewrite (reindex_onto (fun y => fgraph y) (fun p => [ffun x => tnth p (enum_rank x)])); last first.
  move=> t _; by apply/eq_from_tnth => i; rewrite tnth_fgraph ffunE enum_valK.
apply eq_big.
  move=> f /=; apply/eqP/ffunP => i; by rewrite ffunE tnth_fgraph enum_rankK.
move=> f _.
congr (G (F _) _).
  apply/ffunP => i; by rewrite ffunE tnth_fgraph enum_rankK.
have @zero : 'I_#|I| by exists O; apply/card_gt0P; exists idef.
transitivity (tnth (fgraph f) zero).
  apply set_nth_default; by rewrite size_tuple ltn_ord.
rewrite tnth_fgraph; congr (f _).
by apply enum_val_nth.
Qed.

End big_tuple_ffun.

(* This is from master branch of math-comp bigop. *)
Lemma big_ord1_eq (R: Type) (idx: R) (op: Monoid.law idx) (F : nat -> R) (i : nat) n :
  \big[op/idx]_(j < n | j == i :> nat) F j = if i < n then F i else idx.
Proof.
case: ltnP => [i_lt|i_ge]; first by rewrite (big_pred1_eq _ (Ordinal _)).
by rewrite big_pred0// => j; apply: contra_leqF i_ge => /eqP <-.
Qed.


Require Import Util
               Program
               refocusing_semantics.

Module MiniML_PreRefSem <: PRE_REF_SEM.

  Parameter var : Set.

  Inductive expr :=
  | Z   : expr
  | S   : expr -> expr
  | App : expr -> expr -> expr
  | Var : var -> expr
  | Lam : var -> expr -> expr
  | Case: expr -> expr -> var -> expr -> expr
  | Letv: var -> expr -> expr -> expr
  | Fix : var -> expr -> expr
  | Pair: expr -> expr -> expr
  | Fst : expr -> expr
  | Snd : expr -> expr.
  Definition term := expr.
  Hint Unfold term.

  Definition ckind := unit.
  Hint Unfold ckind.

  Inductive val :=
  | vZ    : val
  | vS    : val -> val
  | vLam  : var -> term -> val
  | vVar  : var -> val
  | vPair : val -> val -> val.
  Definition value (_ : ckind) := val.
  Hint Unfold value.
  Notation value' := (value ()).

  Coercion val_value := id : val -> value'.

  Inductive red :=
  | rApp : value' -> value' -> red
  | rLet : var -> value' -> term -> red
  | rFix : var -> term -> red
  | rFst : value' -> red
  | rSnd : value' -> red
  | rCase: value' -> term -> var -> term -> red.
  Definition redex (_ : ckind) := red.
  Hint Unfold redex.
  Notation redex' := (redex ()).

  Coercion red_redex := id : red -> redex'.

  Inductive ec :=
  | s_c    : ec
  | ap_l   : value' -> ec
  | ap_r   : term -> ec
  | case_c : term -> var -> term -> ec
  | let_c  : var -> term -> ec
  | fst_c  : ec
  | snd_c  : ec
  | pair_l : value' -> ec
  | pair_r : term -> ec.
  Definition elem_context := ec.
  Hint Unfold elem_context.

  Definition ckind_trans (_ : ckind) (_ : elem_context) : ckind := ().
  Notation "k +> ec" := (ckind_trans k ec) (at level 50, left associativity).

  Definition init_ckind : ckind     := ().
  Definition dead_ckind (_ : ckind) := False.
  Hint Unfold init_ckind dead_ckind.

  Inductive context (k1 : ckind) : ckind -> Set :=
  | empty : context k1 k1
  | ccons : forall (ec : elem_context) {k2}, context k1 k2 -> context k1 (k2+>ec).
  Arguments empty {k1}. Arguments ccons {k1} _ {k2} _.
  Notation context' := (context () ()).

  Notation "[.]"      := empty.
  Notation "[.]( k )" := (@empty k).
  Infix "=:"          := ccons (at level 60, right associativity).

  Fixpoint value_to_term {k} (v : value k) : term :=
      match v with
      | vZ          => Z
      | vS v        => S (@value_to_term () v)
      | vLam x t    => Lam x t
      | vVar x      => Var x
      | vPair v1 v2 => Pair (@value_to_term () v1) (@value_to_term () v2)
      end.
  Coercion value_to_term : value >-> term.

  Definition redex_to_term {k} (r : redex k) : term :=
      match r with
      | rApp v1 v2      => App (v1 : term) (v2 : term)
      | rLet x v t      => Letv x (v : term) t
      | rFix x v        => Fix x (v : term)
      | rFst v          => Fst (v : term)
      | rSnd v          => Snd (v : term)
      | rCase v t1 x t2 => Case (v : term) t1 x t2
      end.
  Coercion redex_to_term : redex >-> term.

  Instance dead_is_comp : CompPred ckind dead_ckind.
      split. auto.
  Defined.

  Lemma init_ckind_alive :
      ~dead_ckind init_ckind.

  Proof. auto. Qed.

  Lemma value_to_term_injective :
      forall {k} (v v' : value k), value_to_term v = value_to_term v' -> v = v'.

  Proof.
    destruct k.

    induction v;
        destruct v';
        inversion 1;

    solve [ f_equal; auto ].
  Qed.

  Lemma redex_to_term_injective :
      forall {k} (r r' : redex k), redex_to_term r = redex_to_term r' -> r = r'.

  Proof.
    destruct k.

    destruct r; destruct r';
    inversion 1;

    solve [ f_equal; auto using value_to_term_injective ].
  Qed.

  Fixpoint compose {k1 k2} (c0 : context k1 k2)
                      {k3} (c1 : context k3 k1) : context k3 k2 :=
      match c0 in context _ k2' return context k3 k2' with
      | [.]     => c1
      | ec=:c0' => ec =: compose c0' c1
      end.
  Infix "~+" := compose (at level 60, right associativity).

  Definition elem_plug (t : term) (ec : elem_context) : term :=
      match ec with
      | s_c            => S t
      | ap_l v         => App (v : term) t
      | ap_r t2        => App t t2
      | case_c t1 x t2 => Case t t1 x t2
      | let_c x t2     => Letv x t t2
      | fst_c          => Fst t
      | snd_c          => Snd t
      | pair_l v       => Pair (v : term) t
      | pair_r t2      => Pair t t2
      end.
  Notation "ec :[ t ]" := (elem_plug t ec) (at level 0).

  Lemma elem_plug_injective1 : forall ec {t0 t1}, ec:[t0] = ec:[t1] -> t0 = t1.

  Proof.
    intros ec ? ? H.

    destruct ec;
    inversion H;
    solve [ trivial ].
  Qed.

  Fixpoint plug t {k1 k2} (c : context k1 k2) : term :=
      match c with
      | [.]    => t
      | ec=:c' => plug ec:[t] c'
      end.
  Notation "c [ t ]" := (plug t c) (at level 0).

  Definition immediate_ec ec t := exists t', ec:[t'] = t.

  Lemma death_propagation :
      forall k ec, dead_ckind k -> dead_ckind (k+>ec).

  Proof. auto. Qed.

  Lemma proper_death : forall k, dead_ckind k -> ~ exists (r : redex k), True.
  Proof. auto. Qed.

  Parameter subst : var -> term -> term -> term.

  Definition contract {k} (r : redex k) : option term :=
      match r with
      | rApp (vLam x t) v  => Some (subst x v t)
      | rLet x v t         => Some (subst x v t)
      | rFix x e           => Some (subst x (Fix x e) e)
      | rFst (vPair v1 v2) => Some ((v1 : value') : term)
      | rSnd (vPair v1 v2) => Some ((v2 : value') : term)
      | rCase vZ t _ _     => Some t
      | rCase (vS v) _ x t => Some (subst x (v : value') t)
      | _                  => None
      end.
  Notation contract' := (@contract ()).

  Lemma value_trivial1 : forall {k} ec t,
      forall (v : value k), ~dead_ckind (k+>ec) -> ec:[t] = v ->
          exists (v' : value (k+>ec)), t = v'.

  Proof.
    intros k ec t v H H0.
    destruct ec, v; inversion H0; subst;
    eauto.
  Qed.

  Lemma value_redex : forall {k} (v : value k) (r : redex k),
                          value_to_term v <> redex_to_term r.
  Proof.
    destruct v; destruct r; intro H; discriminate H.
  Qed.

  Lemma redex_trivial1 : forall {k} (r : redex k) ec t,
                             ~dead_ckind (k+>ec) -> ec:[t] = r ->
                                 exists (v : value (k+>ec)), t = v.
  Proof with auto.
    intros k r ec t H H0.
    destruct ec, r; inversion H0; subst;
    eauto.
  Qed.

  Definition immediate_subterm t0 t := exists ec, t = ec:[t0].

  Lemma wf_immediate_subterm : well_founded immediate_subterm.
  Proof. REF_LANG_Help.prove_st_wf. Qed.

  Definition subterm_order := clos_trans_1n term immediate_subterm.
  Notation "t1 <| t2" := (subterm_order t1 t2) (at level 70, no associativity).
  Definition wf_subterm_order : well_founded subterm_order
      := wf_clos_trans_l _ _ wf_immediate_subterm.

  Inductive decomp k : Set :=
  | d_red : forall {k'}, redex k' -> context k k' -> decomp k
  | d_val : value k -> decomp k.
  Arguments d_val {k} _. Arguments d_red {k} {k'} _ _.

  Definition decomp_to_term {k} (d : decomp k) :=
      match d with
      | d_val v     => value_to_term v
      | d_red r c => c[r]
      end.
  Coercion decomp_to_term : decomp >-> term.
  Notation decomp'   := (@decomp ()).

  Definition dec (t : term) k (d : decomp k) : Prop :=
      ~dead_ckind k /\ t = d.

  Definition reduce k t1 t2 :=
      exists {k'} (c : context k k') (r : redex k') t,  dec t1 k (d_red r c) /\
          contract r = Some t /\ t2 = c[t].

  Instance lrws : LABELED_REWRITING_SYSTEM ckind term :=
  { ltransition := reduce }.
  Instance rws : REWRITING_SYSTEM term :=
  { transition := reduce init_ckind }.

  Class SafeKRegion (k : ckind) (P : term -> Prop) :=
  {
      preservation :                                                        forall t1 t2,
          P t1  ->  k |~ t1 → t2  ->  P t2;
      progress :                                                               forall t1,
          P t1  ->  (exists (v : value k), t1 = v) \/ (exists t2, k |~ t1 → t2)
  }.

End MiniML_PreRefSem.

Module MiniML_Strategy <: REF_STRATEGY MiniML_PreRefSem.

  Import MiniML_PreRefSem.

  Inductive interm_dec k : Set :=
  | in_red  : redex k -> interm_dec k
  | in_term : term -> elem_context -> interm_dec k
  | in_val  : value k -> interm_dec k
  | in_dead : interm_dec k.
  Arguments in_red {k} _.    Arguments in_val {k} _.
  Arguments in_term {k} _ _. Arguments in_dead {k}.

  Definition dec_term (t : term) k : interm_dec k :=
      match t with
      | Z              =>  in_val vZ
      | S t            =>  in_term t s_c
      | App t1 t2      =>  in_term t1 (ap_r t2)
      | Var x          =>  in_val (vVar x)
      | Lam x t        =>  in_val (vLam x t)
      | Case t ez x es =>  in_term t (case_c ez x es)
      | Letv x t e     =>  in_term t (let_c x e)
      | Fix x t        =>  in_red (rFix x t)
      | Pair t1 t2     =>  in_term t1 (pair_r t2)
      | Fst t          =>  in_term t fst_c
      | Snd t          =>  in_term t snd_c
      end.

  Definition dec_context (ec : elem_context) k (v : value (k+>ec)) : interm_dec k :=
      match ec with
      | s_c             =>  in_val (vS v)
      | ap_l v0         =>  in_red (rApp v0 v)
      | ap_r t          =>  in_term t (ap_l v)
      | case_c ez x es  =>  in_red (rCase v ez x es)
      | let_c x e       =>  in_red (rLet x v e)
      | fst_c           =>  in_red (rFst v)
      | snd_c           =>  in_red (rSnd v)
      | pair_l v0       =>  in_val (vPair v0 v)
      | pair_r t        =>  in_term t (pair_l v)
      end.

  Lemma dec_term_correct :
      forall (t : term) k, match dec_term t k with
      | in_red r      => t = r
      | in_val v      => t = v
      | in_term t' ec => t = ec:[t']
      | in_dead       => dead_ckind k
      end.

  Proof with auto.
    destruct k, t; simpl...
  Qed.

  Lemma dec_term_from_dead : forall t k, dead_ckind k -> dec_term t k = in_dead.
  Proof.
    intros ? k H.
    inversion H.
  Qed.

  Lemma dec_term_next_alive :
      forall t k t0 ec0, dec_term t k = in_term t0 ec0 -> ~ dead_ckind (k+>ec0).

  Proof. auto. Qed.

  Lemma dec_context_correct :
      forall ec k v, match dec_context ec k v with
      | in_red r      => ec:[v] = r
      | in_val v'     => ec:[v] = v'
      | in_term t ec' => ec:[v] = ec':[t]
      | in_dead       => dead_ckind (k+>ec)
      end.

  Proof.
    intros ec k v.
    destruct k, ec; dependent destruction v;
    simpl;
    solve [ auto ].
  Qed.

  Lemma dec_context_from_dead : forall ec k v,
      dead_ckind (k+>ec) -> dec_context ec k v = in_dead.

  Proof.
    intros ec k v H.
    inversion H.
  Qed.

  Lemma dec_context_next_alive : forall ec k v {t ec0},
      dec_context ec k v = in_term t ec0 -> ~ dead_ckind (k+>ec0).

  Proof. auto. Qed.

  Definition search_order (k : ckind) (t : term) (ec ec0 : elem_context) : Prop :=
      match ec, ec0 with
      | ap_l _,   ap_r _   => immediate_ec ec t /\ immediate_ec ec0 t
      | pair_l _, pair_r _ => immediate_ec ec t /\ immediate_ec ec0 t
      | _, _               => False
      end.
  Notation "k , t |~  ec1 << ec2 " := (search_order k t ec1 ec2)
               (at level 70, t, ec1, ec2 at level 50, no associativity).

  Lemma wf_search : forall k t, well_founded (search_order k t).
  Proof. REF_LANG_Help.prove_ec_wf. Qed.

  Lemma search_order_trans :  forall k t ec0 ec1 ec2,
      k,t |~ ec0 << ec1 -> k,t |~ ec1 << ec2 ->
      k,t |~ ec0 << ec2.

  Proof.
    intros k ? ec0 ec1 ec2 ? ?.
    destruct k, ec0, ec1;
    solve [ autof ].
  Qed.

  Ltac inj_vr := match goal with
  | [Hv : value_to_term _ = value_to_term _ |- _] => apply value_to_term_injective in Hv
  | [Hr : redex_to_term _ = redex_to_term _ |- _] => apply redex_to_term_injective in Hr
  | [ |- _] => idtac
  end.

  Lemma search_order_comp_if : forall t ec0 ec1 k,
      immediate_ec ec0 t -> immediate_ec ec1 t ->
      ~ dead_ckind (k+>ec0) -> ~ dead_ckind (k+>ec1) ->
          k,t |~ ec0 << ec1 \/ k,t |~ ec1 << ec0 \/ ec0 = ec1.

  Proof.
    intros t ec0 ec1 k H H0 _ _.

    assert (G := H); assert (G0 := H0).
    destruct H as (t0, H), H0 as (t1, H0).
    rewrite <- H in H0.

    destruct ec0, ec1; inversion H0;

    solve
    [ auto
    | inj_vr; subst;
      simpl;
      intuition constructor ].
  Qed.

  Lemma search_order_comp_fi :
      forall k t ec0 ec1,  k, t |~ ec0 << ec1 ->
          immediate_ec ec0 t /\ immediate_ec ec1 t /\
          ~ dead_ckind (k+>ec0) /\ ~ dead_ckind (k+>ec1).

  Proof.
    intros k t ec0 ec1 H.

    destruct ec0, ec1; inversion H;
    solve [ intuition ].
  Qed.

  Lemma dec_term_term_top : forall t k {t' ec},
            dec_term t k = in_term t' ec -> forall ec', ~ k, t |~ ec << ec'.

  Proof.
    intros t t' k ec H ec' H0.
    destruct t;
        inversion H; subst;
        destruct ec';
    solve [ inversion H0 ].
  Qed.

  Lemma dec_context_red_bot :
      forall k ec v r, dec_context ec k v = in_red r ->
          forall ec', ~ k, ec:[v] |~ ec' << ec.

  Proof.
    intros k ec v r H ec' H0.
    destruct ec, ec';
    solve [ inversion H; inversion H0 ].
  Qed.

  Lemma dec_context_val_bot :
      forall k ec v v', dec_context ec k v = in_val v' ->
          forall ec', ~ k, ec:[v] |~ ec' << ec.
  Proof.
    intros k ec v v' H ec' H0.
    destruct ec, ec';
    solve [ inversion H; inversion H0 ].
  Qed.

  Lemma dec_context_term_next :
      forall ec0 k v t ec1, dec_context ec0 k v = in_term t ec1 ->
      (*1*) k, ec0:[v] |~ ec1 << ec0 /\
      (*2*) forall ec,  k, ec0:[v] |~ ec << ec0  ->  ~ k, ec0:[v] |~ ec1 << ec.

  Proof.
    intros ec0 k v t ec1 H.

    assert (H0 := dec_context_correct ec0 k v); rewrite H in H0.

    destruct ec0;
        inversion H; subst;

    solve [
        split;
        [ constructor;
          eexists; eauto
        | intros ec'' H1 H2;
          destruct ec'';
          inversion H1; inversion H2
    ]   ].
  Qed.

  Lemma elem_context_det :
      forall t ec {t'}, t = ec:[t'] ->
          forall k ec',  k, t |~ ec' << ec ->
              exists (v : value (k+>ec)), t' = v.
  Proof.
    intros t ec t' H k ec' H0.

    destruct ec, ec';
        inversion H0; subst;

    solve
    [ destruct H1 as (t1, H1);
      inversion H1;
      eexists; trivial ].
  Qed.

End MiniML_Strategy.

Module MiniML_RefSem := PreciseRedRefSem MiniML_PreRefSem MiniML_Strategy.

Require Import refocusing_machine.

Module MiniML_EAM := RefEvalApplyMachine MiniML_RefSem.


(* TLC in Coq
 *
 * Module: tlc.semantics.pattern
 * Purpose: Contains the pattern matching algorithm.
 *)

Require Import mathcomp.ssreflect.eqtype.
Require Import mathcomp.ssreflect.path.
Require Import mathcomp.ssreflect.seq.
Require Import mathcomp.ssreflect.ssrbool.
Require Import mathcomp.ssreflect.ssreflect.
Require Import mathcomp.ssreflect.ssrfun.
Require Import mathcomp.ssreflect.ssrnat.
Require Import tlc.syntax.all_syntax.
Require Import tlc.utility.monad.
Require Import tlc.utility.option.
Require Import tlc.utility.partial_map.
Require Import tlc.utility.result.
Require Import tlc.utility.seq.

Set Implicit Arguments.
Unset Strict Implicit.
Unset Printing Implicit Defensive.

(* Errors that may be encountered while operating on patterns *)
Inductive pattern_error :=
| PEIll (p : pattern)
| PEParameter (j : nat) (t t' : term)
| PEConstructor (cp ct : constructor)
| PEStructure (p : pattern) (t : term).

(* Compute the set of parameters appearing in a pattern *)
Fixpoint pattern_params p :=
  match p with
  | PWildcard => [::]
  | PParameter j => [:: j]
  | PConstructor _ => [::]
  | PApplication f a => union (pattern_params f) (pattern_params a)
  end.

(* Compute the arity (number of parameters) in a pattern *)
Definition pattern_arity p :=
  let xs := pattern_params p in
  if nilp xs then 0 else (foldr maxn 0 xs).+1.

(* Determine whether a pattern is well-formed
 * Well-formed patterns are those whose set of used parameters is equal to the
 * set {j | j < |p|}.
 *)
Definition pattern_wf p :=
  set_eq (pattern_params p) (iota 0 (pattern_arity p)).

(* Match a term t against a pattern p and return the map of terms bound by p *)
Fixpoint match_pattern_rec p t (m : partial_map [eqType of nat] term) :=
  match p, t with
  | PWildcard, _ => pure m
  | PParameter j, t =>
    if m{@j} is Some t'
    then if t == t' then pure m else Failure (PEParameter j t t')
    else pure m{= j := t}
  | PConstructor cp, TConstructor ct =>
    if cp == ct then pure m
    else Failure (PEConstructor cp ct)
  | PApplication pf pa, TApplication tf ta =>
    m' <- match_pattern_rec pf tf m;
    match_pattern_rec pa ta m'
  | _, _ => Failure (PEStructure p t)
  end.

(* Match a term t against a pattern p and return the list of terms bound by p *)
Definition match_pattern p t :=
  if pattern_wf p then
    m <- match_pattern_rec p t [::];
    pure (map snd (sort (fun '(j1, _) '(j2, _) => j1 <= j2) m))
  else Failure (PEIll p).

(* Construct a term that will match a pattern
 * Each pattern wildcard and parameter is represented as a parameter in the
 * resulting term.  The term parameters are each produced on a separate binding
 * level.  For pattern parameters, the resulting term parameter is bound at the
 * level of the index of the pattern parameter.  Pattern wildcards are bound at
 * the levels above the arity of the pattern, in the order in which they are
 * encountered.
 *)
Fixpoint pattern_term_rec p n :=
  match p with
  | PWildcard => (TParameter (P n 0), n.+1)
  | PParameter j => (TParameter (P j 0), n)
  | PConstructor c => (TConstructor c, n)
  | PApplication pf pa =>
    let '(tf, n) := pattern_term_rec pf n in
    let '(ta, n) := pattern_term_rec pa n in
    (TApplication tf ta, n)
  end.
Definition pattern_term p :=
  if pattern_wf p
  then Some (pattern_term_rec p (pattern_arity p))
  else None.
Example pattern_term_1 :
  pattern_term {-p ($1, %, $0, $1) -} =
  Some ({-t ($(1, 0), $(2, 0), $(0, 0), $(1, 0)) -}, 3).
by []. Defined.

(* Destruct a case of an argument term into an assertion about that term when
 * it satisfies the case
 * A term that matches the case pattern is constructed using pattern_term.
 * That term is then equated to the argument term and the equality is wrapped
 * in N existential quantifiers, where N is the number of bindings in the
 * pattern term.
 *)
Definition destruct_case a (c : match_case) :=
  let '(p, b) := c in
  ts <- pattern_term p;
  let '(t, n) := ts in
  pure (iter n AExists {-A a = t -}).

(* Destruct a list of cases of an argument term into an assertion about that
 * term when it satisfies at least one of the cases
 *)
Fixpoint destruct_cases a (cs : match_cases) :=
  match cs with
  | nil => pure ATrue
  | [:: c] => destruct_case a c
  | c :: cs =>
    A <- destruct_case a c;
    As <- destruct_cases a cs;
    pure (AOr A As)
  end.


From discprob.basic Require Import base order nify.
From discprob.basic Require Export monad.
From discprob.prob Require Import prob countable finite stochastic_order.
Require Import Reals Fourier ProofIrrelevance FunctionalExtensionality Psatz.
From mathcomp Require Import ssreflect ssrbool ssrfun eqtype ssrnat seq div choice fintype finset finfun bigop.

Record ldist (A: Type) := mklDist {
  outcomes :> list (R * A);
  nonneg : all  (λ r, Rle_dec 0 r) [seq r.1 | r <- outcomes];
  sum1 : \big[Rplus/0]_(a <- map fst outcomes) a == R1
 }.

Arguments outcomes {_}.
Arguments sum1 {_}.
Arguments mklDist {_}.

Lemma ldist_irrel A (l1 l2: ldist A) :
  outcomes l1 = outcomes l2 → l1 = l2.
Proof.
  destruct l1 as [l1 pf1a pf1b].
  destruct l2 as [l2 pf2a pf2b] => //.
  rewrite /outcomes => ?. subst.
  f_equal; auto; exact: bool_irrelevance.
Qed.

Global Instance dist_ret: stdpp.base.MRet ldist.
  refine (λ A x, mklDist [::(R1, x)] _ _).
Proof.
{ abstract (rewrite //= andb_true_r; destruct (Rle_dec) => //=; nra). }
{ abstract (rewrite //= big_cons big_nil; apply /eqP => //=; field). }
Defined.

Definition dist_bind_aux :=
  λ {A B} (f: A → ldist B),
  fix go (l: seq (R * A)) :=
    match l with
      [::] => [::]
    | (r, x) :: l => map (λ py, (r * py.1, py.2)%R) (f x) ++ go l
    end.

Lemma dist_bind_pf1 {A B} (f: A → ldist B) d:
  all [eta (Rle_dec 0)] [seq r.1 | r <- dist_bind_aux f (outcomes d)].
Proof.
  apply /allP => r.
  destruct d as (l, pf1, pf2).
  rewrite //=; clear pf2.
  induction l as [| a l] => //=.
  destruct a as (r'&x) => //=.
  rewrite map_cat mem_cat.
  move /orP => [Hin|Hin].
  - move /allP in pf1. specialize (pf1 r').
    destruct (f x) as [l' pf1' pf2'].
    rewrite //= in Hin; clear pf2'.
    induction l' as [|? l'] => //=.
    rewrite //= in_cons in Hin.
    move /orP in Hin. destruct Hin as [Heq|Hin'].
    * move /eqP in Heq. rewrite Heq. destruct (Rle_dec 0 (_ * _)) as [|Hn] => //=; try nra.
      exfalso; apply Hn.
      apply Rmult_le_0_compat.
      ** cut (is_true (Rle_dec 0 r')); first by destruct (Rle_dec 0 r') => //=.
         apply pf1. rewrite in_cons. apply /orP; by left.
      ** cut (is_true (Rle_dec 0 (fst a))); first by destruct (Rle_dec 0 (fst a)) => //=.
         move /allP in pf1'. apply pf1'. rewrite in_cons. apply /orP; by left.
    * eapply IHl'; eauto; intros.
      rewrite //= in pf1'. move /andP in pf1'. destruct pf1'. auto.
  - eapply IHl; eauto.
      rewrite //= in pf1. move /andP in pf1. destruct pf1. auto.
Qed.

Lemma dist_bind_pf2 {A B} (f: A → ldist B) d:
  \big[Rplus/0]_(a<-[seq i.1 | i <- dist_bind_aux f (outcomes d)]) a == 1.
Proof.
 intros. apply /eqP.  destruct d as (l, pf1, pf2).
 rewrite //=; clear pf1.
 revert pf2.  move /eqP. replace 1 with R1 by auto. generalize R1 as t.
 induction l as [| a l] => t.
 - rewrite big_nil //= => ?.
 - destruct a as (r, x); rewrite //= big_cons => Hsum.
   rewrite map_cat -map_comp big_cat //= big_map //= (IHl (t - r)%R).
   * rewrite -big_distrr //=. specialize (sum1 (f x)). move /eqP. rewrite big_map => ->. field.
   * rewrite -Hsum. ring_simplify. done.
Qed.

Global Instance dist_bind: stdpp.base.MBind ldist :=
  λ A B f d, mklDist (dist_bind_aux f (outcomes d)) (dist_bind_pf1 f d) (dist_bind_pf2 f d).

Global Program Instance dist_join: stdpp.base.MJoin ldist :=
  λ A (lld: ldist (ldist A)),
  let fix go (ls : seq (R * ldist A)) :=
  match ls with [::] => [::] | (r, l) :: ls => [seq (r * i.1, i.2) | i <- l] ++ go ls end in
  mklDist (go lld) _ _.
Next Obligation.
  rewrite //= => A lld.
  apply /allP => r.
  destruct lld as (l, pf1, pf2).
  rewrite //=; clear pf2.
  induction l as [| a l] => //=.
  destruct a as (r'&x) => //=.
  rewrite map_cat mem_cat.
  move /orP => [Hin|Hin].
  - assert (0 ≤ r').
    { cut (is_true (Rle_dec 0 r')); first by destruct (Rle_dec 0 r') => //=.
      move /allP in pf1. apply pf1. rewrite in_cons //= eq_refl //. }
    clear pf1.
    destruct x as (l', pf1', pf2').
    rewrite //= in Hin. clear pf2'. induction l' as [| ? l'].
    { rewrite in_nil in Hin; done. }
    rewrite //= in_cons in Hin.
    move /orP in Hin. destruct Hin as [Heq|Hin'].
    * move /eqP in Heq. rewrite Heq.
      destruct (Rle_dec 0 (_ * _)) as [|Hn] => //=.
      contradiction Hn; apply Rmult_le_0_compat.
      ** auto.
      ** cut (is_true (Rle_dec 0 (fst a))); first by destruct (Rle_dec 0 (fst a)) => //=.
         move /allP in pf1'. apply pf1'. rewrite in_cons. apply /orP; by left.
    * eapply IHl'; eauto; intros.
      rewrite //= in pf1'. move /andP in pf1'. destruct pf1'; done.
  - eapply IHl; eauto.
    rewrite //= in pf1. move /andP in pf1. destruct pf1; done.
Qed.
Next Obligation.
 destruct lld as (l, pf1, pf2).
 apply /eqP.
 rewrite //=; clear pf1.
 revert pf2. move /eqP. generalize R1 as t.
 induction l as [| a l] => t.
 - intros. rewrite //=.
 - destruct a as (r, x); rewrite //= big_cons => Hsum.
   rewrite map_cat -map_comp big_cat //= big_map //= (IHl (t - r)%R).
   * rewrite -big_distrr //=. specialize (sum1 (x)). move /eqP. rewrite big_map => ->. field.
   * rewrite -Hsum; ring_simplify; done.
Qed.

Global Program Instance dist_fmap: stdpp.base.FMap ldist :=
  λ A B (f: A → B) (ld: ldist A),
  let fix go (l : seq (R * A)) := match l with [::] => [::] | (r, x) :: l => (r, f x) :: go l end in
  mklDist (go ld) _ _.
Next Obligation.
  intros A B f ld.
 apply /allP => r.
 destruct ld as (l, pf1, pf2).
 rewrite //=. clear pf2.
 induction l as [|a l] => //=.
 destruct a as (r', x) => //=.
 rewrite in_cons.
 move /orP => [Heq|Hin].
 - move /eqP in Heq. rewrite Heq.
   move /andP in pf1. apply pf1.
 - apply IHl, Hin. intros. rewrite //= in pf1. move /andP in pf1.
   destruct pf1; auto.
Qed.
Next Obligation.
 intros A B f ld.
 apply /eqP.
 destruct ld as (l, pf1, pf2).
 rewrite //=; clear pf1.
 revert pf2. move /eqP. generalize R1 as t.
 induction l as [| a l] => t.
 - intros. rewrite //=.
 - destruct a as (r, x); rewrite //= big_cons => Hsum.
   rewrite big_cons. rewrite (IHl (t - r)%R); first field.
   rewrite -Hsum; ring_simplify; done.
Qed.

Lemma bind_join_fmap {A B} (m: ldist A) (g: A → ldist B):
  outcomes (mbind g m) = outcomes (mjoin (fmap g m)).
Proof.
  rewrite /mbind/mjoin/fmap/dist_bind/dist_join/dist_fmap.
  destruct m as (lm, pf1m, pf2m) => //=.
  clear pf1m pf2m.
  induction lm as [|a l] => //.
  destruct a as (r, x) => //=; f_equal; done.
Qed.

Lemma in_ldist_bind {A B : eqType} (f: A → ldist B) r (x : B) (m: ldist A) :
  (r, x) \in outcomes (x ← m; f x) →
   ∃ r'' r' y, (r'', x) \in outcomes (f y)
               ∧ (r', y) \in outcomes m
               ∧ r = r' * r''.
Proof.
  rewrite /mbind/dist_bind//=.
  destruct m as (l, pf1, pf2) => //=.
  clear pf1 pf2.
  induction l as [|p l IH] => //.
  rewrite //=. destruct p as (r'&y) => //=.
  rewrite mem_cat. move /orP => [Hin1|Hintl].
  - move /mapP in Hin1. destruct Hin1 as [[r'' y'] Hin Heq].
    rewrite //= in Heq. inversion Heq; subst.
    exists r'', r', y; repeat split; auto.
    rewrite in_cons eq_refl //.
  - destruct IH as (r1&r2&y'&?&?&?); auto.
    exists r1, r2, y'; repeat split; auto.
    rewrite in_cons. apply /orP; by right.
Qed.

Lemma in_ldist_bind2 {A B : eqType} (f: A → ldist B) r' r'' (x : B) y (m: ldist A) :
   (r'', x) \in outcomes (f y) →
   (r', y) \in outcomes m →
   (r' * r'', x) \in outcomes (m ≫= f).
Proof.
  rewrite /mbind/dist_bind//=.
  destruct m as (l, pf1, pf2) => //=.
  clear pf1 pf2.
  induction l as [|p l IH] => //.
  rewrite //= => Hin1 Hin2'. destruct p as (?&?) => //=.
  rewrite mem_cat.
  rewrite in_cons in Hin2'. move /orP in Hin2'. destruct Hin2' as [Heq|Htl]; apply /orP.
  - left. move /eqP in Heq. inversion Heq; subst.
    apply /mapP. exists (r'',  x); auto.
  - right. apply IH; eauto.
Qed.

Lemma nth_lt_size {A: eqType} (d: A) m l (Hsize: (m < size l)%N):
  nth d l m \in l.
Proof.
  apply /all_nthP => //.
  apply /allP => x Hin //=.
Qed.

(* Every ldist can be regarded as a random variable *)
Definition dist_of_ldist {A: eqType} (ld: ldist A) : distrib [finType of ordinal (size ld)].
  apply (@mkDistrib _ (λ n, (nth R0 (map fst ld) (nat_of_ord n)))).
  - destruct a as [m Hlt]; destruct ld as (l, pf1, pf2); rewrite //=.
    rewrite //= in Hlt. move /allP in pf1.
    apply Rle_ge.
    cut (is_true (Rle_dec 0 (nth 0 [seq (fst i) | i <- l] m)));
      first by (destruct (Rle_dec) => //=).
    apply /pf1/nth_lt_size. rewrite size_map //.
  - destruct ld as (l, pf1, pf2); rewrite //=. move /eqP in pf2.
    replace 1 with R1 by auto.
    rewrite -pf2 -sum_index_ordinal -SeriesF_big_op.
    apply SeriesF_is_seriesC.
Defined.

Definition ld_def {A} (ld: seq (R *A) ) (n: 'I_(size ld)) : A.
  induction ld as [|(r, x)].
  - destruct n as (n&Hlt). inversion Hlt.
  - exact x.
Qed.

Lemma sum_index_ordinal_Fpair {A: eqType} (l: seq (R * A)) F:
  \big[Rplus/0]_(a in 'I_(size l))
   F (nth 0 [seq i.1 | i <- l] a, nth (ld_def l a) [seq i.2 | i <- l] a) =
  \big[Rplus/0]_(a<- l) F a.
Proof.
    induction l as [| a l].
    * rewrite //= big_nil big_ord0 //.
    * transitivity (\big[Rplus/0%R]_(0<= a' < (size l).+1)
                     F (nth 0 ([seq i.1 | i <- a :: l]) a', (nth (snd a) ([seq i.2 | i <- a :: l]) a'))).
      { rewrite big_mkord. apply eq_bigr => ??. rewrite (set_nth_default (snd a)); first done.
        rewrite size_map. auto. }
      rewrite //= big_cons big_ltn; last auto.
      destruct a; f_equal; rewrite big_add1 //=. rewrite -IHl big_mkord.
      apply eq_bigr => i H. rewrite (set_nth_default (ld_def l i)); first done.
      rewrite size_map. done.
Qed.

Definition rvar_of_ldist {A: eqType} (ld: ldist A) : rvar (dist_of_ldist ld) A :=
  mkRvar (dist_of_ldist ld) (λ n, (nth (ld_def ld n) (map snd ld) (nat_of_ord n))).

(* There must be some nice way to do this... I really need to learn ssreflect better *)
Lemma setP_help {A: finType} (P: A -> bool) (i: A):
  (pred_of_set [set x | P x]) i = P i.
Proof.
  rewrite -in_set.
  transitivity (i \in [set x | P x]); first by rewrite in_set.
  rewrite in_set. done.
Qed.

Lemma pr_rvar_ldist {A: eqType} (ld: ldist A) r :
  pr_eq (rvar_of_ldist ld) r = \big[Rplus/0%R]_(a <- outcomes ld | snd a == r) fst a.
Proof.
  rewrite /pr_eq/pr.
  rewrite SeriesC_fin_big -big_mkcondr.
  rewrite /dist_of_ldist //=.
  apply (sum_index_ordinal_P _ _ (R0, r)) => //= i.
  rewrite (nth_map (nth (0, r) ld i)) //.
  rewrite (set_nth_default (0, r)) //.
Qed.

Lemma pr_gt_rvar_ldist (ld: ldist R) r :
  pr_gt (rvar_of_ldist ld) r = \big[Rplus/0%R]_(a <- outcomes ld | Rgt_dec (snd a) r) fst a.
Proof.
  rewrite /pr_gt/pr.
  rewrite SeriesC_fin_big -big_mkcondr.
  rewrite /dist_of_ldist //=.
  apply (sum_index_ordinal_P _ _ (R0, r)) => //= i.
  rewrite (nth_map (nth (0, r) ld i)) //.
  rewrite (set_nth_default (0, r)) //.
Qed.

Lemma sum_mjoin {B: eqType} (lld: ldist (ldist B)) P:
  \big[Rplus/0]_(a <- outcomes (mjoin lld) | P (snd a)) a.1 =
  \big[Rplus/0]_(ld <- lld) \big[Rplus/0]_(a <- ld.2 | P (snd a)) (ld.1 * a.1).
Proof.
  destruct lld as (lld, pf1, pf2) => //=.
  clear pf1 pf2.
  induction lld as [|[r ld] lld] => //=.
  { rewrite ?big_nil //. }
  rewrite big_cons big_cat //=.
  f_equal; eauto.
  rewrite big_map //=.
Qed.

Lemma sum_fmap {A B} (f: A → B) (ld: ldist A) F:
  \big[Rplus/0]_(x <- outcomes (f <$> ld)) (F (x.1) (x.2)) =
  \big[Rplus/0]_(x <- ld) (F (x.1) (f x.2)).
Proof.
  destruct ld as (ld, pf1, pf2) => //=.
  clear pf1 pf2.
  induction ld as [|[r x] ld] => //=.
  { rewrite ?big_nil //. }
  rewrite ?big_cons //=.
  f_equal; eauto.
Qed.

Lemma nat_of_ord_map_iota {A} n (F: nat → A):
  [seq (F (nat_of_ord i)) | i <- enum 'I_n] =
  [seq (F i) | i <- iota 0%N n].
Proof.
    by rewrite -val_enum_ord map_comp.
Qed.

Lemma nat_of_ord_map_iota' {A} n (F: nat → A):
  [seq (F (nat_of_ord i)) | i <- Finite.enum [finType of 'I_n]] =
  [seq (F i) | i <- iota 0%N n].
Proof.
  rewrite -nat_of_ord_map_iota. f_equal.
  rewrite -enumT//=.
Qed.

Lemma map_nth_iota {A: eqType} (l: seq A) a:
      [seq nth a l i | i <- iota 0 (size l)] = l.
Proof.
  induction l => //=.
  replace 1%nat with (1 + 0)%nat by auto with *.
  rewrite iota_addl -map_comp; f_equal => //=.
Qed.

Lemma img_rvar_of_ldist {A: eqType} (h: ldist A):
  map sval (Finite.enum [finType of (imgT (rvar_of_ldist h))]) = undup [seq i.2 | i <- h].
Proof.
  rewrite {1}[@Finite.enum]unlock //=. rewrite img_fin_enum_sval.
  assert (a: A).
  { destruct h as (l, pf1, pf2). destruct l.
    - rewrite big_nil in pf2. move /eqP in pf2. exfalso.
      rewrite //= in pf2. fourier.
    - exact (snd p).
  }
  f_equal.
  etransitivity.
  - apply (@eq_map _ _ _ (λ n, nth a [seq i.2 | i <- h] (nat_of_ord n))).
    { intro i. erewrite set_nth_default; auto. rewrite size_map. done. }
  - rewrite -enumT. rewrite (nat_of_ord_map_iota (size h) (λ n, nth a [seq (snd i) | i <- h] n)).
  destruct h as (l, pf1, pf2) => //=; clear pf1 pf2.
  rewrite -(size_map snd l) map_nth_iota //.
Qed.

Lemma img_rvar_of_ldist' {A: eqType} (h: ldist A) a:
  a \in img (rvar_of_ldist h) = (a \in undup [seq (snd i) | i <- h]).
Proof.
  apply (introTF (exCP _)); case: ifP.
  - rewrite mem_undup. move /mapP. intros [x Hin Heq]; subst.
    rewrite //=.
    assert (index (snd x) (map snd h) < size h)%nat as Hsize.
    { rewrite -(size_map snd) index_mem. apply /mapP. eexists; eauto. }
    exists (Ordinal Hsize) => //=.
    rewrite nth_index => //=. apply /mapP. eexists; eauto.
  - rewrite mem_undup => Hfalse. intros (x&Heq).
    move /negP in Hfalse. apply Hfalse. move /eqP in Heq.
    rewrite -Heq. rewrite //=. apply nth_lt_size. rewrite size_map. destruct x. auto.
Qed.

(* Lemmas for simplifying monadic expressions *)
Lemma pr_mbind_ldist1 {A B: eqType} (h: ldist A) (f: A → ldist B) (b: B) :
  pr_eq (rvar_of_ldist (mbind f h)) b =
  \big[Rplus/0]_(a : imgT (rvar_of_ldist h))
   (pr_eq (rvar_of_ldist h) (sval a) * pr_eq (rvar_of_ldist (f (sval a))) b).
Proof.
  transitivity (Ex (rvar_comp (rvar_of_ldist h) (λ a, pr_eq (rvar_of_ldist (f a)) b))); last first.
  { rewrite Ex_fin_comp. eapply eq_bigr => ??. field. }
  rewrite pr_rvar_ldist Ex_fin_pt.
  symmetry. etransitivity.
  { apply eq_bigr => ?? /=. rewrite !pr_rvar_ldist. done. }
  rewrite bind_join_fmap.
  rewrite (sum_mjoin _ (λ x, x == b)).
  rewrite (sum_fmap f h (λ x y, \big[Rplus/0]_(a <- y | a.2 == b) (x * a.1))).
  rewrite /index_enum.
  rewrite //= (sum_index_ordinal_Fpair _ (λ xy,
     \big[Rplus/0]_(a <- f (xy.2) | a.2 == b) a.1 * (xy.1))).
  eapply eq_bigr => x _.
  rewrite Rmult_comm.
  rewrite big_distrr //.
Qed.

Lemma pr_mbind_ldist2 {A B: eqType} (h: ldist A) (f: A → ldist B) (b: B) :
  pr_eq (rvar_of_ldist (mbind f h)) b =
  \big[Rplus/0]_(a <- undup [seq i.2 | i <- h])
   (pr_eq (rvar_of_ldist h) a * pr_eq (rvar_of_ldist (f a)) b).
Proof.
  rewrite -img_rvar_of_ldist pr_mbind_ldist1 //.
  rewrite big_map //=.
Qed.

Lemma pr_mret_simpl {A: eqType} (x y: A):
  pr_eq (rvar_of_ldist (mret x)) y = if (x == y) then 1 else 0.
Proof.
  rewrite pr_rvar_ldist big_cons big_nil //=. case: ifP => ?; field.
Qed.

Lemma pr_mbind_mret {A B: eqType} (h: ldist A) (f: A → B) (b: B) :
  pr_eq (rvar_of_ldist (mbind (λ x, mret (f x)) h)) b =
  pr_eq (rvar_comp (rvar_of_ldist h) f) b.
Proof.
  rewrite pr_mbind_ldist1 //.
  symmetry. rewrite {1}/pr_eq pr_eq_alt_comp.
  eapply eq_bigr => a _.
  rewrite pr_mret_simpl.
  case: ifP => ?; rewrite //= ?Rmult_1_r ?Rmult_0_r; done.
Qed.

Lemma pr_mbind_mret_inj {A B: eqType} (h: ldist A) (f: A → B) (a: A) :
  injective f →
  pr_eq (rvar_of_ldist (mbind (λ x, mret (f x)) h)) (f a)=
  pr_eq (rvar_of_ldist h) a.
Proof.
  intros Hinj.
  rewrite pr_mbind_mret //= {1}/pr_eq pr_eq_alt_comp.
  rewrite {2}/pr_eq pr_eq_alt.
  eapply eq_bigr => i _. symmetry.
  case: ifP; move /eqP.
  * intros Heq1. rewrite Heq1 eq_refl. done.
  * intros Hneq. case: ifP; auto. move /eqP. intros Heq.
    apply Hinj in Heq. congruence.
Qed.

Lemma pr_mbind_mret_inj0 {A B: eqType} (h: ldist A) (f: A → B) (b: B) :
  (∀ a, f a ≠ b) →
  pr_eq (rvar_of_ldist (mbind (λ x, mret (f x)) h)) b = 0.
Proof.
  intros Hneq. rewrite pr_mbind_mret {1}/pr_eq pr_eq_alt_comp.
  apply big1 => i _. case: ifP; auto. move /eqP => Heq. exfalso.
  apply (Hneq (sval i)); auto.
Qed.

Lemma pr_gt_mbind_ldist1 {A: eqType} (h: ldist A) (f: A → ldist R) (b: R) :
  pr_gt (rvar_of_ldist (mbind f h)) b =
  \big[Rplus/0]_(a : imgT (rvar_of_ldist h))
   (pr_eq (rvar_of_ldist h) (sval a) * pr_gt (rvar_of_ldist (f (sval a))) b).
Proof.
  transitivity (Ex (rvar_comp (rvar_of_ldist h) (λ a, pr_gt (rvar_of_ldist (f a)) b))); last first.
  { rewrite Ex_fin_comp. eapply eq_bigr => ??. field. }
  rewrite pr_gt_rvar_ldist Ex_fin_pt.
  symmetry. etransitivity.
  { apply eq_bigr => ?? /=. rewrite !pr_gt_rvar_ldist. done. }
  rewrite bind_join_fmap.
  rewrite (sum_mjoin _ (λ x, Rgt_dec x b)).
  rewrite (sum_fmap f h (λ x y, \big[Rplus/0]_(a <- y | Rgt_dec a.2 b) (x * a.1))).
  rewrite /index_enum.
  rewrite //= (sum_index_ordinal_Fpair _ (λ xy,
     \big[Rplus/0]_(a <- f (xy.2) | Rgt_dec a.2 b) a.1 * (xy.1))).
  eapply eq_bigr => x _.
  rewrite Rmult_comm.
  rewrite big_distrr //.
Qed.

Lemma pr_gt_mbind_ldist2 {A: eqType} (h: ldist A) (f: A → ldist R) (b: R) :
  pr_gt (rvar_of_ldist (mbind f h)) b =
  \big[Rplus/0]_(a <- undup [seq i.2 | i <- h])
   (pr_eq (rvar_of_ldist h) a * pr_gt (rvar_of_ldist (f a)) b).
Proof.
  rewrite -img_rvar_of_ldist pr_gt_mbind_ldist1 //.
  rewrite big_map //=.
Qed.

Lemma pr_gt_mret_simpl (x y: R):
  pr_gt (rvar_of_ldist (mret x)) y = if (Rgt_dec x y) then 1 else 0.
Proof.
  rewrite pr_gt_rvar_ldist big_cons big_nil //=.
  destruct Rgt_dec => //=; try fourier; try nra.
Qed.

Lemma pr_gt_mbind_mret {A: eqType} (h: ldist A) (f: A → R) (b: R) :
  pr_gt (rvar_of_ldist (mbind (λ x, mret (f x)) h)) b =
  pr_gt (rvar_comp (rvar_of_ldist h) f) b.
Proof.
  rewrite pr_gt_mbind_ldist1 //.
  symmetry. rewrite {1}/pr_gt pr_gt_alt_comp.
  eapply eq_bigr => a _.
  rewrite pr_gt_mret_simpl.
  destruct Rgt_dec => //=; rewrite ?Rmult_1_r ?Rmult_0_r //=.
Qed.

Lemma Ex_mbind_ldist1 {A: eqType} (h: ldist A) (f: A → ldist R) :
  Ex (rvar_of_ldist (mbind f h)) =
  \big[Rplus/0]_(a : imgT (rvar_of_ldist h))
   (pr_eq (rvar_of_ldist h) (sval a) * Ex (rvar_of_ldist (f (sval a)))).
Proof.
  rewrite /Ex. rewrite SeriesC_fin_big.
  f_equal; etransitivity.
  {eapply eq_bigr => x _. rewrite pr_mbind_ldist1. rewrite big_distrl //. }
  rewrite exchange_big.
  apply eq_bigr => x _.
  rewrite SeriesC_fin_big.
  rewrite big_distrr //=.
  rewrite /index_enum//=.
  rewrite -(big_map sval (λ x, true) (λ i, _ * pr_eq _ i * i)).
  rewrite -(big_map sval (λ x, true) (λ i, _ * (pr_eq _ i * i))).
  rewrite img_rvar_of_ldist.
  rewrite (img_rvar_of_ldist (dist_bind _ _ f h)).
  eapply sum_reidx_map_sym with (h := id).
  - intros; rewrite //=. rewrite //=. rewrite Rmult_assoc. done.
  - intros ? Hin' ? Hnin. rewrite (pr_img_nin (rvar_of_ldist (f (sval x)))); first by ring.
    rewrite img_rvar_of_ldist'. auto.
  - intros ? Hin' ? Hnin. exfalso. apply Hnin.
    destruct x as (x&Hin).
    rewrite //= in Hin'. rewrite img_rvar_of_ldist' in Hin.
    rewrite ?mem_undup in Hin Hin'.
    move /mapP in Hin.
    move /mapP in Hin'.
    destruct Hin as [(r'&x') Hin1 ?]; subst.
    destruct Hin' as [(r''&y'') Hin2 ?]; subst.
    specialize (in_ldist_bind2 f _ _ _ _ h Hin2 Hin1).
    rewrite //= => Hin_comb.
    exists y''.  rewrite mem_undup; split; auto.
    apply /mapP; auto.
    exists (r' * r'', y''); auto.
  - apply undup_uniq.
  - apply undup_uniq.
  - intros ??; done.
Qed.

Lemma Ex_mbind_ldist2 {A: eqType} (h: ldist A) (f: A → ldist R) :
  Ex (rvar_of_ldist (mbind f h)) =
  \big[Rplus/0]_(a <- undup [seq i.2 | i <- h]) (pr_eq (rvar_of_ldist h) a * Ex (rvar_of_ldist (f a))).
Proof.
  rewrite -img_rvar_of_ldist Ex_mbind_ldist1 //.
  rewrite big_map //=.
Qed.

Lemma Ex_mret_simpl (x: R):
  Ex (rvar_of_ldist (mret x)) = x.
Proof.
  rewrite /Ex SeriesC_fin_big /index_enum.
  rewrite -(big_map sval (λ x, true) (λ i, pr_eq _ i * i)).
  rewrite img_rvar_of_ldist //= big_cons big_nil.
  rewrite pr_mret_simpl eq_refl. f_equal; field.
Qed.

Lemma Ex_mbind_mret {A: eqType} (h: ldist A) (f: A → R) :
  Ex (rvar_of_ldist (mbind (λ x, mret (f x)) h)) =
  Ex (rvar_comp (rvar_of_ldist h) f).
Proof.
  rewrite Ex_mbind_ldist1 //.
  rewrite Ex_fin_comp.
  eapply eq_bigr => a _.
  rewrite Ex_mret_simpl.
  f_equal.
Qed.

Lemma outcomes_eq_dist {A: eqType} (r1 r2: ldist A) :
  outcomes r1 = outcomes r2 →
  eq_dist (rvar_of_ldist r1) (rvar_of_ldist r2).
Proof.
  rewrite /eq_dist => Heq x.
  rewrite ?pr_rvar_ldist Heq. done.
Qed.

Lemma ldist_left_id {A B: eqType} (x: A) (f: A → ldist B):
  mbind f (mret x) = f x.
Proof.
  apply ldist_irrel => //=. rewrite cats0.
  rewrite (@eq_map _ _ _ id); first by apply map_id.
  intros (?&?) => //=; f_equal; auto; ring_simplify; done.
Qed.

Lemma ldist_right_id {A: eqType} (m: ldist A) :
  mbind mret m = m.
Proof.
  apply ldist_irrel => //=. induction (outcomes m) as [| a l] => //=.
  destruct a as (?&?) => //=; f_equal; auto.
  rewrite Rmult_1_r. done.
Qed.

Lemma ldist_assoc {A B C: eqType} (m: ldist A) (f: A → ldist B) (g: B → ldist C) :
  mbind g (mbind f m) = mbind (λ x, mbind g (f x)) m.
Proof.
  apply ldist_irrel.
  destruct m as [m pf1 pf2] => //=.
  clear pf1 pf2.
  assert (Haux: ∀ A B (f: A → ldist B) l1 l2, dist_bind_aux f (l1 ++ l2) =
                               dist_bind_aux  f l1 ++ dist_bind_aux f l2).
  { induction l1 as [| a l] => //=.
    intros. destruct a as (r&x).
    rewrite IHl catA. done.
  }
  induction m as [| a l]; first by rewrite //=.
  destruct a as (r&x).
  rewrite //= Haux; eauto.
  f_equal; eauto.
  clear -Haux. induction (outcomes (f x)) as [| a' l'] => //=.
  destruct a' as (r'&x'). rewrite map_cat; f_equal; eauto.
  rewrite -map_comp => //=.
  apply eq_map => c.
  destruct c as (r''&x'') => //=.
  f_equal; auto; field.
Qed.

Lemma ldist_bind_ext {A B: eqType} m (f g: A → ldist B):
  (∀ a, f a = g a) →
  mbind f m = mbind g m.
Proof.
  destruct m as [l pf1 pf2] => Hext //=.
  rewrite /dist_bind.
  apply ldist_irrel => //=.
  clear pf1 pf2. induction l as [| a l] => //=.
  destruct a as (r, x). by rewrite IHl Hext.
Qed.

Lemma ldist_fmap_bind {A B C: eqType} m (h: A → B) (f: B → ldist C):
  mbind f (x ← m; mret (h x)) = mbind (λ x, f (h x)) m.
Proof.
  rewrite ldist_assoc.
  apply ldist_bind_ext => a. rewrite ldist_left_id. done.
Qed.

Lemma eq_dist_ldist_bind_ext {A B: eqType} m (f g: A → ldist B):
  (∀ a, eq_dist (rvar_of_ldist (f a)) (rvar_of_ldist (g a))) →
  eq_dist (rvar_of_ldist (mbind f m)) (rvar_of_ldist (mbind g m)).
Proof.
  rewrite /eq_dist => Heq b. rewrite ?pr_mbind_ldist2.
  eapply eq_bigr => a _. rewrite Heq. done.
Qed.

Lemma le_dist_ldist_bind_ext {A: eqType} m (f g: A → ldist R):
  (∀ a k, pr_gt (rvar_of_ldist (f a)) k <= pr_gt (rvar_of_ldist (g a)) k) →
  (∀ k, pr_gt (rvar_of_ldist (mbind f m)) k <= pr_gt (rvar_of_ldist (mbind g m)) k).
Proof.
  rewrite /eq_dist => Heq b. rewrite ?pr_gt_mbind_ldist2.
  apply Rle_bigr => a _. apply Rmult_le_compat_l.
  * apply Rge_le, pr_eq_ge_0.
  * done.
Qed.

Lemma eq_dist_ldist_bind_congr {A B: eqType} m m' (f g: A → ldist B):
  eq_dist (rvar_of_ldist m) (rvar_of_ldist m') →
  (∀ a, eq_dist (rvar_of_ldist (f a)) (rvar_of_ldist (g a))) →
  eq_dist (rvar_of_ldist (mbind f m)) (rvar_of_ldist (mbind g m')).
Proof.
  rewrite /eq_dist => Hdist Heq b. rewrite ?pr_mbind_ldist2.
  rewrite bigop_cond_non0 [a in _ = a]bigop_cond_non0.
  eapply (sum_reidx_map _ _ _ _ id).
  - intros. rewrite Hdist Heq. done.
  - rewrite //=; intros a Hin Hnon0. rewrite -?img_rvar_of_ldist'; split.
    * apply pr_img. rewrite -Hdist. move /eqP in Hnon0.
      edestruct (ge_pr_0 (rvar_dist (rvar_of_ldist m))) as [|Heq0]; first by eauto.
      rewrite /pr_eq//= in Hnon0. rewrite Heq0 in Hnon0. nra.
    * rewrite -Heq -Hdist. done.
  - rewrite //= => a Hin. rewrite -?img_rvar_of_ldist' => Hnon0 Hfalse.
    exfalso; apply Hfalse. exists a.
    rewrite Heq Hdist; split; auto.
    rewrite -?img_rvar_of_ldist'.
    apply pr_img. rewrite Hdist. move /eqP in Hnon0.
    edestruct (ge_pr_0 (rvar_dist (rvar_of_ldist m'))) as [|Heq0]; first by eauto.
    rewrite /pr_eq in Hnon0. rewrite Heq0 in Hnon0. nra.
  - apply undup_uniq.
  - apply undup_uniq.
  - done.
Qed.

Lemma ldist_bind_swap {A B C: eqType} (m1: ldist A) (m2: ldist B) (f: A → B → ldist C):
  eq_dist (rvar_of_ldist (a ← m1;
                            b ← m2;
                            f a b))
          (rvar_of_ldist (b ← m2;
                            a ← m1;
                            f a b)).
Proof.
  intros c.
  rewrite ?pr_mbind_ldist2.
  etransitivity.
  { eapply eq_bigr. intros. rewrite pr_mbind_ldist2. rewrite big_distrr //. }
  symmetry.
  etransitivity.
  { eapply eq_bigr. intros. rewrite pr_mbind_ldist2. rewrite big_distrr //. }
  rewrite exchange_big.
  do 2 (eapply eq_bigr => ??).
  rewrite //=. nra.
Qed.

Lemma pr_eq_bind_pair {A B: eqType} (m1: ldist A) (m2: ldist B) (a: A) (b: B):
  pr_eq (rvar_of_ldist (a ← m1; b ← m2; mret (a, b))) (a, b) =
  pr_eq (rvar_of_ldist m1) a * pr_eq (rvar_of_ldist m2) b.
Proof.
  rewrite ?pr_mbind_ldist2.
  etransitivity.
  { eapply eq_bigr. intros a' _.
    rewrite pr_mbind_ldist2. apply Rmult_eq_compat_l.
    eapply eq_bigr. intros b' _.
    rewrite pr_mret_simpl.
    transitivity (pr_eq (rvar_of_ldist (mret a')) a *
                     (pr_eq (rvar_of_ldist m2) b' *
                      pr_eq (rvar_of_ldist (mret b')) b)); last done.
    rewrite ?pr_mret_simpl.
    case: ifP; move /eqP => Heq.
    - inversion Heq as [[Heqa Heqb]].
      subst. rewrite ?eq_refl. nra.
    - case: ifP; move /eqP => Heqa.
      { case: ifP; move /eqP => Heqb.
        * subst. congruence.
        * nra.
      }
      nra.
  }
  etransitivity.
  {
    eapply eq_bigr. intros. rewrite -big_distrr.
    rewrite -pr_mbind_ldist2. rewrite -Rmult_assoc. done.
  }
  rewrite -big_distrl. rewrite ldist_right_id //=; f_equal.
  specialize (pr_mbind_ldist2 m1 (λ x, mret x)) => //=. intros <-.
  (* For some reason I cannot just do the appropriate rewrite... *)
  assert (Hfinal: ∀ {A: eqType} (m1: ldist A) a,
             pr_eq (rvar_of_ldist (x ← m1; mret x)) a = pr_eq (rvar_of_ldist m1) a).
  { intros.  rewrite ldist_right_id. done. }
  eapply Hfinal.
Qed.

Lemma eq_dist_bind_rvar_pair {A B: eqType} (m1: ldist A) (m2: ldist B):
  eq_dist (rvar_of_ldist (a ← m1; b ← m2; mret (a, b)))
          (rvar_pair (rvar_of_ldist m1) (rvar_of_ldist m2)).
Proof.
  intros (a&b). rewrite pr_eq_bind_pair pr_eq_rvar_pair //.
Qed.

Program Definition bernoulli (p: R) (Hrange: (0 <= p <= 1)%R) : ldist bool :=
  mklDist [:: (p, true) ; (1 - p, false)%R] _ _.
Next Obligation.
  rewrite //=; intros p (?&?);
  repeat (apply /andP; split; auto); destruct (Rle_dec) => //. nra.
Qed.
Next Obligation.
  rewrite //=; intros p (?&?);
  rewrite ?big_cons ?big_nil; apply /eqP => //=; field.
Qed.

Lemma pr_bernoulli_true (p: R) (Hrange: 0 <= p <= 1):
  pr_eq (rvar_of_ldist (bernoulli p Hrange)) true = p.
Proof.
  rewrite pr_rvar_ldist //= ?big_cons ?big_nil //=. field.
Qed.

Lemma pr_bernoulli_false (p: R) (Hrange: 0 <= p <= 1):
  pr_eq (rvar_of_ldist (bernoulli p Hrange)) false = 1 - p.
Proof.
  rewrite pr_rvar_ldist //= ?big_cons ?big_nil //=. field.
Qed.

Module test.

  Program Definition test1 : ldist (bool * bool) :=
    x ← bernoulli (1/2) _ ;
    y ← bernoulli (1/2) _;
    mret (x, y).
  Next Obligation. split; fourier. Qed.
  Next Obligation. split; fourier. Qed.

  (* Eval compute in (outcomes test1). *)

  Definition von_neumann_coin (p: R) (Hrange: 0 <= p <= 1): ldist (option bool) :=
    x ← bernoulli p Hrange;
    y ← bernoulli p Hrange;
    match x, y with
      | true, false => mret (Some true)
      | false, true => mret (Some false)
      | _, _ => mret None
    end.

  Fixpoint von_neumann_repeat (p: R) (Hrange: 0 <= p <= 1) (n: nat): ldist (option bool) :=
    match n with
      | 0 => mret None
      | S n' =>
        res ← von_neumann_coin p Hrange;
        match res with
          | None => von_neumann_repeat p Hrange n'
          | _ => mret res
        end
    end.

  Lemma range_three_fourths: 0 <= 3/4 <= 1. Proof. split; fourier. Qed.
  (* Eval compute in (outcomes (von_neumann_coin (3/4) range_three_fourths)). *)

  Lemma von_neumann_fair (p: R) (Hrange: 0 <= p <= 1) :
    pr_eq (rvar_of_ldist (von_neumann_coin p Hrange)) (Some true) =
    pr_eq (rvar_of_ldist (von_neumann_coin p Hrange)) (Some false).
  Proof.
    rewrite ?pr_rvar_ldist //= ?big_cons ?big_nil //=. field.
  Qed.

  Lemma von_neumann_fail (p: R) (Hrange: 0 <= p <= 1) :
    pr_eq (rvar_of_ldist (von_neumann_coin p Hrange)) None =
    1 - 2 * (p * (1 - p)).
  Proof.
    rewrite ?pr_rvar_ldist //= ?big_cons ?big_nil //=. field.
  Qed.

  (* These are not yet the most efficient ways to prove these;
     mainly ``brute-force'' computation; see below for a way that uses
     facts about monadic operations and a simple 'coupling' argument *)
  Lemma von_neumann_repeat_fair (p: R) (Hrange: 0 <= p <= 1) n :
    pr_eq (rvar_of_ldist (von_neumann_repeat p Hrange n)) (Some true) =
    pr_eq (rvar_of_ldist (von_neumann_repeat p Hrange n)) (Some false).
  Proof.
    induction n.
    - rewrite ?pr_rvar_ldist //= ?big_cons ?big_nil //=.
    - rewrite ?pr_rvar_ldist //= ?big_cons ?big_nil //=.
      rewrite ?big_cat //=. f_equal.
      * rewrite ?big_map //= -?big_distrr; f_equal.
        by rewrite ?pr_rvar_ldist in IHn.
      * rewrite ?cats0.
        rewrite ?big_cons //=; f_equal; first field.
        rewrite ?big_map //= -?big_distrr; f_equal.
        by rewrite ?pr_rvar_ldist in IHn.
  Qed.

  Lemma von_neumann_repeat_fail n (p: R) (Hrange: 0 <= p <= 1) :
    pr_eq (rvar_of_ldist (von_neumann_repeat p Hrange n)) None =
    (1 - 2 * (p * (1 - p)))^n.
  Proof.
    induction n.
    - rewrite ?pr_rvar_ldist //= ?big_cons ?big_nil //=. field.
    - rewrite ?pr_rvar_ldist //= ?big_cons ?big_nil //=.
      rewrite ?cats0 ?big_cat //=.
      rewrite ?pr_rvar_ldist in IHn.
      rewrite ?big_map -?big_distrr //= IHn.
      rewrite ?big_cons //=.
      rewrite ?big_map //= -?big_distrr //= IHn.
      field.
  Qed.

  Lemma von_neumann_repeat_fair' (p: R) (Hrange: 0 <= p <= 1) n :
    pr_eq (rvar_of_ldist (von_neumann_repeat p Hrange n)) (Some true) =
    pr_eq (rvar_of_ldist (von_neumann_repeat p Hrange n)) (Some false).
  Proof.
    induction n.
    - rewrite ?pr_rvar_ldist //= ?big_cons ?big_nil //=.
    - rewrite !pr_mbind_ldist2.
      (* Coupling *)
      eapply (sum_reidx_surj2 _ _ _ _ _ _
                            (λ a, match a with | Some b => Some (negb b) | None => None end )).
      * destruct a => ?.
        ** f_equal; destruct s; rewrite ?von_neumann_fair //=;
          rewrite (pr_rvar_ldist (dist_ret _ (Some true)));
          rewrite (pr_rvar_ldist (dist_ret _ (Some false)));
          rewrite //= ?big_cons ?big_nil //=; try field.
        ** f_equal. eauto.
      * intros [[|]|]; rewrite //=.
      * intros [[|]|]; rewrite //=.
      * apply undup_uniq.
      * apply undup_uniq.
      * intros [[|]|] [[|]|] => //=.
  Qed.

  Lemma von_neumann_repeat_fail' n (p: R) (Hrange: 0 <= p <= 1) :
    pr_eq (rvar_of_ldist (von_neumann_repeat p Hrange n)) None =
    (1 - 2 * (p * (1 - p)))^n.
  Proof.
    induction n.
    - rewrite ?pr_rvar_ldist //= ?big_cons ?big_nil //=. field.
    - rewrite pr_mbind_ldist2.
      rewrite {1}/von_neumann_coin.
      rewrite ?big_cons ?big_nil IHn //=.
      rewrite von_neumann_fail //=.
      rewrite ?pr_mret_simpl //=; field.
  Qed.

End test.