Require Import Misc Ccpo.

Set Implicit Arguments.

From mathcomp Require Import ssreflect ssrfun eqtype choice.
From mathcomp Require Import ssrbool ssrnat seq order fintype finfun.
From mathcomp Require Import bigop ssralg ssrnum ssrint rat.
Import GRing.
Import Order.Theory.
Import Num.Theory.

Delimit Scope order_scope with Omc.
#[local] Open Scope O_scope.
Delimit Scope O_scope with O.
#[local] Open Scope ring_scope.

#[export] Program Instance ratO : ord rat :=
     { Oeq := fun n m : rat => n = m;
       Ole := fun n m : rat => (n <= m)%R}.

Definition MF (A:Type) := A -> rat.

#[export] Instance MFO (A:Type) : ord (MF A) := ford A rat.

Definition M A := MF A -m> rat.

#[export] Instance MO (A:Type) : ord (M A) := fmono (MF A) rat.

#[export] Instance app_mon (A:Type) (x:A) : monotonic (fun (f:MF A) => f x).
#[export] Hint Resolve app_mon : core.

Definition munit (A:Type) (x:A) : M A := mon (fun (f:MF A) => f x).

Definition mstar (A B:Type) : M A -> (A -> M B) -> M B.
Defined.

Lemma star_simpl (A B:Type) (a:M A) (F:A -> M B) (f:MF B) :
  mstar a F f = a (fun x => F x f).

Lemma law1 (A B:Type) (x:A) (F:A -> M B) (f:MF B) :
  mstar (munit x) F f == F x f.

Lemma law2 (A:Type) (a:M A) (f:MF A) :
  mstar a (fun x:A => munit x) f == a (fun x:A => f x).

Lemma law3 (A B C:Type) (a:M A) (F:A -> M B) (G:B -> M C) (f:MF C) :
  mstar (mstar a F) G f == mstar a (fun x:A => mstar (F x) G) f.

Definition stable_add (A:Type) (m:M A) : Prop :=
  forall f g:MF A, m (f \+ g) = (m f) + (m g).

Definition stable_sub (A:Type) (m:M A) : Prop :=
  forall f g:MF A, m (f \- g) = (m f) - (m g).

Definition stable_mull (A:Type) (m:M A) : Prop :=
  forall (k:rat) (f:MF A), m (k \*o f) = k * (m f).

Definition stable_opp (A:Type) (m:M A) : Prop :=
  forall (f:MF A), m (fun x => - (f x)) = - (m f).

Section StabilityProperties.

Variable A : Type.
Variable m : M A.
Hypothesis Mstable_sub : stable_sub m.

Implicit Types (f g : MF A).

Lemma Mstable_eq : stable m.
Hint Resolve Mstable_eq : core.

Lemma Mstable0 : m \0 = 0.
Hint Resolve Mstable0 : core.

Lemma Mstable_opp : stable_opp m.

Lemma Mstable_add : stable_add m.

Lemma Mstable_addn f (n : nat) : m (fun x => f x *+ n) = (m f) *+ n.

Lemma Mstable_subn f (n : nat) : m (fun x => f x *- n) = (m f) *- n.

Lemma Mstable_divn f (n : nat) :
  m (fun x => f x / n%:~R) = (m f) / (n%:~R).

Lemma Mstable_addi f (n : int) : m (fun x => f x *~ n) = (m f) *~ n.

Lemma Mstable_muli f (n : int) : m (fun x => n%:~R * f x) = (n%:~R) * (m f).

Lemma Mstable_divi f (n : int) : m (fun x => f x / n%:~R) = (m f) / (n%:~R).

Lemma Mstable_mull : stable_mull m.

Lemma Mstable_linear (p q : rat) f g :
  m ((p \*o f) \+ (q \*o g)) = p * (m f) + q * (m g).

End StabilityProperties.
#[export] Hint Resolve Mstable_eq : core.
#[export] Hint Resolve Mstable0 : core.
#[export] Hint Resolve Mstable_opp : core.
#[export] Hint Resolve Mstable_add : core.

Lemma unit_stable_eq (A:Type) (x:A) : stable (munit x).

Lemma star_stable_eq (A B:Type) (m:M A) (F:A -> M B) : stable (mstar m F).

Lemma unit_monotonic A (x:A) (f g : MF A) :
  (f <= g)%O -> munit x f <= munit x g.

