From discprob.basic Require Import base order nify monad.
From discprob.prob Require Import prob countable finite stochastic_order.
From discprob.monad.finite Require Import monad.
From mathcomp Require Import ssreflect ssrbool ssrfun eqtype ssrnat seq div choice fintype.
From mathcomp Require Import tuple finfun bigop prime binomial finset.
Require Import Fourier Psatz Omega.
Require Import Arith.

Section cost.
Definition cost A := (nat * A)%type.
Definition ldist_cost A := ldist (cost A).

Local Open Scope stdpp_scope.
Local Open Scope nat_scope.

Global Instance cost_bind: stdpp.base.MBind cost :=
  λ A B (f: A → cost B) x,
  (x.1 + (f (x.2)).1, (f (x.2)).2).
Global Instance cost_ret: stdpp.base.MRet cost :=
  λ A x, (0, x).

Lemma cost_left_id {A B: eqType} (x: A) (f: A → cost B):
  mbind f (mret x) = f x.
Proof.
  rewrite /cost_bind//=; destruct (f x); f_equal => //=.
Qed.

Lemma cost_right_id {A: eqType} (m: cost A) :
  mbind mret m = m.
Proof.
  destruct m => //=. rewrite /cost_bind; f_equal => //=. nify; lia.
Qed.

Lemma cost_assoc {A B C: eqType} (m: cost A) (f: A → cost B) (g: B → cost C) :
  mbind g (mbind f m) = mbind (λ x, mbind g (f x)) m.
Proof.
  destruct m => //=. rewrite /cost_bind; f_equal => //=. nify; lia.
Qed.

Notation MBind := stdpp.base.MBind.
Global Instance ldist_cost_bind: MBind ldist_cost :=
  λ A B f x,
  '(c1, a) ← x;
  '(c2, b) ← f a;
  mret (c1 + c2, b).
Global Instance ldist_cost_ret: MRet ldist_cost :=
  λ A x, mret (0, x).

Lemma ldist_cost_bind_fold {A B} (f: A → ldist_cost B) (x: ldist_cost A):
  '(c1, a) ← x;
  '(c2, b) ← f a;
  mret (c1 + c2, b) = mbind f x.
Proof.
  symmetry. rewrite {1}/mbind/ldist_cost_bind. done.
Qed.

Lemma ldist_cost_left_id {A B: eqType} (x: A) (f: A → ldist_cost B):
  mbind f (mret x) = f x.
Proof.
  specialize (@ldist_left_id).
  rewrite /mbind/ldist_cost_bind/ldist_cost//= => ->.
  rewrite -[a in _ = a]ldist_right_id /mbind.
  apply ldist_bind_ext => a.
  destruct a as (c, a) => //=.
Qed.

End cost.

Require Import Reals Fourier FunctionalExtensionality.

Lemma ldist_cost_right_id {A: eqType} (m: ldist_cost A) :
  mbind mret m = m.
Proof.
  rewrite /mbind/ldist_cost_bind.
  apply ldist_irrel=>//=.
  destruct m as [l pf1 pf2] => //=; clear pf1 pf2.
  induction l => //=. destruct a as (r, x) => //=.
  destruct x as (c, a) => //=.
  rewrite ?Rmult_1_r; repeat f_equal => //; nify; lia.
Qed.

Lemma ldist_cost_assoc {A B C: eqType} (m: ldist_cost A)
      (f: A → ldist_cost B) (g: B → ldist_cost C) :
  mbind g (mbind f m) = mbind (λ x, mbind g (f x)) m.
