From discprob.basic Require Import base order.
From discprob.prob Require Import prob countable finite.
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 Reals Fourier ProofIrrelevance FunctionalExtensionality Psatz.
Require Import Coq.omega.Omega.
From Coquelicot Require Import Rcomplements Rbar Series Lim_seq Hierarchy Markov.
Global Set Bullet Behavior "Strict Subproofs".

Notation "x ≤ y" := (Rle x y) (at level 70, no associativity).
Notation "x ≥ y" := (Rge x y) (at level 70, no associativity).

Lemma Finite_inj x y: Finite x = Finite y → x = y.
Proof.
  inversion 1. done.
Qed.
Lemma big_cons_Rplus:
  ∀ (I : Type)
  (i : I) (r : seq I) (P : pred I) (F : I → R)
  (x:=\big[Rplus/0]_(j <- r | P j) F j),
  \big[Rplus/0]_(j <- (i :: r) | P j) F j = (if P i then (F i) else 0) + x.
Proof.
  intros. rewrite big_cons. case: ifP => //. by rewrite Rplus_0_l.
Qed.

Lemma Rlt_not_eq': ∀ r1 r2, r1 < r2 → r2 ≠ r1.
Proof. intros ??? Hfalse. symmetry in Hfalse. eapply Rlt_not_eq; last exact Hfalse. done. Qed.