Lemma star_monotonic A B (m:M A) (F:A -> M B) (f g : MF B) :
  (f <= g)%O -> mstar m F f <= mstar m F g.

Lemma star_le_compat A B (m1 m2:M A) (F1 F2:A -> M B) :
  (m1 <= m2)%O -> (F1 <= F2)%O -> (mstar m1 F1 <= mstar m2 F2)%O.
#[export] Hint Resolve star_le_compat : core.

Lemma unit_stable_sub (A:Type) (x:A) : stable_sub (munit x).

Lemma star_stable_sub (A B:Type) (m:M A) (F:A -> M B) :
  stable_sub m -> (forall a:A, stable_sub (F a)) -> stable_sub (mstar m F).

Record distr (A:Type) : Type := {
       mu : M A;
       mu_stable_sub : stable_sub mu;
       mu_prob : mu (fun x => 1) = 1
   }.

#[export] Hint Resolve mu_stable_sub mu_prob : core.

Section MeasureProp.

Context {A : Type} (m: distr A).
Implicit Types (f g : MF A).

Lemma mu_monotonic : monotonic (mu m).
Hint Resolve mu_monotonic : core.

Lemma mu_stable_eq : stable (mu m).
Hint Resolve mu_stable_eq : core.

Lemma mu_zero : mu m \0 = 0.
Hint Resolve mu_zero : core.

Lemma mu_zero_eq f : (forall x, f x = 0) -> mu m f = 0.

Lemma mu_one_eq f : (forall x, f x = 1) -> mu m f = 1.

Lemma mu_stable_inv f : mu m (fun x => 1 - f x) = 1 - (mu m f).

Lemma mu_stable_add f g : mu m (f \+ g) = mu m f + mu m g.

Lemma mu_stable_mull (q : rat) f : mu m (q \*o f) = q * mu m f.

Lemma mu_add_zero f g : mu m f = 0 -> mu m g = 0 -> mu m (f \+ g) = 0.

Lemma mu_stable_pos f : (\0 <= f)%O -> 0 <= mu m f.

Lemma mu_stable_le1 f : (forall x, f x <= 1) -> mu m f <= 1.

Lemma mu_cte (c:rat) : mu m (fun x => c) = c.

Lemma mu_stable_mulr (c:rat) f : mu m (c \o* f) = (mu m f) * c.

Lemma mu_stable_inv_inv f : mu m f = 1 - mu m (fun x => 1 - f x).

End MeasureProp.
#[export] Hint Resolve mu_monotonic : core.
#[export] Hint Resolve mu_stable_eq : core.
#[export] Hint Resolve mu_zero : core.
#[export] Hint Immediate mu_zero_eq : core.
#[export] Hint Immediate mu_one_eq : core.
#[export] Hint Resolve mu_stable_inv : core.
#[export] Hint Resolve mu_stable_add : core.
#[export] Hint Resolve mu_stable_mull : core.
#[export] Hint Resolve mu_add_zero : core.
#[export] Hint Resolve mu_cte : core.
#[export] Hint Resolve mu_stable_inv_inv : core.

#[export] Program Instance Odistr (A:Type) : ord (distr A) :=
    {Ole := fun (f g : distr A) => (mu f <= mu g)%O;
     Oeq := fun (f g : distr A) => Oeq (mu f) (mu g)}.

Section MonDistrib.

Variables (A B : Type).

Definition Munit : A -> distr A.

Definition Mlet : distr A -> (A -> distr B) -> distr B.

Lemma Munit_simpl (q : A -> rat) x : mu (Munit x) q = q x.

Lemma Munit_simpl_eq (q : A -> rat) x : mu (Munit x) q == q x.

Lemma Mlet_simpl (m : distr A) (M : A -> distr B) (f : B -> rat) :
  mu (Mlet m M) f = mu m (fun x => (mu (M x) f)).

Lemma Mlet_simpl_eq (m : distr A) (M : A -> distr B) (f : B -> rat) :
mu (Mlet m M) f == mu m (fun x => (mu (M x) f)).

End MonDistrib.

Section OperDistr.

Variables (A B C : Type).

Lemma Munit_eq_compat (x y : A) : x = y -> Munit x == Munit y.