Proof.
  rewrite /mbind/base.mbind/ldist_cost_bind.
  rewrite !ldist_assoc.
  apply ldist_bind_ext => x.
  destruct x as (c, a).
  rewrite !ldist_assoc.
  apply ldist_bind_ext => x.
  destruct x as (c', b).
  rewrite !ldist_assoc.
  rewrite ldist_left_id.
  apply ldist_bind_ext => x.
  destruct x as (c'', z).
  rewrite ldist_left_id.
  do 2 f_equal. rewrite addnA. done.
Qed.

Lemma ldist_cost_bind_semi {B C} (m: ldist_cost B) (g: nat → C) c1:
  (x ← '(c2, b) ← m;
  mret (c1 + c2, b)%nat;
  mret (g (fst x))) =
  x ← m;
  mret (g (fst x + c1)%nat).
Proof.
  apply ldist_irrel=>//=.
  destruct m as [l pf1 pf2] => //=; clear pf1 pf2.
  induction l => //=. destruct a as (r, x) => //=.
  destruct x as (c, a) => //=.
  rewrite ?Rmult_1_r /=.
  f_equal; auto.
  f_equal; auto.
  f_equal; rewrite addnC //.
Qed.

Lemma ldist_cost_bind_drop {B C} (m: ldist_cost B) (h: B → ldist C) (g: nat → C):
  (x ← (x ← m; mret (h x)) ;
  mret (g (fst x))) =
  x ← m;
  mret (g (fst x)%nat).
Proof.
  apply ldist_irrel=>//=.
  destruct m as [l pf1 pf2] => //=; clear pf1 pf2.
  induction l => //=. destruct a as (r, x) => //=.
  destruct x as (c, a) => //=.
  rewrite ?Rmult_1_r /=.
  f_equal; auto.
  rewrite addn0.
  f_equal; auto.
Qed.

Lemma cost_bind_const {A B C: eqType} d (h: ldist_cost A) (f: A → ldist_cost B) (g: nat → C) (c: C):
  (∀ d', d' \in [seq i.2.1 | i <- outcomes h] → d' = d) →
  pr_eq (rvar_comp (rvar_of_ldist (mbind f h)) (λ x, g (fst x))) c =
  \big[Rplus/0]_(a <- undup [seq i.2.2 | i <- h])
   (pr_eq (rvar_comp (rvar_of_ldist h) snd) a
    * pr_eq (rvar_comp (rvar_of_ldist (f a)) (λ x, g (fst x + d)%nat)) c).
Proof.
  intros Hconst.
  rewrite -pr_mbind_mret ldist_assoc.
  rewrite pr_mbind_ldist2. symmetry.
  eapply sum_reidx_surj1 with (h := λ x, (d, x)).
  - intros a0 Hin. symmetry.
    rewrite -(pr_mbind_mret (f a0)).
    rewrite ldist_cost_bind_semi; f_equal.
    rewrite /pr_eq pr_eq_alt_comp. rewrite pr_eq_alt.
    (* rewrite img_rvar_of_ldist. *)
    rewrite -?big_mkcondr.
    apply eq_bigl.
    rewrite /=. intros x.
    destruct x as ((d'&a)&Hin') => //=.
    rewrite img_rvar_of_ldist' in Hin'.
    rewrite mem_undup in Hin'. move /mapP in Hin'.
    destruct Hin' as [(r&d'''&a'') ? Heq]; subst.
    assert (d' = d) as ->.
    { apply Hconst. apply /mapP. eexists; eauto. rewrite //=. inversion Heq. congruence. }
    rewrite //=.
    cut ((d, a) == (d, a0) = (a == a0)).
    { intros  Heq'. rewrite Heq'. done. }
    (* There must be a better way to rewrite this automatically!!!! *)
    apply /pair_eqP. case: ifP. move /eqP. intros; congruence.
    move /eqP. intros. inversion 1. congruence.
  - intros a0. rewrite !mem_undup.  move /mapP.
    intros [(r&d'&a) ? Heq] _.
    split; auto. apply /mapP. exists (r, (d', a)); auto.
    rewrite //=. f_equal; auto. symmetry; apply Hconst.
    apply /mapP. eexists; eauto.
  - intros (d', a).  rewrite !mem_undup => Hin _.
    exists a. repeat split; auto.
    rewrite mem_undup. rewrite map_comp.
    apply /mapP. eexists; eauto.
    f_equal. symmetry; apply Hconst.
    rewrite map_comp. apply /mapP. eexists; eauto.
  - apply undup_uniq.
  - apply undup_uniq.
  - intros ??.  congruence.
Qed.

Lemma ldist_cost_pair {B1 B2: eqType} {C D: eqType} (m1: ldist_cost B1) (m2: ldist_cost B2)
      (h: B1 → B2 → D) (g: nat → C) (c: C):
  pr_eq (rvar_comp (rvar_of_ldist ((x ← m1; y ← m2; mret (h x y) : ldist_cost D)))
                   (λ x, g (fst x))) c =
  pr_eq (rvar_comp (rvar_pair (rvar_comp (rvar_of_ldist m1) fst)
                              (rvar_comp (rvar_of_ldist m2) fst)) (λ xy, g (fst xy + snd xy)%nat)) c.
Proof.
  rewrite -pr_mbind_mret rvar_pair_comp rvar_comp_comp.
  rewrite ldist_assoc.
  rewrite pr_mbind_ldist2.
  symmetry.
  rewrite {1}/pr_eq pr_eq_alt_comp.
  etransitivity.
  { eapply eq_bigr => i _. rewrite pr_eq_rvar_pair. done. }
  rewrite -(big_map _ (λ x, true) (λ i, if g (fst (fst i) + fst (snd i))%nat == c then
                                pr_eq _ (fst i) * pr_eq _ (snd i)
                                 else
                                   0)).

  rewrite /index_enum.
  rewrite (perm_big _ (img_pair_rv _ _ _ _)).
  rewrite  big_Rplus_allpair'.
  rewrite -(big_map _ (λ x, true) (λ i, \big[Rplus/0]_(i' <- _)
                                         (if (g (fst (fst (i, _)) + fst (snd (i, _))) == c)%nat
                                          then
                                            pr_eq _ (fst (i, _)) *
                                            pr_eq _ (snd (i, _))
                                          else
                                            0))).
  rewrite img_rvar_of_ldist.
  apply eq_bigr => a _.
  etransitivity.
  { eapply eq_bigr => i _.
    rewrite -[a in (if (_ : bool) then _ else a) = _](Rmult_0_r (pr_eq (rvar_of_ldist m1) a)).
    rewrite Rmult_comm -(Rmult_comm 0) -Rmult_if_distrib Rmult_comm. done.
  }
  rewrite -big_distrr. apply Rmult_eq_compat_l.
  destruct a as (c', b).
  rewrite ldist_assoc.
  rewrite ldist_assoc.
  rewrite pr_mbind_ldist2.
  rewrite -(big_map _ (λ x, true) (λ i, (if (g (fst (fst (_, i)) + fst (snd (_, i))) == c)%nat
                                          then
                                            pr_eq _ (snd (_, i))
                                          else
                                            0))).
  rewrite img_rvar_of_ldist.
  eapply eq_bigr => a' _.
    rewrite -[a in (if (_ : bool) then _ else a) = _](Rmult_0_r (pr_eq (rvar_of_ldist m2) a')).
    rewrite -[a in (if (_ : bool) then a else _) = _](Rmult_1_r (pr_eq (rvar_of_ldist m2) a')).
    rewrite Rmult_comm -(Rmult_comm 0) -Rmult_if_distrib Rmult_comm.
  apply Rmult_eq_compat_l.
  destruct a' as (c'', b2).
  rewrite ?ldist_left_id pr_mret_simpl /= addn0 //.
Qed.


(** This file implements cancellation for separation logic.
 **)
Require Coq.Lists.List.
Require ExtLib.Data.Option.
Require Import ExtLib.Tactics.

Set Implicit Arguments.
Set Strict Implicit.

Section iterated.
  Variable T : Type.
  Variable base : T.
  Variable join : T -> T -> T.

  Fixpoint iterated (ls : list T) : option T :=
    match ls with
      | nil => None
      | List.cons l ls => match iterated ls with
                            | None => Some l
                            | Some ls => Some (join l ls)
                          end
    end.

  Definition iterated_base (ls : list T) : T :=
    match iterated ls with
      | None => base
      | Some l => l
    end.

  Lemma iterated_None : forall ls, iterated ls = None -> ls = nil.
  Proof.
    destruct ls; simpl; intros; try congruence.
    destruct (iterated ls); congruence.
  Qed.

  Require Import Relations Morphisms.
  Require Import RelationClasses.

  Hypothesis R : relation T.
  Hypothesis Reflexive_R : Reflexive R.
  Hypothesis Transitive_R : Transitive R.
  Hypothesis join_assoc : forall a b c, R (join a (join b c)) (join (join a b) c).
  Hypothesis join_respects : Proper (R ==> R ==> R)%signature join.

  Definition option_R (a b : option T) : Prop :=
    match a , b with
      | None , None => True
      | Some a , Some b => R a b
      | _ , _ => False
    end.

  Local Instance Reflexive_optionR : Reflexive option_R.
  Proof.
    red. destruct x; simpl; reflexivity.
  Qed.
  Local Instance Transitive_optionR : Transitive option_R.
  Proof.
    red. destruct x; destruct y; destruct z; simpl; intros; auto.
    etransitivity; eauto. intuition.
  Qed.

  Theorem iterated_app
  : forall ls' ls,
      option_R (iterated (ls ++ ls'))
               (match iterated ls' with
                  | None => iterated ls
                  | Some ls' => match iterated ls with
                                  | None => Some ls'
                                  | Some ls => Some (join ls ls')
                                end
                end).
  Proof.
    induction ls; simpl.
    { intros. destruct (iterated ls'); reflexivity. }
    { intros. unfold option_R in *.
      repeat match goal with
               | |- context [ match ?X with _ => _ end ] =>
                 consider X; intros; try congruence
               | H : context [ match ?X with _ => _ end ] |- _ =>
                 consider X; intros; try congruence
               | H : R _ _ |- _ => rewrite H
             end; inv_all; subst; try solve [ congruence | intuition ].
      { rewrite IHls. etransitivity. 2: eapply join_assoc. reflexivity. }
      { rewrite H5. reflexivity. }
      { rewrite H4; reflexivity. }
      { subst. reflexivity. } }
  Qed.

  Hypothesis Rbase_LL : forall a, R (join base a) a.
  Hypothesis Rbase_LR : forall a, R (join a base) a.
  Hypothesis Rbase_RL : forall a, R a (join base a).
  Hypothesis Rbase_RR : forall a, R a (join a base).

  Theorem iterated_base_app
  : forall ls' ls,
      R (iterated_base (ls ++ ls'))
        (join (iterated_base ls) (iterated_base ls')).
  Proof.
    induction ls; simpl.
    { intros. unfold iterated_base. simpl.
      eapply Rbase_RL. }
    { unfold iterated_base in *; simpl in *.
      destruct (iterated (ls ++ ls')).
      { rewrite IHls.
        rewrite join_assoc.
        destruct (iterated ls); try reflexivity.
        rewrite Rbase_LR. reflexivity. }
      { destruct (iterated ls); destruct (iterated ls'); simpl.
        { rewrite <- join_assoc. rewrite <- IHls. eapply Rbase_RR. }
        { rewrite <- join_assoc. rewrite <- IHls. apply Rbase_RR. }
        { rewrite Rbase_LL in IHls.
          rewrite <- IHls. apply Rbase_RR. }
        { eapply Rbase_RR. } } }
  Qed.

  Theorem iterated_base_cons
  : forall l ls,
      R (iterated_base (l :: ls))
        (join l (iterated_base ls)).
  Proof.
    intros.
    change (l :: ls)%list with ((l :: nil) ++ ls)%list.
    etransitivity. eapply iterated_base_app.
    unfold iterated_base at 1. simpl. reflexivity.
  Qed.

  Variable P : T -> Prop.
  Hypothesis Pjoin : forall a b, P a -> P b -> P (join a b).

  Theorem iterated_inv
  : forall ls,
      List.Forall P ls ->
      match iterated ls with
        | None => ls = nil
        | Some e => P e
      end.
  Proof.
    clear - Pjoin. induction 1; simpl; intros; auto.
    destruct (iterated l); auto.
  Qed.

  Hypothesis Pbase : P base.

  Theorem iterated_base_inv
  : forall ls,
      List.Forall P ls ->
      P (iterated_base ls).
  Proof.
    clear - Pjoin Pbase.
    unfold iterated_base; intros.
    eapply iterated_inv in H. destruct (iterated ls); eauto.
  Qed.

End iterated.