Lemma if_sumbool_case_match {A: Type} P Q (b: {P} + {¬ P}) (b': {Q} + {¬ Q}) (x y: A):
  (P ↔ Q) → (if b then x else y) = (if b' then x else y).
Proof.
  intros HPQ; destruct b as [p|np], b' as [q|nq] => //=.
  - contradiction nq. by rewrite -HPQ.
  - contradiction np. by rewrite HPQ.
Qed.

(* TODO: lots of redundancy and duplication here *)

Lemma pr_sum_restrict {A: finType} (Ω: distrib A) (X: rrvar Ω) l P:
  \big[Rplus/0]_(i <- l | P i) pr_eq X i =
  \big[Rplus/0]_(i <- l | P i && (Rgt_dec (pr_eq X i) 0)) pr_eq X i.
Proof.
  symmetry; induction l.
  * rewrite ?big_nil //=.
  * rewrite ?big_cons.
    symmetry. case: ifP.
    **  intros HP. rewrite //=.
        case: ifP.
        *** intros. f_equal; eauto.
        *** rewrite /pr//=. destruct (Rgt_dec) as [|Hnt] => // _. apply Rnot_lt_ge in Hnt.
            rewrite -[a in _ = a]Rplus_0_l; f_equal; eauto.
            move: Hnt.
            rewrite /pr_eq.
            edestruct (ge_pr_0 Ω) as [Hge|]; eauto.
            intros. exfalso. rewrite //= in Hnt. rewrite //= in Hge. fourier.
    ** rewrite //=.
Qed.

Section pr_gt.

Definition pr_gt_indicator {A} {Ω: distrib A} (T: rrvar Ω) r :=
  indicator T (λ n, Rgt_dec n r).

Lemma pr_gt_to_indicator {A} {Ω: distrib A} (T: rrvar Ω) r:
  pr (rvar_dist T) (λ n, Rgt_dec (T n) r) = pr_eq (pr_gt_indicator T r) 1.
Proof.
  rewrite /pr_eq/pr. apply Series_ext => n. rewrite /countable_sum; destruct pickle_inv => //=.
  destruct Rgt_dec; rewrite /rvar_dist//=; eauto; case: ifP => //=; move /eqP; nra.
Qed.

Variable A: finType.
Variable Ω: distrib A.
Variable T: rrvar Ω.

Lemma pr_gt_alt r: pr (rvar_dist T) (λ a, Rgt_dec (T a) r)
                 = \big[Rplus/0]_(v : imgT T) if (Rgt_dec (sval v) r) then pr_eq T (sval v) else 0.
Proof.
  rewrite pr_gt_to_indicator /pr_gt_indicator. apply Finite_inj.
  rewrite -Ex_indicator/indicator//=. rewrite Ex_fin_comp => //=. rewrite /index_enum/id_rvar//=.
  f_equal; eapply eq_bigr => ?? => //=. case: ifP => ?; destruct Rgt_dec => //=; nra.
Qed.

(* There must be a better way to do this kind of splitting on sequences... *)
Lemma pr_gt_alt3 r: pr (rvar_dist T) (λ n, Rgt_dec (T n) r) =
                    \big[Rplus/0]_(i <- filter (λ x, Rgt_dec (sval x) r)
                                     (Finite.enum [finType of imgT T])) (pr_eq T (sval i)).
Proof.
  rewrite pr_gt_alt.
  rewrite big_filter //=.
  symmetry.
  transitivity (\big[Rplus/0]_(i : imgT (A:=A) T | (λ i, true) i && (Rgt_dec (sval i) r))
                 if (Rgt_dec (sval i) r) then pr_eq T (sval i) else 0).
  { apply eq_bigr. rewrite //=.  intros ?; destruct Rgt_dec => //=; done. }
  assert (H0: 0 = (\big[Rplus/0]_(i : imgT (A:=A) T | (λ i, true) i && ~~ (Rgt_dec (sval i) r))
                 if (Rgt_dec (sval i) r) then pr_eq T (sval i) else 0)).
  {
    symmetry; apply big1 => i.  move /negP. destruct (Rgt_dec) => //=.
  }
  rewrite -[a in a = _]Rplus_0_r [a in _ + a = _]H0 -bigID.
  rewrite //=.
Qed.
End pr_gt.

Lemma pr_gt_alt_comp {A: finType} {B: eqType} {Ω: distrib A} (T: rvar Ω B) (f: B → R) r:
 pr (rvar_dist (rvar_comp T f)) (λ n, Rgt_dec ((rvar_comp T f) n) r)
                 = \big[Rplus/0]_(v : imgT T) if (Rgt_dec (f (sval v)) r) then pr_eq T (sval v) else 0.
Proof.
  rewrite pr_gt_to_indicator /pr_gt_indicator. apply Finite_inj.
  rewrite -Ex_indicator/indicator//=. rewrite rvar_comp_assoc.
  rewrite Ex_fin_comp => //=. rewrite /index_enum/id_rvar//=.
  f_equal; eapply eq_bigr => ??; case: ifP; destruct (Rgt_dec) => //=; nra.
Qed.

Section pr_ge.

Definition pr_ge_indicator {A} {Ω: distrib A} (T: rrvar Ω) r :=
  indicator T (λ n, Rge_dec n r).

Lemma pr_ge_to_indicator {A} {Ω: distrib A} (T: rrvar Ω) r:
  pr (rvar_dist T) (λ n, Rge_dec (T n) r) = pr_eq (pr_ge_indicator T r) 1.
Proof.
  rewrite /pr_eq/pr. apply Series_ext => n. rewrite /countable_sum; destruct pickle_inv => //=.
  destruct Rge_dec; rewrite /rvar_dist//=; eauto; case: ifP => //=; move /eqP; nra.
Qed.

Variable A: finType.
Variable Ω: distrib A.
Variable T: rrvar Ω.

Lemma pr_ge_alt r: pr (rvar_dist T) (λ a, Rge_dec (T a) r)
                 = \big[Rplus/0]_(v : imgT T) if (Rge_dec (sval v) r) then pr_eq T (sval v) else 0.
Proof.
  rewrite pr_ge_to_indicator /pr_ge_indicator. apply Finite_inj.
  rewrite -Ex_indicator/indicator//=. rewrite Ex_fin_comp => //=. rewrite /index_enum/id_rvar//=.
  f_equal; eapply eq_bigr => ??; case: ifP; destruct Rge_dec => //=; nra.
Qed.

(* There must be a better way to do this kind of splitting on sequences... *)
Lemma pr_ge_alt3 r: pr (rvar_dist T) (λ n, Rge_dec (T n) r) =
                    \big[Rplus/0]_(i <- filter (λ x, Rge_dec (sval x) r)
                                     (Finite.enum [finType of imgT T])) (pr_eq T (sval i)).
Proof.
  rewrite pr_ge_alt.
  rewrite big_filter //=.
  symmetry.
  transitivity (\big[Rplus/0]_(i : imgT (A:=A) T | (λ i, true) i && (Rge_dec (sval i) r))
                 if (Rge_dec (sval i) r) then pr_eq T (sval i) else 0).
  { apply eq_bigr. rewrite //= => ?; destruct Rge_dec => //=. }
  assert (H0: 0 = (\big[Rplus/0]_(i : imgT (A:=A) T | (λ i, true) i && ~~ (Rge_dec (sval i) r))
                 if (Rge_dec (sval i) r) then pr_eq T (sval i) else 0)).
  {
    symmetry; apply big1 => i.  move /negP. destruct (Rge_dec) => //=.
  }
  rewrite -[a in a = _]Rplus_0_r [a in _ + a = _]H0 -bigID.
  rewrite //=.
Qed.
End pr_ge.

Lemma pr_ge_alt_comp {A: finType} {B: eqType} {Ω: distrib A} (T: rvar Ω B) (f: B → R) r:
 pr (rvar_dist (rvar_comp T f)) (λ n, Rge_dec ((rvar_comp T f) n) r)
                 = \big[Rplus/0]_(v : imgT T) if (Rge_dec (f (sval v)) r) then pr_eq T (sval v) else 0.
Proof.
  rewrite pr_ge_to_indicator /pr_ge_indicator. apply Finite_inj.
  rewrite -Ex_indicator/indicator//=. rewrite rvar_comp_assoc.
  rewrite Ex_fin_comp => //=. rewrite /index_enum/id_rvar//=.
  f_equal; eapply eq_bigr => ??; case: ifP; destruct Rge_dec => //=; nra.
Qed.

Section pr_eq.

Definition pr_eq_indicator {A} {B: eqType} {Ω: distrib A} (T: rvar Ω B) r :=
  indicator T (λ n, n == r).

Lemma pr_eq_to_indicator {A} {B} {Ω: distrib A} (T: rvar Ω B) r:
  pr (rvar_dist T) (λ n, (T n) == r) = pr_eq (pr_eq_indicator T r) 1.
Proof.
  rewrite /pr_eq/pr. apply Series_ext => n. rewrite /countable_sum; destruct pickle_inv => //=.
  case: ifP => ?; rewrite /rvar_dist//=; eauto; case: ifP => //=; move /eqP; nra.
Qed.

Variable A: finType.
Variable B: eqType.
Variable Ω: distrib A.
Variable T: rvar Ω B.

Lemma pr_eq_alt r: pr (rvar_dist T) (λ a, (T a) == r)
                 = \big[Rplus/0]_(v : imgT T) if (sval v == r) then pr_eq T (sval v) else 0.
Proof.
  rewrite pr_eq_to_indicator /pr_gt_indicator. apply Finite_inj.
  rewrite -Ex_indicator/indicator//=. rewrite Ex_fin_comp => //=. rewrite /index_enum/id_rvar//=.
  f_equal; eapply eq_bigr => ??; case: ifP; nra.
Qed.

End pr_eq.

Lemma pr_eq_alt_comp {A: finType} {B C: eqType} {Ω: distrib A} (T: rvar Ω B) (f: B → C) r:
  pr (rvar_dist (rvar_comp T f)) (λ n, (rvar_comp T f) n == r)
  = \big[Rplus/0]_(v : imgT T) if (f (sval v) == r) then pr_eq T (sval v) else 0.
Proof.
  rewrite pr_eq_to_indicator /pr_eq_indicator. apply Finite_inj.
  rewrite -Ex_indicator/indicator//=. rewrite rvar_comp_assoc.
  rewrite Ex_fin_comp => //=. rewrite /index_enum/id_rvar//=.
  f_equal; eapply eq_bigr => ??; case: ifP; nra.
Qed.

Section convert.
Lemma pr_gt_to_geq_eps {A: finType} {Ω: distrib A} (X: rrvar Ω) (x eps: R):
  (eps > 0 ∧ ∀ y, y \in img X → y ≠ x → Rabs (x - y) ≥ eps) →
    pr_gt X x = pr_ge X (x + eps/2).
Proof.
  intros (Hgt0&HX).
  rewrite /pr_ge /pr_gt !pr_ge_alt !pr_gt_alt.
  rewrite big_seq_cond.
  rewrite [a in _ = a]big_seq_cond.
  apply eq_bigr. intros (y&?) Hin => //=. rewrite Bool.andb_true_r in Hin.
  destruct (Rgt_dec) as [Hgt|Hngt] => //=.
  - destruct (Rge_dec) as [Hge|Hnge]; first done.
    exfalso; eapply Rlt_not_ge; last eapply (HX y); eauto.
    * rewrite Rabs_minus_sym Rabs_right; nra.
    * apply Rgt_not_eq. done.
  - destruct (Rge_dec); last done.
    exfalso; eapply Rlt_not_ge; last eapply (HX y); eauto.
    * rewrite Rabs_minus_sym Rabs_right; nra.
    * apply Rlt_not_eq. nra.
Qed.

Lemma pr_ge_to_gt_eps {A: finType} {Ω: distrib A} (X: rrvar Ω) (x eps: R):
  (eps > 0 ∧ ∀ y, y \in img X → y ≠ x → Rabs (x - y) ≥ eps) →
    pr_ge X x = pr_gt X (x - eps/2).
Proof.
  intros (Hgt0&HX).
  rewrite /pr_ge /pr_gt ?pr_ge_alt ?pr_gt_alt.
  rewrite big_seq_cond.
  rewrite [a in _ = a]big_seq_cond.
  apply eq_bigr. intros (y&?) Hin => //=. rewrite Bool.andb_true_r in Hin.
  destruct (Rge_dec).
  - destruct (Rgt_dec); first done.
    exfalso; eapply Rlt_not_ge; last eapply (HX y); eauto.
    * rewrite Rabs_minus_sym Rabs_right; nra.
    * apply Rlt_not_eq. nra.
  - destruct (Rgt_dec); last done.
    exfalso; eapply Rlt_not_ge; last eapply (HX y); eauto.
    * rewrite Rabs_minus_sym Rabs_left; nra.
    * apply Rlt_not_eq. nra.
Qed.

Lemma pr_ge_split_eps {A: finType} {Ω: distrib A} (X: rrvar Ω) (x eps: R):
  (eps > 0 ∧ ∀ y, y \in img X → y ≠ x → Rabs (x - y) ≥ eps) →
    pr_eq X x = pr_ge X (x - eps/2) - pr_ge X (x + eps/2).
Proof.
  intros (Hgt0&HX).
  rewrite /pr_ge ?pr_ge_alt.
  rewrite /Rminus.
  rewrite -[a in _ = _ + - a](Rmult_1_l) Ropp_mult_distr_l.
  rewrite big_distrr //=.
  rewrite -big_split //=.
  rewrite {1}/pr_eq ?pr_eq_alt.
  rewrite big_seq_cond.
  rewrite [a in _ = a]big_seq_cond.
  apply eq_bigr. intros (y&?) Hin => //=. rewrite Bool.andb_true_r in Hin.
  case: ifP.
  - move /eqP => ->.
    destruct (Rge_dec) as [Hle|Hle]; last nra.
    destruct (Rge_dec) as [Hle'|Hle']; first nra.
    rewrite Rmult_0_r Rplus_0_r //.
  - move /eqP => Hneq.
    destruct (Rge_dec) as [Hle|Hle] => //=; try nra;
    destruct (Rge_dec) as [Hle'|Hle'] => //=; try nra.
    exfalso.
    eapply Rlt_not_ge; last eapply (HX y); eauto.
    assert (y - x ≤ eps / 2) by nra.
    assert (x - y ≤ eps / 2) by fourier.
    destruct (Rle_dec 0 (x - y)) as [?|Hgt%Rnot_le_gt].
    ** rewrite Rabs_right; fourier.
    ** rewrite Rabs_minus_sym Rabs_right; fourier.
Qed.

Lemma Rmin_dist_list (l: seq R) (x: R):
  ∃ eps, eps > 0 ∧ ∀ y, y \in l → y ≠ x → Rabs (x - y) >= eps.
Proof.
  exists (foldl (Rmin) 1 (map (λ y, Rabs (x - y)) (filter (predC1 x) l))).
  split.
  * assert (1 > 0) as Hgt0 by fourier.
    revert Hgt0; generalize 1.
    induction l.
    ** rewrite //=; fourier.
    ** rewrite //=. case: ifP.
       *** rewrite map_cons //=.
           move /eqP => Hneq r Hgt.
           eapply IHl.
           apply Rmin_case => //.
           apply Rabs_pos_lt. intros Heq.
           apply Hneq; symmetry; by apply Rminus_diag_uniq.
       *** intros; eauto.
  * intros y Hin Hneq.
    generalize 1. induction l.
    ** rewrite //=.
    ** rewrite in_cons in Hin.
       move /orP in Hin; destruct Hin as [Heqa|Hinl].
       *** move /eqP in Heqa.
           subst. move /eqP /negbTE in Hneq. rewrite //= Hneq /=.
           intros. apply Rle_ge.  eapply Rle_trans; first eapply seq_ext.foldl_Rmin.
           apply Rmin_r.
       *** rewrite //=. case: ifP; intros; eapply IHl; eauto.
Qed.

(* Because we are working with finite probability spaces, the following is true: *)

Lemma pr_ge_split2 {A B: finType} {Ω: distrib A} {Ω': distrib B} (X: rrvar Ω) (Y: rrvar Ω')
      (x: R): ∃ x1 x2,
    pr_eq X x = pr_ge X x1 - pr_ge X x2 ∧
    pr_eq Y x = pr_ge Y x1 - pr_ge Y x2.
Proof.
  edestruct (Rmin_dist_list (map sval (Finite.enum [finType of imgT X])
                            ++ map sval (Finite.enum [finType of imgT Y])) x) as (eps&Hgt0&Hdist).
  exists (x - eps/2). exists (x + eps/2).
  split; eapply pr_ge_split_eps.
  - split; auto. intros. eapply Hdist; auto.
    rewrite mem_cat. apply /orP; left.
    rewrite /index_enum [Finite.enum]unlock //=. rewrite img_fin_enum_sval.
    by apply img_alt', img_alt.
  - split; auto. intros. eapply Hdist; auto.
    rewrite mem_cat. apply /orP; right.
    rewrite /index_enum [Finite.enum]unlock //=. rewrite img_fin_enum_sval.
    by apply img_alt', img_alt.
Qed.

Lemma pr_gt_to_geq2 {A B: finType} {Ω: distrib A} {Ω': distrib B} (X: rrvar Ω) (Y: rrvar Ω')
      (x: R): ∃ x',
    pr_gt X x = pr_ge X x' ∧
    pr_gt Y x = pr_ge Y x'.
Proof.
  edestruct (Rmin_dist_list (map sval (Finite.enum [finType of imgT X])
                            ++ map sval (Finite.enum [finType of imgT Y])) x) as (eps&Hgt0&Hdist).
  exists (x + eps/2).
  split; eapply pr_gt_to_geq_eps.
  - split; auto. intros. eapply Hdist; auto.
    rewrite mem_cat. apply /orP; left.
    rewrite /index_enum [Finite.enum]unlock //=. rewrite img_fin_enum_sval.
    by apply img_alt', img_alt.
  - split; auto. intros. eapply Hdist; auto.
    rewrite mem_cat. apply /orP; right.
    rewrite /index_enum [Finite.enum]unlock //=. rewrite img_fin_enum_sval.
    by apply img_alt', img_alt.
Qed.

Lemma pr_ge_to_gt2 {A B: finType} {Ω: distrib A} {Ω': distrib B} (X: rrvar Ω) (Y: rrvar Ω')
      (x: R): ∃ x',
    pr_ge X x = pr_gt X x' ∧
    pr_ge Y x = pr_gt Y x'.
Proof.
  edestruct (Rmin_dist_list (map sval (Finite.enum [finType of imgT X])
                            ++ map sval (Finite.enum [finType of imgT Y])) x) as (eps&Hgt0&Hdist).
  exists (x - eps/2).
  split; eapply pr_ge_to_gt_eps.
  - split; auto. intros. eapply Hdist; auto.
    rewrite mem_cat. apply /orP; left.
    rewrite /index_enum [Finite.enum]unlock //=. rewrite img_fin_enum_sval.
    by apply img_alt', img_alt.
  - split; auto. intros. eapply Hdist; auto.
    rewrite mem_cat. apply /orP; right.
    rewrite /index_enum [Finite.enum]unlock //=. rewrite img_fin_enum_sval.
    by apply img_alt', img_alt.
Qed.
End convert.

Definition eq_dist {A B: countType} {C: eqType} {Ω1: distrib A} {Ω2: distrib B}
           (X: rvar Ω1 C) (Y: rvar Ω2 C) := ∀ x, pr_eq X x = pr_eq Y x.

Lemma eq_dist_ext {A: countType} {B: eqType} {Ω: distrib A} (X Y: rvar Ω B) :
  (∀ x: A, X x = Y x) →
  eq_dist X Y.
Proof.
  intros Hext v. rewrite /pr_eq. apply pr_eq_pred' => i; by rewrite Hext.
Qed.

Lemma eq_dist_refl {A: countType} {C: eqType} {Ω: distrib A} (X: rvar Ω C) :
  eq_dist X X.
Proof. rewrite /eq_dist => x; done. Qed.

Lemma eq_dist_sym {A B: countType} {C: eqType} {Ω: distrib A} {Ω': distrib B}
      (X: rvar Ω C) (Y: rvar Ω' C) :
  eq_dist X Y → eq_dist Y X.
Proof. rewrite /eq_dist => ??; done. Qed.

Lemma eq_dist_trans {A B C: countType} (D: eqType) {ΩA: distrib A} {ΩB: distrib B} {ΩC: distrib C}
      (X: rvar ΩA D) (Y: rvar ΩB D) (Z: rvar ΩC D) :
  eq_dist X Y → eq_dist Y Z → eq_dist X Z.
Proof. rewrite /eq_dist => Hxy Hyz x. etransitivity; eauto. Qed.

Global Instance eq_dist_Reflexive {A} {B} {Ω}: Reflexive (@eq_dist A A B Ω Ω).
Proof. intros ?. apply eq_dist_refl. Qed.

Lemma eq_dist_comp {A1 A2: finType} {B C: eqType} {Ω1: distrib A1} {Ω2: distrib A2}
      (X1: rvar Ω1 B) (X2: rvar Ω2 B) (f: B → C) :
  eq_dist X1 X2 → eq_dist (rvar_comp X1 f) (rvar_comp X2 f).
Proof.
  intros Heq c. rewrite /pr_eq ?pr_eq_alt_comp.
  rewrite -(big_map sval (λ x, true) (λ v, if f v == c then pr_eq X1 v else 0)).
  rewrite -(big_map sval (λ x, true) (λ v, if f v == c then pr_eq X2 v else 0)).
  rewrite bigop_cond_non0.
  eapply (sum_reidx_map _ _ _ _ (λ x, x)).
  - intros. by rewrite Heq.
  - rewrite //=. intros a ?. case: ifP; last by (intros ?; move /eqP; nra).
    intros _.  move /eqP. rewrite Heq => Hnon0.
    split; auto.
    assert (a \in img X2) as Hin.
    { apply pr_img. move: Hnon0. rewrite /pr_eq. edestruct (ge_pr_0 Ω2); eauto; nra. }
    destruct (proj1 (img_alt X2 a) Hin) as (x&Heq').
    subst. apply /mapP.
    exists (exist _ (X2 x) Hin) => //=.
    rewrite /index_enum [Finite.enum]unlock //=. apply img_fin_mem.
  - intros b Hin _ Hfalse.
    rewrite -Heq. apply /eqP.
    case_eq ((if f b == c then pr_eq X1 b else 0) == 0); auto.
    move /eqP => Hneq.
    exfalso; apply Hfalse.
    exists b; split; auto.
    * move: Hneq. case: ifP; last (intros ?; nra).
      move /eqP => _  Himg.
      assert (b \in img X1) as Hin'.
      { apply pr_img. move: Himg. rewrite /pr_eq.  edestruct (ge_pr_0 Ω1); eauto; nra. }
      destruct (proj1 (img_alt X1 b) Hin') as (x&Heq').
      subst. apply /mapP.
      exists (exist _ (X1 x) Hin') => //=.
      rewrite /index_enum [Finite.enum]unlock //=. apply img_fin_mem.
    * split; auto. rewrite /=. apply /eqP. done.
  - rewrite /index_enum [Finite.enum]unlock //=. rewrite img_fin_enum_sval. apply undup_uniq.
  - rewrite /index_enum [Finite.enum]unlock //=. rewrite img_fin_enum_sval. apply undup_uniq.
  - auto.
Qed.

Lemma big_mkcond_sumbool
     : ∀ (R : Type) (idx : R) (op : Monoid.law idx) (I : Type) (r : seq I) (Q S: I → Prop)
       (P : ∀ i, {Q i} + {S i}) (F : I → R),
       \big[op/idx]_(i <- r | is_left (P i)) F i = \big[op/idx]_(i <- r) (if P i then F i else idx).
Proof.
  intros. rewrite big_mkcond. apply eq_bigr => i ?. destruct (P i) => //=.
Qed.

Lemma eq_dist_pr_ge {A B: finType} {Ω: distrib A} {Ω': distrib B} (X: rrvar Ω)
      (Y: rrvar Ω'):
  eq_dist X Y ↔ (∀ x, pr_ge X x = pr_ge Y x).
Proof.
  rewrite /eq_dist; split.
  - intros Heqd x. rewrite /pr_ge. rewrite ?pr_ge_alt.
    rewrite -?big_mkcond_sumbool.
    rewrite -!(big_map sval (λ y, Rge_dec y x)).
    rewrite //=.
    rewrite (pr_sum_restrict).
    apply sum_reidx_map with (h := λ x, x); auto.
    * intros a Hin. move /andP => [Hge Hgt].
      destruct (Rgt_dec) => //=.
      split; auto.
      rewrite /index_enum [Finite.enum]unlock //=. rewrite img_fin_enum_sval.
      apply img_alt', img_alt.
      apply pr_img. rewrite -Heqd. done.
    * intros ??? Hfalse. edestruct (pr_eq_ge_0 Y); eauto.
      exfalso; apply Hfalse.
      eexists. split; eauto.
      ** rewrite /index_enum [Finite.enum]unlock //=. rewrite img_fin_enum_sval.
         apply img_alt', img_alt, pr_img. rewrite Heqd /pr. eauto.
      ** split; auto. apply /andP; split; auto.
         rewrite Heqd /pr. destruct (Rgt_dec) => //=.
    * rewrite /index_enum [Finite.enum]unlock //=. rewrite img_fin_enum_sval. apply undup_uniq.
    * rewrite /index_enum [Finite.enum]unlock //=. rewrite img_fin_enum_sval. apply undup_uniq.
  - intros Heqd x.
    edestruct (pr_ge_split2 X Y x) as (x1&x2&HeqX&HeqY).
    rewrite HeqX HeqY ?Heqd. done.
Qed.

Lemma eq_dist_pr_gt {A B: finType} {Ω: distrib A} {Ω': distrib B}
      (X: rrvar Ω) (Y: rrvar Ω'):
  eq_dist X Y ↔ (∀ x, pr_gt X x = pr_gt Y x).
Proof.
  rewrite eq_dist_pr_ge; split.
  - intros Heqd x.
    edestruct (pr_gt_to_geq2 X Y x) as (x1&HeqX&HeqY).
    rewrite HeqX HeqY ?Heqd. done.
  - intros Heqd x.
    edestruct (pr_ge_to_gt2 X Y x) as (x1&HeqX&HeqY).
    rewrite HeqX HeqY ?Heqd. done.
Qed.

Module eq_dist_Ex.
Section eq_dist_Ex.

  Variable (A1 A2: countType).
  Variable (B: eqType).
  Variable (Ω1: distrib A1).
  Variable (Ω2: distrib A2).
  Variable (X1: rvar Ω1 B).
  Variable (X2: rvar Ω2 B).
  Variable (f: B → R).

  Variable (EQDIST: eq_dist X1 X2).
  Variable (EX: ex_Ex (rvar_comp X1 f)).

  Definition a := countable_sum (λ r: imgT X1, pr_eq X1 (sval r) * f (sval r)).
  Definition a' :=
         λ n, match n with
              | O => 0
              | S n' => countable_sum (λ r : imgT X1, pr_eq X1 (sval r) * f (sval r)) n'
              end.
  Definition σ := λ n,
            match pickle_inv [countType of imgT X2] n with
            | None => O
            | Some yb =>
              match Rgt_dec (pr_eq X1 (sval yb)) 0 with
              | left Hpf => S (@pickle [countType of imgT X1] (exist _ (sval yb) (pr_img _ _ Hpf)))
              | _ => O
              end
            end.

  Lemma σinj: ∀ n n', a' (σ n) <> 0 → σ n = σ n' → n = n'.
  Proof.
    intros n n'.
    rewrite /σ.
    case_eq (pickle_inv [countType of imgT X2] n); rewrite //=.
    intros s Heqs.
    case_eq (pickle_inv [countType of imgT X2] n'); rewrite //=.
    - intros s' Heqs'. rewrite Heqs'.
      do 2 destruct (Rgt_dec) => //=.
      * intros ? Heq. inversion Heq as [Heq']. apply pickle_inj in Heq'.
        eapply pickle_inv_some_inj; eauto.
        rewrite Heqs Heqs'. f_equal. inversion Heq'.
        by apply sval_inj_pred.
    - intros Hnone.
      destruct (Rgt_dec) => //=.
      rewrite Hnone.
      intros Heqa' Heq0. inversion Heq0.
  Qed.

  Lemma σcov: ∀ n, a' n <> 0 → ∃ m, σ m = n.
  Proof.
    intros n. rewrite /a'. destruct n => //=.
    rewrite /countable_sum.
    case_eq (pickle_inv (img_countType X1) n).
    - intros s Heq => //= Hneq0.
      rewrite EQDIST in Hneq0.
      assert (Hneq0': pr_eq X2 (sval s) <> 0).
      { intros Heq0. rewrite Heq0 Rmult_0_l //= in Hneq0. }
      exists (@pickle [countType of imgT X2] (exist _ (sval s) (pr_img0 _ _ Hneq0'))).
      rewrite /σ pickleK_inv //=.
      destruct (Rgt_dec) as [|Hngt].
      * f_equal. apply pickle_inv_some_inv. rewrite Heq.
        f_equal. apply sval_inj_pred => //=.
      * exfalso. rewrite EQDIST in Hngt. destruct (pr_eq_ge_0 X2 (sval s)); nra.
    - rewrite //=.
  Qed.

  Lemma ex_Ex_alt:
    ex_series (Rabs \o a').
  Proof.
    apply ex_series_incr_1 => //=.
    eapply (ex_series_ext); last by eapply (Ex_comp_pt.ex_series_Ex_comp_pt_abs _ _ _ X1 f).
    intros n. rewrite /countable_sum.
    destruct (pickle_inv) => //=.
    * rewrite Rabs_mult. symmetry. rewrite Rabs_right; last by apply pr_eq_ge_0.
      done.
    * by rewrite Rabs_R0.
  Qed.

  Lemma rearrange:
    ex_series (λ n, Rabs (a' (σ n))) ∧
    is_series (λ n, a' (σ n)) (Series a').
  Proof.
    destruct (ex_Ex_alt) as (v&Hisv).
    edestruct (rearrange.series_rearrange_covering a' σ).
    - apply σinj.
    - apply σcov.
    - eauto.
    - split.
      ** exists v; eauto.
      ** eauto.
  Qed.

  Lemma a'_equal_ex_X1:
    Ex (rvar_comp X1 f) = Series a'.
  Proof.
    rewrite /a'.
    rewrite Series_incr_1; last first.
    { apply ex_series_Rabs, ex_Ex_alt. }
    rewrite Rplus_0_l //.
    by apply Ex_comp_pt.Ex_comp.
  Qed.

  Lemma ex_Ex_eq_dist:
    ex_Ex (rvar_comp X2 f).
  Proof.
    apply ex_Ex_comp_by_alt.
    destruct rearrange as (Hex&_).
    eapply ex_series_ext; last eapply Hex.
    intros n => //=.
    rewrite /a'/σ/countable_sum.
    destruct (pickle_inv) => //=; last by apply Rabs_R0.
    destruct Rgt_dec => //=.
    * rewrite pickleK_inv //= Rabs_mult.
      rewrite Rabs_right; last by nra.
      rewrite EQDIST //.
    * destruct (pr_eq_ge_0 X1 (sval s)) as [?|Heq0]; first by nra.
      rewrite -EQDIST Heq0 Rabs_R0; nra.
  Qed.

  Lemma Ex_eq_dist:
    Ex (rvar_comp X1 f) = Ex (rvar_comp X2 f).
  Proof.
    rewrite a'_equal_ex_X1.
    transitivity (Series (a' \o σ)).
    { symmetry.  apply is_series_unique. destruct rearrange; eauto. }
    rewrite Ex_comp_pt.Ex_comp; last apply ex_Ex_eq_dist.
    apply Series_ext.
    intros n => //=.
    rewrite /a'/σ/countable_sum.
    destruct (pickle_inv) => //=.
    destruct Rgt_dec => //=.
    - rewrite pickleK_inv //= EQDIST //.
    - destruct (pr_eq_ge_0 X1 (sval s)) as [?|Heq0]; first by nra.
      rewrite -EQDIST Heq0; nra.
  Qed.
End eq_dist_Ex.
End eq_dist_Ex.

Definition ex_Ex_eq_dist := eq_dist_Ex.ex_Ex_eq_dist.
Definition Ex_eq_dist := eq_dist_Ex.Ex_eq_dist.
Arguments ex_Ex_eq_dist {_ _ _ _ _}.
Arguments Ex_eq_dist {_ _ _ _ _}.

Lemma Ex_eq_dist_fin {A: finType} {B: countType}
      {Ω: distrib A} {Ω': distrib B} (X: rrvar Ω) (Y: rrvar Ω'):
  eq_dist X Y → Ex X = Ex Y.
Proof.
  intros. apply: (Ex_eq_dist _ _ id) => //=.
Qed.


Require Import String.
Require Import Names.
Require Import PNames.
Require Import Arith.
Require Import Bool.
Require Import Lambda.
Require Import Arith_Lambda.
Require Import Bool_Lambda.
Require Import Functors.
(* Require Import MonadLib. *)

Open Scope string_scope.

Section Type_Test_Section.
  (* Type Testing, of course. *)
  Definition D := AType :+: LType :+: BType.

  Definition tex2' : DType D := tnat' _.
  Definition tex2 := proj1_sig tex2'.
  Definition tex3' : DType D := tarrow' _ tex2' tex2'. (* Nat -> Nat *)
  Definition tex3 := proj1_sig tex3'.
  Definition tex4' : DType D := tarrow' _ (tbool' _) tex2'. (* Bool -> Nat *)
  Definition tex4 := proj1_sig tex4'.

  Global Instance AType_eqTArrow T : FAlgebra eq_TArrowName T (eq_TArrowR D) AType :=
    {| f_algebra := @Default_eq_TArrow _ _ _ _ _ T |}.

  Eval compute in ("Evaluating 'tnat == tnat -> tnat' as an 'TNat :+: TArrow :+: TBool'.").
  Eval compute in (eq_DType D tex2 tex3').
  Eval compute in ("Evaluating 'tnat == tnat' as an 'TNat :+: TArrow :+: TBool'.").
  Eval compute in (eq_DType D tex2 tex2').
  Eval compute in ("Evaluating 'tnat -> tnat == tnat' as an 'TNat :+: TArrow :+: TBool'.").
  Eval compute in (eq_DType D tex3 tex2').
  Eval compute in ("Evaluating 'tnat -> tnat == tnat -> tnat' as an 'TNat :+: TArrow :+: TBool'.").
  Eval compute in (eq_DType D tex3 (tarrow' _ tex2' tex2')).
  Eval compute in ("Evaluating 'tbool -> tnat == tnat -> tnat' as an 'TNat :+: TArrow :+: TBool'.").
  Eval compute in (eq_DType D tex4 (tarrow' _ tex2' tex2')).

  Eval compute in ("Evaluating 'tbool == tnat' as an 'TNat :+: TArrow :+: TBool'.").
  Eval compute in (eq_DType D (tbool _) (tnat' _)).

  Global Instance eq_TArrow_eq_D :
    PAlgebra eq_TArrow_eqName (sig (UP'_P (eq_TArrow_eq_P D))) D.
  Proof.
    eauto 150 with typeclass_instances.
  Defined.

  Global Instance eq_DType_eq_D :
    PAlgebra eq_DType_eqName (sig (UP'_P (eq_DType_eq_P D))) D.
    eauto 150 with typeclass_instances.
  Defined.

End Type_Test_Section.

Section Test_Section.

  Definition E (A : Set) := Arith :+: (Lambda D A) :+: Bool.

  Global Instance Fun_E : forall (A : Set),
    Functor (E A).
  Proof.
    eauto with typeclass_instances.
  Defined.

  Definition ex_3 A : Names.Exp (E A) := (lit' _ 1). (* ex_3 := 1 *)

  Definition ex_1 (A : Set) t1 t2 : Names.Exp (E A) := (* ex_1 := (\x : t1. x) (\y : t2 . y) *)
    app' D E (lam' D E t1 (fun e : A => var' D E e)) (lam' D E t2 (fun e : A => var' D E e)).
  Definition ex_4 A : Names.Exp (E A) := (* ex_4 := \x : tnat . x + x*)
    lam' D E (tnat' _) (fun x => @add' _ _ _ (var' D _ x) (var' D _ x)).
  Example ex_id A : Names.Exp (E A) :=  (* ex_id := \x. x *)
    lam' D _ (tnat' _) (fun x => (var' D _ x)).
  Example ex_5 A : Names.Exp (E A) := (* ex_5 := \x. (x (\y. y)) 1*)
    lam' D _ (tnat' _ ) (fun x => app' D _ (app' D _ (var' D _ x) (ex_id _)) (ex_3 _)).
  Example ex_6 A : Names.Exp (E A) := (* ex_6 := \x. \y. x y*)
    lam' D _ (tnat' _) (fun x => lam' D _ (tnat' _) (fun y => app' D _ (var' D _ x) (var' D _ y))).
  Example ex_7 A : Names.Exp (E A):= (* ex_7 := (\x. (x (\y. y)) 1) (\x. \y. x y)*)
    (app' D _ (ex_5 _) (ex_6 _)).

  Instance D_typeof T : FAlgebra TypeofName T (typeofR D) (E (typeofR D)).
    eauto with typeclass_instances.
  Defined.

  Definition test_typeof e :=
    match (typeof D (E _) e) with
      | Some t1 => DTypePrint _ (proj1_sig t1)
      | None => "Type Error"
    end.

  Eval compute in ("The type of '(\x:tnat->tnat. x) (\y:tnat.y)' should be tnat->tnat").
  Eval compute in (test_typeof (proj1_sig (ex_1 _ tex3' tex2'))).
  Eval compute in ("The type of '1' as a 'Arith :+: Lambda' should be tnat").
  Eval compute in (test_typeof (proj1_sig (ex_3 _))).
  Eval compute in ("The type of '\x.x+x' as a 'Arith :+: Lambda' should be 'tnat -> tnat'").
  Eval compute in (test_typeof (proj1_sig (ex_4 _))).

  Eval compute in (ExpPrint _ (proj1_sig (ex_1 _ tex3' tex2'))).
  Eval compute in (ExpPrint _ (proj1_sig (ex_3 _))).
  Eval compute in (ExpPrint _ (proj1_sig (ex_4 _))).
  Eval compute in (ExpPrint _ (proj1_sig (ex_7 _))).

  Definition V := StuckValue :+: BotValue :+: NatValue :+: (ClosValue E) :+: (BoolValue).

  Instance V_eval : FAlgebra EvalName (Exp E nat) (evalR V) (E nat).
    eauto 150 with typeclass_instances.
  Defined.

  Eval compute in ("Evaluating '(\x:tnat->tnat. x) (\y:tnat.y)'").
  Eval compute in (ValuePrint V (proj1_sig (beval V _ 3 (ex_1 _ tex3' tex2') nil))).
  Eval compute in ("Evaluating '((\x:tnat->tnat. x) (\y:tnat.y)) 1'").
  Eval compute in (ValuePrint V (proj1_sig (beval V _ 3 (app' D E (ex_1 _ tex3' tex2') (lit' _ 1)) nil))).
  Eval compute in ("Evaluating '\x.x+x'").
  Eval compute in (ValuePrint V (proj1_sig (beval V _ 2 (ex_4 _) nil))).
  Eval compute in ("Evaluating '(\x.x+x) 1'").
  Eval compute in (ValuePrint V (proj1_sig (beval V _ 4 (app' D E (ex_4 _) (lit' _ 1)) nil))).
  Eval compute in ("Evaluating '(\x.x+x) 3'").
  Eval compute in (ValuePrint V (proj1_sig (beval V _ 4 (app' D E (ex_4 _) (lit' _ 3)) nil))).
  Eval compute in ("Evaluating '(\x. (x (\y. y)) 1) (\x. \y. x y)'").
  Eval compute in (ValuePrint V (proj1_sig (beval V _ 5 (ex_7 _) nil))).
  Eval compute in ("Evaluating '\x. x'").
  Eval compute in (ValuePrint V (proj1_sig (beval V _ 5 (ex_id _) nil))).

  Definition Eqv (A B : Set) :=
    (NP_Functor_eqv E Arith A B) ::+:: (NP_Functor_eqv E Bool A B) ::+::
    (Lambda_eqv D E A B).

  Definition WFV := (WFValue_Clos D E V Eqv ((fun e => typeof _ _ (proj1_sig e)))) ::+:: (WFValue_Bot D V) ::+::
    (WFValue_VI D V) ::+:: (WFValue_VB D V).

  Definition SV := (SubValue_refl V) ::+:: (SubValue_Bot V) ::+:: (SubValue_Clos E V).

  Definition EV_Alg : PAlgebra EC_ExpName (sig (UP'_P (eval_continuous_Exp_P V (E nat) SV))) (E nat).
    eauto 200 with typeclass_instances.
  Defined.

  Lemma eval_continuous : forall m,
    forall (e : Exp E nat) (gamma gamma' : Env _),
      forall n (Sub_G_G' : Sub_Environment V SV gamma gamma'),
        m <= n ->
        SubValueC _ SV (beval _ _ m e gamma) (beval _ _ n e gamma').
    eapply beval_continuous with (eval_continuous_Exp_E := EV_Alg);
      eauto with typeclass_instances.
  Qed.

  Global Instance Bool_sound P_bind P pb typeof_rec eval_rec :
    PAlgebra eval_Soundness_alg_Name
    (sig (UP'_P2 (eval_alg_Soundness_P D V (E nat) WFV P_bind P
      (E (typeofR D)) (Fun_E (typeofR D)) pb typeof_rec
      eval_rec f_algebra f_algebra))) Bool.
  eauto 250 with typeclass_instances.
  Defined.

  Eval compute in ("Continuity of Evaluation Proven!").

  Hint Extern 0 (id _) => unfold id; eauto with typeclass_instaces : typeclass_instances.

  Theorem soundness : forall n gamma gamma' gamma'' e' e'',
    E_eqvC _ Eqv gamma gamma' e' e'' ->
    forall (WF_gamma : forall n b, lookup gamma' n = Some b ->
      exists T, lookup gamma b = Some T)
    (WF_gamma2 : List.length gamma = List.length gamma')
    (WF_gamma' : forall n b, lookup gamma' n = Some b -> b = n)
    (WF_gamma'' : WF_Environment _ _ WFV gamma'' gamma) T,
    typeof D (E _) (proj1_sig e') = Some T -> WFValueC _ _ WFV (beval _ _ n e'' gamma'') T.
  Proof.
    apply soundness_X with (eval_continuous_Exp_E := EV_Alg);
      eauto 250 with typeclass_instances.
  Qed.

  Eval compute in ("Type Soundness for Arith :+: Lambda :+: Boolean Proven!").

End Test_Section.

(*
*** Local Variables: ***
*** coq-prog-args: ("-emacs-U" "-impredicative-set") ***
*** End: ***
*)


Require Import Lib.

(** printing |-- %\ensuremath{\vdash}% *)
(** printing \in %\ensuremath{\in}% *)
(** printing Gamma %\ensuremath{\Gamma}% *)

(** * Syntax

      Here is the syntax for ReWire. *)

(** ** Monads and Types *)

(** Here is the syntax for layered state monad. *)
Inductive label : Type :=
  | LNil : label
  | LRead : label
  | LWrite : label
  | LReadWrite : label.

Inductive mo : Type :=
  | MIdentity : mo  (* zero-layer state monad *)
  | MStateT   : ty -> label -> mo -> mo

(** Here is the syntax for ``ordinary'' types. *)

with ty : Type :=
  | TArrow   : ty -> ty -> ty
  | TProd    : ty -> ty -> ty
  | TSum     : ty -> ty -> ty
  | TNil     : ty
  | TMonadic : mo -> ty -> ty.

(** Mutual induction schemes for types and monads will be useful. *)
Scheme ty_mo_ind := Induction for ty Sort Prop
  with mo_ty_ind := Induction for mo Sort Prop.

Combined Scheme ty_mo_mutind from ty_mo_ind, mo_ty_ind.

(** *** Relating monads by permissiveness *)

(** Partial order on monads. M1 less permissive than M2 means that M1 has an identical
    store shape but fewer or the same permissions. *)

Inductive monad_less_permissive : mo -> mo -> Prop :=
  | LP_StateT_00 : forall T M1 M2,
                     monad_less_permissive M1 M2 ->
                     monad_less_permissive (MStateT T LNil M1)
                                           (MStateT T LNil M2)
  | LP_StateT_0R : forall T M1 M2,
                     monad_less_permissive M1 M2 ->
                     monad_less_permissive (MStateT T LNil M1)
                                           (MStateT T LRead M2)
  | LP_StateT_0W : forall T M1 M2,
                     monad_less_permissive M1 M2 ->
                     monad_less_permissive (MStateT T LNil M1)
                                           (MStateT T LWrite M2)
  | LP_StateT_0RW : forall T M1 M2,
                     monad_less_permissive M1 M2 ->
                     monad_less_permissive (MStateT T LNil M1)
                                           (MStateT T LReadWrite M2)
  | LP_StateT_RR : forall T M1 M2,
                     monad_less_permissive M1 M2 ->
                     monad_less_permissive (MStateT T LRead M1)
                                           (MStateT T LRead M2)
  | LP_StateT_WW : forall T M1 M2,
                     monad_less_permissive M1 M2 ->
                     monad_less_permissive (MStateT T LWrite M1)
                                           (MStateT T LWrite M2)
  | LP_StateT_RRW : forall T M1 M2,
                     monad_less_permissive M1 M2 ->
                     monad_less_permissive (MStateT T LRead M1)
                                           (MStateT T LReadWrite M2)
  | LP_StateT_WRW : forall T M1 M2,
                     monad_less_permissive M1 M2 ->
                     monad_less_permissive (MStateT T LWrite M1)
                                           (MStateT T LReadWrite M2)
  | LP_StateT_RWRW : forall T M1 M2,
                     monad_less_permissive M1 M2 ->
                     monad_less_permissive (MStateT T LReadWrite M1)
                                           (MStateT T LReadWrite M2)
  | LP_Identity  : monad_less_permissive MIdentity MIdentity.

Hint Constructors monad_less_permissive.

Theorem less_permissive_refl : forall M, monad_less_permissive M M.
Proof.
  intros; induction M; [| destruct l]; auto.
Qed.

Theorem less_permissive_antisymm : forall M1 M2, monad_less_permissive M1 M2 ->
                                                 monad_less_permissive M2 M1 ->
                                                 M1 = M2.
Proof.
  intros M1 M2 Hw12; induction Hw12; intros Hw21; auto; inversion Hw21; subst; rewrite (IHHw12 H0); auto.
Qed.

Theorem less_permissive_trans: forall M1 M2 M3, monad_less_permissive M1 M2 ->
                                                monad_less_permissive M2 M3 ->
                                                monad_less_permissive M1 M3.
Proof.
  induction M1; intros.
    inversion H; subst. inversion H0; subst; auto.
    inversion H; subst; inversion H0; subst; constructor; eapply IHM1; eauto.
Qed.

(* Relation on monads. "M1 disjoint with M2" means that M1 and M2 have the
   same store shape and their permission sets are disjoint. *)
Inductive monad_disjoint : mo -> mo -> Prop :=
  | Disjoint_StateT_00 : forall T M1 M2,
                           monad_disjoint M1 M2 ->
                           monad_disjoint (MStateT T LNil M1)
                                          (MStateT T LNil M2)
  | Disjoint_StateT_0R : forall T M1 M2,
                           monad_disjoint M1 M2 ->
                           monad_disjoint (MStateT T LNil M1)
                                          (MStateT T LRead M2)
  | Disjoint_StateT_0W : forall T M1 M2,
                           monad_disjoint M1 M2 ->
                           monad_disjoint (MStateT T LNil M1)
                                          (MStateT T LWrite M2)
  | Disjoint_StateT_0RW : forall T M1 M2,
                           monad_disjoint M1 M2 ->
                           monad_disjoint (MStateT T LNil M1)
                                          (MStateT T LReadWrite M2)
  | Disjoint_StateT_R0 : forall T M1 M2,
                           monad_disjoint M1 M2 ->
                           monad_disjoint (MStateT T LRead M1)
                                          (MStateT T LNil M2)
  | Disjoint_StateT_W0 : forall T M1 M2,
                           monad_disjoint M1 M2 ->
                           monad_disjoint (MStateT T LWrite M1)
                                          (MStateT T LNil M2)
  | Disjoint_StateT_RW0 : forall T M1 M2,
                           monad_disjoint M1 M2 ->
                           monad_disjoint (MStateT T LReadWrite M1)
                                          (MStateT T LNil M2)
  | Disjoint_Identity  : monad_disjoint MIdentity MIdentity.

Theorem monad_disjoint_symm : forall M1 M2, monad_disjoint M1 M2 ->
                                            monad_disjoint M2 M1.
Proof.
  induction M1; intros ? ?H; inversion H; subst; auto; constructor; auto.
Qed.

(** ** Terms and Configurations *)

Inductive tm : Type :=
  (* Lambda calculus *)
  | tvar     : id -> tm
  | tapp     : tm -> tm -> tm
  | tabs     : id -> ty -> tm -> tm  (* ty is type of binder *)

  (* Data *)
  | tunit    : tm
  | tpair    : tm -> tm -> tm
  | tpi1     : tm -> tm
  | tpi2     : tm -> tm
  | tinl     : tm -> ty -> tm        (* ty is type on right of sum *)
  | tinr     : tm -> ty -> tm        (* ty is type on left of sum *)
  | tcase    : tm -> tm -> tm -> tm

  (* State monad operations *)
  | treturn  : tm -> mo -> tm
  | tbind    : tm -> tm -> tm
  | tlift    : tm -> ty -> tm        (* ty is type of state for new layer *)
  | televate : tm -> mo -> tm        (* mo is monad we are elevating to *)
  | tget     : ty -> mo -> tm        (* ty is type of state for new layer, mo is base monad *)
  | tput     : tm -> ty -> mo -> tm  (* ty is type of state for new layer, mo is base monad *)
  .

Tactic Notation "tm_cases" tactic(first) ident(c) :=
  first;
  [ Case_aux c "tvar" | Case_aux c "tapp" | Case_aux c "tabs"
  | Case_aux c "tunit" | Case_aux c "tpair"
  | Case_aux c "tpi1" | Case_aux c "tpi2"
  | Case_aux c "tinl" | Case_aux c "tinr"
  | Case_aux c "tcase" | Case_aux c "treturn" | Case_aux c "tbind" | Case_aux c "tlift"
  | Case_aux c "televate" | Case_aux c "tget" | Case_aux c "tput" ].

Definition store := list tm.
Definition configuration := (tm*store)%type.

(** * Typing Judgments *)

(** ** For terms *)
Definition context := partial_map ty.

Reserved Notation "Gamma '|--' t '\in' T" (at level 40).

Inductive has_type : context -> tm -> ty -> Prop :=
  | T_Var  : forall Gamma x T,
               Gamma x = Some T -> Gamma |-- tvar x \in T
  | T_Abs  : forall Gamma x T_11 T_12 t_12,
               extend Gamma x T_11 |-- t_12 \in T_12 ->
                 Gamma |-- tabs x T_11 t_12 \in TArrow T_11 T_12
  | T_App  : forall T_11 T_12 Gamma t_1 t_2,
               Gamma |-- t_1 \in TArrow T_11 T_12 ->
               Gamma |-- t_2 \in T_11 ->
                 Gamma |-- tapp t_1 t_2 \in T_12

  | T_Unit : forall Gamma,
               Gamma |-- tunit \in TNil
  | T_Pair : forall Gamma T_1 T_2 t_1 t_2,
               Gamma |-- t_1 \in T_1 ->
               Gamma |-- t_2 \in T_2 ->
                 Gamma |-- tpair t_1 t_2 \in TProd T_1 T_2
  | T_Pi1  : forall Gamma T_1 T_2 t,
               Gamma |-- t \in TProd T_1 T_2 ->
                 Gamma |-- tpi1 t \in T_1
  | T_Pi2  : forall Gamma T_1 T_2 t,
               Gamma |-- t \in TProd T_1 T_2 ->
                 Gamma |-- tpi2 t \in T_2
  | T_Inl  : forall Gamma T_1 T_2 t,
               Gamma |-- t \in T_1 ->
                 Gamma |-- tinl t T_2 \in TSum T_1 T_2
  | T_Inr  : forall Gamma T_1 T_2 t,
               Gamma |-- t \in T_2 ->
                 Gamma |-- tinr t T_1 \in TSum T_1 T_2
  | T_Case : forall Gamma T_l T_r T_res t_s t_l t_r,
               Gamma |-- t_s \in TSum T_l T_r ->
               Gamma |-- t_l \in TArrow T_l T_res ->
               Gamma |-- t_r \in TArrow T_r T_res ->
                 Gamma |-- tcase t_s t_l t_r \in T_res

  | T_Return  : forall Gamma T Mo t,
                  Gamma |-- t \in T ->
                    Gamma |-- treturn t Mo \in TMonadic Mo T
(*  | T_BindM : forall Gamma t t' T_1 T_2 Mo L1 L2,
                   Gamma |-- t \in TMonadic (MStateT T_1 L1 Mo) T_1 ->
                   Gamma |-- t' \in TArrow T_1 (TMonadic (MStateT T_2 L2 Mo) TNil) ->
                   Gamma |-- tbind t t' \in TMonadic (MStateT T_2 (LReadWrite L1 L2) Mo) TNil *)
  | T_Bind : forall Gamma t Mo T_1 t' T_2,
                  Gamma |-- t  \in TMonadic Mo T_1 ->
                  Gamma |-- t' \in TArrow T_1 (TMonadic Mo T_2) ->
                    Gamma |-- tbind t t' \in TMonadic Mo T_2
  | T_Lift    : forall Gamma t Mo T T',
                  Gamma |-- t \in TMonadic Mo T' ->
                    Gamma |-- tlift t T \in TMonadic (MStateT T LNil Mo) T'
  | T_Elevate : forall Gamma Mo Mo' t T,
                  monad_less_permissive Mo Mo' ->
                  Gamma |-- t \in TMonadic Mo T ->
                    Gamma |-- televate t Mo' \in TMonadic Mo' T
  | T_Get     : forall Gamma Mo T,
                  Gamma |-- tget T Mo \in TMonadic (MStateT T LRead Mo) T
  | T_Put     : forall Gamma Mo t T,
                  Gamma |-- t \in T ->
                   Gamma |-- tput t T Mo \in TMonadic (MStateT T LWrite Mo) TNil

where "Gamma '|--' t '\in' T" := (has_type Gamma t T).

Tactic Notation "has_type_cases" tactic(first) ident(c) :=
  first;
  [ Case_aux c "T_Var" | Case_aux c "T_Abs" | Case_aux c "T_App"
  | Case_aux c "T_Unit" | Case_aux c "T_Pair"
  | Case_aux c "T_Pi1" | Case_aux c "T_Pi2"
  | Case_aux c "T_Inl" | Case_aux c "T_Inr"
  | Case_aux c "T_Case" | Case_aux c "T_Return" | Case_aux c "T_Bind" | Case_aux c "T_Lift"
  | Case_aux c "T_Elevate" | Case_aux c "T_Get" | Case_aux c "T_Put" ].

Lemma type_unique : forall t Gamma T1 T2, Gamma |-- t \in T1 -> Gamma |-- t \in T2 -> T1 = T2.
Proof.
  induction t; intros Gamma T1 T2 Ht1 Ht2; inversion Ht1; subst; inversion Ht2; subst; try auto.
  Case "tvar".
    assert (Some T1=Some T2). rewrite <- H2. rewrite H1. auto.
    inversion H. auto.
  Case "tapp".
    assert (T_11 = T_0). eapply IHt2. eassumption. assumption.
    subst.
    assert (TArrow T_0 T2=TArrow T_0 T1). eapply IHt1. eassumption. assumption.
    inversion H. auto.
  Case "tabs".
    assert (T_12=T_0). eapply IHt; eauto.
    subst; auto.
  Case "tpair".
    assert (T_1=T_0); eauto.
    assert (T_2=T_3); eauto.
    subst; auto.
  Case "tpi1".
    assert (TProd T1 T_2=TProd T2 T_0); eauto.
    inversion H; auto.
  Case "tpi2".
    assert (TProd T_1 T1=TProd T_0 T2); eauto.
    inversion H; auto.
  Case "tinl".
    assert (T_1 = T_0); eauto.
    subst; auto.
  Case "tinr".
    assert (T_2 = T_0); eauto.
    subst; auto.
  Case "case".
    assert (TSum T_l T_r = TSum T_l0 T_r0); eauto.
    inversion H; subst.
    assert (TArrow T_r0 T1 = TArrow T_r0 T2); eauto.
    inversion H0; auto.
  Case "treturn".
    assert (T=T0); eauto.
    subst; auto.
  Case "tbindM1". (*
    pose proof (IHt1 _ _ _ H2 H3).
    inversion H.
    pose proof (IHt2 _ _ _ H4 H6).
    inversion H0.
    subst. auto.
  Case "tbindM2".
    pose proof (IHt1 _ _ _ H2 H3).
       inversion H.
    pose proof (IHt2 _ _ _ H4 H6).
       inversion H0. subst.
     inversion Ht1. subst. inversion Ht2. subst. admit. admit. *) pose proof (IHt2 _ _ _ H4 H6).
     inversion H. auto.
  Case "tlift". subst.
    assert (TMonadic Mo T' = TMonadic Mo0 T'0); eauto.
    inversion H; auto.
  Case "televate".
    assert (TMonadic Mo T = TMonadic Mo0 T0); eauto.
    inversion H; auto.
Qed.

(** ** For configurations *)
Inductive store_matches_monad : store -> mo -> Prop :=
  | matches_monad_empty : store_matches_monad nil MIdentity
  | matches_monad_cons  : forall Mo l T t Sto,
                            \empty |-- t \in T ->
                            store_matches_monad Sto Mo ->
                              store_matches_monad (t::Sto) (MStateT T l Mo).

Reserved Notation "co '|>' T" (at level 40).

Inductive configuration_has_type : configuration -> ty -> Prop :=
  | configuration_has_type_intro : forall t T Mo Sto,
                                     \empty |-- t \in TMonadic Mo T ->
                                     store_matches_monad Sto Mo ->
                                       (t,Sto) |> TMonadic Mo T
where "co |> T" := (configuration_has_type co T).

Lemma less_permissive_store_matches : forall Mo1 Mo2 Sto,
                                         monad_less_permissive Mo1 Mo2 ->
                                         store_matches_monad Sto Mo1 ->
                                         store_matches_monad Sto Mo2.
Proof.
  intros Mo1 Mo2 Sto.
  generalize dependent Mo2. generalize dependent Mo1.
  induction Sto; intros Mo1 Mo2 Hless Hmatches.
    Case "empty".
      inversion Hmatches; subst; inversion Hless; auto.
    Case "cons".
      inversion Hmatches; subst; inversion Hless; subst; constructor.
        assumption. apply (IHSto Mo); auto.
        assumption. apply (IHSto Mo); auto.
        assumption. apply (IHSto Mo); auto.
        assumption. apply (IHSto Mo); auto.
         assumption. apply (IHSto Mo); auto.
        assumption. apply (IHSto Mo); auto.
        assumption. apply (IHSto Mo); auto.
        assumption. apply (IHSto Mo); auto.
         assumption. apply (IHSto Mo); auto.
Qed.

Lemma more_permissive_store_matches : forall Mo1 Mo2 Sto,
                                        monad_less_permissive Mo2 Mo1 ->
                                        store_matches_monad Sto Mo1 ->
                                        store_matches_monad Sto Mo2.
Proof.
  intros Mo1 Mo2 Sto.
  generalize dependent Mo2. generalize dependent Mo1.
  induction Sto; intros Mo1 Mo2 Hless Hmatches.
    Case "empty".
      inversion Hmatches; subst; inversion Hless; auto.
    Case "cons".
      inversion Hmatches; subst; inversion Hless; subst; constructor.
            assumption. apply (IHSto Mo); auto.
        assumption. apply (IHSto Mo); auto.
        assumption. apply (IHSto Mo); auto.
        assumption. apply (IHSto Mo); auto.
         assumption. apply (IHSto Mo); auto.
        assumption. apply (IHSto Mo); auto.
        assumption. apply (IHSto Mo); auto.
        assumption. apply (IHSto Mo); auto.
         assumption. apply (IHSto Mo); auto.
Qed.

(** * Substitution *)

Reserved Notation "'[' x ':=' s ']' t" (at level 20).

(* Substitution function. *)
Fixpoint subst (x:id) (s:tm) (t:tm) : tm :=
  match t with
  | tvar x' =>
      if eq_id_dec x x' then s else t
  | tapp t1 t2 =>
      tapp ([x:=s] t1) ([x:=s] t2)
  | tabs x' T t1 =>
      tabs x' T (if eq_id_dec x x' then t1 else ([x:=s] t1))

  | tunit =>
      tunit
  | tpair t1 t2 =>
      tpair ([x:=s] t1) ([x:=s] t2)
  | tpi1 t1 =>
      tpi1 ([x:=s] t1)
  | tpi2 t1 =>
      tpi2 ([x:=s] t1)
  | tinl t1 T =>
      tinl ([x:=s] t1) T
  | tinr t1 T =>
      tinr ([x:=s] t1) T
  | tcase t1 t2 t3 =>
      tcase ([x:=s] t1) ([x:=s] t2) ([x:=s] t3)

  | treturn t1 M =>
      treturn ([x:=s] t1) M
  | tbind t1 t2 =>
      tbind ([x:=s] t1) ([x:=s] t2)
  | tlift t1 T =>
      tlift ([x:=s] t1) T
  | televate t1 Mo =>
      televate ([x:=s] t1) Mo
  | tget T Mo =>
      tget T Mo
  | tput t1 T Mo =>
      tput ([x:=s] t1) T Mo
  end

where "'[' x ':=' s ']' t" := (subst x s t).

(** A variable _x_ appears free in a term _t_ .*)

Inductive appears_free_in : id -> tm -> Prop :=
  | afi_var : forall x,
      appears_free_in x (tvar x)
  | afi_app1 : forall x t1 t2,
      appears_free_in x t1 ->
      appears_free_in x (tapp t1 t2)
  | afi_app2 : forall x t1 t2,
      appears_free_in x t2 ->
      appears_free_in x (tapp t1 t2)
  | afi_abs : forall x y T11 t12,
      y <> x  ->
      appears_free_in x t12 ->
      appears_free_in x (tabs y T11 t12)
  (* pairs *)
  | afi_pair1 : forall x t1 t2,
      appears_free_in x t1 ->
      appears_free_in x (tpair t1 t2)
  | afi_pair2 : forall x t1 t2,
      appears_free_in x t2 ->
      appears_free_in x (tpair t1 t2)
  | afi_pi1 : forall x t,
      appears_free_in x t ->
      appears_free_in x (tpi1 t)
  | afi_pi2 : forall x t,
      appears_free_in x t ->
      appears_free_in x (tpi2 t)
  (* copairs *)
  | afi_tinl : forall x t TR,
      appears_free_in x t ->
      appears_free_in x (tinl t TR)
  | afi_tinr : forall x t TL,
      appears_free_in x t ->
      appears_free_in x (tinr t TL)
  | afi_tcase1 : forall x t1 t2 t3,
      appears_free_in x t1 ->
      appears_free_in x (tcase t1 t2 t3)
  | afi_tcase2 : forall x t1 t2 t3,
      appears_free_in x t2 ->
      appears_free_in x (tcase t1 t2 t3)
  | afi_tcase3 : forall x t1 t2 t3,
      appears_free_in x t3 ->
      appears_free_in x (tcase t1 t2 t3)
  (* monadic constructors *)
  | afi_treturn : forall x t M,
      appears_free_in x t ->
      appears_free_in x (treturn t M)
  | afi_tbind1 : forall x t1 t2,
      appears_free_in x t1 ->
      appears_free_in x (tbind t1 t2)
  | afi_tbind2 : forall x t1 t2,
      appears_free_in x t2 ->
      appears_free_in x (tbind t1 t2)
  | afi_tlift : forall x t T,
      appears_free_in x t ->
      appears_free_in x (tlift t T)
  | afi_televate : forall x t M,
      appears_free_in x t ->
      appears_free_in x (televate t M)
  | afi_tput : forall x t T M,
      appears_free_in x t ->
      appears_free_in x (tput t T M)
.

Hint Constructors appears_free_in.

Definition closed (t:tm) := forall x, ~ appears_free_in x t.

Lemma context_invariance : forall Gamma Gamma' t S,
      Gamma |-- t \in S  ->
      (forall x, appears_free_in x t -> Gamma x = Gamma' x)  ->
      Gamma' |-- t \in S.
Proof with eauto.
  intros. generalize dependent Gamma'.
  induction H; intros Gamma' Heqv; auto; try (econstructor;eauto;fail)...
  Case "tvar".
    apply T_Var... rewrite <- Heqv...
  Case "tabs".
    apply T_Abs... apply IHhas_type. intros y Hafi.
    unfold extend. destruct (eq_id_dec x y)...
Qed.

Lemma free_in_context : forall x t T Gamma,
      appears_free_in x t ->
        Gamma |-- t \in T ->
        (exists T', Gamma x = Some T').
Proof with eauto.
  intros x t T Gamma Hafi Htyp.
  induction Htyp; inversion Hafi; subst...
  Case "T_Abs".
    destruct IHHtyp as [T' Hctx]... exists T'.
    unfold extend in Hctx.
    rewrite neq_id in Hctx...
Qed.

Lemma subst_preserves_typing : forall Gamma x U v t T,
      (extend Gamma x U) |-- t \in T ->
                  \empty |-- v \in U ->
                   Gamma |-- ([x:=v] t) \in T.
Proof with eauto.
  intros Gamma x U v t S Htypt Htypv.
  generalize dependent Gamma. generalize dependent S.
  induction t; intros S Gamma Htypt; simpl; inversion Htypt; subst; try (econstructor;eauto;fail).
  Case "tvar".
    simpl. rename i into y.
    destruct (eq_id_dec x y).
    SCase "x=y".
      subst.
      unfold extend in H1. rewrite eq_id in H1.
      inversion H1; subst. clear H1.
      eapply context_invariance...
      intros x Hcontra.
      destruct (free_in_context _ _ S empty Hcontra) as [T' HT']...
      inversion HT'.
    SCase "x<>y".
      apply T_Var... unfold extend in H1. rewrite neq_id in H1...
  Case "tabs".
      rename i into y. rename t into T11.
      apply T_Abs...
      destruct (eq_id_dec x y).
      SCase "x=y".
         eapply context_invariance...
         subst.
         intros x Hafi. unfold extend.
         destruct (eq_id_dec y x)...
      SCase "x<>y".
         apply IHt.
         eapply context_invariance...
         intros z Hafi.
         unfold extend.
         destruct (eq_id_dec y z)...
         subst. rewrite neq_id...
Qed.

(** * Values *)

(* Values. Note that computations are *always* values (much like functional
   abstractions), i.e. they do not reduce via the lambda calculus reduction
   relation (but they *will* reduce under the monadic reduction relation,
   defined later). *)
Inductive value : tm -> Prop :=
  | v_abs : forall x T t,
      value (tabs x T t)
  | v_unit : value tunit
  | v_pair : forall t1 t2,
      value t1 -> value t2 -> value (tpair t1 t2)
  | v_inl : forall t T,
      value t -> value (tinl t T)
  | v_inr : forall t T,
      value t -> value (tinr t T)
  | v_return : forall t Mo,
      value (treturn t Mo)
  | v_bind : forall t1 t2,
      value (tbind t1 t2)
  | v_lift : forall t T,
      value (tlift t T)
  | v_elevate : forall t Mo,
      value (televate t Mo)
  | v_get : forall T Mo,
      value (tget T Mo)
  | v_put : forall t T Mo,
      value (tput t T Mo).

Theorem value_dec : forall (t:tm), value t \/ ~ (value t).
  intro.
  induction t; try (right; intro H; inversion H; fail);
               try (left; constructor; fail).

  Case "T_Pair".
    destruct IHt1.
    SCase "t1 is a value".
      destruct IHt2.
      SSCase "t2 is a value".
        left; constructor; assumption.
      SSCase "t2 is not a value".
        right; intro H1; inversion H1; contradiction.
    SCase "t1 is not a value".
      right. intro H1. inversion H1; contradiction.

  Case "T_InL".
    destruct IHt.
      left; apply v_inl; assumption.
      right; intro Hval; inversion Hval; contradiction.

  Case "T_InR".
    destruct IHt.
      left; apply v_inr; assumption.
      right; intro Hval; inversion Hval; contradiction.
Qed.

Ltac decide_value e := destruct (value_dec e) as [?Hval | ?Hnotval].

(** ** Done computations *)

Inductive done_mo : configuration -> Prop :=
  | done_return : forall v Mo Sto, value v -> done_mo (treturn v Mo,Sto).

Theorem done_mo_dec (co:configuration) : done_mo co \/ ~ (done_mo co).
  destruct co as [t Sto].
  induction t; try (right; intro H; inversion H; fail).
  Case "treturn".
    destruct (value_dec t).
    SCase "t is a value".
      left; constructor; assumption.
    SCase "t is not a value".
      right; intro Hdone; inversion Hdone; contradiction.
Qed.


From stdpp Require Import prelude.
From VLSM.Lib Require Import EquationsExtras.
From VLSM.Lib Require Import Preamble.

(** * Natural number utility definitions and lemmas

  Given a decidable property on naturals and a bound, finds the
  largest natural (not larger the than bound) for which the property holds.
*)
Equations find_largest_nat_with_property_bounded
  (P : nat -> Prop) `{forall n, Decision (P n)}
  (bound : nat)
  : option nat :=
| P, 0 => None
| P, S n => if decide (P n) then Some n else find_largest_nat_with_property_bounded P n.

Lemma find_largest_nat_with_property_bounded_Some
  (P : nat -> Prop) `{forall n, Decision (P n)}
  (bound : nat)
  : forall n,
      find_largest_nat_with_property_bounded P bound = Some n
        <->
      maximal_among lt (fun n => n < bound /\ P n) n.
Proof.
  induction bound; simp find_largest_nat_with_property_bounded.
  - by split; [| intros []; lia].
  - case_decide as HP; split.
    + by intros [= ->]; repeat split; cbn; [lia | | lia].
    + intros [[Hn HPn] Hmax]; f_equal; cbn in Hmax.
      destruct (decide (n < bound)); [| lia].
      cut (bound < n); [lia |].
      by apply Hmax; [split; [lia |] |].
    + rewrite IHbound.
      intros [[] Hmax]; split; [by split; [lia |] |].
      intros m [] ?.
      apply Hmax; [split; [| done] | done].
      by destruct (decide (m = bound)); [subst | lia].
    + intros [[] Hmax]; apply IHbound; repeat split; [| done |].
      * by destruct (decide (n = bound)); [subst | lia].
      * by intros m [] ?; apply Hmax; [split; [lia |] |].
Qed.

Lemma find_largest_nat_with_property_bounded_None
  (P : nat -> Prop) `{forall n, Decision (P n)}
  (bound : nat)
  : find_largest_nat_with_property_bounded P bound = None
      <->
    forall n, n < bound -> ~ P n.
Proof.
  induction bound; simp find_largest_nat_with_property_bounded; [by split; [lia |] |].
  case_decide as HP; split; [done | | |].
  - by intro Hmax; contradict HP; apply Hmax; lia.
  - intros.
    destruct (decide (n < bound)); [by eapply IHbound |].
    by replace n with bound; [| lia].
  - intros Hmax; apply IHbound.
    by intros; apply Hmax; lia.
Qed.

(**
  Given a predicate on indices and naturals, a bound for each index, and a
  list of indices, finds first index in the list and largest natural (less than
  bound) corresponding to it for which the predicate holds.
*)
Equations find_first_indexed_largest_nat_with_propery_bounded {index}
  (P : index -> nat -> Prop) `{!RelDecision P}
  (bound : index -> nat)
  (indices : list index)
  : option (index * nat) :=
| _, _, [] => None
| P, bound, i :: indices' with inspect (find_largest_nat_with_property_bounded (P i) (bound i)) =>
  | None eq: Hdec => find_first_indexed_largest_nat_with_propery_bounded P bound indices';
  | (Some n) eq: Hdec => Some (i, n).

Lemma find_first_indexed_largest_nat_with_propery_bounded_Some {index}
  (P : index -> nat -> Prop) `{!RelDecision P}
  (bound : index -> nat)
  (indices : list index) (Hindices : NoDup indices)
  : forall (i : index) (n : nat),
      find_first_indexed_largest_nat_with_propery_bounded P bound indices = Some (i, n)
        <->
      i ∈ indices /\
      find_largest_nat_with_property_bounded (P i) (bound i) = Some n /\
      forall (prefix suffix : list index),
        indices = prefix ++ [i] ++ suffix ->
        forall (j : index), j ∈ prefix ->
        find_largest_nat_with_property_bounded (P j) (bound j) = None.
Proof.
  induction indices; simp find_first_indexed_largest_nat_with_propery_bounded.
  - by split; [inversion 1 | intros [Hcontra]; inversion Hcontra].
  - apply NoDup_cons in Hindices as [Ha Hindices].
    cbn; split;
      destruct (find_largest_nat_with_property_bounded  (P a) (bound a)) as [_n |] eqn: Hpos.
    + inversion 1; subst a _n; split_and!; [by left | done |].
      intros [| _i prefix] ? Heq; [by inversion 1 |].
      contradict Ha; simplify_list_eq.
      by apply elem_of_app; right; left.
    + rewrite IHindices by done.
      intros (Hi & Hposition & Hfirst).
      split_and!; [by right | done |].
      intros [| _a prefix] ? Heq; [by inversion 1 |].
      simplify_list_eq.
      by inversion 1; [| eapply Hfirst].
    + intros (Hi & Hposi & Hfirst).
      apply elem_of_cons in Hi as [<- | Hi]; [by congruence |].
      apply elem_of_list_split in Hi as (prefix' & suffix & ->).
      by erewrite Hfirst in Hpos; [| rewrite app_comm_cons | left].
    + intros (Hi & Hposi & Hfirst).
      apply elem_of_cons in Hi as [<- | Hi]; [by congruence |].
      apply IHindices; [done |].
      split_and!; [done.. |].
      intros prefix suffix -> j Hj.
      eapply Hfirst; [| by right].
      by rewrite app_comm_cons.
Qed.

Lemma find_first_indexed_largest_nat_with_propery_bounded_None {index}
  (P : index -> nat -> Prop) `{!RelDecision P}
  (bound : index -> nat)
  (indices : list index) (Hindices : NoDup indices)
  : find_first_indexed_largest_nat_with_propery_bounded P bound indices = None
      <->
    forall (i : index), i ∈ indices ->
      find_largest_nat_with_property_bounded (P i) (bound i) = None.
Proof.
  induction indices; simp find_first_indexed_largest_nat_with_propery_bounded;
    [by split; [inversion 2 |] |].
  apply NoDup_cons in Hindices as [Ha Hindices].
  cbn; split; destruct find_largest_nat_with_property_bounded eqn: Hp; [done | | |].
  - by intros Hnone i Hi; apply elem_of_cons in Hi as  [-> | Hi]; [| eapply IHindices].
  - intro Hmin.
    replace (find_largest_nat_with_property_bounded (P a) (bound a))
      with (@None nat) in Hp; [done |].
    by symmetry; apply Hmin; left.
  - intros Hmin.
    apply IHindices; [done |].
    by intros; apply Hmin; right.
Qed.


(** * ImpParser: Lexing and Parsing in Coq *)

(** The development of the Imp language in [Imp.v] completely ignores
    issues of concrete syntax -- how an ascii string that a programmer
    might write gets translated into abstract syntax trees defined by
    the datatypes [aexp], [bexp], and [com].  In this chapter, we
    illustrate how the rest of the story can be filled in by building
    a simple lexical analyzer and parser using Coq's functional
    programming facilities.

    It is not important to understand all the details here (and
    accordingly, the explanations are fairly terse and there are no
    exercises).  The main point is simply to demonstrate that it can
    be done.  You are invited to look through the code -- most of it
    is not very complicated, though the parser relies on some
    "monadic" programming idioms that may require a little work to
    make out -- but most readers will probably want to just skim down
    to the Examples section at the very end to get the punchline. *)

(* ################################################################# *)
(** * Internals *)

Require Import Coq.Strings.String.
Require Import Coq.Strings.Ascii.
Require Import Coq.Arith.Arith.
Require Import Coq.Arith.EqNat.
Require Import Coq.Lists.List.
Import ListNotations.
Require Import Maps.
Require Import Imp.

(* ================================================================= *)
(** ** Lexical Analysis *)

Definition isWhite (c : ascii) : bool :=
  let n := nat_of_ascii c in
  orb (orb (beq_nat n 32)   (* space *)
           (beq_nat n 9))   (* tab *)
      (orb (beq_nat n 10)   (* linefeed *)
           (beq_nat n 13)). (* Carriage return. *)

Notation "x '<=?' y" := (leb x y)
  (at level 70, no associativity) : nat_scope.

Definition isLowerAlpha (c : ascii) : bool :=
  let n := nat_of_ascii c in
    andb (97 <=? n) (n <=? 122).

Definition isAlpha (c : ascii) : bool :=
  let n := nat_of_ascii c in
    orb (andb (65 <=? n) (n <=? 90))
        (andb (97 <=? n) (n <=? 122)).

Definition isDigit (c : ascii) : bool :=
  let n := nat_of_ascii c in
     andb (48 <=? n) (n <=? 57).

Inductive chartype := white | alpha | digit | other.

Definition classifyChar (c : ascii) : chartype :=
  if isWhite c then
    white
  else if isAlpha c then
    alpha
  else if isDigit c then
    digit
  else
    other.

Fixpoint list_of_string (s : string) : list ascii :=
  match s with
  | EmptyString => []
  | String c s => c :: (list_of_string s)
  end.

Fixpoint string_of_list (xs : list ascii) : string :=
  fold_right String EmptyString xs.

Definition token := string.

Fixpoint tokenize_helper (cls : chartype) (acc xs : list ascii)
                       : list (list ascii) :=
  let tk := match acc with [] => [] | _::_ => [rev acc] end in
  match xs with
  | [] => tk
  | (x::xs') =>
    match cls, classifyChar x, x with
    | _, _, "("      =>
      tk ++ ["("]::(tokenize_helper other [] xs')
    | _, _, ")"      =>
      tk ++ [")"]::(tokenize_helper other [] xs')
    | _, white, _    =>
      tk ++ (tokenize_helper white [] xs')
    | alpha,alpha,x  =>
      tokenize_helper alpha (x::acc) xs'
    | digit,digit,x  =>
      tokenize_helper digit (x::acc) xs'
    | other,other,x  =>
      tokenize_helper other (x::acc) xs'
    | _,tp,x         =>
      tk ++ (tokenize_helper tp [x] xs')
    end
  end %char.

Definition tokenize (s : string) : list string :=
  map string_of_list (tokenize_helper white [] (list_of_string s)).

Example tokenize_ex1 :
    tokenize "abc12==3  223*(3+(a+c))" %string
  = ["abc"; "12"; "=="; "3"; "223";
       "*"; "("; "3"; "+"; "(";
       "a"; "+"; "c"; ")"; ")"]%string.
Proof. reflexivity. Qed.

(* ================================================================= *)
(** ** Parsing *)

(* ----------------------------------------------------------------- *)
(** *** Options With Errors *)

(** An [option] type with error messages: *)

Inductive optionE (X:Type) : Type :=
  | SomeE : X -> optionE X
  | NoneE : string -> optionE X.

Implicit Arguments SomeE [[X]].
Implicit Arguments NoneE [[X]].

(** Some syntactic sugar to make writing nested match-expressions on
    optionE more convenient. *)

Notation "'DO' ( x , y ) <== e1 ; e2"
   := (match e1 with
         | SomeE (x,y) => e2
         | NoneE err => NoneE err
       end)
   (right associativity, at level 60).

Notation "'DO' ( x , y ) <-- e1 ; e2 'OR' e3"
   := (match e1 with
         | SomeE (x,y) => e2
         | NoneE err => e3
       end)
   (right associativity, at level 60, e2 at next level).

(* ----------------------------------------------------------------- *)
(** *** Symbol Table *)

(** Build a mapping from [tokens] to [nats].  A real parser would do
    this incrementally as it encountered new symbols, but passing
    around the symbol table inside the parsing functions is a bit
    inconvenient, so instead we do it as a first pass. *)

Fixpoint build_symtable (xs : list token) (n : nat)
                      : (token -> nat) :=
  match xs with
  | [] => (fun s => n)
  | x::xs =>
    if (forallb isLowerAlpha (list_of_string x))
    then (fun s => if string_dec s x then n
                   else (build_symtable xs (S n) s))
    else build_symtable xs n
  end.

(* ----------------------------------------------------------------- *)
(** *** Generic Combinators for Building Parsers *)

Open Scope string_scope.

Definition parser (T : Type) :=
  list token -> optionE (T * list token).

Fixpoint many_helper {T} (p : parser T) acc steps xs :=
  match steps, p xs with
  | 0, _ =>
      NoneE "Too many recursive calls"
  | _, NoneE _ =>
      SomeE ((rev acc), xs)
  | S steps', SomeE (t, xs') =>
      many_helper p (t::acc) steps' xs'
  end.

(** A (step-indexed) parser that expects zero or more [p]s: *)

Fixpoint many {T} (p : parser T) (steps : nat) : parser (list T) :=
  many_helper p [] steps.

(** A parser that expects a given token, followed by [p]: *)

Definition firstExpect {T} (t : token) (p : parser T)
                     : parser T :=
  fun xs => match xs with
            | x::xs' =>
              if string_dec x t
              then p xs'
              else NoneE ("expected '" ++ t ++ "'.")
            | [] =>
              NoneE ("expected '" ++ t ++ "'.")
            end.

(** A parser that expects a particular token: *)

Definition expect (t : token) : parser unit :=
  firstExpect t (fun xs => SomeE(tt, xs)).

(* ----------------------------------------------------------------- *)
(** *** A Recursive-Descent Parser for Imp *)

(** Identifiers: *)

Definition parseIdentifier (symtable :string->nat)
                           (xs : list token)
                         : optionE (id * list token) :=
match xs with
| [] => NoneE "Expected identifier"
| x::xs' =>
    if forallb isLowerAlpha (list_of_string x) then
      SomeE (Id (symtable x), xs')
    else
      NoneE ("Illegal identifier:'" ++ x ++ "'")
end.

(** Numbers: *)

Definition parseNumber (xs : list token)
                     : optionE (nat * list token) :=
match xs with
| [] => NoneE "Expected number"
| x::xs' =>
    if forallb isDigit (list_of_string x) then
      SomeE (fold_left
               (fun n d =>
                  10 * n + (nat_of_ascii d -
                            nat_of_ascii "0"%char))
               (list_of_string x)
               0,
             xs')
    else
      NoneE "Expected number"
end.

(** Parse arithmetic expressions *)

Fixpoint parsePrimaryExp (steps:nat) symtable
                         (xs : list token)
                       : optionE (aexp * list token) :=
  match steps with
  | 0 => NoneE "Too many recursive calls"
  | S steps' =>
      DO (i, rest) <-- parseIdentifier symtable xs ;
          SomeE (AId i, rest)
      OR DO (n, rest) <-- parseNumber xs ;
          SomeE (ANum n, rest)
                OR (DO (e, rest) <== firstExpect "("
                       (parseSumExp steps' symtable) xs;
          DO (u, rest') <== expect ")" rest ;
          SomeE(e,rest'))
  end

with parseProductExp (steps:nat) symtable
                     (xs : list token) :=
  match steps with
  | 0 => NoneE "Too many recursive calls"
  | S steps' =>
    DO (e, rest) <==
      parsePrimaryExp steps' symtable xs ;
    DO (es, rest') <==
       many (firstExpect "*" (parsePrimaryExp steps' symtable))
            steps' rest;
    SomeE (fold_left AMult es e, rest')
  end

with parseSumExp (steps:nat) symtable (xs : list token)  :=
  match steps with
  | 0 => NoneE "Too many recursive calls"
  | S steps' =>
    DO (e, rest) <==
      parseProductExp steps' symtable xs ;
    DO (es, rest') <==
      many (fun xs =>
        DO (e,rest') <--
           firstExpect "+"
             (parseProductExp steps' symtable) xs;
           SomeE ( (true, e), rest')
        OR DO (e,rest') <==
        firstExpect "-"
           (parseProductExp steps' symtable) xs;
            SomeE ( (false, e), rest'))
        steps' rest;
      SomeE (fold_left (fun e0 term =>
                          match term with
                            (true,  e) => APlus e0 e
                          | (false, e) => AMinus e0 e
                          end)
                       es e,
             rest')
  end.

Definition parseAExp := parseSumExp.

(** Parsing boolean expressions: *)

Fixpoint parseAtomicExp (steps:nat) (symtable : string->nat)
                        (xs : list token)  :=
match steps with
  | 0 => NoneE "Too many recursive calls"
  | S steps' =>
     DO    (u,rest) <-- expect "true" xs;
         SomeE (BTrue,rest)
     OR DO (u,rest) <-- expect "false" xs;
         SomeE (BFalse,rest)
     OR DO (e,rest) <--
            firstExpect "not"
               (parseAtomicExp steps' symtable)
               xs;
         SomeE (BNot e, rest)
     OR DO (e,rest) <--
              firstExpect "("
                (parseConjunctionExp steps' symtable) xs;
          (DO (u,rest') <== expect ")" rest;
              SomeE (e, rest'))
     OR DO (e, rest) <== parseProductExp steps' symtable xs;
            (DO (e', rest') <--
              firstExpect "=="
                (parseAExp steps' symtable) rest;
              SomeE (BEq e e', rest')
             OR DO (e', rest') <--
               firstExpect "<="
                 (parseAExp steps' symtable) rest;
               SomeE (BLe e e', rest')
             OR
               NoneE
      "Expected '==' or '<=' after arithmetic expression")
end

with parseConjunctionExp (steps:nat)
                         (symtable : string->nat)
                         (xs : list token) :=
  match steps with
  | 0 => NoneE "Too many recursive calls"
  | S steps' =>
    DO (e, rest) <==
      parseAtomicExp steps' symtable xs ;
    DO (es, rest') <==
       many (firstExpect "&&"
               (parseAtomicExp steps' symtable))
            steps' rest;
    SomeE (fold_left BAnd es e, rest')
  end.

Definition parseBExp := parseConjunctionExp.

Check parseConjunctionExp.

Definition testParsing {X : Type}
           (p : nat -> (string -> nat) ->
                list token ->
                optionE (X * list token))
           (s : string) :=
  let t := tokenize s in
  p 100 (build_symtable t 0) t.

(*
Eval compute in
  testParsing parseProductExp "x*y*(x*x)*x".

Eval compute in
  testParsing parseConjunctionExp "not((x==x||x*x<=(x*x)*x)&&x==x".
*)

(** Parsing commands: *)

Fixpoint parseSimpleCommand (steps:nat)
                            (symtable:string->nat)
                            (xs : list token) :=
  match steps with
  | 0 => NoneE "Too many recursive calls"
  | S steps' =>
    DO (u, rest) <-- expect "SKIP" xs;
      SomeE (SKIP, rest)
    OR DO (e,rest) <--
         firstExpect "IF" (parseBExp steps' symtable) xs;
       DO (c,rest')  <==
         firstExpect "THEN"
           (parseSequencedCommand steps' symtable) rest;
       DO (c',rest'') <==
         firstExpect "ELSE"
           (parseSequencedCommand steps' symtable) rest';
       DO (u,rest''') <==
         expect "END" rest'';
       SomeE(IFB e THEN c ELSE c' FI, rest''')
    OR DO (e,rest) <--
         firstExpect "WHILE"
           (parseBExp steps' symtable) xs;
       DO (c,rest') <==
         firstExpect "DO"
           (parseSequencedCommand steps' symtable) rest;
       DO (u,rest'') <==
         expect "END" rest';
       SomeE(WHILE e DO c END, rest'')
    OR DO (i, rest) <==
         parseIdentifier symtable xs;
       DO (e, rest') <==
         firstExpect ":=" (parseAExp steps' symtable) rest;
       SomeE(i ::= e, rest')
  end

with parseSequencedCommand (steps:nat)
                           (symtable:string->nat)
                           (xs : list token) :=
  match steps with
  | 0 => NoneE "Too many recursive calls"
  | S steps' =>
      DO (c, rest) <==
        parseSimpleCommand steps' symtable xs;
      DO (c', rest') <--
        firstExpect ";"
          (parseSequencedCommand steps' symtable) rest;
        SomeE(c ;; c', rest')
      OR
        SomeE(c, rest)
  end.

Definition bignumber := 1000.

Definition parse (str : string) : optionE (com * list token) :=
  let tokens := tokenize str in
  parseSequencedCommand bignumber
                        (build_symtable tokens 0) tokens.

(* ################################################################# *)
(** * Examples *)

(*
Compute parse "
  IF x == y + 1 + 2 - y * 6 + 3 THEN
    x := x * 1;;
    y := 0
  ELSE
    SKIP
  END  ".
====>
  SomeE
     (IFB BEq (AId (Id 0))
              (APlus
                 (AMinus (APlus (APlus (AId (Id 1)) (ANum 1)) (ANum 2))
                    (AMult (AId (Id 1)) (ANum 6)))
                 (ANum 3))
      THEN Id 0 ::= AMult (AId (Id 0)) (ANum 1);; Id 1 ::= ANum 0
      ELSE SKIP FI, [])
*)

(*
Compute parse "
  SKIP;;
  z:=x*y*(x*x);;
  WHILE x==x DO
    IF z <= z*z && not x == 2 THEN
      x := z;;
      y := z
    ELSE
      SKIP
    END;;
    SKIP
  END;;
  x:=z  ".
====>
  SomeE
     (SKIP;;
      Id 0 ::= AMult (AMult (AId (Id 1)) (AId (Id 2)))
                     (AMult (AId (Id 1)) (AId (Id 1)));;
      WHILE BEq (AId (Id 1)) (AId (Id 1)) DO
        IFB BAnd (BLe (AId (Id 0)) (AMult (AId (Id 0)) (AId (Id 0))))
                  (BNot (BEq (AId (Id 1)) (ANum 2)))
           THEN Id 1 ::= AId (Id 0);; Id 2 ::= AId (Id 0)
           ELSE SKIP FI;;
        SKIP
      END;;
      Id 1 ::= AId (Id 0),
     [])
*)

(*
Compute parse "
  SKIP;;
  z:=x*y*(x*x);;
  WHILE x==x DO
    IF z <= z*z && not x == 2 THEN
      x := z;;
      y := z
    ELSE
      SKIP
    END;;
    SKIP
  END;;
  x:=z  ".
=====>
  SomeE
     (SKIP;;
      Id 0 ::= AMult (AMult (AId (Id 1)) (AId (Id 2)))
            (AMult (AId (Id 1)) (AId (Id 1)));;
      WHILE BEq (AId (Id 1)) (AId (Id 1)) DO
        IFB BAnd (BLe (AId (Id 0)) (AMult (AId (Id 0)) (AId (Id 0))))
                 (BNot (BEq (AId (Id 1)) (ANum 2)))
          THEN Id 1 ::= AId (Id 0);;
               Id 2 ::= AId (Id 0)
          ELSE SKIP
        FI;;
        SKIP
      END;;
      Id 1 ::= AId (Id 0),
     []).
*)

(** $Date: 2016-07-14 17:02:35 -0400 (Thu, 14 Jul 2016) $ *)


Require Import List Arith Bool.
Require Import Expr Env.
Require Import EquivDec EqdepClass.
Require Import DepList.
Require Import Word Prover.

Set Implicit Arguments.
Set Strict Implicit.

Local Notation "[ x , .. , y ]" := (cons x .. (cons y nil) ..).

(** * The Transitivity Prover **)

(* Algorithm for grouping expressions by equalities. Terribly inefficient... *)
Section Grouper.
  Variable A : Type.
  Variable A_seq : A -> A -> bool.

  Fixpoint in_seq (ls : list A) (a : A) : bool :=
    match ls with
      | nil => false
      | x :: ls' => A_seq x a || in_seq ls' a
    end.

  Fixpoint groupWith (grps : list (list A)) (g : list A) (a : A) :=
    match grps with
      | nil => [a :: g]
      | g' :: grps' => if in_seq g' a
                       then (g' ++ a :: g) :: grps'
                       else g' :: groupWith grps' g a
    end.

  Fixpoint addEquality (ls : list (list A)) (a : A) (b : A) : list (list A) :=
    match ls with
      | nil => [[a, b]] (* matched nothing *)
      | grp :: ls' => if in_seq grp a
                        then groupWith ls' grp b
                        else if in_seq grp b
                               then groupWith ls' grp a
                               else grp :: addEquality ls' a b
    end.

  Fixpoint inSameGroup (grps : list (list A)) (a : A) (b : A) :=
    match grps with
      | nil => false
      | g :: grps' =>
        if in_seq g a then
          if in_seq g b
            then true
            else inSameGroup grps' a b
        else inSameGroup grps' a b
    end.

  Variable R : A -> A -> Prop.
  (* An arbitrary partial equivalence relation *)

  Hypothesis Rsym : forall x y, R x y -> R y x.
  Hypothesis Rtrans : forall x y z, R x y -> R y z -> R x z.
  Hypothesis A_seq_correct : forall x y, A_seq x y = true -> R x y.

  Fixpoint InR (x : A) (ls : list A) : Prop :=
    match ls with
      | nil => False
      | y :: ls' => R y x \/ InR x ls'
    end.

  Definition groupEqualTo (a : A) := Forall (R a).

  Definition groupEqual (g : list A) :=
    match g with
      | nil => True
      | a' :: g' => groupEqualTo a' g'
    end.

  Definition groupsEqual := Forall groupEqual.

  Hint Extern 1 (groupEqual _) => hnf.

  Hint Resolve Rsym Rtrans.

  Lemma Rweaken : forall x y l,
    Forall (R x) l
    -> R x y
    -> Forall (R y) l.
    induction 1; t.
  Qed.

  Hint Resolve Rweaken.

  Lemma groupEqualTo_groupEqual : forall x xs,
    Forall (R x) xs
    -> groupEqual xs.
    induction 1; t.
  Qed.

  Hint Resolve groupEqualTo_groupEqual.

  Hint Resolve Folds.Forall_app.

  Lemma groupEqualTo_In : forall x y g,
    InR y g
    -> Forall (R x) g
    -> R x y.
    induction 2; t.
  Qed.

  Hint Immediate groupEqualTo_In.

  Hint Extern 1 (Forall _ _) => progress hnf.

  Lemma in_seq_correct : forall (a : A) (ls : list A),
    in_seq ls a = true -> InR a ls.
  Proof.
    induction ls; t.
  Qed.

  Hint Resolve in_seq_correct A_seq_correct.

  Lemma groupWith_sound : forall x xs grps,
    Forall groupEqual grps
    -> Forall (R x) xs
    -> Forall groupEqual (groupWith grps xs x).
    induction 1; t. eauto 10.
      apply in_seq_correct in H3. eauto 7.
  Qed.

  Hint Resolve groupWith_sound.

  Theorem addEquality_sound : forall x y grps,
    groupsEqual grps
    -> R x y
    -> groupsEqual (addEquality grps x y).
    induction 1; t;
      match goal with
        | [ H : _ |- _ ] => apply A_seq_correct in H || apply in_seq_correct in H
      end; eauto 7.
  Qed.

  Theorem inSameGroup_sound : forall grps, groupsEqual grps
    -> forall x y, inSameGroup grps x y = true
      -> R x y.
    induction 1; t;
      repeat match goal with
        | [ H : _ |- _ ] => apply A_seq_correct in H || apply in_seq_correct in H
      end; eauto 7.
  Qed.
End Grouper.

Section TransitivityProver.
  Variable types : list type.
  Variable fs : functions types.
  Section with_vars.
  Variables uvars vars : env types.

  Definition transitivity_summary : Type := list (list (expr types)).

  Coercion typeof_env : env >-> tenv.

  Definition eqD (e1 e2 : expr types) : Prop :=
    match typeof (typeof_funcs fs) uvars vars e1 with
      | None => False
      | Some t =>
        match exprD fs uvars vars e1 t, exprD fs uvars vars e2 t with
          | Some v1, Some v2 => v1 = v2
          | _, _ => False
        end
    end.

  Theorem eqD_refl : forall e1 e2, e1 = e2
    -> forall t, typeof (typeof_funcs fs) uvars vars e1 = Some t
      -> forall v, exprD fs uvars vars e1 t = Some v
        -> eqD e1 e2.
    t.
  Qed.

  Lemma nth_error_nil : forall T n,
    nth_error (@nil T) n = None.
  Proof.
    destruct n; simpl; auto.
  Qed.

(*
  Theorem eqD_refl_wt : forall e1 e2, e1 = e2 ->
    match is_well_typed (typeof_funcs fs) uvars vars e1 (typeof (typeof_funcs fs) uvars vars e1) return Prop with
      | None => True
      | Some t => eqD e1 e2
    end.
  Proof.
    intros; subst.
    case_eq (well_typed fs uvars vars e2); intros; auto.
    generalize is_well_typed_typeof.
    generalize well_typed_is_well_typed. intros.
    generalize H. apply H0 in H.
    generalize H.
    apply is_well_typed_correct in H.
    intros.
    apply H1 in H2. destruct H.
    eapply eqD_refl; eauto.
  Qed.
*)

  Fixpoint transitivityLearn (sum : transitivity_summary) (hyps : list (expr types)) : transitivity_summary :=
    match hyps with
      | nil => sum
      | h :: hyps' =>
        let grps := transitivityLearn sum hyps' in
          match h with
            | Equal t x y => addEquality (@expr_seq_dec  _) grps x y
            | _ => grps
          end
    end.
  Definition groupsOf := transitivityLearn nil.

  Definition transitivityEqProver (groups : transitivity_summary)
    (x y : expr types) := inSameGroup (@expr_seq_dec _) groups x y.

  Fixpoint proveEqual (groups : transitivity_summary) (e1 e2 : expr types) {struct e1} :=
    expr_seq_dec e1 e2 ||
      (inSameGroup (@expr_seq_dec _) groups e1 e2
        || match e1, e2 with
             | Func f1 args1, Func f2 args2 =>
               if eq_nat_dec f1 f2
                 then (fix proveEquals (es1 es2 : list (expr types)) :=
                   match es1, es2 with
                     | nil, nil => true
                     | e1 :: es1', e2 :: es2' => proveEqual groups e1 e2 && proveEquals es1' es2'
                     | _, _ => false
                   end) args1 args2
                 else false
             | _, _ => false
           end
    ).

  Definition transitivityProve (groups : transitivity_summary)
    (goal : expr types) :=
    match goal with
      | Equal _ x y => proveEqual groups x y
      | _ => false
    end.

  Hint Resolve addEquality_sound.
  Hint Immediate lookupAs_det.

  Ltac foundTypeof E := generalize (@exprD_principal _ fs uvars vars (typeof_funcs fs) uvars vars
    (typeof_env_WellTyped_env _) (typeof_env_WellTyped_env _) (typeof_funcs_WellTyped_funcs _) E);
  destruct (typeof (typeof_funcs fs) uvars vars E); intuition.

  Ltac foundDup E T1 T2 := match T1 with
                             | T2 => fail 1
                             | _ =>
                               assert (T1 = T2) by (apply (exprD_det fs uvars vars E);
                                 try match goal with
                                       | [ H : _ |- _ ] => solve [ intro; apply H with T1; t ]
                                     end); subst
                           end.

  Ltac eqD1 :=
    match goal with
      | [ _ : context[typeof _ _ _ ?E] |- _ ] => foundTypeof E
      | [ |- context[typeof _ _ _ ?E] ] => foundTypeof E
      | [ _ : context[exprD fs uvars vars ?E ?T1] |- context[exprD fs uvars vars ?E ?T2] ] => foundDup E T1 T2
      | [ _ : context[exprD fs uvars vars ?E ?T1], _ : context[exprD fs uvars vars ?E ?T2] |- _ ] => foundDup E T1 T2
      | [ x : tvar, H1 : forall y : tvar, _ |- False ] => apply H1 with x; t
    end.

  Ltac eqD := unfold eqD; intros; repeat eqD1; t.

  Theorem eqD_sym : forall x y : expr types, eqD x y -> eqD y x.
    unfold eqD; intros.
    eqD.
  Qed.

  Theorem eqD_trans : forall x y z : expr types, eqD x y -> eqD y z -> eqD x z.
    eqD.
  Qed.

  Hint Immediate eqD_sym eqD_trans.

  Theorem groupsOf_sound : forall hyps,
    AllProvable fs uvars vars hyps
    -> groupsEqual eqD (groupsOf hyps).
    induction hyps. intuition. simpl in *. constructor.

    intro. simpl in H. destruct H.
    specialize (IHhyps H0).

    t1.
    destruct a; auto.
    revert H; case_eq (exprD fs uvars vars (Equal t a1 a2) tvProp); intros; try contradiction.
    simpl in *. apply addEquality_sound; eauto.

    Focus 2.

    revert H.
    unfold eqD.
    case_eq (exprD fs uvars vars a1 t); try congruence; intros.
(*
    rewrite (exprD_typeof _ _ _ _ _ H). rewrite H. destruct (exprD fs uvars vars a2 t); try congruence.
    inversion H2. subst; auto.
*)
    admit. (** TODO : this isn't true in general, but the fact that everything is provable makes it true **)
    admit.
  Qed.
  End with_vars.

  Definition transitivityValid (uvars vars : env types) (sum : transitivity_summary) : Prop :=
    (forall ls, In ls sum -> forall e, In e ls -> ValidExp fs uvars vars e) ->
    groupsEqual (eqD uvars vars) sum.

  Definition transitivitySummarize := groupsOf.

  Theorem transitivitySummarizeCorrect : forall uvars vars hyps,
    AllProvable fs uvars vars hyps ->
    transitivityValid uvars vars (transitivitySummarize hyps).
  Proof.
  Admitted.

  Theorem transitivityLearnCorrect : forall uvars vars sum,
    transitivityValid uvars vars sum -> forall hyps,
    AllProvable fs uvars vars hyps ->
    transitivityValid uvars vars (transitivityLearn sum hyps).
  Proof.
  Admitted.

  Theorem transitivityProverCorrect : ProverCorrect fs transitivityValid transitivityProve.
    admit.
(*
    unfold transitivityProver; hnf; intros;
      destruct goal; try discriminate;
        match goal with
          | [ H1 : _, H2 : _ |- _ ] =>
            apply (inSameGroup_sound eqD_sym eqD_trans expr_seq_dec
              (groupsOf_sound _ H1)) in H2
        end; hnf in *; simpl in *; eqD.
*)
  Qed.

  Theorem transitivityValid_extensible : forall (u g : env types) (f : transitivity_summary)
     (ue ge : list (sigT (tvarD types))),
   transitivityValid u g f -> transitivityValid (u ++ ue) (g ++ ge) f.
  Proof.
  Admitted.

  Definition transitivityProver : ProverT types :=
  {| Facts := transitivity_summary
   ; Summarize := transitivitySummarize
   ; Learn := transitivityLearn
   ; Prove := transitivityProve
   |}.

  Definition transitivityProver_correct : ProverT_correct transitivityProver fs.
  eapply Build_ProverT_correct with (Valid := transitivityValid);
    eauto using transitivityValid_extensible, transitivitySummarizeCorrect, transitivityLearnCorrect, transitivityProverCorrect.
  Qed.

End TransitivityProver.

Definition TransitivityProver : ProverPackage :=
{| ProverTypes := nil_Repr EmptySet_type
 ; ProverFuncs := fun ts => nil_Repr (Default_signature ts)
 ; Prover_correct := fun ts fs => transitivityProver_correct fs
|}.

Ltac unfold_transitivityProver H :=
  match H with
    | tt =>
      cbv delta [
        transitivityProver proveEqual
        transitivitySummarize transitivityLearn transitivityProve

        groupsOf addEquality
        transitivityLearn
        inSameGroup
        expr_seq_dec
        in_seq
        groupWith
      ]
    | _ =>
      cbv delta [
        transitivityProver proveEqual
        transitivitySummarize transitivityLearn transitivityProve

        groupsOf addEquality
        transitivityLearn
        inSameGroup
        expr_seq_dec
        in_seq
        groupWith
      ] in H
  end.


Require Import Crypto.Arithmetic.PrimeFieldTheorems.
Require Import Crypto.Specific.solinas64_2e416m2e208m1_9limbs.Synthesis.

(* TODO : change this to field once field isomorphism happens *)
Definition mul :
  { mul : feBW_loose -> feBW_loose -> feBW_tight
  | forall a b, phiBW_tight (mul a b) = F.mul (phiBW_loose a) (phiBW_loose b) }.
Proof.
  Set Ltac Profiling.
  Time synthesize_mul ().
  Show Ltac Profile.
Time Defined.

Print Assumptions mul.


From iris.base_logic.lib Require Import invariants.
From iris.program_logic Require Import weakestpre.
From iris.heap_lang Require Import lang.
From iris.heap_lang Require Import notation proofmode.

Section invariants.
  Context `{!heapG Σ}. (* Ghost stated required for HeapLang *)

  Definition prog1 : expr :=
    let: "x" := ref #0 in
    Fork ("x" <- #1);;
    !"x".

  Definition N := nroot.@"myinvariant".

  Lemma prog1_spec :
    ⊢ WP prog1 {{ v, ⌜v = #0⌝ ∨ ⌜v = #1⌝ }}.
  Proof.
    unfold prog1. wp_alloc x as "Hx".
    iMod (inv_alloc N _ (x ↦ #0 ∨ x ↦ #1)%I with "[Hx]") as "#H".
    { iNext. iLeft; eauto. }
    wp_pures. wp_apply wp_fork.
    - iInv N as "[Hx|Hx]"; wp_store; eauto.
    - wp_pures. iInv N as "[Hx|Hx]"; wp_load; iModIntro; eauto.
  Qed.

End invariants.


Require Import
  CMorphisms
  Coq.Program.Basics
  FormTopC.FormTop
  FormTopC.Cont
  FormTopC.FormalSpace
  Algebra.OrderC
  Algebra.PreOrder
  Algebra.SetsC
  Prob.StdLib.

Set Universe Polymorphism.
Set Asymmetric Patterns.

Local Open Scope Subset.
Local Open Scope FT.

(** Information bases, which are the predicative versions of
    Scott domains. Perhaps, see Definition 1.9 of [2].
    Though my formulation is a little different; I build off
    of a pre-existing notion of a meet semi-lattice.

    Also, I directly show that this formal topology is
    inductively generated by generating it with an axiom set. *)
Module InfoBase.
Section InfoBase.

Set Printing Universes.
Universes A P AP.
Variable (S : PreOrder@{A P}).

Context {PO : PreO.t@{A P} (le S)}.

(** The axiom set essentially says that if [s <= t], then
    [s] is covered by the singleton set [{t}]. *)
Inductive Ix@{} {s : S} : Type@{P} := .

Arguments Ix : clear implicits.

Definition C@{} (s : S) (s' : Ix s) : Subset@{A P} S := match s' with
  end.

Definition IBInd@{} : PreISpace.t@{A P P} :=
  {| PreISpace.S := S
   ; PreISpace.Ix := Ix
   ; PreISpace.C := C
  |}.

Definition Cov (s : S) (U : Subset@{A P} S) : Type@{P} :=
  In (⇓ U) s.

Definition IB@{} : PreSpace.t@{A P P} :=
  {| PreSpace.S := S
   ; PreSpace.Cov := Cov |}.

(** This axiom set is localized. *)
Local Instance loc@{} : FormTop.localized@{A P P} IBInd.
Proof.
unfold FormTop.localized. intros. induction i.
Qed.

Theorem CovEquiv : PreSpace.Cov IB ==== GCovL IBInd.
Proof.
intros a U. simpl. unfold Cov. split; intros.
- destruct X as [t Ut st].
  apply FormTop.glle_left with t. assumption.
  apply FormTop.glrefl. assumption.
- induction X.
  + exists a. assumption. reflexivity.
  + destruct IHX as [t Ut bt].
    exists t. assumption. etransitivity; eassumption.
  + destruct i.
Qed.

(** The proof that [Cov] is a valid formal topology. *)
Local Instance isCovG : FormTop.t (toPSL IBInd) :=
  FormTop.GCovL_formtop _.

(** Should prove this via homeomorphism with IBInd. *)
Local Instance isCov@{} : FormTop.t IB.
Proof.
Admitted.

Local Instance Pos : FormTop.gtPos IBInd.
Proof.
apply FormTop.gall_Pos.
intros. destruct i.
Qed.

Definition InfoBase@{} : IGt@{A P P AP} :=
  {| IGS := IBInd |}.

End InfoBase.
End InfoBase.

Arguments InfoBase.Ix : clear implicits.

Module InfoBaseCont.
Section InfoBaseCont.

Generalizable All Variables.

Context {S : PreSpace.t} {POS : PreO.t (le (PreSpace.S S))}.
Context {T : PreOrder} {POT : PreO.t (le T)}.

Record ptNM {F : Subset T} : Type :=
  { ptNM_local : forall {a b}, F a -> F b ->
     Inhabited (F ∩ (eq a ↓ eq b))
  ; ptNM_le_right : forall a b, a <=[T] b -> F a -> F b
  ; ptNM_here : Inhabited F
  }.

Arguments ptNM : clear implicits.

Instance ptNM_proper : Proper ((eq ==> iffT) ==> iffT) ptNM.
Proof.
Admitted.

(** I have no idea whether this is in fact
    a good definition *)
Record tNM {F_ : Cont.map S (InfoBase.IB T)} :=
  { NMle_left : forall a b c, a <=[PreSpace.S S] b -> F_ c b -> F_ c a
  ; NMle_right :  forall a b c, F_ b a -> b <=[T] c -> F_ c a
  ; NMlocal : forall {a b c}, F_ b a -> F_ c a ->
     Inhabited ((fun t => F_ t a) ∩ (eq b ↓ eq c))
  ; NMhere : forall s : S, In (union (fun _ => True) F_) s
  }.

Arguments tNM : clear implicits.

Hypothesis FTS : FormTop.t S.

Theorem contNM : forall (F : Cont.map S (InfoBase.IB T)),
  tNM F
  -> Cont.t S (InfoBase.IB T) F.
Proof.
intros. constructor; intros.
- unfold InfoBase.Cov. apply FormTop.refl.
  apply (NMhere X).
- eapply (NMle_left X); eassumption.
- unfold InfoBase.Cov. apply FormTop.refl.
  pose proof (NMlocal X X0 X1).
  destruct X2. destruct i.
  econstructor; eassumption.
- simpl in *. unfold InfoBase.Cov in *.
  destruct X1 as [t0 Vt0 bt0].
  apply FormTop.refl. exists t0. assumption.
  apply (NMle_right X) with b; assumption.
Qed.

End InfoBaseCont.

Arguments tNM : clear implicits.
Arguments ptNM : clear implicits.

(*
Section InfoBaseML.

Context {S : PreSpace.t} {POS : PreO.t (le S)}.
Context {T} `{MeetLat.t T}.

Record pt {F : Subset T} : Type :=
  { pt_local : forall {a b}, F a -> F b -> F (MeetLat.min a b)
  ; pt_le_right : forall a b, MeetLat.le a b -> F a -> F b
  ; pt_here : Inhabited F
  }.

Arguments pt : clear implicits.

Instance pt_proper : Proper ((eq ==> iffT) ==> iffT) pt.
Proof.
Admitted.

Lemma down_min : forall a b,
 In (FormTop.down MeetLat.le a b) (MeetLat.min a b).
Proof.
intros. constructor; apply MeetLat.min_ok.
Qed.

Theorem pt_ptNM : forall F, pt F -> ptNM MeetLat.le F.
Proof.
intros F H0. destruct H0. constructor; eauto.
intros. constructor 1 with (MeetLat.min a b).
econstructor. unfold In. eauto. apply down_min.
Qed.

(** I have no idea whether this is in fact
    a good definition *)
Record t {F_ : Cont.map S T} :=
  { le_left : forall a b c, leS a b -> F_ c b -> F_ c a
  ; le_right :  forall a b c, F_ b a -> MeetLat.le b c -> F_ c a
  ; local : forall {a b c}, F_ b a -> F_ c a ->
     F_ (MeetLat.min b c) a
  ; here : forall s : S, In (union (fun _ => True) F_) s
  }.

Arguments t : clear implicits.

Variable CovS : S -> Subset S -> Type.
Hypothesis FTS : FormTop.t leS CovS.
Let CovT : T -> Subset T -> Type := @InfoBase.Cov _ MeetLat.le.

Theorem cont : forall (F : Cont.map S T),
  t F -> Cont.t leS MeetLat.le CovS CovT F.
Proof.
intros. apply contNM. assumption.
destruct X. constructor; eauto.
intros. specialize (local0 _ _ _ X X0).
constructor 1 with (MeetLat.min b c).
constructor; eauto. apply down_min.
Qed.

Definition above_pt (x : T) : pt (MeetLat.le x).
Proof.
constructor; intros.
- apply MeetLat.min_ok; assumption.
- etransitivity; eassumption.
- econstructor. unfold In. reflexivity.
Qed.

Definition lift_op (f : S -> T) (y : T) (x : S) : Type :=
  MeetLat.le (f x) y.

Definition lift_monotone (f : S -> T)
  (fmono : forall x y, leS x y -> MeetLat.le (f x) (f y))
  : t (lift_op f).
Proof.
constructor; unfold lift_op; intros.
- etransitivity. apply fmono. eassumption. assumption.
- etransitivity; eassumption.
- apply MeetLat.min_ok; assumption.
- econstructor. unfold In. constructor. reflexivity.
Qed.

End InfoBaseML.

Arguments t {_} leS {_ _} F_.
Arguments pt {_} {_} F.
Arguments lift_op {_ _ _} f y x.

Section Product.

Context {S} `{MeetLat.t S}.
Context {T} `{MeetLat.t T}.
Context {U} `{MeetLat.t U}.

Definition lift_binop (f : S -> T -> U)
  (result : U) (args : S * T) : Type :=
  let (l, r) := args in MeetLat.le (f l r) result.

Existing Instances MeetLat.product_ops MeetLat.product.

Theorem lift_binop_monotone : forall (f : S -> T -> U)
  (fmono : forall x x' y y', MeetLat.le x x' -> MeetLat.le y y'
     -> MeetLat.le (f x y) (f x' y'))
  , t MeetLat.le (lift_binop f).
Proof.
intros. unfold lift_binop. constructor; intros.
- destruct a, b. simpl in *. unfold prod_op in *.
  destruct X. rewrite <- X0. apply fmono; assumption.
- destruct a. rewrite <- X0.  assumption.
- destruct a. apply MeetLat.min_ok; assumption.
- destruct s. econstructor. constructor.
  reflexivity.
Qed.

End Product.

Section Compose.

Context {S} {leS : crelation S} {SOps} {MLS : MeetLat.t S SOps}.

Instance OneOps : MeetLat.Ops True := MeetLat.one_ops.

Theorem to_pt : forall (F : Cont.map True S), t MeetLat.le F ->
  pt (fun s => F s I).
Proof.
intros F H. constructor; intros.
- apply (local H); assumption.
- eapply (le_right H); eassumption.
- pose proof (here H I) as X. destruct X.
  econstructor; eauto.
Qed.

Theorem from_pt : forall (F : Subset S), pt F -> t MeetLat.le (fun t' _ => F t').
Proof.
intros F H. constructor; intros.
- assumption.
- eapply (pt_le_right H); eassumption.
- apply (pt_local H); assumption.
- pose proof (pt_here H) as X. destruct X.
  repeat (econstructor || eauto).
Qed.

Context {T TOps} {MLT : MeetLat.t T TOps}.
Context {U UOps} {MLU : MeetLat.t U UOps}.

Theorem t_compose (F : Cont.map S T) (G : Cont.map T U)
  : t MeetLat.le F -> t MeetLat.le G
  -> t MeetLat.le (compose G F).
Proof.
intros HF HG.
constructor; unfold compose; intros.
- destruct X0 as (t & Fbt & Gtc).
  exists t. split.
  + assumption.
  + eapply (le_left HF); eassumption.
- destruct X as (t & Fat & Gtb).
  exists t. split. eapply (le_right HG); eassumption.
  assumption.
- destruct X as (t & Fat & Gtb).
  destruct X0 as (t' & Fat' & Gt'c).
  exists (MeetLat.min t t'). split.
  + apply (local HG); eapply (le_left HG).
    apply MeetLat.min_l. assumption.
    apply MeetLat.min_r. assumption.
  + apply (local HF); assumption.
- destruct (here HF s). destruct (here HG a).
  exists a0. constructor. exists a. auto.
Qed.

End Compose.

Section EvalPt.

Context {S SOps} {MLS : MeetLat.t S SOps}.
Context {T TOps} {MLT : MeetLat.t T TOps}.

Definition eval (F : Cont.map S T) (x : Subset S) (t : T) : Type :=
  Inhabited (x ∩ F t).

Theorem eval_pt (F : Cont.map S T) (x : Subset S)
  : pt x -> t MeetLat.le F -> pt (eval F x).
Proof.
intros Hx HF.
pose proof (t_compose (fun t _ => x t) F (from_pt _ Hx) HF) as H.
apply to_pt in H.
eapply pt_proper. 2: eassumption. simpl_crelation.
unfold eval. split; intros.
- destruct X. destruct i. econstructor; eauto.
- destruct X. destruct p. repeat (econstructor || eauto).
Qed.

End EvalPt.
*)

End InfoBaseCont.

(*
Arguments InfoBaseCont.t {S} leS {T} {TOps} F : rename, clear implicits.
*)

Module One.
Section One.

Definition OnePO : PreOrder :=
  {| PO_car := True
   ; le := fun _ _ => True
  |}.

Definition One : PreISpace.t := InfoBase.IBInd OnePO.

Section One_intro.

Context {S : PreSpace.t} {FTS : FormTop.t S}.

Definition One_intro : Cont.map S (toPSL One) :=
   fun (_ : True) (s : S) => True.

Theorem One_intro_cont :
  Cont.t S (toPSL One) One_intro.
Proof.
constructor; unfold One_intro; intros; simpl; try auto.
- apply FormTop.refl. unfold In; simpl. constructor 1 with I.
  unfold In; simpl; constructor. constructor.
- apply FormTop.refl. unfold In; simpl.
  exists I. destruct b, c. unfold In.
  apply (down_eq (A := toPSL One)). auto. auto.
- apply FormTop.refl. constructor 1 with I.
  induction X. destruct a0. assumption.
  assumption.  destruct i. auto.
Qed.

End One_intro.

Context {S : PreSpace.t} {POS : PreO.t (le (PreSpace.S S))}.

Definition Point (f : Subset S) := Cont.t (toPSL One) S (fun t _ => f t).

End One.
End One.

Module Sierpinski.

Definition SierpPO : PreOrder :=
  {| PO_car := bool
   ; le := Bool.leb |}.

Definition Sierp := InfoBase.IBInd SierpPO.

(*
Definition sand : Cont.map (InfoBase.IBInd (ProdPO SierpPO SierpPO))
  Sierp :=
  InfoBaseCont.lift_binop andb.

Existing Instances MeetLat.product MeetLat.product_ops.

Theorem sand_cont : InfoBaseCont.t MeetLat.le sand.
Proof.
apply InfoBaseCont.lift_binop_monotone.
simpl. intros. destruct x, x', y, y'; auto.
Qed.

Definition sor : Cont.map (bool * bool) bool :=
  InfoBaseCont.lift_binop orb.

Theorem sor_cont : InfoBaseCont.t MeetLat.le sor.
Proof.
apply InfoBaseCont.lift_binop_monotone.
simpl. intros. destruct x, x', y, y'; auto; congruence.
Qed.

Definition const_cont (b : bool) : InfoBaseCont.pt (MeetLat.le (negb b)).
Proof.
apply InfoBaseCont.above_pt.
Qed.
*)

End Sierpinski.


Require Import Crypto.Arithmetic.PrimeFieldTheorems.
Require Import Crypto.Specific.solinas32_2e266m3_11limbs.Synthesis.

(* TODO : change this to field once field isomorphism happens *)
Definition sub :
  { sub : feBW_tight -> feBW_tight -> feBW_loose
  | forall a b, phiBW_loose (sub a b) = F.sub (phiBW_tight a) (phiBW_tight b) }.
Proof.
  Set Ltac Profiling.
  Time synthesize_sub ().
  Show Ltac Profile.
Time Defined.

Print Assumptions sub.