Lemma Mlet_le_compat (m1 m2 : distr A) (M1 M2 : A -> distr B) :
  (m1 <= m2 -> M1 <= M2 -> Mlet m1 M1 <= Mlet m2 M2)%O.
Hint Resolve Mlet_le_compat : core.

Add Parametric Morphism : (Mlet (A:=A) (B:=B))
    with signature Ole ==> Ole ==> Ole
      as Mlet_le_morphism.

Add Parametric Morphism : (Mlet (A:=A) (B:=B))
    with signature Ole ==> (@pointwise_relation A (distr B) (@Ole _ _)) ==> Ole
      as Mlet_le_pointwise_morphism.

#[export] Instance Mlet_mon2 : monotonic2 (@Mlet A B).

Definition MLet : distr A -m> (A -> distr B) -m> distr B
  := mon2 (@Mlet A B).

Lemma MLet_simpl (m : distr A) (M : A -> distr B) (f : B -> rat) :
  mu (MLet m M) f = mu m (fun x => mu (M x) f).

Lemma Mlet_eq_compat (m1 m2 : distr A) (M1 M2 : A -> distr B) :
  (m1 == m2 -> M1 == M2 -> Mlet m1 M1 == Mlet m2 M2)%type.
Hint Resolve Mlet_eq_compat : core.

Add Parametric Morphism : (Mlet (A:=A) (B:=B))
    with signature Oeq ==> Oeq ==> Oeq
      as Mlet_eq_morphism.

Add Parametric Morphism : (Mlet (A:=A) (B:=B))
    with signature Oeq ==> (@pointwise_relation A (distr B) (@Oeq _ _)) ==> Oeq
      as Mlet_Oeq_pointwise_morphism.

Lemma mu_le_compat (m1 m2 : distr A) :
  (m1 <= m2 -> forall f g : A -> rat, f <= g -> mu m1 f <= mu m2 g)%O.

Lemma mu_eq_compat (m1 m2 : distr A) :
  (m1 == m2 -> forall f g : A -> rat, f == g -> mu m1 f = mu m2 g)%type.
Hint Immediate mu_le_compat mu_eq_compat : core.

Add Parametric Morphism : (mu (A:=A))
    with signature Ole ==> Ole
      as mu_le_morphism.

Add Parametric Morphism : (mu (A:=A))
    with signature Oeq ==> Oeq
      as mu_eq_morphism.

Add Parametric Morphism (a : distr A) : (@mu A a)
    with signature (@pointwise_relation A rat (@eq _) ==> Oeq)
      as mu_distr_eq_morphism.

Add Parametric Morphism (a : distr A) : (@mu A a)
    with signature (@pointwise_relation A rat (@Oeq _ _) ==> Oeq)
      as mu_distr_Oeq_morphism.

Add Parametric Morphism (a : distr A) : (@mu A a)
    with signature (@pointwise_relation _ _ (@Ole _ _) ==> Ole)
      as mu_distr_le_morphism.

Add Parametric Morphism : (@Mlet A B)
    with signature (Ole ==> @pointwise_relation _ _ (@Ole _ _) ==> Ole)
      as mlet_distr_le_morphism.

Add Parametric Morphism : (@Mlet A B)
    with signature (Oeq ==> @pointwise_relation _ _ (@Oeq _ _) ==> Oeq)
      as mlet_distr_eq_morphism.

Lemma Mlet_unit (x : A) (m : A -> distr B) : Mlet (Munit x) m == m x.

Lemma Mlet_ext (m : distr A) : Mlet m (fun x => Munit x) == m.

Lemma Mlet_assoc (m1 : distr A) (m2 : A -> distr B) (m3 : B -> distr C) :
  Mlet (Mlet m1 m2) m3 == Mlet m1 (fun x:A => Mlet (m2 x) m3).

Lemma let_indep (m1 : distr A) (m2 : distr B) (f : MF B) :
  mu m1 (fun => mu m2 f) = mu m2 f.

Lemma let_indep_distr (m1 : distr A) (m2 : distr B) :
  Mlet m1 (fun => m2) == m2.

Section MuBool.
Context {m : distr A}.
Implicit Types (f g : A -> bool).

Lemma mu_bool_le1 {f} : mu m (fun x => (f x)%:~R) <= 1%:~R.
Hint Resolve mu_bool_le1 : core.

Lemma mu_bool_0le {f} : 0%:~R <= mu m (fun x => (f x)%:~R).
Hint Resolve mu_bool_0le : core.

