Require Import Coq.Lists.List Coq.Sorting.Permutation.

From Coq Require Export Morphisms RelationClasses Setoid Program.Equality.
From ITree Require Export ITree ITreeFacts Eq.Rutt Props.Divergence Props.Finite.
From Paco Require Import paco.
From Coq Require Export Eqdep EqdepFacts Logic.Classical.
Require Import Lia.
Require Export Refinement Merge ConvergentRefinement ConcreteRefinement.

Import Monads.
Import MonadNotation.
Local Open Scope monad_scope.

Open Scope list_scope.

Lemma or_spec_l E1 E2 RE REAns R1 R2 RR (t1 t2 : itree_spec E1 R1) (t3 : itree_spec E2 R2) :
  @padded_refines E1 E2 R1 R2 RE REAns RR t1 t3 ->
  @padded_refines E1 E2 R1 R2 RE REAns RR t2 t3 ->
  @padded_refines E1 E2 R1 R2 RE REAns RR (or_spec t1 t2) t3.
Proof.
  intros. unfold or_spec. unfold padded_refines in *.
  rewrite pad_vis. pstep. constructor. intros [ | ]; constructor; pstep_reverse.
Qed.

Lemma or_spec_r E1 E2 RE REAns R1 R2 RR (t1 : itree_spec E1 R1) (t2 t3: itree_spec E2 R2) :
  (@padded_refines E1 E2 R1 R2 RE REAns RR t1 t2 \/
  @padded_refines E1 E2 R1 R2 RE REAns RR t1 t3) ->
  @padded_refines E1 E2 R1 R2 RE REAns RR t1 (or_spec t2 t3).
Proof.
  intros. unfold or_spec. unfold padded_refines in *.
  rewrite pad_vis. pstep.
  destruct H.
  - econstructor. constructor. Unshelve. all : try (apply true). pstep_reverse.
  - econstructor. constructor. Unshelve. all : try (apply false). pstep_reverse.
Qed.

Lemma spin_bind E R S k :  (@ITree.spin E R) >>= k ≅ @ITree.spin E S.
Proof.
  ginit. gcofix CIH. gstep. red. cbn. constructor.
  gfinal. eauto.
Qed.

Lemma and_or_spec E R (t1 t2 : itree_spec E R) :
  padded_refines_eq eq (and_spec t1 t2) (or_spec t1 t2).
Proof.
  unfold and_spec, or_spec, padded_refines_eq, padded_refines.
  setoid_rewrite pad_vis. pstep. red. cbn. econstructor. constructor.
  Unshelve. all : try apply true. econstructor. cbn. constructor.
  Unshelve. all : try apply true. pstep_reverse.
  enough (padded_refines_eq eq t1 t1). auto. reflexivity.
Qed.

Fixpoint trepeat {E R} (n : nat) (t : itree E R) : itree E unit :=
  match n with
  | 0 => Ret tt
  | S m => t;; trepeat m t end.

Lemma trepeat_Sn E R (n : nat) (t : itree E R) :
  trepeat (S n) t ≅ t;; trepeat n t.
Proof.
  apply eqit_bind; reflexivity.
Qed.

Lemma trepeat_0 E R (t : itree E R) :
  trepeat 0 t ≅ Ret tt.
Proof. reflexivity. Qed.

Definition halve_spec_fix {E} : list nat -> itree_spec E (list nat * list nat) :=
  rec_fix_spec (fun halve_spec_rec l =>
                  n <- spec_exists nat;;
                  trepeat n ( l' <- spec_exists (list nat);; halve_spec_rec l' );;
                  halve_spec l
               ).

Theorem padded_refines_ret E1 E2 R1 R2 RE REAns RR r1 r2 :
  (RR r1 r2 : Prop) ->
  @padded_refines E1 E2 R1 R2 RE REAns RR (Ret r1) (Ret r2).
Proof.
  intros Hr. red. repeat rewrite pad_ret. pstep. constructor. auto.
Qed.

Theorem padded_refines_vis E1 E2 R1 R2 A B
        RE REAns RR (e1 : E1 A) (e2 : E2 B)
        (k1 : A -> itree_spec E1 R1) (k2 : B -> itree_spec E2 R2) :
  (RE A B e1 e2 : Prop) ->
  (forall a b, (REAns A B e1 e2 a b : Prop) -> padded_refines RE REAns RR (k1 a) (k2 b) ) ->
  padded_refines RE REAns RR (Vis (Spec_vis e1) k1) (Vis (Spec_vis e2) k2).
Proof.
  unfold padded_refines. setoid_rewrite pad_vis. intros.
  pstep. constructor; auto. intros. left. pstep. constructor.
  left. eapply H0; auto.
Qed.

Variant halve_call : forall A, callE (list nat) (list nat * list nat) A -> A -> Prop :=
  | halve_call_i (l l1 l2: list nat) :
    Permutation l (l1 ++ l2) ->
    (length l1 >= length l2 /\ (length l > length l1 \/ length l <= 1)) ->
    halve_call _ (Call l) (l1,l2).

Lemma halve_correct' E l : padded_refines_eq eq (@halve E l) (halve_spec_fix l).
Proof.
  eapply padded_refines_monot with (RE1 := eqE)
                                   (REAns1 := eqEAns)
                                   (RR1 := sub_eqEAns halve_call _ _ (Call l) (Call l) ); eauto.
  { intros [l1 l2] [l3 l4]. intros. inv PR. inj_existT. subst. auto. }
  eapply padded_refines_mrec with (REInv := eqE)
                                  (* in this post condition need to include *)
                                  (REAnsInv := sub_eqEAns halve_call).
  2 : constructor. intros.
  eqE_inv H. clear l. destruct d2 as [l]. cbn.
  destruct l as [ | a [ | b l] ].
  - existssr 0. rewrite trepeat_0. rewrite bind_ret_l.
    existssr (@nil nat). existssr (@nil nat). assertsr. reflexivity.
    assertsr. cbn. lia.
    apply padded_refines_ret. do 2 constructor. reflexivity.
    cbn. lia. (*note reflexivity is not enough this is a subreflexive rel*)
  - existssr 0. rewrite trepeat_0. rewrite bind_ret_l.
    existssr (a :: nil). existssr (@nil nat). assertsr. reflexivity.
    assertsr. cbn. lia. apply padded_refines_ret. do 2 constructor.
    reflexivity. cbn. lia.
  - cbn. existssr 1. rewrite trepeat_Sn. setoid_rewrite trepeat_0.
    cbn. repeat rewrite bind_bind. existssr l.
    eapply padded_refines_bind with (RR :=
                                    sub_eqEAns halve_call (list nat * list nat) (list nat * list nat) (Call (l)) (Call (l))).
    + unfold call_spec. apply padded_refines_vis. repeat constructor.
      intros [l1 l2] r Hl12. inv Hl12. inj_existT. subst.
      inv H5. inj_existT. subst. apply padded_refines_ret.
      constructor. auto.
    + intros [l1 l2] r2 Heq. inv Heq. inj_existT. subst.
      inv H4. inj_existT. injection H5. intros. subst.
      rewrite bind_ret_l.
      existssr (a :: l1). existssr (b :: l2). assertsr.
      { cbn. rewrite Permutation_app_comm. cbn.
      rewrite H0. rewrite Permutation_app_comm. reflexivity. }
      assertsr.
      { cbn. rewrite H0. repeat rewrite app_length. lia. }
      apply padded_refines_ret. do 2 constructor.
      { cbn. rewrite Permutation_app_comm. cbn.
      rewrite H0. rewrite Permutation_app_comm. reflexivity. }
      cbn. repeat rewrite H0. repeat rewrite app_length. lia.