Lemma mu_bool_impl {f g} :
  (forall x, (f x) ==> (g x)%B) ->
  (mu m (fun x => (f x)%:~R) <= mu m (fun x => (g x)%:~R)).

Lemma mu_bool_impl1 {f g} :
  (forall x, (f x) ==> (g x)%B) ->
  mu m (fun x => (f x)%:~R) = 1 -> mu m (fun x => (g x)%:~R) = 1.

Lemma mu_bool_negb0 {f g} :
  (forall x, (f x) ==> ~~ (g x)%B) ->
  mu m (fun x => (f x)%:~R) = 1 ->
  mu m (fun x => (g x)%:~R) = 0.

Lemma mu_bool_negb {f} :
  mu m (fun x => (~~ f x)%:~R) = 1 - mu m (fun x => (f x)%:~R).

Lemma mu_bool_negb1 {f g} :
  (forall x, (~~ (f x) ==> g x)%B) ->
  mu m (fun x => (f x)%:~R) = 0 ->
  mu m (fun x => (g x)%:~R) = 1.

End MuBool.

End OperDistr.

#[export] Hint Resolve mu_bool_le1 : core.
#[export] Hint Resolve mu_bool_0le : core.

Notation "[1/2]" := (2%:~R)^-1.

#[export] Instance flip_mon :
  monotonic (fun (f : bool -> rat) => [1/2] * (f true) + [1/2] * (f false)).

Definition flip : M bool :=
  mon (fun (f : bool -> rat) => [1/2] * (f true) + [1/2] * (f false)).

Lemma flip_stable_sub : stable_sub flip.

Lemma flip_prob : flip (fun x => 1) = 1.

Lemma flip_true : flip (fun b => (b%:~R)) = [1/2].

Lemma flip_false : flip (fun b => (~~b)%:~R) = [1/2].

#[export] Hint Resolve flip_true flip_false : core.

Definition Flip : distr bool.

Lemma Flip_simpl f : mu Flip f = [1/2] * (f true) + [1/2] * (f false).

Require Arith.


Section FiniteDistributions.

Variable A : Type.
Variable p : seq A.

Definition weight (c : A -> rat) : rat := \sum_(i <- p | 0 < c i) (c i).

Lemma weight_nonneg c : 0 <= weight c.
Hint Resolve weight_nonneg : core.

Lemma weight_case c : (weight c = 0) \/ 0 < (weight c)^-1.

#[export] Instance finite_mon (c : A -> rat) :
  monotonic (fun f => (\sum_(i <- p | 0 < c i) (c i * f i)) / weight c).

Definition mfinite (c : A -> rat) : M A :=
  mon (fun f => (\sum_(i <- p | 0 < c i) (c i * f i))/weight c).

Lemma finite_simpl (c : A -> rat) f :
  mfinite c f = (\sum_(i <- p | 0 < c i) (c i * f i))/weight c.

Lemma finite_stable_sub (c : A -> rat) : stable_sub (mfinite c).

End FiniteDistributions.

Record fin (A : Type) : Type :=
  mkfin { points : seq A;
          coeff : A -> rat;
          weight_pos : 0 < weight points coeff
        }.
#[export] Hint Resolve weight_pos : core.

Lemma inv_weight_pos A (d : fin A) : 0 < (weight (points d) (coeff d))^-1.
#[export] Hint Resolve inv_weight_pos : core.

Lemma weight_is_unit A (d : fin A) : (weight (points d) (coeff d)) \is a unit.
#[export] Hint Resolve weight_is_unit : core.

Definition fprob A (d : fin A) (a : A) : rat :=
  coeff d a / weight (points d) (coeff d).

Definition Finite A : fin A -> distr A.

Lemma Finite_simpl A (d : fin A) :
  mu (Finite d) = mfinite (points d) (coeff d).

Lemma Finite_eq_in (A : eqType) (d : fin A) (a : A) :
  uniq (points d) -> (a \in points d)%SEQ -> 0 < coeff d a ->
  mu (Finite d) (fun x => (x==a)%:~R) = fprob d a.

Lemma Finite_eq_out (A : eqType) (d : fin A) (a : A) :
  (a \notin points d)%SEQ \/ (coeff d a <= 0)%R ->
  mu (Finite d) (fun x => (x==a)%:~R) = 0.

Lemma Finite_in_seq (A : eqType) (d : fin A) :
  mu (Finite d) (fun x => (x \in points d)%:~R) = 1%R.

Record unif (A : Type) : Type :=
  mkunif { upoints :> seq A; _ : size upoints != O }.

Definition usize A (p : unif A) : nat := size (upoints p).

Lemma usize_pos A (p : unif A) : usize p != O.

Definition unif2fin A (p:unif A) : fin A.

Definition Uniform A (d : unif A) : distr A := Finite (unif2fin d).

Lemma Uniform_simpl A (d : unif A) :
  mu (Uniform d) = mfinite (upoints d) (fun A => 1%R).

Lemma weight1_size A (d : seq A) : weight d (fun x => 1) = (size d)%:~R.

Lemma mu_uniform_sum A (d : unif A) (f : A -> rat) :
  mu (Uniform d) f = (\sum_(i <- d) f i) / (size d)%:~R.

Lemma Uniform_eq_in (A : eqType) (d : unif A) (a : A) :
  uniq d -> (a \in upoints d)%SEQ ->
  mu (Uniform d) (fun x => (x==a)%:~R) = 1 / (usize d)%:~R.

Lemma Uniform_eq_out (A : eqType) (d : unif A) (a : A) :
  (a \notin upoints d)%SEQ ->
  mu (Uniform d) (fun x => (x==a)%:~R) = 0%:~R.

Lemma Uniform_in_seq (A : eqType) (d : unif A) :
  mu (Uniform d) (fun x => (x \in upoints d)%:~R) = 1%:~R.

Fact succ_neq0 (n m : nat) : (n==m.+1)%N -> (n!=0)%N.

Lemma Uniform_unif_seq_eq (A : eqType) (d1 d2 : unif A) :
  (d1 == d2 :> seq A) -> Uniform d1 == Uniform d2.

Definition unif_def A (d : A) (s : seq A) : unif A.

Lemma Uniform_def_ne A (d : A) (s : seq A) :
  forall (Hs : (size s != 0)%N), Uniform (unif_def d s) = Uniform (mkunif s Hs).

Section UnifNat.

Implicit Types (n a : nat).

Definition unifnat n : unif nat := mkunif (iota 0 (n.+1)) (eq_refl true).
Definition Random n : distr nat := Uniform (unifnat n).

Lemma Random_simpl n : mu (Random n) = mfinite (iota 0 (n.+1)) (fun x => 1).

Lemma Random_eq_in n a :
  (a <= n)%N -> mu (Random n) (fun x => (x==a)%:~R) = 1 / (n.+1)%:~R.

Lemma Random_eq_out n a :
  (n < a)%N -> mu (Random n) (fun x => (x==a)%:~R) = 0.

Lemma mu_random_sum n (f : nat -> rat) :
  mu (Random n) f = (\sum_(0 <= i < n.+1) f i) / (n.+1)%:~R.

Lemma Random_in_range n :
  mu (Random n) (fun x => (x <= n)%N%:~R) = 1%:~R.

End UnifNat.

Lemma mu_stable_sum (A : Type) (m : distr A)
      (I : Type) (s : seq I) (f : I -> A -> rat) :
  mu m (fun a => \sum_(i <- s) f i a) = \sum_(i <- s) (mu m (f i)).

Section Bigsums.

Variable A : eqType.
Implicit Types (x : A) (s : seq A) (m : distr A).

Lemma in_seq_sum s x :
  uniq s -> (x \in s)%:~R = \sum_(i <- s) (x == i)%:~R :> rat.

Lemma mu_in_seq m s :
  uniq s ->
  mu m (fun x => (x \in s)%:~R) = \sum_(a <- s) mu m (fun x => (x == a)%:~R).

Lemma mu_bool_cond m (f g : A -> bool) :
  mu m (fun x => (f x)%:~R) = 1 ->
  mu m (fun x => (g x)%:~R) = mu m (fun x => (f x && g x)%:~R).

Lemma mu_pos_cond (m : distr A) (f : A -> bool) (g : A -> rat) :
  (forall x, 0 <= g x <= 1) ->
  mu m (fun x => (f x)%:~R) = 1 ->
  mu m (fun x => (g x)) = mu m (fun x => ((f x)%:~R * g x)).

End Bigsums.