Qed.
(*
Section ConvRef.
Context (E D : Type -> Type).
Context (Pre : forall A, D A -> Prop).
Context (bodies : D ~> itree_spec (D +' E) ).
Context (A : Type) (init : D A) (Hinit : Pre A init).
Context (Hconv : forall A (d : D A), Pre A d -> )

Context {A B: Type} {E : Type -> Type}.
Context (body : ((A -> itree_spec (callE A B +' E) B) -> A -> itree_spec (callE A B +' E) B)).
Context (Inv : A -> Prop).

Lemma rec_fix_conv_rev :
  (* (forall a, Inv a -> conv_ref (body a) ) -> *)
  forall a, Inv a -> conv_ref (rec_fix_spec body a).

End ConvRef. *)

Global Instance padded_refines_eq_itree E1 E2 R1 R2 RE REAns RR
  : Proper (eq_itree eq ==> eq_itree eq ==> impl)
           (@padded_refines E1 E2 R1 R2 RE REAns RR).
Proof.
  repeat intro. assert (x ≈ y). rewrite H. reflexivity.
  rewrite <- H2. assert (x0 ≈ y0). rewrite H0. reflexivity.
  rewrite <- H3. auto.
Qed.

Lemma interp_mrec_spec_or  (E : Type -> Type) (R : Type) (D : Type -> Type) (ctx : forall T : Type, D T -> itree_spec (D +' E) T)
  (t1 t2 : itree_spec _ R)
  : interp_mrec_spec ctx (or_spec t1 t2) ≅ or_spec (interp_mrec_spec ctx t1) (interp_mrec_spec ctx t2).
Proof.
  setoid_rewrite interp_mrec_spec'_exists. apply eqit_Vis. intros [ | ]; reflexivity.
Qed.

Lemma interp_mrec_spec_spin (E : Type -> Type) (R : Type) (D : Type -> Type) (ctx : forall T : Type, D T -> itree_spec (D +' E) T) :
  interp_mrec_spec ctx ITree.spin ≅ @ITree.spin _ R.
Proof.
  pcofix CIH. pstep. red. cbn. constructor.
  right. eauto.
Qed.

Lemma conv_ref_mrec_spin (E D : Type -> Type) (R : Type) (P : forall A : Type, D A -> Prop) :
  @conv_ref_mrec E D R P ITree.spin.
Proof.
  pcofix CIH. pstep. red. cbn. constructor.
  eauto.
Qed.

Lemma conv_ref_mrec_ret (E D : Type -> Type) (R : Type) (P : forall A : Type, D A -> Prop) r:
  @conv_ref_mrec E D R P (Ret r).
Proof.
  pstep. constructor.
Qed.

Lemma conv_ref_mrec_ex (E D : Type -> Type) (R A : Type) (P : forall A : Type, D A -> Prop) k :
  (forall a, conv_ref_mrec P (k a)) ->
  @conv_ref_mrec E D R P (Vis (@Spec_exists _ A) k ).
Proof.
  intros. pstep. constructor. left. apply H.
Qed.

(*should be able to generalize this to cany concrete tree *)
Lemma or_spec_r_Ret_inv E1 E2 R1 R2 RE REAns RR r t2 t3 :
  @padded_refines E1 E2 R1 R2 RE REAns RR (Ret r) (or_spec t2 t3) ->
  padded_refines RE REAns RR (Ret r) t2 \/ padded_refines RE REAns RR (Ret r) t3.
Proof.
  intros Hor. unfold padded_refines in *. setoid_rewrite pad_ret in Hor.
  setoid_rewrite pad_ret. setoid_rewrite pad_vis in Hor.
  pinversion Hor. inj_existT. subst. inv H1. destruct a.
  - left. pstep. auto.
  - right. pstep. auto.
Qed.

Lemma conv_ref_spec_exists E R :
  @conv_ref E R (spec_exists R).
Proof.
  pstep. red. cbn. constructor.
  intros. left. pstep. constructor.
Qed.

Section SpecFix.
  Context {A B: Type} {E : Type -> Type}.
  Context (Pre : A -> Prop) (Post : A -> B -> Prop).
  Context (prog : A -> itree_spec E B).

  Context (Rdec : A -> A -> Prop).
  Context (HRdec : well_founded Rdec).

  Context (Hprog : forall a, isConcrete (prog a)).
  Variant pre_call : forall C, callE A B C -> Prop :=
    | pre_call_intro (a : A) : Pre a -> pre_call B (Call a).
  Variant post_call : forall C, callE A B C -> C -> Prop :=
    | post_call_intro a b : Post a b -> post_call B (Call a) b.

  Definition total_spec (a : A) : itree_spec E B :=
    assume_spec (Pre a);;
    b <- spec_exists B;;
    assert_spec (Post a b);;
    Ret b.

  (* better for using as a final spec *)
  Definition partial_spec (a : A) : itree_spec E B :=
    or_spec (total_spec a) ITree.spin.
  (* better for syntax driven solving *)
  (**)

  Definition partial_spec_fix : A -> itree_spec E B :=
    rec_fix_spec (fun rec a =>
                    assume_spec (Pre a);;
                    (
                      n <- spec_exists nat;;
                      trepeat n (a' <- spec_exists A;; assert_spec (Pre a');; rec a' )
                    );;
                    b <- spec_exists B;;
                    assert_spec (Post a b);;
                    Ret b
                 ).

  Definition partial_spec_fix' : A -> itree_spec E B :=
    rec_fix_spec (fun rec a =>
                        assume_spec (Pre a);;
                        (
                          n <- spec_exists nat;;
                          trepeat n (a' <- spec_exists A;; assert_spec (Pre a');; rec a' )
                        );;
                        or_spec (
                            b <- spec_exists B;;
                            assert_spec (Post a b);;
                            Ret b)
                        ITree.spin
                 ).

  Definition total_spec_fix : A -> itree_spec E B :=
    rec_fix_spec (fun rec a =>
                    assume_spec (Pre a);;
                    (
                      n <- spec_exists nat;;
                      trepeat n (a' <- spec_exists A;; assert_spec (Pre a' /\ Rdec a' a);; rec a' )
                    );;
                    b <- spec_exists B;;
                    assert_spec (Post a b);;
                    Ret b
                 ).

  (* here would be a good place for proofs that partial_spec_fix' is greater than any of these
     from there I should be able to prove that total_spec_fix
   *)

  Lemma total_spec_fix_total a : padded_refines_eq eq (total_spec_fix a) (total_spec a).
  Proof.
    revert a.
    apply well_founded_ind with (R := Rdec). auto.
    intros a Hind. unfold total_spec_fix. simpl.
    setoid_rewrite interp_mrec_spec_bind.
    match goal with |- context [ interp_mrec_spec ?k _ ] => set k as body end.
    unfold total_spec. assumesr. intros Ha. setoid_rewrite interp_mrec_spec_assume.
    assumesl. auto. rewrite bind_bind.  setoid_rewrite interp_mrec_spec_bind.
    setoid_rewrite interp_mrec_spec_exists. existssl. intros n.
    induction n.
    - simpl. rewrite bind_ret_l. setoid_rewrite interp_mrec_spec_bind.
      setoid_rewrite interp_mrec_spec_exists.
      setoid_rewrite interp_mrec_spec_bind. setoid_rewrite interp_mrec_spec_assert.
      setoid_rewrite interp_mrec_spec_ret. reflexivity.
    - simpl. rewrite bind_bind. setoid_rewrite interp_mrec_spec_bind.
      setoid_rewrite IHn. setoid_rewrite interp_mrec_spec_bind.
      setoid_rewrite interp_mrec_spec_exists. rewrite bind_bind. existssl.
      intros a'. setoid_rewrite interp_mrec_spec_bind.
      setoid_rewrite interp_mrec_spec_assert. rewrite bind_bind. assertsl.
      intros [Hprea' Hdec]. simpl.
      assert (Htot : interp_mrec_spec body (call_spec a') ≈ total_spec_fix a' ).
      { unfold total_spec_fix, rec_fix_spec, rec_spec, mrec_spec .
        setoid_rewrite interp_mrec_spec_inl. rewrite tau_eutt. rewrite bind_ret_r.
        reflexivity. }
      rewrite Htot, Hind; auto. unfold total_spec.
      setoid_rewrite bind_bind at 1. assumesl. auto.
      setoid_rewrite bind_bind at 1. existssl. intros b.
      setoid_rewrite bind_bind at 1. assertsl. intros. rewrite bind_ret_l. reflexivity.
   Qed.

    Lemma partial_spec_fix_conv_ref :
    forall a, Pre a -> conv_ref (partial_spec_fix' a).
  Proof.
    intros a Ha.
    eapply conv_ref_mrec_conv_ref with (Pre := pre_call).
    constructor. auto.
    intros. inv H. inj_existT. subst.
    clear a Ha.
    rename a0 into a. rename H2 into Ha. clear H0. cbn.
    apply conv_ref_mrec_bind. rewrite bind_trigger. apply conv_ref_mrec_forall.
    auto. intros. apply conv_ref_mrec_ret.
    intros []. rewrite bind_bind. apply conv_ref_mrec_bind.
    apply conv_ref_mrec_ex. intros. apply conv_ref_mrec_ret.
    intros n. induction n; cbn.
    - rewrite bind_ret_l. apply conv_ref_mrec_ex. intros [ | ].
      2 : apply conv_ref_mrec_spin. apply conv_ref_mrec_bind.
      + apply conv_ref_mrec_ex. intros. apply conv_ref_mrec_ret.
      + intros. rewrite bind_bind. setoid_rewrite bind_vis.
        apply conv_ref_mrec_ex. intros. repeat rewrite bind_ret_l.
        apply conv_ref_mrec_ret.
    - rewrite bind_bind. apply conv_ref_mrec_bind; intros; auto.
      apply conv_ref_mrec_bind.
      + apply conv_ref_mrec_ex. intros. apply conv_ref_mrec_ret.
      + intros. rewrite bind_bind. setoid_rewrite bind_vis.
        apply conv_ref_mrec_ex.
        intros. repeat rewrite bind_ret_l. pstep. red. cbn.
        constructor. constructor. auto. left.
        apply conv_ref_mrec_ret.
  Qed.

  Lemma prog_ret_spin : forall a, Pre a ->
                             padded_refines_eq eq (prog a) (partial_spec_fix' a) ->
                             (exists b, prog a ≈ Ret b) \/ (prog a ≈ ITree.spin).
  Proof.
    intros a Href Ha. apply no_ev_to_dec.
    eapply concrete_conv_ref_padded_to_no_ev; eauto.
    apply partial_spec_fix_conv_ref. auto.
  Qed.
    (* need a padded_refines version *)

(*   Lemma partial_spec_fix_partial_spec' a :
    (Pre a -> padded_refines_eq eq (prog a) (partial_spec_fix' a)) ->
    padded_refines_eq eq (prog a) (partial_spec a). *)

  Lemma partial_spec_fix_partial_spec_assume_succeed':
    forall (a : A) (b : B),
      Pre a ->
      padded_refines_eq eq (Ret b) (partial_spec_fix' a) ->
      padded_refines_eq eq (Ret b) (partial_spec a).
  Proof.
    intros a b Ha Hfix. unfold partial_spec_fix' in Hfix. cbn in Hfix.
    repeat setoid_rewrite interp_mrec_spec_or in Hfix.

    (*need an inversion principle for or spec *)
    (*could use rewrite for interp_mrec_spec and or_spec *)
    do 2 setoid_rewrite interp_mrec_spec_bind in Hfix. revert Hfix.
    match goal with | |- context [ interp_mrec_spec ?f _ ] => set f as body end.
    intro Hfix.
    assert (
        padded_refines_eq eq (interp_mrec_spec body (or_spec (b <- spec_exists B;; _ <- assert_spec (Post a b);; Ret b) ITree.spin)) (partial_spec a)).
    {
      rewrite interp_mrec_spec_or. unfold partial_spec.
      apply or_spec_l.
      - apply or_spec_r. left. unfold total_spec.
        assumesr. intros. repeat setoid_rewrite interp_mrec_spec_bind.
        normalize_interp_mrec_spec. setoid_rewrite interp_mrec_spec_ret.
        setoid_rewrite bind_bind. setoid_rewrite bind_ret_l.
        eapply padded_refines_bind. reflexivity.
        intros. subst. eapply padded_refines_bind. reflexivity.
        intros. subst. reflexivity.
      - apply or_spec_r. right. rewrite interp_mrec_spec_spin. reflexivity.
    } setoid_rewrite H in Hfix.
    setoid_rewrite interp_mrec_spec_ret in Hfix.
    match type of Hfix with padded_refines_eq eq _ ?t => set t as t1 end.
    fold t1 in Hfix.

    assert (t1 ≈ (interp_mrec_spec body (trigger (@Spec_forall _ (Pre a) ) ) >>
                 Ret tt >>
                  (interp_mrec_spec body
           (ITree.bind (spec_exists nat)
              (fun n : nat => trepeat n (ITree.bind (spec_exists A) (fun a' : A =>
                       trigger (@Spec_exists _ (Pre a') ) >> Ret tt >> call_spec a'))))) >> partial_spec a)
                 ).
    {
      unfold t1. repeat rewrite bind_bind.
      apply eqit_bind. reflexivity. intro.
      reflexivity.
    }
    rewrite H0 in Hfix. clear H0.
    (*might have dropped an important assumption *)
    eapply padded_conv_ref_ret_bind; eauto.
    normalize_interp_mrec_spec. cbn.
    apply conv_ref_bind.
    { pstep. red. cbn. constructor; auto. intros a0.
          left. pstep. red. cbn. constructor. }
    intros _.  setoid_rewrite bind_ret_l.
    eapply conv_ref_interp_mrec_conv_ref.
    Unshelve. 3 : { intros C c. destruct c. apply (Pre a0). }
    - cbn.
      intros. destruct d. cbn.
      clear - H0 Ha. setoid_rewrite bind_bind.
      apply conv_ref_mrec_bind.
      { pstep. red. cbn. constructor. auto. left.
        pstep. constructor. }
      intros _. rewrite bind_ret_l.
      repeat rewrite bind_bind. unfold or_spec. setoid_rewrite bind_vis.
      apply conv_ref_mrec_ex. intros n.
      rewrite bind_ret_l.
      eapply conv_ref_mrec_bind.
      + induction n. cbn. pstep. constructor.
        cbn. apply conv_ref_mrec_bind; intros; eauto.
        apply conv_ref_mrec_bind. apply conv_ref_mrec_ex.
        intros. pstep. constructor. intros.
        rewrite bind_bind.
        setoid_rewrite bind_vis.
        apply conv_ref_mrec_ex. intros Hr.
        repeat rewrite bind_ret_l. apply conv_ref_mrec_inl.
        auto. intros. pstep. constructor.
      + intros []. apply conv_ref_mrec_ex.
        intros [ | ]; try apply conv_ref_mrec_spin.
        setoid_rewrite bind_vis. apply conv_ref_mrec_ex.
        intros. setoid_rewrite bind_ret_l. rewrite bind_bind. setoid_rewrite bind_vis.
        apply conv_ref_mrec_ex. intros.
        repeat rewrite bind_ret_l. pstep. constructor.
    - setoid_rewrite bind_vis. setoid_rewrite bind_ret_l.
      apply conv_ref_mrec_exists. intros n. clear - n Ha.
      induction n. cbn. pstep. constructor.
      cbn. eapply conv_ref_mrec_bind; eauto.
      setoid_rewrite bind_bind. setoid_rewrite bind_ret_l.
      setoid_rewrite bind_vis.
      apply conv_ref_mrec_exists. intros a0. rewrite bind_ret_l.
      setoid_rewrite bind_vis. apply conv_ref_mrec_exists. intros Ha0.
      rewrite bind_ret_l. apply conv_ref_mrec_inl. auto. intros. pstep. constructor.
  Qed.

  Lemma partial_spec_fix_partial_spec_assume_fail a :
    ~ Pre a -> padded_refines_eq eq (prog a) (partial_spec a).
  Proof.
    intros HPre. apply or_spec_r.
    left. assumesr. intros. contradiction.
  Qed.

  Lemma partial_spec_fix_partial_spec_assume_succeed:
    forall (a : A) (b : B),
      padded_refines_eq eq (Ret b) (partial_spec_fix a) ->
      padded_refines_eq eq (Ret b) (partial_spec a).
  Proof.
    intros a b Hfix.
    apply or_spec_r. left. unfold total_spec.
    assumesr. intros Ha.
    unfold partial_spec_fix in Hfix.
    do 2 setoid_rewrite interp_mrec_spec_bind in Hfix. revert Hfix.
    match goal with | |- context [ interp_mrec_spec ?f _ ] => set f as body end.
    intro Hfix.
    assert (interp_mrec_spec body (b <- spec_exists B;; _ <- assert_spec (Post a b);; Ret b) ≈
                             (b <- spec_exists B;; _ <- assert_spec (Post a b);; Ret b)).
    {
      repeat setoid_rewrite interp_mrec_spec_bind.
      normalize_interp_mrec_spec.
      setoid_rewrite interp_mrec_spec_ret. cbn. reflexivity.
    } setoid_rewrite H in Hfix.
    (* need to manage diff between padded_refines and refines *)
    cbn in Hfix. setoid_rewrite <- bind_bind in Hfix.
    eapply padded_conv_ref_ret_bind; eauto.
    setoid_rewrite interp_mrec_spec_ret.
    apply conv_ref_bind.
    { pstep. red. cbn. constructor; auto. intros a0.
          left. pstep. red. cbn. constructor. }
    intros _.
    eapply conv_ref_interp_mrec_conv_ref.
    Unshelve. 3 : { intros C c. destruct c. apply (Pre a0). }
    - cbn.
      intros. destruct d. cbn.
      clear - H0 Ha.
      repeat rewrite bind_bind. setoid_rewrite bind_ret_l.
      apply conv_ref_mrec_bind.
      apply conv_ref_mrec_forall. auto. intros. pstep. constructor.
      intros. rewrite <- bind_bind.
      apply conv_ref_mrec_bind.
      + rewrite bind_bind. apply conv_ref_mrec_bind.
        apply conv_ref_mrec_exists. intros. pstep. constructor.
        intros n.
        induction n; cbn. rewrite bind_ret_l. apply conv_ref_mrec_exists.
        intros. pstep. constructor. rewrite bind_bind.
        eapply conv_ref_mrec_bind; intros; eauto.
        setoid_rewrite bind_vis. apply conv_ref_mrec_exists.
        intros. setoid_rewrite bind_ret_l. rewrite bind_bind. setoid_rewrite bind_vis.
        apply conv_ref_mrec_exists. intros. repeat rewrite bind_ret_l.
        apply conv_ref_mrec_inl. auto. intros. pstep. constructor.
      + intros b. rewrite bind_bind. setoid_rewrite bind_vis.
        apply conv_ref_mrec_exists. intros HPost. repeat rewrite bind_ret_l.
        pstep. constructor.
    - setoid_rewrite bind_vis. setoid_rewrite bind_ret_l.
      apply conv_ref_mrec_exists. intros n. clear - n Ha.
      induction n. cbn. pstep. constructor.
      cbn. eapply conv_ref_mrec_bind; eauto.
      setoid_rewrite bind_bind. setoid_rewrite bind_ret_l.
      setoid_rewrite bind_vis.
      apply conv_ref_mrec_exists. intros a0. rewrite bind_ret_l.
      setoid_rewrite bind_vis. apply conv_ref_mrec_exists. intros Ha0.
      rewrite bind_ret_l. apply conv_ref_mrec_inl. auto. intros. pstep. constructor.
  Qed.

  Lemma partial_spec_fix_partial_spec' a :
    (Pre a -> padded_refines_eq eq (prog a) (partial_spec_fix' a)) ->
    padded_refines_eq eq (prog a) (partial_spec a).
  Proof.
    intros Hfix. destruct (classic (Pre a) ).
    - specialize (Hfix H). apply prog_ret_spin in H as Hdec; auto.
      destruct Hdec as [ [b Hb] | Hspin ].
      + rewrite Hb. rewrite Hb in Hfix.
        apply partial_spec_fix_partial_spec_assume_succeed'; auto.
      + rewrite Hspin. apply or_spec_r. right. reflexivity.
    - apply or_spec_r. left. assumesr. intro. contradiction.
  Qed.

  Lemma partial_spec_fix_partial_spec a :
    (Pre a -> padded_refines_eq eq (prog a) (partial_spec_fix a)) ->
    padded_refines_eq eq (prog a) (partial_spec a).
  Proof.
    intros Hfix. apply partial_spec_fix_partial_spec'.
    intros Ha. specialize (Hfix Ha). etransitivity; eauto.
    eapply padded_refines_monot with (RE1 := eqE) (REAns1 := eqEAns)
    (RR1 := eqEAns _ _ (Call a) (Call a) ); auto.
    { intros. eqEAns_inv PR. auto. }
    eapply padded_refines_mrec with (REInv := eqE).
    2 : constructor.
    intros. eqE_inv H. clear Ha Hfix a. destruct d2.
    simpl. assumesr. intros Ha. assumesl. auto.
    repeat rewrite bind_bind.
    existssl. intros n. existssr n.
    apply padded_refines_bind with (RR := eq).
    + apply refl_refines. red. intros. destruct e; do 2 constructor.
      red. intros. inv H. inj_existT. subst. eqEAns_inv  H6. auto.
      inj_existT. subst. eqEAns_inv H6. auto.
      auto. apply pad_is_padded.
    + intros [] [] _ .
      apply or_spec_r. left. existssl. intros b.
      existssr b. assertsl. intros. assertsr. auto. apply padded_refines_ret.
      constructor.
   Qed.

  (*
  (* TODO : restructure this section so something like this can work *)
  Lemma partial_spec_fix_partial_spec_inv a :
    padded_refines (sub_eqE pre_call) (sub_eqEAns post_call) (sub_eq (Post a) )
                   (prog a) (partial_spec_fix a).
  *)
  (* should work as long as Pre is inhabitted that allows *)
  Lemma partial_spec_fix_spin a :
    padded_refines_eq eq ITree.spin (partial_spec_fix a).
  Proof.
    unfold partial_spec_fix. setoid_rewrite interp_mrec_spec_bind.
    setoid_rewrite interp_mrec_spec_assume. assumesr.
    intros Ha. cbn. rewrite <-  bind_bind.
    match goal with | |- padded_refines _ _ _ _ (interp_mrec_spec _ ( ITree.bind ?t1 ?t2) ) =>
                        remember t2 as t end.
    clear Heqt. generalize dependent t.
    (* this seems like a good coind hyp*)
    (* better solution might be to create a divergent refinement,
       showing that there is a path that you can keep getting taus in
       then give that some kind of mrec rule
     *)
    pcofix CIH.
  Admitted.

  Lemma partial_spec_fix_spin' a :
    padded_refines_eq eq ITree.spin (partial_spec_fix' a).
  Proof.
    unfold partial_spec_fix'. setoid_rewrite interp_mrec_spec_bind.
    repeat normalize_interp_mrec_spec. assumesr. intros.
    setoid_rewrite interp_mrec_spec_bind.
    setoid_rewrite interp_mrec_spec_bind.
    normalize_interp_mrec_spec.
    match goal with |- context [ interp_mrec_spec ?body_ _  ] => set body_ as body end.
    setoid_rewrite interp_mrec_spec_exists. existssr 0. cbn.
    rewrite interp_mrec_spec_ret, bind_ret_l. rewrite interp_mrec_spec_or.
    apply or_spec_r. right. rewrite interp_mrec_spec_spin. reflexivity.
  Qed.

End SpecFix.

Opaque assume_spec.
Opaque assert_spec.
Opaque spec_forall.
Opaque spec_exists.

Lemma halve_spec_fix_correct' E l :
  padded_refines_eq eq (@halve E l) (partial_spec_fix (fun _ => True)
                                                   (fun l '(l1,l2) => Permutation l (l1 ++ l2) /\
                               (length l1 >= length l2 /\ (length l > length l1 \/ length l <= 1)) ) l ).
Proof.
  eapply padded_refines_monot with (RE1 := eqE)
                                   (REAns1 := eqEAns)
                                   (RR1 := sub_eqEAns halve_call _ _ (Call l) (Call l) ); eauto.
  { intros [l1 l2] [l3 l4]. intros. inv PR. inj_existT. subst. auto. }
    eapply padded_refines_mrec with (REInv := eqE)
                                  (* in this post condition need to include *)
                                  (REAnsInv := sub_eqEAns halve_call).
  2 : constructor. intros.
  eqE_inv H. clear l. destruct d2 as [l]. cbn.
  destruct l as [ | a [ | b l] ].
  - assumesr. intros _. rewrite bind_bind. existssr 0.
    cbn. rewrite bind_ret_l. existssr (@nil nat,@nil nat).
    assertsr. split. reflexivity. lia.
    apply padded_refines_ret. do 2 constructor. reflexivity. cbn. lia.
  - assumesr. intros _. rewrite bind_bind. existssr 0.
    cbn. rewrite bind_ret_l. existssr (a :: nil, @nil nat).
    assertsr. split. reflexivity. cbn. lia.
    apply padded_refines_ret. do 2 constructor. reflexivity. cbn. lia.
  - assumesr. intros _. rewrite bind_bind. existssr 1.
    cbn. repeat rewrite bind_bind. existssr l.
    repeat rewrite bind_bind.
    assertsr. auto. setoid_rewrite bind_ret_l.
    eapply padded_refines_bind with  (RR :=
                                    sub_eqEAns halve_call (list nat * list nat) (list nat * list nat) (Call (l)) (Call (l))).
    { apply padded_refines_vis. constructor. constructor. intros [l1 l2] b0 Hl12.
      inv Hl12. inj_existT. subst. inv H5. inj_existT. subst.
      apply padded_refines_ret. constructor. auto. }
    intros [l1 l2] r Hl12. inv Hl12. inj_existT. subst. inv H4. inj_existT.
    injection H5; intros; subst. existssr (a :: l1,b :: l2). assertsr.
    split.
    { cbn. rewrite Permutation_app_comm. cbn.
      rewrite H0. rewrite Permutation_app_comm. reflexivity. }
    { cbn. repeat rewrite H0. repeat rewrite app_length. lia. }
    apply padded_refines_ret. constructor. constructor.
    { cbn. rewrite Permutation_app_comm. cbn.
      rewrite H0. rewrite Permutation_app_comm. reflexivity. }
    { cbn. repeat rewrite H0. repeat rewrite app_length. lia. }
 Qed.
(* this requires lem *)
Lemma halve_ret_spin E l : (exists l', @halve E l ≈ Ret l') \/ (@halve E l ≈ ITree.spin).
Admitted.

(* ok so there is a bunch of small things I still need to do but I think I have solved the problem *)
Theorem halve_correct'' E l : padded_refines_eq eq (@halve E l) (partial_spec (fun _ => True)
                                                   (fun l '(l1,l2) => Permutation l (l1 ++ l2) /\
                               (length l1 >= length l2 /\ (length l > length l1 \/ length l <= 1)) ) l).
Proof.
  eapply  partial_spec_fix_partial_spec. intros. apply halve_ret_spin.
  intros _.
  apply halve_spec_fix_correct'.
Qed.

Lemma sub_eqE_eq_type E P A B (ea : E A) (eb : E B) : sub_eqE P A B ea eb -> A = B.
Proof.
  intros. inv H. auto.
Qed.

Lemma sub_eqE_eq E P A (e1 : E A) (e2 : E A) : sub_eqE P A A e1 e2 -> e1 = e2 /\ P A e1.
Proof.
  intros. inv H. inj_existT. subst. auto.
Qed.

Lemma sub_eqEAns_eq_type E P A B (ea : E A) (eb : E B) a b : sub_eqEAns P A B ea eb a b -> A = B.
Proof.
  intros. inv H. auto.
Qed.

Lemma sub_eqEAns_eq E P A (e1 : E A) (e2 : E A) a1 a2 : sub_eqEAns P A A e1 e2 a1 a2 -> e1 = e2 /\ a1 = a2
                                                                                /\ P A e1 a1.
Proof.
  intros. inv H. inj_existT. subst. auto.
Qed.

Ltac sub_eqE_inv H := eapply sub_eqE_eq_type in H as ?H; subst; eapply sub_eqE_eq in H as [?H ?H] ; subst.
Ltac sub_eqEAns_inv H := apply sub_eqEAns_eq_type in H as ?H; subst; apply sub_eqEAns_eq in H as [?H [?H ?H] ]; subst.

Section merge_correct.
  Definition merge_pre '(l1,l2) := sorted l1 /\ sorted l2.
  Definition merge_post '(l1,l2) l := sorted l /\ Permutation l (l1 ++ l2).

  Theorem merge_correct'' E l1 l2 : padded_refines_eq eq (@Merge.merge E (l1,l2))
                                                  (partial_spec merge_pre merge_post (l1,l2)).
  Proof.
    eapply partial_spec_fix_partial_spec. admit. (* proof that merge has no event *)
    (* I should be able to assume here that Pre holds, because if it does not,
       then I can easily *)
    intros Hl12.
    eapply padded_refines_monot with (RE1 := eqE )
                                   (REAns1 := eqEAns)
                                   (RR1 := sub_eqEAns (post_call (merge_post) ) _ _
                                                      (Call (l1,l2)) (Call (l1,l2)) ); eauto.
    { intros. inv PR. inj_existT. subst. auto. }
    eapply padded_refines_mrec with (REInv := sub_eqE (pre_call merge_pre) )
                                    (REAnsInv := sub_eqEAns (post_call (merge_post) ) ).
    2 : { do 2 constructor. auto. }
    intros. sub_eqE_inv H.
    clear Hl12 l1 l2. destruct d2 as [ [ l1 l2] ]. inv H0. destruct H1 as [Hl1 Hl2].
    cbn.
    destruct l1; destruct l2.
    - assumesr. intros _. repeat rewrite bind_bind. existssr (0). setoid_rewrite bind_ret_l.
      existssr (@nil nat). assertsr. split; auto. apply padded_refines_ret.
      repeat constructor.
    - assumesr. intros _. repeat rewrite bind_bind. existssr 0. setoid_rewrite bind_ret_l.
      existssr (n :: l2). assertsr. auto.  apply padded_refines_ret.
      do 2 constructor. split; auto.
    - assumesr. intros _. repeat rewrite bind_bind. existssr 0. setoid_rewrite bind_ret_l.
      existssr (n :: l1). assertsr. split; auto. rewrite app_nil_r. auto. apply padded_refines_ret.
      do 2 constructor. split. auto. rewrite app_nil_r. auto.
    - rename n0 into m. destruct (Nat.leb n m) eqn : Hnm.
      + cbn. assumesr. intros _. repeat rewrite bind_bind. existssr 1.
        cbn. repeat rewrite bind_bind. existssr (l1, m :: l2). setoid_rewrite bind_bind. assertsr.
        split; auto. eapply sorted_tail. eauto. apply padded_refines_bind with (RR := sub_eq (merge_post (l1, m :: l2) ) ).
        { apply padded_refines_vis. do 3 constructor. split; auto. eapply sorted_tail. eauto.
          intros. inv H. inj_existT. subst. sub_eqEAns_inv  H6.
          inv H2. inj_existT. subst. apply padded_refines_ret. constructor. auto.
        }
        intros ? lm Hlm. inv Hlm. setoid_rewrite bind_ret_l. existssr (n :: lm).
        assertsr.
        { destruct H. split.
          - admit. (*fact about sorted list *)
          - rewrite H0. reflexivity.
        }
        apply padded_refines_ret. destruct H. do 3 constructor.
        * admit.
        * setoid_rewrite H0. reflexivity.
      + cbn. assumesr. intros _. repeat rewrite bind_bind. existssr 1.
        cbn. repeat rewrite bind_bind. existssr (n :: l1, l2). setoid_rewrite bind_bind. assertsr.
        split; auto. eapply sorted_tail. eauto. apply padded_refines_bind with (RR := sub_eq (merge_post (n :: l1,  l2) ) ).
        { apply padded_refines_vis. do 3 constructor. split; auto. eapply sorted_tail. eauto.
          intros. inv H. inj_existT. subst. sub_eqEAns_inv  H6.
          inv H2. inj_existT. subst. apply padded_refines_ret. constructor. auto.
        }
        intros ? lm Hlm. inv Hlm. setoid_rewrite bind_ret_l. existssr (m :: lm).
        assertsr.
        { destruct H. split.
          - admit. (*fact about sorted list *)
          - rewrite H0. cbn. rewrite (Permutation_app_comm (l1) (m :: l2) ) .  cbn. rewrite Permutation_app_comm.
            constructor.
        }
        apply padded_refines_ret. destruct H. do 3 constructor.
        * admit.
        * rewrite H0. cbn. rewrite (Permutation_app_comm (l1) (m :: l2) ) .  cbn. rewrite Permutation_app_comm.
          constructor.
 Admitted.

End merge_correct.

(* this should be contravariant, you can weaken the first and strengthen the second *)

(*
Instance test_le2 n : Proper (le ==> flip impl) (fun m => le m n).
Proof.
  repeat intro. subst. lia.
Qed.

Instance test_le1 : Proper (le ==> (flip le) ==> flip impl) le.
Proof.
  repeat intro. red in H0. lia.
Qed.

Instance test_le1' : Proper (le ==> (flip le) ==> impl) (flip le).
Proof.
  repeat intro. red in H0. red. red in H1. lia.
Qed. *)
(*
Instance le_trans : Transitive le.
Proof.
  repeat intro. lia.
Qed.

Goal 1 <= 2 -> 2 <= 3 -> 1 <= 3.
*)

Global Instance padded_refines_sub_eq E R (RR : R -> R -> Prop) (P1 : forall A, E A -> Prop) (P2 : forall A, E A -> A -> Prop) :
  Transitive RR ->
  Transitive (padded_refines (sub_eqE P1) (sub_eqEAns P2) RR ).
Proof.
  intros HRR t1 t2 t3 Ht12 Ht23. unfold padded_refines in *.
  eapply refines_monot; try eapply refinesTrans; eauto with solve_padded; try apply pad_is_padded.
  - intros A B ea eb PR. inv PR. inj_existT. subst. sub_eqE_inv H3. auto.
  - intros. sub_eqEAns_inv  PR. constructor. intros. destruct H. sub_eqE_inv H.
    exists x1. split; constructor; auto.
  - intros. inv PR. etransitivity; eauto.
Qed.

Global Instance sub_eqEAns_trans E (P : forall A, E A -> A -> Prop) A (e : E A) : Transitive (sub_eqEAns P A A e e).
Proof.
  intros x y z Hxy Hyz. sub_eqEAns_inv Hxy. auto.
Qed.

(* I think I can save the rest for tomorrow*)

(*
Global Instance refines_proper_bind_RE E R RE REAns S :
  Proper (@padded_refines E  R R eq ==> pointwise_relation _ (padded_refines_eq eq) ==> @padded_refines_eq E S S  eq ) ITree.bind.
Proof.
  repeat intro. subst. unfold padded_refines_eq in *. eapply padded_refines_bind; eauto.
  intros; subst. auto.
Qed.
*)

(*do I need the proper instances? Probably *)

Global Instance  padded_refines_eq_gen E1 E2 R1 R2 RE REAns RR : Proper (padded_refines_eq eq ==> flip (padded_refines_eq eq) ==> flip impl )
                                                                (@padded_refines E1 E2 R1 R2 RE REAns RR).
Proof.
  intros t1 t2 Ht12 t3 t4 Ht43. red in Ht43. intro.
  unfold padded_refines_eq, padded_refines in *. rename H into H24.
  specialize @refinesTrans with (E1 := E1) (E2 := E1) (E3 := E2) (RE1 := eqE) (RE2 := RE) as Htrans.
  eapply Htrans in Ht12; eauto; try apply pad_is_padded. clear Htrans.
  specialize @refinesTrans with (E1 := E1) (E2 := E2) (E3 := E2) (RE1 := rcomposeE eqE RE) (RE2 := eqE) as Htrans.
  eapply Htrans in Ht12; eauto; try apply pad_is_padded. eapply refines_monot; try eapply Ht12.
  { intros. destruct PR. destruct H. eqE_inv H. eqE_inv H0. auto. }
  { intros. econstructor. intros. destruct H. eqE_inv H0.
    inv H. inj_existT. subst. eqE_inv H4. exists x1. split. 2 : constructor.
    constructor. intros. destruct H. eqE_inv H. exists x0. split; auto. constructor.
  }
  { intros. inv PR. inv REL1. auto. }
Qed.
  (* this is maybe a tad more subtle then I realized, does it require transitivity?,
     that would not be problematic for these relations because there seem to be PER's

   *)

Lemma or_spec_bind_l  E1 E2 R1 S R2 RE REAns RR (t1 t2 : itree_spec E1 S) (k : S -> itree_spec E1 R1) (t3 : itree_spec E2 R2) :
  padded_refines RE REAns RR (ITree.bind t1 k) t3 ->
  padded_refines RE REAns RR (ITree.bind t2 k) t3 ->
  padded_refines RE REAns RR (ITree.bind (or_spec t1 t2) k) t3.
Proof.
  unfold padded_refines. intros. setoid_rewrite bind_vis. setoid_rewrite pad_vis.
  pstep. constructor. intros [ | ]; constructor; pstep_reverse.
Qed.

Section sort_correct.
  Definition sort_pre (l : list nat) := True.
  Definition sort_post (l l' : list nat) := Permutation l l' /\ sorted l'.

  Theorem sort_correct' E l : padded_refines_eq eq (@sort E l) (partial_spec sort_pre sort_post l).
  Proof.
    eapply partial_spec_fix_partial_spec'. admit.
    intros _.
    eapply padded_refines_monot with (RE1 := eqE )
                                   (REAns1 := eqEAns)
                                   (RR1 := sub_eqEAns (post_call (sort_post) ) _ _
                                                      (Call (l)) (Call (l)) ); eauto.
    { intros. sub_eqEAns_inv  PR. auto. }
    eapply padded_refines_mrec with (REInv := sub_eqE (pre_call sort_pre) )
                                    (REAnsInv := sub_eqEAns (post_call (sort_post) ) ).
    2 : { repeat constructor. }
    intros. clear l. sub_eqE_inv  H. destruct d2. inv H0.
    cbn. destruct (Nat.leb (length l) 1 ) eqn : Hlen.
    - assumesr. intros _. repeat rewrite bind_bind. existssr 0.
      cbn. rewrite bind_ret_l. apply or_spec_r. left.
      existssr l.
      assert (sort_post l l).
      { destruct l. split; auto. constructor. destruct l. split; auto.
        constructor. discriminate. }
      assertsr. auto. apply padded_refines_ret.
      do 2 constructor. auto.
    - rewrite halve_correct''.
      (* here I really want to push the or_spec left, probably only works if only quantifiers,
         if I can get past this then it should work *)

      unfold partial_spec, total_spec. rewrite or_spec_bind.
      apply or_spec_l.
      2 : { setoid_rewrite spin_bind.
            assumesr. intros _ . repeat rewrite bind_bind. existssr 0.
            cbn. rewrite bind_ret_l. apply or_spec_r. right.
            apply padded_refines_spin.
          }
      setoid_rewrite bind_bind at 1. assumesl. auto.  setoid_rewrite bind_bind at 1.
      existssl. intros [l1 l2]. setoid_rewrite bind_bind at 1. assertsl.
      intros Hl12. rewrite bind_ret_l.
      (* now I really need to push the or_spec further because in order to get to the
         merge call (which itself might diverge , unless I actually  I need to just add more off ramps with o) *)

      setoid_rewrite bind_bind. assumesr. intros _ .
      existssr 2. cbn. repeat rewrite bind_bind.
      existssr l1. rewrite bind_bind. assertsr.
      auto. apply padded_refines_bind with (RR := sub_eq (sort_post l1) ).
      { apply padded_refines_vis. repeat constructor.
        intros. inv H. inj_existT. subst. sub_eqEAns_inv H7.
        subst. apply padded_refines_ret. constructor. inv H3. inj_existT.
        subst. auto.
      }
      intros ? l1s Hl1s. inv Hl1s.
      repeat rewrite bind_bind.
      existssr l2. repeat rewrite bind_bind. assertsr.
      auto. apply padded_refines_bind with (RR := sub_eq (sort_post l2) ).
      {
        apply padded_refines_vis. repeat constructor. intros.
        inv H0. inj_existT. subst. sub_eqEAns_inv H8.
        subst. apply padded_refines_ret. constructor. inv H3. inj_existT.
        subst. auto. }
      intros ? l2s Hl2s. inv Hl2s. rewrite bind_ret_l.
      rewrite merge_correct''. unfold partial_spec.
      apply or_spec_l.
      2 : { apply or_spec_r. right. apply padded_refines_spin. } (*same lemma as before*)
      apply or_spec_r. left. unfold total_spec. cbn. red in H. red in H0.
      assumesl. tauto.
      existssl. intros ls. assertsl. intros Hls. destruct H. destruct H0.
      destruct Hls. existssr ls. assertsr.
      split; auto.
      destruct Hl12. rewrite H6, H, H0, H5. reflexivity.
      apply padded_refines_ret.
      constructor. split; auto.
      destruct Hl12.
      split; auto.
      rewrite H6, H, H0, H5. reflexivity.
Admitted.

End sort_correct.

(*
Theorem merge
*)


Add LoadPath "../Coq".

Require Import Contractible.

Close Scope nat_scope.

Set Implicit Arguments.

Local Notation " f === g " := (forall x, f x == g x) (at level 62).

Definition pr20 { A } { P : A -> Type } { Q : A -> Type } : sigT2 P Q -> A.
  intro x.
  destruct x; auto.
Defined.

Definition pr21 { A } { P : A -> Type } { Q : A -> Type } :
  forall x : sigT2 P Q, P (pr20 x).
Proof.
    intro x.
    destruct x; auto.
Defined.

Definition pr22 { A } { P : A -> Type } { Q : A -> Type } :
  forall x : sigT2 P Q, Q (pr20 x).
Proof.
    intro x.
    destruct x; auto.
Defined.

Notation " ( a ; b ; c ) " := (existT2 _ _ a b c).

Definition is_sect { A B : Type } (f : A -> B) ( g : B -> A )
 := forall a, f (g a) == a.

Section HIso.

    Hypotheses (A B : Type) (f : A -> B).

    Definition is_equiv := forall b : B, is_contr ( hfiber f b ).

    Definition is_hiso : Type :=
     { g : B -> A & is_sect f g & { h : A -> B & is_sect g h }}.

    Section HIsos.

        Hypotheses g : B -> A.
        Hypotheses h : A -> B.
        Hypotheses sect : is_sect f g.
        Hypotheses retr : is_sect g h.

        Definition mk_iso : is_hiso := ( g ; sect ; ( h ; retr )).

        Lemma same_obverse : forall a, f a == h a.
        Proof.
            intro. unfold is_sect in sect.
            path_via (f ( g (h a))).
        Defined.

        Lemma same_obv_nat : forall x y, forall p : x == y,
            map f p @ same_obverse y == same_obverse x @ map h p.
            intros.
            induction p.
            cancel_units.
        Defined.

        Lemma sect_natural :
          forall x y, forall p : x == y, sect x @ p == map f (map g p) @ sect y.
        Proof.
            intros.
            induction p.
            cancel_units.
        Defined.

        Lemma sect_absb : forall x, sect (f (g x)) == map f (map g (sect x)).
            intros.
            apply @concat_cancel_right with (r := sect x).
            apply sect_natural.
        Defined.

        (** Coq's automatic parameter detection, in the end, doesn't distinguish
            the next two lemmata from the previous two; they differ locally, though
            in that the local instances refer to distinct local parameters. *)

        Lemma retr_natural :
          forall x y, forall q : x == y, retr x @ q == map g (map h q) @ retr y.
        Proof.
            intros.
            induction q.
            cancel_units.
        Defined.

        Lemma retr_absb : forall x, retr (g (h x)) == map g (map h (retr x)).
        Proof.
            intro.
            apply @concat_cancel_right with (r := retr x).
            apply retr_natural.
        Defined.

    End HIsos.

    Definition is_hequiv := { g : B -> A & is_sect f g & is_sect g f }.

    (** These really aren't needed *)
    (*
    Section IsoVsHequiv.

        Lemma is_hiso_to_is_hequiv : is_hiso -> is_hequiv.
        Proof.
            intro H.
            destruct H as [ g s [h r] ].
            exists g.
            auto.
            intro.
            path_via (g (h a)).
            apply @same_obverse with g; auto.
        Defined.

        Lemma is_hequiv_to_is_hiso : is_hequiv -> is_hiso.
        Proof.
            intro.
            destruct X as [g  p q ].
            exists g; auto.
            exists f; auto.
        Defined.

    End IsoVsHequiv.
    *)
End HIso.

Section HIsoCompose.

    Hypotheses A B C : Type.
    Hypotheses (f : A -> B) ( g : B -> C).

    Section Compose.

        Hypotheses hf : is_hiso f.
        Hypotheses hg : is_hiso g.

        Lemma compose_hiso : is_hiso (g o f).
        Proof.
            remember hf as Hf.
            destruct Hf as [ g1 s1 [ g2 s2 ]].
            remember hg as Hg.
            destruct Hg as [ h1 r1 [ h2 r2 ]].
            exists ( g1 o h1 ).
            intro.
              path_via (g (h1 a)).
              unfold compose. apply map.
              apply s1.
            exists ( h2 o g2 ).
            intro.
              path_via (g1 (g2 a)).
              unfold compose.
              apply map.
              apply r2.
        Defined.

    End Compose.

    Section Cancel.

        Hypotheses hgf : is_hiso (g o f).

        Section CancelRight.

            Hypotheses hf : is_hiso f.
            Lemma cancel_right_hiso : is_hiso g.
            Proof.
                remember hgf as Hgf.
                destruct Hgf as [ Lgf s [ Rgf r ]].
                remember hf as Hf.
                destruct Hf as [ Lf u [ Rf v ]].
                exists ( f o Lgf ).
                intro.
                path_via ( (g o f) (Lgf a)).
                exists ( Rgf o Lf ).
                intro.
                path_via ( f (Lf a)).
                unfold compose.
                apply map.
                apply r.
            Defined.

        End CancelRight.

        (** This is the tricky one. It makes much more sense in pasting figures *)

        Section CancelLeft.

            Hypotheses hg : is_hiso g.
            Lemma cancel_left_hiso : is_hiso f.
            Proof.
                remember hgf as Hgf.
                destruct Hgf as [ Lgf s [ Rgf r ]].
                remember hg as Hg.
                destruct Hg as [ Lg u [ Rg v ]].
                exists ( Lgf o g ).
                intro.
                path_via (Lg (g ( f ( Lgf (g a))))).
                path_via (Lg (Rg (f (Lgf (g a))))).
                refine (! same_obverse u v _).
                path_via (Lg ((g o f) (Lgf (g a)))).
                path_via (Lg (g a)).
                path_via (Lg (Rg a)).
                refine ( same_obverse u v _ ).
                exists f.
                intro.
                path_via (Lgf ( (g o f) a)).
                path_via ( Lgf ( Rgf a )).
                refine ( same_obverse s r _ ).
            Defined.
        End CancelLeft.

    End Cancel.

End HIsoCompose.

Section HIsEquiv.

    Hypotheses (A B : Type)
     (f : A -> B) ( W : is_hiso f ).

    Lemma h_is_equiv :
     is_equiv f.
     remember W as V.
     destruct V as [ g s [ h r ]].
     set ( R := same_obverse s r ).
     intro a.
     exists ( g a ; s a ).
     intro.
     destruct y as [ b p ].
     apply total_path with (! r b @ ! map g (R b) @ map g p ).
     path_via (! map f (! r b @ ! map g (R b) @ map g p) @ p ).
     refine ((trans_paths A B f (fun _ => a) _ _
              ( ! r b @ ! map g (R b) @ map g p )  p ) @ _ ).
     cancel_units.
     unfold R.
     do_concat_map.
     do_opposite_map.
     moveright_onleft.
     associate_right.
     moveleft_onleft.
     moveleft_onleft.
     unfold same_obverse.
     do_concat_map.
     do_opposite_map.
     associate_right.
     moveright_onleft.
     path_via (s (f (g (h b))) @ (map f (r b) @ p)).
     apply opposite.

       apply sect_absb.

     path_via (map f (map g (map f (r b) @ p)) @ s a).

       apply sect_natural.

     do_concat_map.
     associate_left.
   Defined.

End HIsEquiv.


Require Import IOA.
Require Import Misc.
Require Import ListOps.
Require Import OptionOps.
Require Import Simulation.
Require Import RefinementFacts.
Require Import Automation.
Require Import Channel.

Require Import List.
Import ListNotations.
Require Import Peano_dec.
Require Import Arith.Even.

Inductive API :=
| Put (k : nat) (v : nat)
| PutOk
| Get (k : nat)
| GetResult (v : option nat).

Module Spec.

  Inductive Request :=
  | Req_Put (k : nat) (v : nat)
  | Req_Get (k : nat).
  Inductive Response :=
  | Resp_Ok
  | Resp_Value (v : option nat).
  Inductive Internal := Execute.

  Record state : Type := mkState {
                             requests : list Request;
                             responses : list Response;
                             data : nat -> option nat;
                           }.
  Definition start (st : state) : Prop := st = mkState [] [] (fun _ => None).

  Definition step (st : state) (act : Internal + API) (st' : state) : Prop :=
    let (req, res, d) := st in
    match act with
    | inr (Put k v) => st' = mkState (req ++ [Req_Put k v]) res d
    | inr (Get k) => st' = mkState (req ++ [Req_Get k]) res d
    | inl Execute =>
      exists hd tl, req = hd :: tl /\
                    st' = mkState tl
                                  (res ++ [match hd with
                                           | Req_Get k => Resp_Value (d k)
                                           | _ => Resp_Ok
                                           end])
                                  (match hd with
                                   | Req_Put k v =>
                                     fun k' => if eq_nat_dec k k' then
                                                 Some v else
                                                 d k'
                                   | _ => d
                                   end)
    | inr PutOk => exists tl, res = Resp_Ok :: tl /\ st' = mkState req tl d
    | inr (GetResult v) => exists tl, res = (Resp_Value v) :: tl /\ st' = mkState req tl d
    end.

  Definition KVStore : AutomatonDef API :=
    mkAutomatonDef API state Internal start step.

  Ltac break_state :=
    match goal with
    | [ |- context[ ?x ] ] =>
      is_evar x;
      let T := type of x in
      lazymatch eval simpl in T with
      | state =>
        let x1 := fresh "x" in
        let x2 := fresh "y" in
        let x3 := fresh "z" in
        evar (x1 : list Request);
        evar (x2 : list Response);
        evar (x3 : nat -> option nat);
        unify x (mkState x1 x2 x3);
        clear x1 x2 x3
      end
    end.

  Tactic Notation "ext" :=
    estep_ext_break; [solve [break_state; simpl; eauto] |].

  Tactic Notation "int" :=
    estep_int_break Execute; [
      solve [break_state; simpl; eauto] |].

  Goal in_traces KVStore [
         Put 3 4;
           Get 3;
           PutOk;
           GetResult (Some 4);
           Put 3 5;
           PutOk;
           Get 3;
           GetResult (Some 5);
           Get 4;
           GetResult None
       ].
  Proof.
    unfold in_traces.
    eexists; exists_lift; [cbv; eauto |].
    eexists.
    repeat (ext || int); eauto.
  Qed.

End Spec.

Module ClientServer.

  Section Main.

    Import Channel.
    Import Spec.

    Definition ServerAPI : Type :=
      (ChannelReceiverAPI Request) + (ChannelSenderAPI Response).

    Definition server_state : Type := (list Response) * (nat -> option nat).

    Definition server_start (st : server_state) : Prop :=
      st = ([], fun _ => None).

    Definition server_step (st : server_state) (act : EmptySet + ServerAPI) (st' : server_state) :=
      let (resp, m) := st in
      match act with
      | inl e => match e with end
      | inr (inl (Recv (Req_Put k v))) =>
        st' = (resp ++ [Resp_Ok], fun k' => if eq_nat_dec k k' then Some v else m k')
      | inr (inl (Recv (Req_Get k))) =>
        st' = (resp ++ [Resp_Value (m k)], m)
      | inr (inr (Send msg)) =>
        exists tl, resp = msg :: tl /\ st' = (tl, m)
      end.

    Definition Server : AutomatonDef ServerAPI :=
      mkAutomatonDef ServerAPI server_state EmptySet server_start server_step.

    Definition ClientAPI : Type :=
      API + (ChannelSenderAPI Request + ChannelReceiverAPI Response).

    Definition client_state : Type := (list Request * list Response).

    Definition client_start (st : client_state) := st = ([], []).

    Definition client_step (st : client_state) (act : EmptySet + ClientAPI) (st' : client_state) :=
      let (req, res) := st in
      match act with
      | inl e => match e with end
      | inr (inl (Put k v)) =>
        st' = (req ++ [Req_Put k v], res)
      | inr (inl (Get k)) =>
        st' = (req ++ [Req_Get k], res)
      | inr (inr (inl (Send m))) =>
        exists tl, req = m :: tl /\ st' = (tl, res)
      | inr (inr (inr (Recv m))) =>
        st' = (req, res ++ [m])
      | inr (inl (PutOk)) =>
        exists tl, res = Resp_Ok :: tl /\ st' = (req, tl)
      | inr (inl (GetResult v)) =>
        exists tl, res = (Resp_Value v) :: tl /\ st' = (req, tl)
      end.

    Definition Client : AutomatonDef ClientAPI :=
      mkAutomatonDef ClientAPI client_state EmptySet client_start client_step.

    Definition ServerAndChannels :=
      compose
        (compose
           Server
           (Channel.MediatedReliableReordering Request)
           (ChannelAPI Request + ChannelSenderAPI Response)
           (fun act => match act with
                       | inl (inl _) => None
                       | inl (inr act') => Some (inl act')
                       | inr act' => Some (inr act')
                       end)
           (fun act => match act with
                       | inl act' => Some act'
                       | inr _ => None
                       end)
        )
        (Channel.MediatedReliableReordering Response)
        (ChannelAPI Request + ChannelAPI Response)
        (fun act => match act with
                    | inl act' => Some (inl act')
                    | inr (inl act') => Some (inr act')
                    | inr (inr _) => None
                    end)
        (fun act => match act with
                    | inl _ => None
                    | inr act' => Some act'
                    end).

    Definition ServerAndFIFOChannels :=
      compose
        (compose
           Server
           (Channel.Reliable Request)
           (ChannelAPI Request + ChannelSenderAPI Response)
           (fun act => match act with
                       | inl (inl _) => None
                       | inl (inr act') => Some (inl act')
                       | inr act' => Some (inr act')
                       end)
           (fun act => match act with
                       | inl act' => Some act'
                       | inr _ => None
                       end)
        )
        (Channel.Reliable Response)
        (ChannelAPI Request + ChannelAPI Response)
        (fun act => match act with
                    | inl act' => Some (inl act')
                    | inr (inl act') => Some (inr act')
                    | inr (inr _) => None
                    end)
        (fun act => match act with
                    | inl _ => None
                    | inr act' => Some act'
                    end).

    Definition System' :=
      compose
        ServerAndChannels
        Client
        (API + (ChannelAPI Request + ChannelAPI Response))
        (fun act => match act with
                    | inl _ => None
                    | inr act' => Some act'
                    end)
        (fun act => match act with
                    | inl act' => Some (inl act')
                    | inr (inl (inl act')) => Some (inr (inl act'))
                    | inr (inl (inr _)) => None
                    | inr (inr (inl _)) => None
                    | inr (inr (inr act')) => Some (inr (inr act'))
                    end).

    Definition SystemFIFO' :=
      compose
        ServerAndFIFOChannels
        Client
        (API + (ChannelAPI Request + ChannelAPI Response))
        (fun act => match act with
                    | inl _ => None
                    | inr act' => Some act'
                    end)
        (fun act => match act with
                    | inl act' => Some (inl act')
                    | inr (inl (inl act')) => Some (inr (inl act'))
                    | inr (inl (inr _)) => None
                    | inr (inr (inl _)) => None
                    | inr (inr (inr act')) => Some (inr (inr act'))
                    end).

    Definition System :=
      rename
        System'
        API
        (fun act => inl act)
        (ChannelAPI Request + ChannelAPI Response)
        (fun act => inr act).

    Definition SystemFIFO :=
      rename
        SystemFIFO'
        API
        (fun act => inl act)
        (ChannelAPI Request + ChannelAPI Response)
        (fun act => inr act).

  End Main.

  Section SystemCorrect.

    Import Spec.

    Theorem System_correct : refines System Spec.KVStore.
    Proof.
      eapply refines_trans.
      - apply refines_rename.
        apply refines_comp_subst.
        eapply refines_trans.
        + apply refines_comp_subst.
          apply refines_comp_subst_r.
          apply Channel.reordering_mediator_correct.
        + apply refines_comp_subst_r.
          apply Channel.reordering_mediator_correct.
      - apply forward_simulation with
        (fun st_impl st_spec =>
           let s_req := requests st_spec in
           let s_res := responses st_spec in
           let s_d := data st_spec in
           let i_reqbuf := snd (fst (fst st_impl)) in
           let i_req := fst (snd st_impl) in
           let i_resbuf := snd (fst st_impl) in
           let i_res := snd (snd st_impl) in
           let i_ress := fst (fst (fst (fst st_impl))) in
           let i_d := snd (fst (fst (fst st_impl))) in

           s_req = i_reqbuf ++ i_req /\
           s_res = i_res ++ i_resbuf ++ i_ress /\
           (*
         (forall k, s_d k = i_d k)
            *)
           s_d = i_d
        ).
        split.
        + intros s Hs.
          repeat match goal with
                 | [ H : IOA.start _ _ |- _ ] => inversion H; clear H
                 end.
          destruct s as [[[[? ?] ?] ?] ?].
          simpl in *.
          cleanup; subst.
          unfold start; intuition; eauto.
        + intros s1' s1 act s2' Hstep Hrel.
          destruct act as [act | act].
          * (* channel stuff *)
            destruct act as [[[[[] | []] | []] | []] | act].
            destruct act as [[[] | [req]] | [[] | []]];
              try solve [
                    destruct s1 as [[[[? ?] ?] ?] ?];
                    destruct s1' as [[[[? ?] ?] ?] ?];
                    exists s2';
                    destruct s2';
                    simpl in *; intuition;
                    try match goal with
                        | [ H : client_step ?x _ _ |- _ ] =>
                          destruct x; simpl in H; inv_clear H
                        end;
                    cleanup; subst; simpl;
                    try crush_list; auto
                  ].
            destruct s1 as [[[[? ?] ?] ?] ?];
              destruct s1' as [[[[? ?] ?] ?] ?];
              destruct s2'.
            simpl in *.
            eexists (mkState _ _ _); intuition.
            subst.
            estep_int_break Execute; eauto.
            simpl; repeat eexists; eauto;
              destruct req; cleanup; subst;
                f_equal; crush_list.
          * (* API *)
            destruct act;
              destruct s1 as [[[[? ?] ?] ?] ?];
              destruct s1' as [[[[? ?] ?] ?] ?];
              destruct s2';
              simpl in *;
              eexists (mkState _ _ _); intuition;
                subst;
                estep_ext_break; eauto;
                  cleanup; subst; simpl in *;
                    try match goal with
                        | [ H : client_step ?x _ _ |- _ ] =>
                          destruct x; simpl in H; inv_clear H
                        end;
                    cleanup; subst; simpl;
                      try solve [eexists; eauto];
                      try solve [f_equal; crush_list].
    Qed.

  End SystemCorrect.

End ClientServer.

Module Sharded.

  Section Main.

    Import Channel.

    (* note: this doesn't refine Spec because it reorders stuff
       between the shards *)

    Inductive ShardRequest :=
    | Req_Put (k : nat) (v : nat)
    | Req_Get (k : nat)
    .

    Inductive ShardResponse :=
    | Resp_Ok
    | Resp_Value (v : option nat)
    .

    Definition ShardAPI : Type :=
      (ChannelReceiverAPI ShardRequest) + (ChannelSenderAPI ShardResponse).

    Definition shard_state : Type := (list ShardResponse) * (nat -> option nat).

    Definition shard_start (st : shard_state) : Prop :=
      st = ([], fun _ => None).

    Definition shard_step (st : shard_state) (act : EmptySet + ShardAPI) (st' : shard_state) :=
      let (resp, m) := st in
      match act with
      | inl e => match e with end
      | inr (inl (Recv (Req_Put k v))) =>
        st' = (resp ++ [Resp_Ok], fun k' => if eq_nat_dec k k' then Some v else m k')
      | inr (inl (Recv (Req_Get k))) =>
        st' = (resp ++ [Resp_Value (m k)], m)
      | inr (inr (Send msg)) =>
        exists tl, resp = msg :: tl /\ st' = (tl, m)
      end.

    Definition Shard : AutomatonDef ShardAPI :=
      mkAutomatonDef ShardAPI shard_state EmptySet shard_start shard_step.

    (* coordinator coordinates between 2 shards. overall API is:
    channel to S_even, channel from S_even, channel to S_odd, channel
    from S_odd, API *)
    Definition CoordinatorChannels : Type :=
      (ChannelSenderAPI ShardRequest + ChannelReceiverAPI ShardResponse) +
      (ChannelSenderAPI ShardRequest + ChannelReceiverAPI ShardResponse).
    Definition CoordinatorAPI : Type := API + CoordinatorChannels.

    Record coord_state : Type := mkCoordState {
                                     req_even : list ShardRequest;
                                     req_odd : list ShardRequest;
                                     resp : list ShardResponse
                                   }.

    Definition coord_start (st : coord_state) :=
      st = mkCoordState [] [] [].

    Definition coord_step (st : coord_state) (act : EmptySet + CoordinatorAPI) (st' : coord_state) : Prop :=
      let (re, ro, r) := st in
      match act with
      | inl e => match e with end

      (* input *)
      | inr (inl (Put k v)) =>
        if even_odd_dec k then
          st' = mkCoordState (re ++ [Req_Put k v]) ro r else
          st' = mkCoordState re (ro ++ [Req_Put k v]) r
      | inr (inl (Get k)) =>
        if even_odd_dec k then
          st' = mkCoordState (re ++ [Req_Get k]) ro r else
          st' = mkCoordState re (ro ++ [Req_Get k]) r

      (* even shard *)
      | inr (inr (inl (inl (Send m)))) =>
        exists tl, re = m :: tl /\ st' = mkCoordState tl ro r
      | inr (inr (inl (inr (Recv m)))) =>
        st' = mkCoordState re ro (r ++ [m])

      (* odd shard *)
      | inr (inr (inr (inl (Send m)))) =>
        exists tl, ro = m :: tl /\ st' = mkCoordState re tl r
      | inr (inr (inr (inr (Recv m)))) =>
        st' = mkCoordState re ro (r ++ [m])

      (* output *)
      | inr (inl PutOk) =>
        exists tl, r = Resp_Ok :: tl /\ st' = mkCoordState re ro tl
      | inr (inl (GetResult v)) =>
        exists tl, r = (Resp_Value v) :: tl /\ st' = mkCoordState re ro tl
      end.

    Definition Coordinator : AutomatonDef CoordinatorAPI :=
      mkAutomatonDef CoordinatorAPI coord_state EmptySet coord_start coord_step.

    Definition ShardAndChannels :=
      compose
        (compose
           Shard
           (Channel.Reliable ShardRequest)
           (ChannelAPI ShardRequest + ChannelSenderAPI ShardResponse)
           (fun act => match act with
                       | inl (inl _) => None
                       | inl (inr act') => Some (inl act')
                       | inr act' => Some (inr act')
                       end)
           (fun act => match act with
                       | inl act' => Some act'
                       | inr _ => None
                       end)
        )
        (Channel.Reliable ShardResponse)
        (ChannelAPI ShardRequest + ChannelAPI ShardResponse)
        (fun act => match act with
                    | inl act' => Some (inl act')
                    | inr (inl act') => Some (inr act')
                    | inr (inr _) => None
                    end)
        (fun act => match act with
                    | inl _ => None
                    | inr act' => Some act'
                    end).

  End Main.

End Sharded.