Require Import HoTT.Basics HoTT.Types.
Require Import MyTacs MyLemmas Colimits.Diagram Colimits.Colimit Colimits.CoEqualizer Colimits.Colimit_Sigma.

Local Open Scope path_scope.
Generalizable All Variables.

Context `{Funext}.

Module Export KP.
  Private Inductive KP' {A B:Type} (f: A -> B) : Type :=
  | kp : A -> (KP' f).

  Axiom kp_eq : forall {A B} {f: A -> B} (a b: A), f a = f b -> kp f a = kp f b.

  Axiom kp_eq2 : forall {A B f} (a: A), @kp_eq A B f a a 1 = 1.

  Arguments kp {A B f} a.

  Definition KP_ind {A B: Type} {f: A -> B} (P: KP' f -> Type)
             (kp': forall a, P (kp a))
             (kp_eq': forall a b p, transport P (kp_eq a b p) (kp' a) = kp' b)
             (kp_eq2': forall a, transport2 P (kp_eq2 a) (kp' a) = kp_eq' a a 1)
    : forall w, P w
    := fun w => match w with
                |kp a => fun _ => kp' a
                end kp_eq2'.

  Axiom KP_ind_beta_kp_eq : forall {A B: Type} {f: A -> B} (P: KP' f -> Type)
                                   (kp': forall a, P (kp a))
                                   (kp_eq': forall a b p, transport P (kp_eq a b p) (kp' a) = kp' b)
                                   (kp_eq2': forall a, transport2 P (kp_eq2 a) (kp' a) = kp_eq' a a 1)
                                   a b p,
      apD (KP_ind P kp' kp_eq' kp_eq2') (kp_eq a b p) = kp_eq' a b p.

  Axiom KP_ind_beta_kp_eq2 : forall {A B:Type} {f:A -> B} (P : KP' f -> Type)
             (kp': forall a, P (kp a))
             (kp_eq': forall a b p, transport P (kp_eq a b p) (kp' a) = kp' b)
             (kp_eq2': forall a, transport2 P (kp_eq2 a) (kp' a) = kp_eq' a a 1)
             a,
      apD02 (KP_ind P kp' kp_eq' kp_eq2') (kp_eq2 a) @ (concat_p1 _) @ (kp_eq2' a) = KP_ind_beta_kp_eq P kp' kp_eq' kp_eq2' a a 1.

  Definition KP_rec {A B: Type} {f: A -> B} (P: Type)
             (kp': A -> P)
             (kp_eq': forall (a b: A) (p: f a = f b), kp' a = kp' b)
             (kp_eq2': forall a, kp_eq' a a 1 = 1)
    : KP' f -> P.
  Proof.
    simple refine (KP_ind _ kp' (fun a b p => transport_const _ _ @ kp_eq' a b p)  _).
    intro a.
    exact ((whiskerL (transport2 (λ _ : KP' f, P) (kp_eq2 a) (kp' a)) (kp_eq2' a) @ concat_p1 _)^
           @ (whiskerR (transport2_const (A:=KP' f) (B:= P) (kp_eq2 a) (kp' a) @ concat_p1 _)^ (kp_eq' a a 1))).
  Defined.

  Definition KP_rec_beta_kp_eq {A B: Type} {f: A -> B} (P: Type)
             (kp': A -> P)
             (kp_eq': forall (a b: A) (p: f a = f b), kp' a = kp' b)
             (kp_eq2': forall a, kp_eq' a a 1 = 1)
             a b p
    : ap (KP_rec P kp' kp_eq' kp_eq2') (kp_eq a b p) = kp_eq' a b p.
  Proof.
  Admitted.

  Definition KP_rec_beta_kp_eq2 {A B: Type} {f: A -> B} (P: Type)
             (kp': A -> P)
             (kp_eq': forall (a b: A) (p: f a = f b), kp' a = kp' b)
             (kp_eq2': forall a, kp_eq' a a 1 = 1)
             a
    : ap02 (KP_rec P kp' kp_eq' kp_eq2') (kp_eq2 a) = KP_rec_beta_kp_eq P kp' kp_eq' kp_eq2' a a 1 @ (kp_eq2' a).
  Proof.
  Admitted.

  Definition KP_eta `{f: A -> B} {Z} (F G: KP' f -> Z)
             (H1: F o kp == G o kp)
             (H2: forall x y p, H1 x @ ap G (kp_eq x y p) = ap F (kp_eq x y p) @ H1 y)
             (H3: forall x, H2 x x 1
                            = transport (λ p, H1 x @ ap G p = ap F p @ H1 x ) (kp_eq2 x)^
                              ((concat_p1 (H1 x)) @ (concat_1p (H1 x))^) )
    : F == G.
  Proof.
    simple refine (KP_ind _ _ _ _).
    - exact H1.
    - intros x y p. etransitivity.
      rapply @transport_paths_FlFr.
      simple refine (concat_pp_p @ _).
      apply moveR_Vp. by symmetry.
    - admit.
  Defined.

  Lemma transport_KP `{f: forall (z: Z), A z -> B z} `(p: z1 = z2 :> Z) (x: A z1)
    : transport (λ z, KP' (f z)) p (kp x) = kp (p # x).
      by destruct p.
  Defined.
End KP.

Definition KP_equiv_fun `{ua: Univalence} {A B C:Type}
           (f: A -> B)
           (g: C -> B)
           (α: A -> C)
           (e: g o α = f)
  : KP' f -> KP' g.
Proof.
  simple refine (KP_rec _ _ _ _).
  intro a; apply kp. exact (α a).
  intros a b p; cbn.
  apply kp_eq. exact (ap10 e a @ p @ (ap10 e b)^).
  intro a; cbn.
  path_via (kp_eq (f:=g) (α a) (α a) 1).
  apply ap.
  simple refine ((concat_p1 _ @@ 1) @ _).
  apply concat_pV.
  apply kp_eq2.
Defined.

Definition isequiv_KP_equiv_fun_path `{ua: Univalence} {A B C:Type}
           (f: A -> B)
           (g: C -> B)
           (α: A = C)
           (e: g o (equiv_path _ _ α) = f)
  : IsEquiv (KP_equiv_fun f g (equiv_path _ _ α) e).
Proof.
  destruct α. cbn. destruct e.
  unfold KP_equiv_fun; cbn.
  assert ((KP_rec (KP' g) (λ a : A, kp a)
        (λ (a b : A) (p : g a = g b), kp_eq a b ((1 @ p) @ 1))
        (λ a : A, 1 @ kp_eq2 a)) == idmap).
  { simple refine (KP_eta _ _ _ _ _).
    intro; reflexivity.
    intros a b p; cbn.
    simple refine (concat_1p _ @ _ @ (concat_p1 _)^).
    simple refine (_ @ (KP_rec_beta_kp_eq _ _ _ _ _ _ _)^).
    simple refine (ap_idmap _ @ _).
    apply ap.
    simple refine (_ @ (concat_p1 _)^). exact (concat_1p _)^.
    intro a; cbn.
    rewrite transport_paths_FlFr. cbn.
    rewrite ap_V. rewrite inv_V.
    rewrite concat_ap_pFq.
    rewrite concat_ap_Fpq.
    apply moveR_pV.
    match goal with
    |[|- _ = ((?P1 @ _) @ ?P2) @ ?P3] =>
     rewrite (concat_p1 P1);
       rewrite (concat_pp_p (r:= P3))
    end.
    match goal with
    |[|- _ = _ @ (whiskerR ?hh 1 @ _) ]
     => pose (rew := whiskerR_p1 hh);
       rewrite concat_pp_p in rew;
       apply moveL_Vp in rew;
       rewrite rew; clear rew
    end.
    rewrite inv_V.
    apply moveR_Mp.
    match goal with
    |[|- _ = ?P1 @ ((whiskerL 1 ?hh) @ ?P3)] =>
     rewrite (concat_p_pp (p := P1));
       pose (rew := whiskerL_1p hh);
       apply moveL_pV in rew;
       rewrite rew; clear rew
    end. cbn.
    repeat rewrite concat_p1; repeat rewrite concat_1p.
    rewrite <- (apD (λ U, ap_idmap U) (kp_eq2 a)^).
    rewrite transport_paths_FlFr. cbn. rewrite ap_V. rewrite inv_V.
    rewrite concat_p1.
    repeat rewrite concat_pp_p. apply whiskerL.
    repeat rewrite ap_V.
    rewrite <- (ap02_is_ap (KP_rec (KP' g) (λ a0 : A, kp a0)
            (λ (a0 b : A) (p : g a0 = g b), kp_eq a0 b ((1 @ p) @ 1))
            (λ a0 : A, 1 @ kp_eq2 a0))).
    rewrite KP_rec_beta_kp_eq2.
    rewrite ap_idmap. rewrite (concat_1p (kp_eq2 a)).
    rewrite inv_pp.  reflexivity. }
  simple refine (isequiv_homotopic idmap  _).
  exact (λ x, (X x)^).
Defined.

Definition isequiv_KP_equiv_fun `{ua: Univalence} {A B C:Type}
           (f: A -> B)
           (g: C -> B)
           (α: A <~> C)
           (e: g o α = f)
  : IsEquiv (KP_equiv_fun f g α e).
Proof.
  assert ((KP_equiv_fun f g α e) = (KP_equiv_fun f g (equiv_path _ _ (path_universe_uncurried α)) ((ap (λ u, g o u) (ap (@equiv_fun A C) (equiv_path_path_universe_uncurried α))) @ e))).
  { cbn.
    pose (ap (@equiv_fun A C) (equiv_path_path_universe_uncurried α)).
    pose (apD (λ U, KP_equiv_fun f g U) p^).
    cbn in p0. rewrite <- p0. cbn.
    rewrite transport_arrow. rewrite transport_const.
    apply ap.
    cbn. rewrite transport_paths_Fl. apply moveL_Vp.
    rewrite inv_V. reflexivity. }
  simple refine (isequiv_homotopic _ (λ x, ap10 X^ x)).
  exact (isequiv_KP_equiv_fun_path f g ((path_universe_uncurried α)) (ap (λ (u : A → C) (x : A), g (u x))
                                                                       (ap (equiv_fun (B:=C)) (equiv_path_path_universe_uncurried α)) @ e)).
Qed.

Section T.
  Definition T (X: Type) := KP' (λ x: X, tt).

  Definition α {X} : X -> T X := kp.

  Definition β {X} : forall x y: X, forall p:tt=tt, α x = α y
    := λ x y p, kp_eq (f:=(λ x: X, tt)) x y p.

  Definition Ɣ {X} : forall x: X, β x x 1 = 1
    := λ x, kp_eq2 x.

  Definition functoriality_T `(f: X -> Y) : T X -> T Y.
  Proof.
    simple refine (KP_rec _ (α o f) _ _).
    - intros a b p; apply β. exact ((1 @ p) @ 1).
    - intros a; exact (1 @ Ɣ (f a)).
  Defined.

  Definition functoriality_equiv_T `{ua: Univalence} `(f: X <~> Y) : T X <~> T Y.
  Proof.
    exists (functoriality_T f).
    exact (isequiv_KP_equiv_fun (λ x: X, tt) (λ x: Y, tt) f 1).
  Defined.
End T.

Section KP_Sigma.

  Lemma KP_pr1_to_sigma (Y:Type) (A:Y -> Type)
  : KP' (λ x : ∃ y : Y, A y, pr1 x) <~> (∃ y : Y, KP' (λ _ : A y, tt)).
  Proof.
    transparent assert (F : ((∃ y : Y, KP' (λ _ : A y, tt)) -> KP' (λ x : ∃ y : Y, A y, pr1 x))).
  { intros [y a].
    simple refine (KP_rec _ _ _ _ a).
    intro x. apply kp.
    exists y.
    exact x.
    intros u v p; cbn.
    apply kp_eq; cbn. reflexivity.
    intro u; cbn.
    match goal with |[|- @kp_eq _ _ ?ff ?aa ?bb _ = _] => exact (kp_eq2 (f:=ff) aa)  end. }

  transparent assert (G: (KP' (λ x : ∃ y : Y, A y, pr1 x) -> (∃ y : Y, KP' (λ _ : A y, tt)))).
  { simple refine (KP_rec _ _ _ _).
    intros x. exists x.1. apply kp. exact x.2.
    intros a b p. cbn.
    simple refine (path_sigma' _ _ _).
    exact p.
    cbn.
    simple refine (transport_KP (f:=(λ y a, tt)) p (pr2 a) @ _).
    apply kp_eq. apply path_ishprop.
    intro a; cbn.
    match goal with
    |[|- ?XX _ = _] => path_via (XX 1)
    end.
    apply ap.

    simple refine (concat_1p _ @ _).
    cbn.
    match goal with
    |[|- kp_eq _ _ ?pp = _] => assert (r: 1 = pp) by apply path_ishprop
    end.
    destruct r.
    match goal with |[|- @kp_eq _ _ ?ff ?aa ?bb _ = _] => exact (kp_eq2 (f:=ff) aa) end. }
  simple refine (equiv_adjointify G F _ _).
  - intros [y x]. revert x.
    unfold F; clear F; unfold G; clear G.
    simple refine (KP_eta _ _ _ _ _).
    intro x; reflexivity.
    intros a b p; cbn.
    simple refine (concat_1p _ @ _ @ (concat_p1 _)^).
    match goal with
    |[|- _ = ap (λ x, ?gg (?ff x)) ?pp] =>
     simple refine (_ @ (ap_compose ff gg pp)^)
    end.
    simple refine (ap_existT _ _ _ _ _ @ _).
    simple refine (_ @ (ap02 _ (KP_rec_beta_kp_eq _ _ _ _ _ _ _)^)).
    simple refine (_ @ (KP_rec_beta_kp_eq _ _ _ _ _ _ _)^).
    cbn.
    apply ap. symmetry; simple refine (concat_1p _ @ _).
    apply ap; apply path_ishprop.

    intro a; cbn. rewrite transport_paths_FlFr.
    rewrite concat_ap_Fpq. rewrite concat_ap_pFq.
    simpl. rewrite (concat_p1 (whiskerL 1 (ap (ap (exist (λ y0 : Y, KP' (λ _ : A y0, tt)) y)) (kp_eq2 a)^))^).
    repeat rewrite ap_V.
    match goal with
    |[|- _ = _ @ whiskerR (ap (ap (λ x, ?gg (?ff x))) ?pp)^ 1]
     => rewrite <- (ap02_is_ap (λ x, gg (ff x)) pp);
       rewrite (ap02_compose ff gg pp)
    end.
    apply moveL_Vp.
    rewrite whiskerR_RV.
    rewrite whiskerR_pp. rewrite inv_pp.
    match goal with
    |[|- _ = (whiskerR (ap02 _ (ap02 (KP_rec ?P1 ?P2 ?P3 ?P4) (kp_eq2 a)) @ _) _)^ @ _]
     => rewrite (KP_rec_beta_kp_eq2 P1 P2 P3 P4 a)
    end.
    rewrite whiskerR_pp. rewrite inv_pp.
    do 3 rewrite <- whiskerR_RV. rewrite inv_V.
    unfold ap_compose at 2. cbn.
    match goal with
    |[|- _ = (_ @ ?P) @ _] => rewrite (concat_1p P)
    end.
    repeat rewrite concat_p_pp.

    match goal with
    |[|- ((((((whiskerL ?P1 ?P2 @ concat_1p _) @_)@_)@_)@_)@_)@_= _] =>
     pose (rew:=  whiskerL_1p P2); rewrite concat_pp_p in rew
    end.
    apply moveL_Vp in rew. rewrite rew; clear rew. cbn.
    rewrite (concat_1p (ap (ap (exist (λ y0 : Y, KP' (λ _ : A y0, tt)) y)) (kp_eq2 a))^).
    apply moveR_pV.
    match goal with
    |[|- _ = (?P1 @ whiskerR ?hh 1) @ _] =>
     rewrite (concat_pp_p (p:=P1));
       pose (rew := whiskerR_p1 hh); rewrite concat_pp_p in rew;
       apply moveL_Vp in rew
    end.
    rewrite rew; clear rew; cbn. rewrite inv_V.
    match goal with
    |[|- _ = whiskerR ?hh _ @ (?P1 @ _)]
     => rewrite (concat_p_pp (q:=P1));
       pose (rew := whiskerR_p1 hh); rewrite concat_pp_p in rew;
       apply moveL_Vp in rew
    end.
    rewrite rew; clear rew; cbn.
    match goal with
    |[|- _ = (1 @ ?P1) @ _] =>
     rewrite (concat_1p P1)
    end.
    apply whiskerR.
    rewrite ap02_pp. rewrite inv_pp. rewrite ap02_V. apply whiskerR.
    match goal with
    |[|- _ = (ap02 (KP_rec ?P1 ?P2 ?P3 ?P4) (kp_eq2 ?P5))^]
     => rewrite (KP_rec_beta_kp_eq2 P1 P2 P3 P4 P5)
    end. cbn.
    rewrite inv_pp. apply whiskerR.
    match goal with
    |[|- _ = (ap _ (_ @ match path_ishprop ?bar ?foo in (_ = y0) return _ with |_=>_ end) @ _)^] =>
     assert (r: 1 = foo) by apply path_ishprop; destruct r;
     assert (r: 1 = path_ishprop bar 1) by apply path_ishprop; destruct r
    end. cbn.
    repeat rewrite ap_pp. repeat rewrite inv_pp. cbn.
    hott_simpl. apply whiskerR.
    rewrite concat_1p.
    apply moveR_Vp.
    pose (p := kp_eq2 (f:=λ _:A y, tt) a).
    pose (r := apD (λ U, ap_existT (λ y0 : Y, KP' (λ _ : A y0, tt)) y (kp a) (kp a) U) p^).
    cbn in r; rewrite <- r; clear r; unfold p; clear p.
    rewrite transport_paths_FlFr.
    rewrite concat_p1. rewrite ap_V. rewrite inv_V. rewrite ap_V. reflexivity.
  - simple refine (KP_eta _ _ _ _ _).
    intro x. reflexivity.
    intros [a x] [b y] p; cbn.
    simple refine (concat_1p _ @ (ap_idmap _) @ _ @ (concat_p1 _)^).

    simple refine (_ @ (ap_compose G F (kp_eq (a;x) (b;y) p))^).
    unfold G.
    simple refine (_ @ (ap02 F (KP_rec_beta_kp_eq _ _ _ _ _ _ _)^)).
    cbn in *.
    destruct p. cbn. unfold F; clear F.
    pose (F := λ x:Y, λ x':KP' (λ _ : A x, tt),
                         KP_rec (KP' (λ x0 : ∃ y : Y, A y, pr1 x0))
                               (λ x0 : A x, kp (x; x0))
                               (λ (u v : A (x)) (_ : tt = tt),
                                kp_eq (x; u) (x; v) (idpath x))
                               (λ u : A (x), kp_eq2 (x; u))
                               x').
    simple refine (_ @ (ap_path_sigma_1 F a (1 @ kp_eq x y (path_ishprop tt tt)))^).
    simple refine ((KP_rec_beta_kp_eq _ _ _ _ _ _ _)^ @ _).
    exact (path_ishprop tt tt).
    simple refine (_ @ (ap_pp _ _ _)^). cbn.
    simple refine (concat_1p _)^.

    intros [a x]; cbn.
    match goal with |[|- ?XX = _] => set (X := XX) end.
    rewrite transport_paths_FlFr.
    rewrite concat_ap_Fpq. rewrite concat_ap_pFq.
    do 2 rewrite ap_V.
    rewrite <- (ap02_is_ap (λ x : KP' (λ x : ∃ y : Y, A y, pr1 x), F (G x))).
    rewrite ap02_compose.
    cbn.
    rewrite whiskerR_RV. rewrite whiskerR_pp.
    match goal with
    |[|- _ = (?P1 @ ?P2) @ (?P3 @ ?P4)^ ] =>
     rewrite (concat_p1 P1)
    end.
    apply moveL_Vp. rewrite inv_pp. apply moveL_pV.
    match goal with
    |[|- _ = (whiskerR (?P1 @ _) _)^ ]
     => rewrite (concat_p1 P1)
    end.
    rewrite <- whiskerR_RV.
    unfold X; clear X.
    match goal with
    |[|- (whiskerL _ ?hh @ (((?P @ _) @ _) @ _)) @ _ = _] =>
     repeat rewrite (concat_pp_p (p:=P));
       rewrite (concat_p_pp (q := P));
       pose (rew := whiskerL_1p hh); rewrite concat_pp_p in rew;
       apply moveL_Vp in rew
    end.
    rewrite rew; clear rew. cbn.
    match goal with
    |[|- ( _ @ (_ @ ?P)) @ (whiskerR ?hh _) = _] =>
     repeat rewrite (concat_p_pp (r := P));
       rewrite (concat_pp_p (q:=P));
       pose (rew := whiskerR_p1 hh);
       apply moveL_pV in rew
    end.
    rewrite rew; clear rew.
    repeat rewrite concat_p_pp. apply moveR_pV.
    match goal with
    |[|- _ = whiskerR ?hh _ @ _] =>
     pose (rew := whiskerR_p1 hh);
       rewrite concat_pp_p in rew;
       apply moveL_Vp in rew
    end.
    rewrite rew; clear rew; cbn.
    rewrite (concat_1p (ap02 F (ap02 G (kp_eq2 (a;x))))^).
    unfold G at 23.
    rewrite KP_rec_beta_kp_eq2.
    rewrite <- ap02_V. rewrite inv_pp.
    rewrite ap02_pp.
    cbn. rewrite concat_pV_p. apply whiskerR.
    rewrite inv_pp. rewrite ap02_pp. cbn.
    match goal with |[|- _ = 1 @ ?P] => rewrite (concat_1p P) end.
    assert (r: 1 = (path_ishprop tt tt)) by apply path_ishprop.
    destruct r.
    match goal with
    |[|- (((((( 1 @ ?P) @ ?Q) @ _) @ _)@ _)@ _) = _] =>
     rewrite (concat_1p P)
    end.

    pose (p :=apD (λ U, ap_idmap U) (kp_eq2 (f:=(λ x : ∃ y : Y, A y, pr1 x)) (a;x))^). cbn in p.
    rewrite <- p; clear p.
    rewrite transport_paths_FlFr. cbn. repeat rewrite ap_V. repeat rewrite inv_V.
    match goal with
    |[|- ((((?P @ ((?Q @ 1) @ _)) @_)@_)@_)@_ = _] =>
     rewrite (concat_p1 Q)
    end.
    rewrite concat_V_pp.

    pose (p := apD (λ U, (ap_path_sigma_1
      (λ (x0 : Y) (x' : KP' (λ _ : A x0, tt)),
       KP_rec (KP' (λ x1 : ∃ y : Y, A y, pr1 x1)) (λ x1 : A x0, kp (x0; x1))
         (λ (u v : A x0) (_ : tt = tt), kp_eq (x0; u) (x0; v) (idpath x0))
         (λ u : A x0, kp_eq2 (x0; u)) x') a (1 @ U))) (kp_eq2 x)^). cbn in p.
    rewrite <- p; clear p.
    rewrite transport_paths_FlFr.
    cbn.
    pose (p := apD (λ U, ap_pp
       (λ x' : KP' (λ _ : A a, tt),
        KP_rec (KP' (λ x0 : ∃ y : Y, A y, pr1 x0)) (λ x0 : A a, kp (a; x0))
          (λ (u v : A a) (_ : tt = tt), kp_eq (a; u) (a; v) (idpath a))
          (λ u : A a, kp_eq2 (a; u)) x') 1 U) (kp_eq2 x)^); cbn in p.
    rewrite <- p; clear p.
    rewrite transport_paths_FlFr.
    repeat rewrite inv_pp. repeat rewrite ap_V.
    repeat rewrite inv_V.
    cbn.
    match goal with
    |[|- _ @ (_ @ (1 @ ?P)) = _ ] => rewrite (concat_1p P)
    end.
    rewrite ap_idmap. repeat rewrite concat_pp_p. do 3  apply moveR_Vp.
    match goal with
    |[|- _ @ (1 @ ?P) = _] => rewrite (concat_1p P)
    end.
    rewrite concat_V_pp.
    rewrite concat_ap_pFq.
    apply moveR_pV. repeat rewrite concat_pp_p. apply moveL_Mp.
    match goal with
    |[|- _ @ whiskerL 1 ?hh = _]
     => pose (rew := whiskerL_1p hh);
       (* rewrite concat_pp_p in rew *)
       apply moveL_pV in rew
    end.
    rewrite rew; clear rew.
    rewrite <- (ap02_is_ap (λ x' : KP' (λ _ : A a, tt),
         KP_rec (KP' (λ x0 : ∃ y : Y, A y, pr1 x0)) (λ x0 : A a, kp (a; x0))
           (λ (u v : A a) (_ : tt = tt), kp_eq (a; u) (a; v) (idpath a))
           (λ u : A a, kp_eq2 (a; u)) x')).
    rewrite (KP_rec_beta_kp_eq2 (KP' (λ x0 : ∃ y : Y, A y, pr1 x0)) (λ x0 : A a, kp (a; x0))
        (λ (u v : A a) (_ : tt = tt), kp_eq (a; u) (a; v) (idpath a))
        (λ u : A a, kp_eq2 (a; u))).
    repeat rewrite concat_pp_p.
    do 2 apply whiskerL.
    rewrite ap02_is_ap.
    cbn.
    rewrite ap_V. apply moveL_Vp. rewrite concat_p1.
    match goal with
    |[|- ap _ (ap _ (_ @ match ?foo in (_=y) return _ with |_ => _ end)) = _]
     => assert (r: 1 = foo) by apply path_ishprop; destruct r; cbn
    end.
    match goal with
    |[|- _ = ap (λ x0, ?ff ((path_sigma' (λ x1 : Y, KP' (λ _ : A x1, tt)) 1 (1 @ x0)))) ?pp ] =>
     rewrite (ap_compose (λ x0, (path_sigma' (λ x1 : Y, KP' (λ _ : A x1, tt)) 1 (1 @ x0))) ff pp)
    end. cbn.
    apply ap.
    rewrite ap_pp.
    pose (p := apD (λ U,  ap (path_sigma' (λ y : Y, KP' (λ _ : A y, tt)) 1)
                             (concat_1p U)) (kp_eq2 x)^); cbn in p.
    rewrite <- p; clear p.
    rewrite transport_paths_FlFr. simpl. rewrite ap_V. rewrite inv_V.
    rewrite concat_p1. rewrite ap_V.
    rewrite concat_pV_p. reflexivity.
Defined.

  Context `(f: A -> B).

  Lemma KP_slicing_fun : KP' f -> {y: B & T (hfiber f y)}.
  Proof.
    simple refine (KP_rec _ _ _ _).
    - intro a. exact (f a ; (kp (a; 1))).
    - intros a b p; cbn.
      simple refine (path_sigma' _ p _).
      etransitivity. simple refine (transport_KP _ _).
        by apply kp_eq.
    - intros a; cbn. exact (ap _ (ap _ (kp_eq2 _))).
  Defined.

  Lemma T_to_sigmahfiber `{ua: Univalence}
    : KP' f <~> {y:B & KP' (λ _:hfiber f y, tt)}.
  Proof.
    etransitivity; [idtac | exact (KP_pr1_to_sigma B (hfiber f))].
    simple refine (BuildEquiv (isequiv_KP_equiv_fun _ _ _ _)).
    simple refine (equiv_adjointify _ _ _ _).
    intro x; exact (f x; (x;1)).
    intros [y [x p]]. exact x.
    intros [y [x p]].
    destruct p; reflexivity.
    intro x; reflexivity.
    reflexivity.
  Defined.

  Lemma KP_slicing_isequiv `{ua: Univalence} : IsEquiv KP_slicing_fun.
  Proof.
    simple refine (isequiv_homotopic (T_to_sigmahfiber) (H := equiv_isequiv (T_to_sigmahfiber)) _).
    intro x; cbn. unfold KP_slicing_fun.
  revert x.
  simple refine (KP_eta _ _ _ _ _).
  - intro; reflexivity.
  - intros a b p; cbn.
    simple refine (concat_1p _ @ _ @ (concat_p1 _)^).
    unfold KP_equiv_fun.
    match goal with
    |[|- _ = ap (λ x, ?gg (?ff x)) ?pp ] =>
     simple refine (_ @ (ap_compose ff gg pp)^);
       simple refine (_ @ (ap02 gg (KP_rec_beta_kp_eq _ _ _ _ _ _ _)^))
    end.
    simple refine (_ @ (KP_rec_beta_kp_eq _ _ _ _ _ _ _)^).
    simple refine ((KP_rec_beta_kp_eq _ _ _ _ a b p) @ _).
    cbn.
    etransitivity.
    2: exact (apD (λ U, path_sigma' (λ y : B, KP' (λ _ : hfiber f y, tt))
                                    U
                                    (transport_KP U (a; 1) @
                                                 kp_eq (transport (λ y : B, hfiber f y) U (a; 1))
                                                 (b; 1) (path_ishprop tt tt))) (concat_p1 (1 @ p) @ (concat_1p p))^).
    simple refine (_ @ (transport_const _ _)^).
    apply ap.
    apply ap. apply ap. apply path_ishprop.
  - intro a; cbn.
    rewrite transport_paths_FlFr. cbn.
    do 2 rewrite concat_ap_pFq.
    match goal with
    |[|- (?P1 @ ?P2) @ ?P3 = (?P4^ @ _) @ _] =>
     rewrite (concat_p1 P4^);
       apply moveL_Vp;
       rewrite (concat_pp_p (p := P1));
       rewrite (concat_p_pp (p := P4))
    end.
    match goal with
    |[|- (whiskerL 1 ?hh @ _) @ _ = _] =>
     pose (rew := whiskerL_1p hh);
       rewrite concat_pp_p in rew; apply moveL_Vp in rew;
       rewrite rew; clear rew
    end.
    rewrite inv_V.
    rewrite concat_ap_Fpq.
    match goal with
    |[|- ?P1 @ (?P2 @ ?P3^) = _ ]
     => rewrite (concat_p_pp (r := P3^));
       apply moveR_pV
    end.
    match goal with
    |[|- _ = whiskerR ?hh 1 @ _] =>
     pose (rew := whiskerR_p1 hh);
       rewrite concat_pp_p in rew; apply moveL_Vp in rew;
       rewrite rew; clear rew
    end.
    rewrite inv_V.
    match goal with
    |[|- _ = _ @ ap (ap (λ x, ?gg (?ff x))) ?pp] =>
     rewrite <- (ap02_is_ap (gg o ff) pp);
       rewrite (ap02_compose ff gg pp)
    end. cbn.
    match goal with
    |[|- ?P1 @ (?P2 @ ?P3^) = 1 @ (1 @ (?P5 @ ?P6))]
     => rewrite (concat_1p (1 @ (P5 @ P6)));
       rewrite (concat_1p (P5 @ P6));
       rewrite (concat_p_pp (r:=P3^))
    end.
    apply whiskerR.
    unfold KP_equiv_fun.
    repeat rewrite ap02_V.
    rewrite (KP_rec_beta_kp_eq2). cbn.
    rewrite (ap02_pp). rewrite inv_pp.
    match goal with
    |[|- (_ @ ?P1) @ (?P2 @ ?P3) = _] =>
     rewrite (concat_p_pp (r:=P3));
       rewrite (concat_1p P1)
    end.
    apply whiskerR.
    rewrite ap02_pp. cbn.
    match goal with |[|- _ = (_ @ ?P)^] => rewrite (concat_1p P) end.
    rewrite <- ap02_is_ap.
    match goal with
    |[|- _ = (ap02 (KP_rec ?P1 ?P2 ?P3 ?P4) (kp_eq2 ?aa))^] =>
     rewrite (KP_rec_beta_kp_eq2 P1 P2 P3 P4 aa)
    end.
    rewrite inv_pp. cbn.
    rewrite concat_p_pp. apply whiskerR.
    match goal with
    |[|- ap (ap ?ff) ?pp @ _ = _] =>
     rewrite <- (ap02_is_ap ff pp);
       rewrite ap02_V
    end.
    rewrite (KP_rec_beta_kp_eq2).
    rewrite inv_pp. rewrite concat_p_pp. rewrite concat_pV_p.
    repeat rewrite <- ap02_is_ap.
    do 3 rewrite concat_p1.
    rewrite <- (ap02_V (exist (λ y : B, KP' (λ _ : hfiber f y, tt)) (f a)) (ap (concat 1) (kp_eq2 (a; 1)))).
    match goal with
    |[|- ap02 ?ff ?pp @ ap02 _ ?qq = (ap02 ?hh ?rr)^] =>
     rewrite <- (ap02_pp ff pp qq);
       rewrite <- (ap02_V hh rr)
    end.
    apply ap.
    assert (r: 1 = path_ishprop tt tt) by apply path_ishprop.
    destruct r; cbn.
    assert (r: 1 = path_ishprop (idpath tt) 1) by apply path_ishprop.
    destruct r.
    rewrite ap_idmap. cbn. rewrite concat_p1. rewrite inv_pp.
    pose (p:= apD (λ U, (concat_1p U)^) (kp_eq2 (f:=(λ _ : hfiber f (f a), tt)) (a;1))^).
    cbn in p; rewrite <- p; clear p.
    rewrite transport_paths_FlFr.
    rewrite ap_idmap. cbn.
    rewrite concat_p1. rewrite inv_V. rewrite concat_V_pp.
    rewrite ap_V.
    reflexivity.
  Defined.

  Definition KP_slicing `{ua : Univalence} : KP' f <~> {y: B & T (hfiber f y)}
    := BuildEquiv KP_slicing_isequiv.

  Definition KP_lift : KP' f -> B.
  Proof.
    simple refine (KP_rec _ f _ _); intros; try reflexivity.
  Defined.

  Definition KP_lift_slicing `{ua : Univalence} : KP_lift == pr1 o KP_slicing.
  Proof.
    simple refine (KP_eta _ _ _ _ _).
    - intro; reflexivity.
    - intros x y p. cbn.
      simple refine (concat_1p _ @ _ @ (concat_p1 _)^).
      unfold KP_slicing, KP_slicing_fun; cbn.
      unfold KP_lift.
      symmetry.
      simple refine (KP_rec_beta_kp_eq _ _ _ _ _ _ _ @ _).
      simple refine (_ @ (ap_compose KP_slicing pr1 (kp_eq _ _ p))^).
      simple refine (_ @ (ap (ap pr1) _)).
      3: exact (KP_rec_beta_kp_eq _ _ _ _ _ _ _)^.
      exact (pr1_path_sigma _ _)^.
    - intro a. cbn beta.
      Opaque pr1_path_sigma.
      rewrite transport_paths_FlFr.
      rewrite concat_ap_pFq. rewrite concat_ap_Fpq.
      cbn.
      apply moveR_pV.
      repeat rewrite concat_pp_p.
      pose (rew := whiskerR_p1 (ap (ap KP_lift) (kp_eq2 a)^)).
      rewrite concat_pp_p in rew.
      apply moveL_Vp in rew.
      rewrite rew; clear rew.
      apply moveL_Vp.
      repeat rewrite concat_p_pp.
      pose (rew := whiskerL_1p (ap
         (ap (λ x : KP' f, let (proj1_sig, _) := KP_slicing_fun x in proj1_sig))
         (kp_eq2 a)^)).
      rewrite concat_pp_p in rew.
      apply moveL_Vp in rew.
      rewrite rew; clear rew.
      repeat rewrite inv_V.
      repeat rewrite ap_V.
      repeat rewrite <- ap02_is_ap.
      unfold KP_lift.
      rewrite KP_rec_beta_kp_eq2. cbn.
      rewrite <- (ap02_is_ap (λ x : KP' f, let (proj1_sig, _) := KP_slicing_fun x in proj1_sig)).
      unfold KP_slicing_fun. cbn.
      rewrite (ap02_compose _ pr1).
      rewrite KP_rec_beta_kp_eq2.
      rewrite <- (ap02_is_ap pr1).
      repeat rewrite concat_1p. repeat rewrite concat_p1.
      repeat rewrite inv_pp. repeat rewrite inv_V.
      repeat rewrite concat_p_pp.
      simple refine (_ @ concat_1p _). apply whiskerR.
      cbn.
      unfold transport_KP.
      match goal with
      |[|- _ @ ?foo (?bar @ _) = _] =>
       pose (p := apD (λ U: kp (a;1) = kp (a;1), foo (bar @ U)) (kp_eq2 (f:=(λ _ : hfiber f (f a), tt)) (a;1))^)
      end.
      cbn in p; rewrite <- p; clear p.
      rewrite transport_paths_Fl.
      Transparent pr1_path_sigma. cbn.
      rewrite concat_p1. rewrite ap_V. rewrite inv_V.
      rewrite concat_pV_p. rewrite ap02_pp. rewrite inv_pp.
      rewrite concat_pV_p.
      apply moveR_Vp. rewrite <- (ap_compose (concat 1) (path_sigma' (λ y : B, KP' (λ _ : hfiber f y, tt)) 1) (kp_eq2 (a; 1))); cbn.
      rewrite ap02_is_ap.
      unfold pr1_path.
      rewrite (ap_compose (λ x, (path_sigma (λ y : B, KP' (λ _ : hfiber f y, tt))
                                            (f a; kp (a; 1)) (f a; kp (a; 1)) 1 (1 @ x))) (ap pr1)).
      rewrite concat_p1.
      apply ap. cbn. reflexivity.
  Defined.

  Definition KP_lift_hfiber `{ua: Univalence} (y: B) : hfiber KP_lift y <~> T (hfiber f y).
  Proof.
    etransitivity. simple refine (equiv_functor_sigma' KP_slicing _).
    exact (λ x, KP_lift (KP_slicing^-1 x) = y).
    intros a. simple refine (equiv_paths _). f_ap.
    exact (eissect KP_slicing a)^.
    etransitivity. simple refine (equiv_functor_sigma_id _).
    exact (λ x, x.1 = y).
    intros w. simple refine (equiv_paths _).
    etransitivity. exact (KP_lift_slicing _).
    apply ap. apply eisretr.
    { simple refine (equiv_adjointify _ _ _ _).
      - intros w. simple refine (w.2 # w.1.2).
      - intros w. simple refine (exist _ _ _). exists y. exact w. reflexivity.
      - intros w. reflexivity.
      - intros [w e]. destruct e. reflexivity. }
  Defined.
End KP_Sigma.

Section KP_lift_equiv.
  Definition hfiber_KP_lift_equiv `{ua: Univalence} `(f: X -> Y) {A: Type} (y: Y) (e: (hfiber f y) <~>  A)
  : {e': hfiber (KP_lift f) y <~> T A & α o e == e' o (λ x: hfiber f y, (kp x.1; x.2))}.
  Proof.
    simple refine (exist _ _ _).
    - etransitivity.
      rapply @KP_lift_hfiber.
      exact (functoriality_equiv_T e).
    - intros [x p].
      pose (Te := functoriality_T e).
      match goal with
      | |- _ = ?AA => change (Te (α (x; p)) = AA)
      end.
      match goal with
      | |- ?AA = _ (transitive_equiv _ _ _ ?ee1 ?ee4) ?xx =>
        change (AA = ee4 (ee1 (kp x; p))); set (e4 := equiv_fun ee4)
      end. unfold KP_lift_hfiber.
      match goal with
      | |- ?AA = e4 (_ (transitive_equiv _ _ _ ?ee1 ?ee2) ?xx) =>
        change (AA = e4 (ee2 (ee1 (kp x; p)))); set (e1 := equiv_fun ee1)
      end.
      match goal with
      | |- ?AA = e4 (_ (transitive_equiv _ _ _ ?ee2 ?ee3) ?xx) =>
        change (AA = e4 (ee3 (ee2 xx))); set (e2 := equiv_fun ee2);
        set (e3 := equiv_fun ee3)
      end.
      set (sl := equiv_fun (KP_slicing f)) in *.
      symmetry. etransitivity (e4 (e3 (e2 (sl (kp x); _)))). f_ap.
      subst e1.
      etransitivity (e4 (e3 (sl (kp x); _))). f_ap.
      subst e2.
      match goal with
      | |- e4 (e3 (_; ?ppp)) = _ => set (pp := ppp)
      end.
      etransitivity (e4 (transport (λ y0 : Y, T (hfiber f y0)) pp (sl (kp x)).2)).
      reflexivity.
      clear e3. f_ap. clear e4.
      change (transport (λ y, T (hfiber f y)) pp (α (x; 1)) = α (x;p)).
      etransitivity.
      rapply @transport_KP.
      apply β. reflexivity.
  Defined.
End KP_lift_equiv.


From stdpp Require Import gmap.
From iris.program_logic Require Import hoare.
From iris_hashtable Require Import hashtable_invariants util.
From iris.heap_lang Require Import proofmode.

Section tests.

  Context Σ `{!heapG Σ}.

  Definition equalf : val := λ: "x" "y", "x" = "y".
  Definition idf : val := λ: "x", "x".
  Definition as_natkey v :=
    match v with
    | #(LitInt z) => if decide (0 ≤ z)%Z then Some (Z.to_nat z) else None
    | _ => None
    end.

  Lemma as_natkey_inj v n:
    as_natkey v = Some n -> v = #n.
  Proof.
    destruct v as [|l| | |]; try discriminate. destruct l as [n'| | |] ; try discriminate.
    simpl ; case_decide ; last discriminate. intro Heq. injection Heq.
    intro. rewrite -(Z2Nat.id n') ; last done.  by do 3 f_equal.
  Qed.

  Lemma equalf_spec k1 k2 v1 v2 :
    as_natkey v1 = Some k1 ->
    as_natkey v2 = Some k2 ->
    (WP (equalf v1) v2 {{ u, ⌜u = #(bool_decide (k1 = k2))⌝ }})%I.
  Proof.
    intros H1 H2. do 2 wp_lam. wp_op ; intro ; iPureIntro.
    - simplify_eq. by rewrite bool_decide_true.
    - rewrite bool_decide_false ; first reflexivity.
      apply as_natkey_inj in H1. apply as_natkey_inj in H2. simplify_eq. by intros <-.
  Qed.

  Lemma idhash_spec k v :
    as_natkey v = Some k ->
    (WP idf v {{ u, ⌜u = #(id k)⌝ }})%I.
  Proof.
    intro. wp_lam. eauto using as_natkey_inj.
  Qed.

  Instance nat_hashable : Hashable Σ nat id :=
    {| equal_spec := equalf_spec ; hash_spec := idhash_spec  |}.

  Parameter (table : table Σ nat id (gmap nat)).

  Lemma insert_nat_spec m (state :table_state table) t (n : nat) x :
    {{{ table_in_state table m state t }}}
      table_insert table t #n x
      {{{ (state' : table_state table), RET #();
          table_in_state table (insert_val m n x) state' t}}}.
  Proof.
    rewrite (_:Lit n = LitV n) ; last reflexivity.
    eapply table_insert_spec.
    simpl. rewrite decide_True. by rewrite Nat2Z.id. lia.
  Qed.

  Definition test_1 : expr :=
    let: "t" := table_create table #10 in
    table_insert table "t" #1 #1 ;;
    table_insert table "t" #2 #2 ;;
    let: "a" := ref #0 in
    table_fold table (λ: "k" "x" <>, "a" <- !"a" + "x") "t" #() ;;
    !"a".

  Definition int_table : gmap nat _ -> iProp Σ :=
    table_inv (fun _ v => ∃ (j : Z), ⌜v = #j⌝)%I.

  Fixpoint int_val_sum (seq : list (val)) :=
    match seq with
    | [] => Some 0%Z
    |  #(LitInt x) :: seq => z ← int_val_sum seq ; Some (z + x)%Z
    | _ => None
    end.

  Instance int_val_sum_proper : Proper ((≡ₚ) ==> (=)) int_val_sum.
  Proof.
    induction 1 as[|x|x y l| ???? Heq1 ? Heq2]. reflexivity.
    - destruct x as [|x| | |]; try reflexivity. destruct x as [n| | |] ; try reflexivity.
      simpl. by f_equal.
    - destruct x as [|x| | |]; cycle 1 ; [destruct x|..] ; (destruct y as [|y| | |] ; cycle 1 ; [destruct y|..]) ; try reflexivity.
      simpl. destruct (int_val_sum l) ; last reflexivity. simpl. f_equal. lia.
    - by rewrite Heq1.
  Qed.

  Lemma int_val_sum_app seq seq':
    int_val_sum (seq ++ seq') = x ← int_val_sum seq ; y ← int_val_sum seq' ; Some (x + y)%Z.
  Proof.
    induction seq as [|x seq IH]. simpl. rewrite /mbind /option_bind ; by destruct (int_val_sum seq').
    destruct x as [|x | | |] ; try reflexivity. destruct x ; try reflexivity.
    simpl. rewrite IH. destruct (int_val_sum seq) ; last reflexivity.
    destruct (int_val_sum seq') ; [simpl; f_equal; lia | reflexivity].
  Qed.

  Lemma int_val_sum_complete m seq:
    complete m seq -> int_val_sum (seq.*2) = int_val_sum ((all_elements m).*2).
  Proof.
    intros HCom. pose proof (complete_all_elements _ _ HCom) as [seq' [Hseq <-]].
    f_equal.
    rewrite (_:forall l1 l2, iff (l1 = l2) (Forall2 (=) (l2) (l1))).
    apply Forall2_fmap, (Forall2_impl _ _ _ _ Hseq).
    by intros [??] [??] [<- _]. intros l1.
    induction l1 as [|?? IH]. split.
    intros <-. done.
    intros. symmetry. by eapply Forall2_nil_inv_r.
    split. intros <-. decompose_Forall.
    intros HF. inversion HF. simplify_eq. f_equal. by apply IH.
  Qed.

  Definition test_1_inv m l (seq: list (val * val)) (_ : val) : iProp Σ :=
    ( int_table m ∗
      ∃ i, ⌜int_val_sum (seq.*2) = Some i⌝
          ∗ l ↦ #i)%I.

  Lemma test_1_inner M l k x seq a:
    permitted M (seq ++ [(k,x)]) ->
    {{test_1_inv M l seq a}} (λ: "k" "x" <>, #l <- !#l + "x") k x a {{v, test_1_inv M l (seq ++ [(k,x)]) v }}.
  Proof.
    intros [M' HRem]. iIntros "!# [HiTab Hl]".
    iDestruct "Hl" as (i) "[% Hl]".
    iDestruct (table_inv_removal _ _ _ _ HRem with "HiTab") as "[Hseq HiTabM']".
    rewrite big_sepL_app big_sepL_singleton.
    iDestruct "Hseq" as "[Hseq Hkx]".
    iDestruct "Hkx" as (k') "[% Hx]".
    iDestruct "Hx" as (j) "%". iSimplifyEq.
    do 3 wp_lam. wp_load. wp_op. wp_store.
    iSplit.
    - iApply table_inv_removal ; first exact HRem.
      iSplit.
      + iApply big_sepL_app. iSplit. iAssumption.
        iApply big_sepL_singleton. iExists k'. eauto.
      + iAssumption.
    - iExists (i + j). iFrame. iPureIntro.
      rewrite fmap_app int_val_sum_app. by rewrite (_:int_val_sum (seq.*2) = Some i).
  Qed.

  Lemma test_1_spec :
    WP test_1 {{v, ⌜v = #3⌝}}%I.
  Proof.
    assert (H10: (10 > 0)%nat) ; first lia.
    unfold test_1. rewrite (_:10%Z = 10%nat) ; last lia.
    wp_bind (table_create _ _). iApply wp_wand.
    iApply (table_create_spec _ _ H10).
    iIntros (t) "HTable".
    iDestruct "HTable" as (state) "HTable".
    wp_lam. wp_apply (insert_nat_spec _ _ _ 1 #1 with "HTable").
    iIntros (state') "HTable".
    wp_lam. wp_apply (insert_nat_spec _ _ _ 2 #2 with "HTable").
    iIntros (state'') "HTable".
    wp_lam. wp_alloc a as "Ha". wp_lam.
    wp_apply (table_fold_spec _ _ _ _ (LamV "k" (λ: "x" <>, #a <- ! #a + "x")) _ #() with "[Ha $HTable]").
    iSplitR.
    - iIntros (????) "!# [% ?]". by iApply (test_1_inner _ _ ).
    - iSplit ; last eauto. rewrite /int_table. do 2 rewrite -table_inv_insert.
      do 2 ( iSplit ; last eauto). iApply (table_inv_empty (Key:=nat)).
    - iIntros (v seq) "[% [HTable [_ HInv]]]".
      assert (HCom:complete (insert_val (insert_val ∅ 1%nat #1) 2%nat #2) seq) by assumption.
      iDestruct "HInv" as (n) "[% Ha]".
      assert (HSum: int_val_sum (seq.*2) = Some n) ; first assumption.
      wp_seq. wp_load. iPureIntro. f_equal.
      rewrite (int_val_sum_complete _ _ HCom) in HSum. vm_compute in HSum.
      injection HSum. by intros <-.
  Qed.

   Definition test_2 (t : val) : expr :=
    table_fold table (λ: <> "x" "a", "x" + "a")%V t #0.

  Lemma test_2_spec m (state : table_state table) t :
    {{{int_table m ∗ table_in_state table m state t}}}
      test_2 t
      {{{x, RET #x ; ⌜int_val_sum ((all_elements m).*2) = Some x⌝ ∗
                     int_table m ∗ table_in_state table m state t}}}.
  Proof.
    iIntros (Φ) "[HiTab HTable] HΦ".
    wp_apply (table_fold_spec _ _ _ (fun seq v => int_table m ∗ ∃ x, ⌜int_val_sum (seq.*2) = Some x⌝ ∗ ⌜v = #x⌝)%I with "[$HTable HiTab]"). iSplit ; last eauto.
    - iIntros (k x seq a) "!# [% [HiTab HInv]]". rename_last HRem. destruct HRem as [? HRem].
      iDestruct "HInv" as (x') "[% %]". iSimplifyEq.
      iDestruct (table_inv_removal _ _ _ _ HRem with "HiTab") as "[Hseq HiTabM']".
      rewrite big_sepL_app big_sepL_singleton.
      iDestruct "Hseq" as "[Hseq Hkx]".
      iDestruct "Hkx" as (k') "[% Hx]".
      iDestruct "Hx" as (j) "%". iSimplifyEq.
      do 3 wp_lam. wp_op. iSplit.
      + iApply table_inv_removal. done.
        iSplit. iApply big_sepL_app. iSplit.
        iAssumption. iApply big_sepL_singleton. eauto.
        iAssumption.
      + iExists (j + x')%Z. iSplit ; iPureIntro.
        rewrite fmap_app int_val_sum_app (_:int_val_sum (seq.*2) = Some x') /=; last assumption.
        f_equal. lia. reflexivity.
    - iIntros (v seq) "[% [HTable [HiTab HInv]]]".
      iDestruct "HInv" as (x) "[% %]". iSimplifyEq. iApply "HΦ".
      iSplit. by erewrite <-int_val_sum_complete. iFrame.
  Qed.

  Definition test_3 : expr :=
    let: "t" := table_create table #10 in
    table_insert table "t" #1 #1 ;;
    table_insert table "t" #2 #2 ;;
    let: "a" := ref #0 in
    (rec: "loop" "c" :=
       match: "c" #() with
         NONE     => !"a"
       | SOME "p" => "a" <- !"a" + Snd (Fst "p") ;;
                     "loop" (Snd "p")
       end) (table_cascade table "t").

  Lemma test_3_spec :
    WP test_3 {{ v, ⌜v = #3⌝ }}%I.
  Proof.
    assert (H10: (10 > 0)%nat) ; first lia.
    unfold test_3. rewrite (_:10%Z = 10%nat) ; last lia.
    wp_bind (table_create _ _). iApply wp_wand.
    iApply (table_create_spec _ _  H10).
    iIntros (t) "HTable".
    iDestruct "HTable" as (state) "HTable".
    wp_lam. wp_apply (insert_nat_spec _ _ _ 1 #1 with "HTable").
    iIntros (state') "HTable".
    wp_lam. wp_apply (insert_nat_spec _ _ _ 2 #2 with "HTable").
    iIntros (state'') "HTable".
    wp_lam. wp_alloc a. wp_lam.
    wp_apply (table_cascade_spec with "HTable").
    set (m:=(insert_val (insert_val ∅ 1%nat #1) 2%nat #2)).
    iIntros (c) "[#HCas HTable]".
    assert (Hloop: forall x seq,
               int_val_sum (seq.*2) = Some x ->
               is_cascade table m c seq state'' t -∗
               table_in_state table m state'' t -∗ a ↦ #x -∗
               WP (rec: "loop" "c" := match: "c" #() with
                                        NONE => ! #a
                                      | SOME "p" => #a <- ! #a + Snd (Fst "p") ;;
                                                    "loop" (Snd "p")
                                      end) c
               {{ v, ⌜exists x seq, complete m seq /\ int_val_sum (seq.*2) = Some x /\ v = #x⌝ }}).
    { iIntros (i seq Hi) "#HCas".
      iLöb as "IH" forall (c i seq Hi) "HCas".
      iIntros "HTable Ha".
      wp_rec. wp_apply (is_cascade_spec with "[$HCas $HTable]").
      iIntros (? k x c') "[HTable [[% %]|[% [% #HCas']]]]".
      - simplify_eq. wp_match. wp_load. iPureIntro. exists i, seq. auto.
      - simplify_eq. simpl.
        assert (permitted m (seq ++ [(k, x)])) as [m' HRem] ; first assumption.
        iAssert (int_table m) as "HiTab".
        { do 2 (iApply (table_inv_insert (Key:=nat)) ; iSplit ; last eauto).
          iApply (table_inv_empty (Key:=nat)).
        }
        iDestruct (table_inv_removal _ _ _ _ HRem with "HiTab") as "[Hseq _]".
        rewrite big_sepL_app big_sepL_singleton.
        iDestruct "Hseq" as "[_ Hkx]". iDestruct "Hkx" as (?) "[_ Hx]".
        iDestruct "Hx" as (j) "%". iSimplifyEq.
        wp_match. wp_load. do 2 wp_proj. wp_op. wp_store. wp_proj.
        iApply ("IH" $! _ (i + j)%Z (seq ++ [(k , #j)]) with "[] [] HTable Ha").
        by rewrite fmap_app int_val_sum_app Hi /=. done.
    }
    iApply (wp_wand with "[-]"). by iApply (Hloop 0 [] with "HCas HTable [-]").
    iIntros (v) "%".
    assert (∃ (x : Z) (seq : list (val * val)), complete m seq ∧ int_val_sum (seq.*2) = Some x ∧ v = #x) as [x [seq [Hcom [Hsum ->]]]] ; first assumption.
    rewrite (int_val_sum_complete _ _ Hcom) in Hsum.
    vm_compute in Hsum. injection Hsum as <-. done.
  Qed.

End tests.


From HoTT Require Import Basics.
From HoTT Require Import Categories.
From HoTT Require Import Spaces.Finite.
Require Import Misc.
Require Import Limits.Graph.

Set Universe Polymorphism.
Set Implicit Arguments.
Generalizable All Variables.
Set Asymmetric Patterns.

Local Open Scope morphism_scope.

Section FunctorLimits.
  Local Open Scope object_scope.

  Context {C D : PreCategory}.
  Context (F : Functor C D).

  Context {Size Arrows : Type}.
  Context (G : Graph C Size Arrows).

  Definition graph_of : Graph D Size Arrows :=
    {| gr_vertex := fun s => (F _0 (gr_vertex G s))%object
    ;  gr_edges  := fun a => {| proj1 := gr_src G a
                          ;  proj2 :=
                            {| proj1 := gr_dst G a
                            ;  proj2 := F _1 (gr_edge G a)
                            |}
                          |}
    |}.

  Definition cone_of (cn : Cone G) : Cone graph_of.
  Proof.
    srapply mkCone; [ exact (F _0 (cn_top cn))%object
                    | intro side; exact (F _1 (cn_side cn side)) | ].
    intro a; unfold graph_of; unfold gr_edge; unfold gr_src; unfold gr_dst; simpl.
    rewrite <- composition_of. f_ap. apply cn_comm.
  Defined.

  Definition FunctorPreservesLimit :=
    forall(L : Limit G), exists(L' : Limit graph_of), lim_cone L' = cone_of (lim_cone L).
End FunctorLimits.

From HoTT Require Import Categories.Adjoint.UnitCounit.
From HoTT Require Import Categories.Adjoint.UnitCounitCoercions.
From HoTT Require Import Categories.Adjoint.Hom.
From HoTT Require Import Categories.Adjoint.HomCoercions.

Section AdjunctsPreserveLimits.
  Local Open Scope object_scope.

  Context {C D : PreCategory}.
  Context {F : Functor C D}.
  Context {G : Functor D C}.
  Context (Adj : F -| G).

  Context {Size Arrows : Type}.
  Context (Gr : Graph D Size Arrows).

  Definition counitCone :
    Cone (graph_of F (graph_of G Gr)) -> Cone Gr.
  Proof.
    intro cn. srapply mkCone; [ exact (cn_top cn) | | ].
    - intro n. refine (_ o cn_side cn n). simpl. apply (counit Adj).
    - intro a; simpl. rewrite <- (cn_comm cn a); simpl.
      unfold gr_edge; unfold gr_src; unfold gr_dst. simpl.
      unfold gr_edge; unfold gr_src; unfold gr_dst. simpl.
      repeat rewrite <- associativity. f_ap.
      rewrite (commutes (counit Adj)); simpl. reflexivity.
  Defined.

  Definition ConeMorphismLToR (cnC : Cone (graph_of G Gr)) (cnD : Cone Gr) :
    ConeMorphism (counitCone (cone_of F cnC)) cnD ->
    ConeMorphism cnC (cone_of G cnD).
  Proof.
    intro mph. srapply mkCnMph.
    - simpl. refine (equiv_fun (equiv_hom_set_adjunction _ _ _)^-1 _).
      + apply adjunction_unit__of__adjunction_unit_counit. exact Adj.
      + apply (cnmph_mph mph).
    - intro n. simpl. rewrite <- associativity. rewrite <- composition_of.
      rewrite (cnmph_comm mph); simpl. rewrite composition_of.
      rewrite associativity. rewrite <- (commutes (unit Adj)). simpl.
      rewrite <- associativity. rewrite unit_counit_equation_2.
      apply left_identity.
  Defined.
  Definition ConeMorphismRToL (cnC : Cone (graph_of G Gr)) (cnD : Cone Gr) :
    ConeMorphism cnC (cone_of G cnD) ->
    ConeMorphism (counitCone (cone_of F cnC)) cnD.
  Proof.
    intro mph. srapply mkCnMph.
    - simpl. refine (equiv_fun (equiv_hom_set_adjunction _ _ _) _).
      + apply adjunction_unit__of__adjunction_unit_counit. exact Adj.
      + apply (cnmph_mph mph).
    - intro n. simpl. rewrite <- (cnmph_comm mph).
      rewrite composition_of. repeat rewrite <- associativity. f_ap.
      simpl. rewrite (commutes (counit Adj)). reflexivity.
  Defined.
  Definition RightAdjunctLimit (L : Limit Gr) : Limit (graph_of G Gr).
  Proof.
    srapply mkLim; [ exact (cone_of G (lim_cone L)) | | ].
    - intro c. apply ConeMorphismLToR. apply lim_ex.
    - intros c m1 m2.
      pose (mph1 := ConeMorphismRToL m1). pose (mph2 := ConeMorphismRToL m2).
      refine (equiv_fun (equiv_ap'
                           (equiv_hom_set_adjunction
                              (adjunction_unit__of__adjunction_unit_counit Adj) _ _) _ _)^-1 _).
      exact (lim_uniq L _ mph1 mph2).
  Defined.

  Theorem RightAdjunctPreservesLimit : FunctorPreservesLimit G Gr.
  Proof.
    intro L. exists(RightAdjunctLimit L). reflexivity.
  Qed.
End AdjunctsPreserveLimits.

From HoTT Require Import Categories.HomFunctor.

Section HomPreservesLimits.
  Context `{Funext}.
  Context {C : PreCategory}.
  Context (A : object (C^op)).
  Context {Size Arrows : Type}.
  Context (Gr : Graph C Size Arrows).
  Definition F := covariant_hom_functor C A.

  Definition mkCConeFromSetCone (cn : Cone (graph_of F Gr)) :
    cn_top cn -> Cone Gr.
  Proof.
    intro x. srapply mkCone; [ exact A | | ].
    - intro n. exact (cn_side cn n x).
    - intro a; simpl. rewrite <- (cn_comm cn a); simpl.
      unfold gr_edge; simpl. rewrite right_identity. reflexivity.
  Defined.
  Definition CHomConeMorphism (cSet : Cone (graph_of F Gr)) (cC : Cone Gr) (x : cn_top cSet) :
    ConeMorphism cSet (cone_of F cC) ->
    ConeMorphism (mkCConeFromSetCone cSet x) cC.
  Proof.
    intro mph. srapply mkCnMph.
    - simpl. exact (cnmph_mph mph x).
    - intro n; simpl. rewrite <- (cnmph_comm mph n); simpl.
      rewrite right_identity. reflexivity.
  Defined.

  Definition CovariantHomLimit (L : Limit Gr) : Limit (graph_of F Gr).
  Proof.
    srapply mkLim; [ exact (cone_of F (lim_cone L)) | | ].
    - intro c. pose (c' := mkCConeFromSetCone c). srapply mkCnMph.
      + simpl. intro x. pose (mph := lim_ex L (c' x)). exact (cnmph_mph mph).
      + intro n. apply path_forall. intro x; simpl. rewrite right_identity.
        pose (mph := lim_ex L (c' x)). rewrite (cnmph_comm mph). reflexivity.
    - intros c m1 m2. apply path_forall. intro x.
      pose (mph1 := CHomConeMorphism x m1). pose (mph2 := CHomConeMorphism x m2).
      pose (Huniq := lim_uniq L _ mph1 mph2). exact Huniq.
  Defined.

  Theorem CovariantHomPreservesLimit : FunctorPreservesLimit (covariant_hom_functor C A) Gr.
  Proof.
    intro L. exists(CovariantHomLimit L). reflexivity.
  Qed.
End HomPreservesLimits.


Set Implicit Arguments.
Unset Strict Implicit.

Require Import Arith ZArith.
Require Import Vector.
Require Import String.
Require Import JMeq.

Require Import Integers.

(* semantics definitions *)
Definition bit := bool.
Definition bvec64 := Int64.int.
Definition bvec64_nil : bvec64 := Int64.zero.

Definition iN (N : nat) := Fin.t N.

Definition arr (N:nat) T := Vector.t T N.
Definition arr_upd N T (a:arr N T) (i:iN N) (v:T) : arr N T :=
  Vector.replace a i v.

Program Fixpoint vector_keymap
         A B n (v : Vector.t A n)
         (f : Fin.t n -> A -> B)
  : Vector.t B n
  := match v with
     | Vector.nil => (@Vector.nil B)
     | Vector.cons a n' v' =>
       Vector.cons
         _ (f (@Fin.of_nat_lt (n-1) n _) a) n'
         (vector_keymap v' (fun i => f (@Fin.of_nat_lt (proj1_sig (Fin.to_nat i)) n _)))
     end.
Next Obligation.
  induction n';
    omega.
Defined.
Next Obligation.
  clear -i.
  induction i; simpl; try omega.
  destruct (Fin.to_nat i); simpl in *; omega.
Defined.

Definition arr_init (N : nat) T (t0 : T) (f : iN N -> T) : arr N T :=
  vector_keymap (Vector.const t0 N) (fun ix _ => f ix).

(* syntax *)
Inductive ty : Type :=
  | TVec64 : ty
  | TArr : nat -> ty -> ty.

Ltac solve_ty_dec :=
  repeat
    (match goal with
     | |- context[not (?X)] =>
       let x := fresh "x" in intros x
     | [ H : (?T /\ ?T1) |- _ ] => destruct H
     | [ H : ~ (?T /\ ?T1) |- _ ] => destruct H
     | [ H : ?P |- _] =>
       match type of P with
         Prop =>
         repeat
           (match goal with
            | H : ?T |- False =>
              solve [ subst; inversion H]
            end)
       end
     end); auto.

Program Fixpoint ty_dec (t1 t2 : ty) : {t1=t2} + {t1<>t2} :=
  match t1, t2 with
  | TVec64, TVec64 => left eq_refl
  | TArr n1 t1', TArr n2 t2' =>
    match eq_nat_dec n1 n2 with
    | left pf =>
      match ty_dec t1' t2' with
      | left pf' => left _
      | right pf' => right _
      end
    | right pf => right _
    end
  | _, _ => right _
  end.
Next Obligation.
  solve_ty_dec.
  destruct t1; destruct t2; solve_ty_dec.
  exfalso.
  eapply H; eauto.
Defined.
Next Obligation.
  split; intros;
  solve_ty_dec.
Defined.
Next Obligation.
  split; intros;
  solve_ty_dec.
Defined.

Fixpoint interp_ty (t : ty) : Type :=
  match t with
  | TVec64 => bvec64
  | TArr N t2 => arr N (interp_ty t2)
  end.

Program Definition cast_interp_ty (t1 t2 : ty) (pf : t1=t2)
  : interp_ty t1 -> interp_ty t2 :=
  fun v =>
    match pf as p in eq _ _ return interp_ty t2 with
    | eq_refl => v
    end.

Definition id (t : ty) := string.

Inductive binop : Type :=
    OAnd (* and *)| OXor (* xor *) | OAdd (* 64-bit unsigned add *)
    | OShr (* shift right *) | OShru (* unsigned shift right *)
    | OShl (* shift left *)
    | OOr
    | OSub
    | OEq
    (*| OIfOr (* or used in IF statements *)*)
    .

Inductive exp : ty -> Type :=
  | EVal : bvec64 -> exp TVec64
  | EVar : forall t, id t -> exp t
  | EDeref : forall N t (i : iN N), id (TArr N t) -> exp t
  | EBinop : binop -> exp TVec64 -> exp TVec64 -> exp TVec64
  | ENot : exp TVec64 -> exp TVec64.

Inductive assign_kind : Type := Blocking | Nonblocking.

Inductive stmt : Type :=
  | SAssign : forall (k:assign_kind) (x:id TVec64) (e:exp TVec64), stmt
  | SUpdate : forall (k:assign_kind) N (x:id (TArr N TVec64)) (i:iN N) (e:exp TVec64), stmt
  | SSeq : stmt -> stmt -> stmt
  | SITE : exp TVec64 -> stmt -> stmt -> stmt
  | SIter : forall (lo hi:nat) (f:iN hi -> stmt), stmt
  | SSkip : stmt.

Definition r_rotate (x : exp TVec64)
           (i : exp TVec64) : exp TVec64 :=
  EBinop OOr (EBinop OShr x i) (EBinop OShl x
           (EBinop OSub (EVal (Int64.repr 64)) i)).

Inductive var_kind : Type := Local | Input | Output.

Inductive prog : Type :=
  | VDecl : var_kind -> id TVec64 -> bvec64 -> prog -> prog
  | ADecl : var_kind -> forall N t, id (TArr N t) -> prog -> prog
  | PStmt : stmt -> prog
  | PSeq : prog -> prog -> prog
  | PDone : prog.

Definition binop_interp (op : binop) (v1 v2 : bvec64) : bvec64 :=
  match op with
  | OAnd => Int64.and v1 v2
  | OSub => Int64.sub v1 v2
  | OXor => Int64.xor v1 v2
  | OAdd => Int64.add v1 v2
  | OShr => Int64.shr v1 v2
  | OShru => Int64.shru v1 v2
  | OShl => Int64.shl v1 v2
  | OOr =>  Int64.or v1 v2
  | OEq => v1 (*BOGUS*)
  (*| OIfOr => v1 (*BOGUS*)*)
  end.

Section state.
  Definition state := forall t, id t -> interp_ty t.

  Definition upd t (x : id t) (v : interp_ty t) (s : state)
    : state :=
    fun t2 (y : id t2) =>
      match ty_dec t t2 as t' in sumbool _ _ with
      | left pf =>
        match string_dec x y with
        | left pf' => cast_interp_ty pf v
        | right pf' => s t2 y
        end
      | right _ => s t2 y
      end.
End state.

Section exp_interp.
  Variable s : forall t, id t -> interp_ty t.

 Fixpoint exp_interp (e : exp TVec64) : interp_ty TVec64 :=
    match e as e' in exp t return interp_ty t with
    | EVal v => v
    | EVar _ x => s x
    | EDeref _ _ i x => Vector.nth (s x) i
    | EBinop op e1 e2 =>
      let v1 := exp_interp e1 in
      let v2 := exp_interp e2 in
      binop_interp op v1 v2
    | ENot e' =>
      Int64.not (exp_interp e')
    end.

End exp_interp.

Definition enum_iN (N : nat) : list (iN N) :=
  Vector.to_list (vector_keymap (Vector.const O N) (fun ix _ => ix)).

Definition itern lo hi T (f : iN hi -> T -> T) (t : T) : T :=
  List.fold_right
    f t
    (List.filter (fun i => let n := proj1_sig (Fin.to_nat i) in
                           (lo <=? n))  (enum_iN hi)).

Fixpoint stmt_interp (s : state) (c : stmt) (cont : state -> state) : state :=
  match c with
  | SAssign k x e =>
    let v := exp_interp s e in
    cont (upd x v s)
  | SUpdate k t x i e =>
    let v := exp_interp s e in
    cont (upd x (arr_upd (s _ x) i v) s)
  | SSeq c1 c2 =>
    stmt_interp s c1 (fun s' => stmt_interp s' c2 cont)
  | SITE e c1 c2 =>
    let v := exp_interp s e in
    if Int64.eq v Int64.zero then stmt_interp s c2 cont
    else stmt_interp s c1 cont
  | SIter lo hi f =>
    cont (itern lo (fun i s => stmt_interp s (f i) (fun s' => s')) s)
  | SSkip => cont s
  end.

Program Fixpoint prog_interp (s : state) (p : prog) : state :=
  match p with
  | VDecl vk x v p' =>
    prog_interp (upd x v s) p'
  | ADecl vk N t x p' =>
    prog_interp s p'
  | PStmt c => stmt_interp s c (fun s => s)
  | PSeq p1 p2 =>
    prog_interp (prog_interp s p1) p2
  | PDone => s
  end.


(**

Syntax of the simply typed lambda calculus as a
multisorted signature.

Written by: Anders Mörtberg, 2021 (adapted from STLC.v)

version for simplified notion of HSS by Ralph Matthes (2022, 2023)
the file is identical to the homonymous file in the parent directory, except for importing files from the present directory

*)
Require Import UniMath.Foundations.PartD.
Require Import UniMath.Foundations.Sets.

Require Import UniMath.MoreFoundations.Tactics.

Require Import UniMath.Combinatorics.Lists.

Require Import UniMath.CategoryTheory.Core.Categories.
Require Import UniMath.CategoryTheory.Core.NaturalTransformations.
Require Import UniMath.CategoryTheory.Core.Functors.
Require Import UniMath.CategoryTheory.FunctorCategory.
Require Import UniMath.CategoryTheory.whiskering.
Require Import UniMath.CategoryTheory.limits.graphs.colimits.
Require Import UniMath.CategoryTheory.limits.binproducts.
Require Import UniMath.CategoryTheory.limits.products.
Require Import UniMath.CategoryTheory.limits.bincoproducts.
Require Import UniMath.CategoryTheory.limits.coproducts.
Require Import UniMath.CategoryTheory.limits.terminal.
Require Import UniMath.CategoryTheory.limits.initial.
Require Import UniMath.CategoryTheory.FunctorAlgebras.
Require Import UniMath.CategoryTheory.exponentials.
Require Import UniMath.CategoryTheory.Adjunctions.Core.
Require Import UniMath.CategoryTheory.Chains.All.
Require Import UniMath.CategoryTheory.Monads.Monads.
Require Import UniMath.CategoryTheory.categories.HSET.Core.
Require Import UniMath.CategoryTheory.categories.HSET.Colimits.
Require Import UniMath.CategoryTheory.categories.HSET.Limits.
Require Import UniMath.CategoryTheory.categories.HSET.Structures.
Require Import UniMath.CategoryTheory.categories.StandardCategories.
Require Import UniMath.CategoryTheory.Groupoids.

Require Import UniMath.SubstitutionSystems.Signatures.
Require Import UniMath.SubstitutionSystems.SumOfSignatures.
Require Import UniMath.SubstitutionSystems.BinProductOfSignatures.
Require Import UniMath.SubstitutionSystems.SimplifiedHSS.SubstitutionSystems.
Require Import UniMath.SubstitutionSystems.SimplifiedHSS.LiftingInitial_alt.
Require Import UniMath.SubstitutionSystems.SimplifiedHSS.MonadsFromSubstitutionSystems.
Require Import UniMath.SubstitutionSystems.SignatureExamples.
Require Import UniMath.SubstitutionSystems.MultiSorted_alt.
Require Import UniMath.SubstitutionSystems.SimplifiedHSS.MultiSortedMonadConstruction_alt.
Require Import UniMath.SubstitutionSystems.MonadsMultiSorted_alt.

Local Open Scope cat.

(** * The simply typed lambda calculus from a multisorted binding signature *)
Section Lam.

  Variable (sort : hSet) (arr : sort → sort → sort).

Local Lemma hsort : isofhlevel 3 sort.
Proof.
exact (isofhlevelssnset 1 sort (setproperty sort)).
Defined.

Let sortToSet : category := [path_pregroupoid sort hsort,HSET].

Local Lemma TerminalSortToSet : Terminal sortToSet.
Proof.
apply Terminal_functor_precat, TerminalHSET.
Defined.

Local Lemma BinCoprodSortToSet : BinCoproducts sortToSet.
Proof.
apply BinCoproducts_functor_precat, BinCoproductsHSET.
Defined.

Local Lemma BinProd : BinProducts [sortToSet,HSET].
Proof.
apply BinProducts_functor_precat, BinProductsHSET.
Defined.

(** Some notations *)
Local Infix "::" := (@cons _).
Local Notation "[]" := (@nil _) (at level 0, format "[]").
Local Notation "a + b" := (setcoprod a b) : set.
Local Notation "s ⇒ t" := (arr s t).
Local Notation "'Id'" := (functor_identity _).
Local Notation "a ⊕ b" := (BinCoproductObject (BinCoprodSortToSet a b)).
Local Notation "'1'" := (TerminalObject TerminalSortToSet).
Local Notation "F ⊗ G" := (BinProduct_of_functors BinProd F G).

Let sortToSet2 := [sortToSet,sortToSet].

Local Lemma BinCoprodSortToSet2 : BinCoproducts sortToSet2.
Proof.
apply BinCoproducts_functor_precat, BinCoprodSortToSet.
Defined.

(** The signature of the simply typed lambda calculus *)
Definition STLC_Sig : MultiSortedSig sort.
Proof.
use make_MultiSortedSig.
- apply ((sort × sort) + (sort × sort))%set.
- intros H; induction H as [st|st]; induction st as [s t].
  + exact ((([],,(s ⇒ t)) :: ([],,s) :: nil),,t).
  + exact (((cons s [],,t) :: []),,(s ⇒ t)).
Defined.

(** The signature with strength for the simply typed lambda calculus *)
Definition STLC_Signature : Signature sortToSet _ _ :=
  MultiSortedSigToSignatureSet sort hsort STLC_Sig.

Definition STLC_Functor : functor sortToSet2 sortToSet2 :=
  Id_H _ BinCoprodSortToSet STLC_Signature.

Lemma STLC_Functor_Initial : Initial (FunctorAlg STLC_Functor).
Proof.
apply SignatureInitialAlgebra.
- apply InitialHSET.
- apply ColimsHSET_of_shape.
- apply is_omega_cocont_MultiSortedSigToSignature.
  + apply ProductsHSET.
  + apply Exponentials_functor_HSET.
  + apply ColimsHSET_of_shape.
Defined.

Definition STLC_Monad : Monad sortToSet :=
  MultiSortedSigToMonadSet sort hsort STLC_Sig.

(** Extract the constructors of the STLC from the initial algebra *)
Definition STLC_M : sortToSet2 :=
  alg_carrier _ (InitialObject STLC_Functor_Initial).

(* The functor parts coincide *)
Lemma STLC_Monad_ok : STLC_M = pr1 (pr1 (pr1 STLC_Monad)).
Proof.
apply idpath.
Qed.

Let STLC_M_mor : sortToSet2⟦STLC_Functor STLC_M,STLC_M⟧ :=
  alg_map _ (InitialObject STLC_Functor_Initial).

Let STLC_M_alg : algebra_ob STLC_Functor :=
  InitialObject STLC_Functor_Initial.

(** The variables *)
Definition var_map : sortToSet2⟦Id,STLC_M⟧ :=
  BinCoproductIn1 (BinCoprodSortToSet2 _ _) · STLC_M_mor.

(** The source of the application constructor *)
Definition app_source (s t : sort) : functor sortToSet2 sortToSet2 :=
    (post_comp_functor (projSortToSet sort hsort (s ⇒ t)) ⊗ post_comp_functor (projSortToSet sort hsort s))
  ∙ (post_comp_functor (hat_functorSet sort hsort t)).

(** The application constructor *)
Definition app_map (s t : sort) : sortToSet2⟦app_source s t STLC_M,STLC_M⟧ :=
    CoproductIn _ _ (Coproducts_functor_precat _ _ _ _ (λ _, _)) (ii1 (s,,t))
  · BinCoproductIn2 (BinCoprodSortToSet2 _ _)
  · STLC_M_mor.

(** The source of the lambda constructor *)
Definition lam_source (s t : sort) : functor sortToSet2 sortToSet2 :=
    pre_comp_functor (sorted_option_functorSet sort hsort s)
  ∙ post_comp_functor (projSortToC sort hsort _ t)
  ∙ post_comp_functor (hat_functorSet sort hsort (s ⇒ t)).

Definition lam_map (s t : sort) : sortToSet2⟦lam_source s t STLC_M,STLC_M⟧ :=
    CoproductIn _ _ (Coproducts_functor_precat _ _ _ _ (λ _, _)) (ii2 (s,,t))
  · BinCoproductIn2 (BinCoprodSortToSet2 _ _)
  · STLC_M_mor.

Definition make_STLC_M_Algebra X (fvar : sortToSet2⟦Id,X⟧)
  (fapp : ∏ s t, sortToSet2⟦app_source s t X,X⟧)
  (flam : ∏ s t, sortToSet2⟦lam_source s t X,X⟧) :
    algebra_ob STLC_Functor.
Proof.
apply (tpair _ X), (BinCoproductArrow _ fvar), CoproductArrow; intros b.
induction b as [st|st]; induction st as [s t].
- exact (fapp s t).
- exact (flam s t).
Defined.

(** The recursor for the stlc *)
Definition foldr_map X (fvar : sortToSet2⟦Id,X⟧)
                       (fapp : ∏ s t, sortToSet2⟦app_source s t X,X⟧)
                       (flam : ∏ s t, sortToSet2⟦lam_source s t X,X⟧) :
                        algebra_mor _ STLC_M_alg (make_STLC_M_Algebra X fvar fapp flam) :=
  InitialArrow STLC_Functor_Initial (make_STLC_M_Algebra X fvar fapp flam).

(** The equation for variables *)
Lemma foldr_var X (fvar : sortToSet2⟦Id,X⟧)
  (fapp : ∏ s t, sortToSet2⟦app_source s t X,X⟧)
  (flam : ∏ s t, sortToSet2⟦lam_source s t X,X⟧) :
  var_map · foldr_map X fvar fapp flam = fvar.
Proof.
unfold var_map.
rewrite <- assoc, (algebra_mor_commutes _ _ _ (foldr_map _ _ _ _)), assoc.
etrans; [eapply cancel_postcomposition, BinCoproductOfArrowsIn1|].
rewrite id_left.
apply BinCoproductIn1Commutes.
Qed.

Lemma foldr_app X (fvar : sortToSet2⟦Id,X⟧)
  (fapp : ∏ s t, sortToSet2⟦app_source s t X,X⟧)
  (flam : ∏ s t, sortToSet2⟦lam_source s t X,X⟧)
  (s t : sort) :
  app_map s t · foldr_map X fvar fapp flam =
  # (pr1 (app_source s t)) (foldr_map X fvar fapp flam) · fapp s t.
Proof.
unfold app_map.
rewrite <- assoc.
etrans; [apply maponpaths, (algebra_mor_commutes _ _ _ (foldr_map X fvar fapp flam))|].
rewrite assoc.
etrans; [eapply cancel_postcomposition; rewrite <- assoc;
         apply maponpaths, BinCoproductOfArrowsIn2|].
rewrite <- !assoc.
etrans; [apply maponpaths, maponpaths, BinCoproductIn2Commutes|].
rewrite assoc.
etrans; [apply cancel_postcomposition; use (CoproductOfArrowsIn _ _ (Coproducts_functor_precat _ _ _ _ (λ _, _)))|].
rewrite <- assoc.
apply maponpaths.
exact (CoproductInCommutes _ _ _ _ _ _ (inl (s,,t))).
Qed.

Lemma foldr_lam X (fvar : sortToSet2⟦Id,X⟧)
  (fapp : ∏ s t, sortToSet2⟦app_source s t X,X⟧)
  (flam : ∏ s t, sortToSet2⟦lam_source s t X,X⟧)
  (s t : sort) :
  lam_map s t · foldr_map X fvar fapp flam =
  # (pr1 (lam_source s t)) (foldr_map X fvar fapp flam) · flam s t.
Proof.
unfold lam_map.
rewrite <- assoc.
etrans; [apply maponpaths, (algebra_mor_commutes _ _ _ (foldr_map X fvar fapp flam))|].
rewrite assoc.
etrans; [eapply cancel_postcomposition; rewrite <- assoc;
         apply maponpaths, BinCoproductOfArrowsIn2|].
rewrite <- !assoc.
etrans; [apply maponpaths, maponpaths, BinCoproductIn2Commutes|].
rewrite assoc.
etrans; [apply cancel_postcomposition; use (CoproductOfArrowsIn _ _ (Coproducts_functor_precat _ _ _ _ (λ _, _)))|].
rewrite <- assoc.
apply maponpaths.
exact (CoproductInCommutes _ _ _ _ _ _ (inr (s,,t))).
Qed.

(* Now substitution *)
Let STLC := STLC_Monad.

(* Parallel substitution *)
Definition psubst {X Y : sortToSet} (f : sortToSet⟦X, STLC Y ⟧) :
  sortToSet⟦ STLC (X ⊕ Y), STLC Y ⟧ := monadSubstGen_instantiated _ _ _ _ f.

(* Substitution of a single variable *)
Definition subst {X : sortToSet} (f : sortToSet⟦ 1, STLC X ⟧) :
  sortToSet⟦ STLC (1 ⊕ X), STLC X ⟧ := monadSubstGen_instantiated _ _ _ _ f.

Definition weak {X Y : sortToSet} : sortToSet⟦STLC Y,STLC (X ⊕ Y)⟧ :=
  mweak_instantiated sort hsort HSET BinCoproductsHSET.

Definition exch {X Y Z : sortToSet} : sortToSet⟦STLC (X ⊕ (Y ⊕ Z)), STLC (Y ⊕ (X ⊕ Z))⟧ :=
  mexch_instantiated sort hsort HSET BinCoproductsHSET.

Lemma psubst_interchange {X Y Z : sortToSet}
        (f : sortToSet⟦X,STLC (Y ⊕ Z)⟧) (g : sortToSet⟦Y, STLC Z⟧) :
        psubst f · psubst g = exch · psubst (g · weak) · psubst (f · psubst g).
Proof.
apply subst_interchange_law_gen_instantiated.
Qed.

Lemma subst_interchange {X : sortToSet}
        (f : sortToSet⟦1,STLC (1 ⊕ X)⟧) (g : sortToSet⟦1,STLC X⟧) :
  subst f · subst g = exch · subst (g · weak) · subst (f · subst g).
Proof.
apply subst_interchange_law_gen_instantiated.
Qed.

(* We could also unfold these as statements about sort-indexed sets, but
   this quickly gets very cumbersome: *)
(* Definition psubst {X Y : sort → hSet} (f : ∏ t, X t → STLC Y t) (t : sort) : *)
(*   STLC (λ t, (X t + Y t)%set) t → STLC Y t. *)
(* Proof. *)
(* intros u. *)
(* transparent assert (X' : (sortToSet)). *)
(* { use (functor_path_pregroupoid _); apply X. } *)
(* transparent assert (Y' : (sortToSet)). *)
(* { use (functor_path_pregroupoid _); apply Y. } *)
(* transparent assert (f' : (sortToSet⟦ X' , STLC_Monad Y' ⟧)). *)
(* { use nat_trans_functor_path_pregroupoid; apply homset_property; use f. } *)
(* use (pr1 (@monadSubstGen_instantiated sort SET BinCoproductsHSET STLC_Monad X' Y' f') t). *)

End Lam.


Require Export ind_potype pcfpo_mod monad_h_module orders.

Set Implicit Arguments.
Unset Strict Implicit.

Unset Automatic Introduction.

Section PCFPO_representation.

(** a lot of notation to make things readable *)

Notation "'TY'" := PCF.TY.
Notation "'IP'" := (IPO TY).
Notation "a '~>' b" := (PCF.arrow a b) (at level 69, right associativity).
(* Notation "a ~> b" := (Module_Hom a b) (at level 50).*)
(*Notation "a 'x' b" := (Prod_Mod _ a b) (at level 30).*)

Notation "a 'x' b" := (product (*C:= MOD _ _*) a b) (at level 30).

(* Notation "P ^ z" := (FIB_MOD _ z P). *)
Notation "P [ z ]" := (FIB_MOD _ z P) (at level 35).

(*Notation "P ''' s" := (DER_MOD _ _ s P ) (at level 25).*)
Notation "'d' P // s" := (DER_MOD _ _ s P) (at level 25).
(*Notation " '*' " := (Term (*C:=MOD _ _*)).*)
Notation "'*'" := (Term (C:=MOD _ PO)).
Notation "'Bool'" := PCF.Bool.
Notation "'Nat'" := PCF.Nat.

Notation "f 'X' g" := (product_mor (*C:= MOD _ _ *) _ f g)(at level 30).

Notation "'FM'" := (#(FIB_MOD _ _ )).
Notation "'FPB'":= (FIB_PB _ _ _ ).
Notation "'PRPB'":= (PROD_PB _ _ _ _ ).
Notation "'PBF'":= (PB_FIB _ _ _ ).
Notation "'PBM'":= (#(PB_MOD _ _ )).
Notation "'DM'":= (#(DER_MOD _ _ _ )).
Notation "'DPB'":= (DER_PB _ _ _ ).
Notation "y [* := z ]":= (IPsubstar z _ y)(at level 55).
(** a monad is a representation if it is accompagnied by
    - a lot of morphisms of modules
    - beta-reduction
*)

Class PCFPO_rep_struct (P : Monad IP) := {
  App : forall r s, (P[r~>s]) x (P[r]) ---> P[s];
  Abs : forall r s, (d P // r)[s] ---> P[r ~> s];
  Rec : forall t, P[t ~> t] ---> P[t];
  ttt :  * ---> P[Bool];
  fff :  * ---> P[Bool];
  nats: forall m:nat, * ---> P[Nat];
  succ: * ---> P[Nat ~> Nat];
  zero : * ---> P[Nat ~> Bool];
  condN: * ---> P[Bool ~> Nat ~> Nat ~> Nat];
  condB: * ---> P[Bool ~> Bool ~> Bool ~> Bool];
  Bottom: forall t, * ---> P[t];

  beta: forall r s V y z,
        App _ _ _ (Abs r s V y, z) <<  (*IPsubstar z _ y*) y[*:= z] ;

  propag_App1: forall r s V y y' z,
        y << y' -> App r s V (y,z) << App _ _ _ (y',z) ;

  propag_App2: forall r s V y z z',
        z << z' -> App r s V (y,z) << App _ _ _ (y,z') ;

  propag_Abs: forall r s V y y',
        y << y' -> Abs r s V y << Abs _ _ _ y' ;

  propag_Rec: forall s V y y',
        y << y' -> Rec s V y << Rec _ _ y'

(*
  these two are not necessary, since we are over PO

  beta_refl : forall r V (y : P V r), y << y ;

  beta_trans : forall r V (a b c : P V r),
          a << b -> b << c -> a << c
*)
}.

(** the type of representations *)

Record PCFPO_rep := {
  pcfpo_rep_monad :> Monad IP;
  pcfpo_rep_struct :> PCFPO_rep_struct pcfpo_rep_monad
}.

Existing Instance pcfpo_rep_struct.

(** morphisms of representations *)

Section PCFPO_rep_Hom.

Variables P R : PCFPO_rep.

Section Rep_Hom_Class.

Variable Sig : Monad_Hom P R.

Notation "'sig'":= (PbMod_ind_Hom Sig).

(** for the constants we need a special morphism of modules

    [* ---> (\Sigma * ) *]

    being the empty product  *)

Lemma id_Term_PB_struct:
   Module_Hom_struct (S:= Term (C:= MOD _ PO))  (T:=(PB_MOD Sig PO Term))
                 (fun r => id PO_TERM).
Proof.
  constructor.
  intros.
  rewrite idl.
  rewrite idr.
  apply hom_refl.
Qed.

Definition PBT : Term ---> (PB_MOD Sig PO Term) :=
      Build_Module_Hom id_Term_PB_struct.

(*
Notation "'FM'" := (!(FIB_MOD _ _ )).
Notation "'FPB'":= (FIB_PB _ _ _ ).
Notation "'PRPB'":= (PROD_PB _ _ _ _ ).
Notation "'PBF'":= (PB_FIB _ _ _ ).
Notation "'sig'":= (PbMod_ind_Hom Sig).
Notation "'PBM'":= (!(PB_MOD _ _ )).
Notation "'DM'":= (!(DER_MOD _ _ _ )).
Notation "'DPB'":= (DER_PB _ _ _ ).
*)

(** Sig : P -> R is a morphism of representations if it makes commute
    all these weird diagrams
*)

Class PCFPO_rep_Hom_struct := {
  App_Hom: forall r s,
        App r s ;; FM sig  ==
          (FM sig ;; FPB) X (FM sig ;; FPB);;
                   PRPB ;; PBM (App r s) ;; PBF        ;

  Abs_Hom: forall r s,
         Abs r s ;; FM sig ==
          FM (DM sig ;; DPB) ;; FPB ;; PBM (Abs r s) ;; PBF  ;

  Rec_Hom: forall t,
         Rec t ;; FM sig ==
            FM sig ;; FPB ;; PBM (Rec t) ;; PBF ;

  ttt_Hom: ttt ;; FM sig ==
          PBT ;; PBM ttt ;; PBF ;

  fff_Hom: fff ;; FM sig ==
          PBT ;; PBM fff ;; PBF ;

  nats_Hom : forall m,
         nats m ;; FM sig ==
            PBT ;; PBM (nats m) ;; PBF ;

  succ_Hom: succ ;; FM sig ==
          PBT ;; PBM succ ;; PBF  ;

  zero_Hom: zero ;; FM sig ==
          PBT ;; PBM zero ;; PBF  ;

  condN_Hom: condN ;; FM sig ==
          PBT ;; PBM condN ;; PBF  ;

  condB_Hom: condB ;; FM sig ==
          PBT ;; PBM condB ;; PBF ;

  Bottom_Hom: forall t,
          Bottom t ;; FM sig ==
          PBT ;; PBM (Bottom t) ;; PBF

}.

End Rep_Hom_Class.

(** the type of morphismes of representations P -> R *)

Record PCFPO_rep_Hom := {
  rep_Hom_monad :> Monad_Hom P R ;
  rep_Hom_monad_struct :> PCFPO_rep_Hom_struct rep_Hom_monad
}.

End PCFPO_rep_Hom.
(*Existing Instance MONAD_struct.*)
Existing Instance rep_Hom_monad_struct.

(** on our way to a category of representations:
    - an equality on morphisms of reps*)

(** two morphisms are equal if their monad homs are equal,
     proof of "equivalence relation" is more or less already done in
           Monad_Hom_equiv *)
(*
Definition eq_Rep (P R: PCFPO_rep) : relation (PCFPO_rep_Hom P R) :=
       fun a c => rep_Hom_monad a == rep_Hom_monad c.
*)

Lemma eq_Rep_equiv (P R: PCFPO_rep) :
     Equivalence (A:=PCFPO_rep_Hom P R)
          (fun a c => rep_Hom_monad a == rep_Hom_monad c).
Proof.
  intros P R.
  assert (H:= Monad_Hom_equiv P R).
  destruct H as [Hr Hs Ht].
  constructor;
  simpl in *.
  unfold Reflexive; intros. apply Hr.
  unfold Symmetric; intros r s; apply Hs.
  unfold Transitive; intros r s t; apply Ht.
Qed.

Definition eq_Rep_oid (P R : PCFPO_rep) := Build_Setoid (eq_Rep_equiv P R).

(** Identity Representation *)

Lemma Rep_id_struct (P : PCFPO_rep) :
         PCFPO_rep_Hom_struct (Monad_Hom_id P).
Proof.
  intro P;
  unfold Monad_Hom_id.
  simpl.
  constructor; intros;
  simpl;
(*Next Obligation.
Proof.*)
  simpl;
  intros;
  try rewrite <- surjective_pairing;
  auto.
Qed.

Definition Rep_id (P: PCFPO_rep) := Build_PCFPO_rep_Hom (Rep_id_struct P).

(** Composition of Representations *)
Section Rep_comp.
Variables P Q R: PCFPO_rep.
Variable M: PCFPO_rep_Hom P Q.
Variable N: PCFPO_rep_Hom Q R.

Program Instance Rep_comp_struct :
      PCFPO_rep_Hom_struct (Monad_Hom_comp M N).
Next Obligation.
Proof.
  set (HM:=App_Hom (Sig := M)).
  set (HN:=App_Hom (Sig := N)).
  simpl in *.
  rewrite HM.
  rewrite HN.
  auto.
Qed.
Next Obligation.
  set (HM:=Abs_Hom (Sig := M)).
  set (HN:=Abs_Hom (Sig := N)).
  simpl in *.

  rewrite HM.
  rewrite HN.
  auto.
Qed.
Next Obligation.
  set (HM:=Rec_Hom (Sig := M)).
  set (HN:=Rec_Hom (Sig := N)).
  simpl in *.

  rewrite HM.
  rewrite HN.
  auto.
Qed.
Next Obligation.
  set (HM:=ttt_Hom (Sig := M)).
  set (HN:=ttt_Hom (Sig := N)).
  simpl in *.

  rewrite HM.
  rewrite HN.
  auto.
Qed.
Next Obligation.
  set (HM:=fff_Hom (Sig := M)).
  set (HN:=fff_Hom (Sig := N)).
  simpl in *.

  rewrite HM.
  rewrite HN.
  auto.
Qed.
Next Obligation.
  set (HM:=nats_Hom (Sig := M)).
  set (HN:=nats_Hom (Sig := N)).
  simpl in *.

  rewrite HM.
  rewrite HN.
  auto.
Qed.
Next Obligation.
  set (HM:=succ_Hom (Sig := M)).
  set (HN:=succ_Hom (Sig := N)).
  simpl in *.

  rewrite HM.
  rewrite HN.
  auto.
Qed.
Next Obligation.
  set (HM:=zero_Hom (Sig := M)).
  set (HN:=zero_Hom (Sig := N)).
  simpl in *.

  rewrite HM.
  rewrite HN.
  auto.
Qed.
Next Obligation.
  set (HM:=condN_Hom (Sig := M)).
  set (HN:=condN_Hom (Sig := N)).
  simpl in *.

  rewrite HM.
  rewrite HN.
  auto.
Qed.
Next Obligation.
  set (HM:=condB_Hom (Sig := M)).
  set (HN:=condB_Hom (Sig := N)).
  simpl in *.

  rewrite HM.
  rewrite HN.
  auto.
Qed.
Next Obligation.
  set (HM:=Bottom_Hom (Sig := M)).
  set (HN:=Bottom_Hom (Sig := N)).
  simpl in *.

  rewrite HM.
  rewrite HN.
  auto.
Qed.

Definition Rep_comp := Build_PCFPO_rep_Hom Rep_comp_struct.

End Rep_comp.

(** category of representations *)

Program Instance PCFPO_REP_struct : Cat (fun a c => PCFPO_rep_Hom a c) := {
  mor_oid a c := eq_Rep_oid a c;
  id a := Rep_id a;
  comp P Q R f g := Rep_comp f g
}.
Next Obligation.
Proof.
  unfold Proper in *; do 2 red.
  simpl.
  intros.
  rewrite H0.
  rewrite H.
  auto.
Qed.

Definition PCFPO_REP := Build_Category PCFPO_REP_struct.

End PCFPO_representation.

Existing Instance pcfpo_rep_struct.


Require Import Coq.Relations.Relation_Operators.
Require Import Logic.GeneralLogic.KripkeModel.
Require Import Logic.SeparationLogic.Model.SeparationAlgebra.

Local Open Scope kripke_model.
Import KripkeModelNotation_Intuitionistic.

Inductive MetaState (state: Type): Type :=
  | Error: MetaState state
  | NonTerminating: MetaState state
  | Terminating: state -> MetaState state.

Arguments Error {_}.
Arguments NonTerminating {_}.
Arguments Terminating {_} _.

Module Type FORWARD.

Parameter forward: forall {state: Type} {state_R: Relation state},
  MetaState state -> MetaState state -> Prop.

End FORWARD.

Module Partial <: FORWARD.

Inductive forward'
          {state: Type}
          {state_R: Relation state}:
  MetaState state -> MetaState state -> Prop :=
| forward_Error:
    forward' Error Error
| forward_NonTerminating:
    forward' NonTerminating NonTerminating
| forward_Terminating_NonTerminating:
    forall x, forward' (Terminating x) NonTerminating
| forward_Terminating:
    forall x y, x <= y -> forward' (Terminating x) (Terminating y).

Definition forward {state: Type} {state_R: Relation state} := forward'.

End Partial.

Module Total <: FORWARD.

Inductive forward'
          {state: Type}
          {state_R: Relation state}:
  MetaState state -> MetaState state -> Prop :=
| forward_Error:
    forward' Error Error
| forward_NonTerminating:
    forward' NonTerminating NonTerminating
| forward_Terminating:
    forall x y, x <= y -> forward' (Terminating x) (Terminating y).

Definition forward {state: Type} {state_R: Relation state} := forward'.

End Total.

Lemma Total2Partial_forward {state: Type} {state_R: Relation state}: forall ms1 ms2,
  Total.forward ms1 ms2 -> Partial.forward ms1 ms2.
Proof.
  intros.
  inversion H; constructor; auto.
Qed.

Inductive lift_relation {state: Type} (R: state -> MetaState state -> Prop):
  MetaState state-> MetaState state -> Prop :=
| lift_relation_Error:
    lift_relation R Error Error
| lift_relation_NonTerminating:
    lift_relation R NonTerminating NonTerminating
| lift_relation_Terminating:
    forall s ms, R s ms -> lift_relation R (Terminating s) ms.

Definition lift_Krelation {state: Type} {state_R: Relation state}: MetaState state -> MetaState state -> Prop := Total.forward.

Inductive lift_join
          {state: Type}
          {J_state: Join state}:
  MetaState state -> MetaState state -> MetaState state -> Prop :=
| lift_join_Error1:
    forall mx my, lift_join Error mx my
| lift_join_Error2:
    forall mx my, lift_join mx Error my
| lift_join_NonTerminating1:
    forall x, lift_join NonTerminating (Terminating x) NonTerminating
| lift_join_NonTerminating2:
    forall x, lift_join (Terminating x) NonTerminating NonTerminating
| lift_join_NonTerminating3:
    lift_join NonTerminating NonTerminating NonTerminating
| lift_join_Terminating:
    forall x y z,
      join x y z ->
      lift_join (Terminating x) (Terminating y) (Terminating z).

Inductive strong_lift_join
          {state: Type}
          {J_state: Join state}:
  MetaState state -> MetaState state -> MetaState state -> Prop :=
| strong_lift_join_Error1:
    forall mx, strong_lift_join Error mx Error
| strong_lift_join_Error2:
    forall mx, strong_lift_join mx Error Error
| strong_lift_join_NonTerminating1:
    forall x, strong_lift_join NonTerminating (Terminating x) NonTerminating
| strong_lift_join_NonTerminating2:
    forall x, strong_lift_join (Terminating x) NonTerminating NonTerminating
| strong_lift_join_NonTerminating3:
    strong_lift_join NonTerminating NonTerminating NonTerminating
| strong_lift_join_Terminating:
    forall x y z,
      join x y z ->
      strong_lift_join (Terminating x) (Terminating y) (Terminating z).

Definition lift_function {A B: Type} (f: A -> B): MetaState A -> MetaState B :=
  fun ma =>
  match ma with
  | NonTerminating => NonTerminating
  | Error => Error
  | Terminating a => Terminating (f a)
  end.

(*
Instance MetaState_SA (state: Type) {SA: SeparationAlgebra state}: SeparationAlgebra (MetaState state).
*)

Lemma lift_relation_forward {state: Type} (R: state -> MetaState state -> Prop):
  forall x y, lift_relation R x y ->
  match x with
  | Error => y = Error
  | NonTerminating => y = NonTerminating
  | _ => True
  end.
Proof.
  intros.
  destruct x, y; inversion H; subst; try congruence; auto.
Qed.

Lemma clos_refl_trans_lift_relation_forward {state: Type} (R: state -> MetaState state ->
Prop):
  forall x y, clos_refl_trans _ (lift_relation R) x y ->
  match x with
  | Error => y = Error
  | NonTerminating => y = NonTerminating
  | _ => True
  end.
Proof.
  intros.
  induction H.
  + apply lift_relation_forward in H; auto.
  + destruct x; auto.
  + destruct x; subst; simpl; subst; auto.
Qed.

Lemma lift_function_rev {A B: Type} (f: A -> B):
  forall ma mb, lift_function f ma = mb ->
  match mb with
  | NonTerminating => ma = NonTerminating
  | Error => ma = Error
  | _ => True
  end.
Proof.
  intros.
  destruct mb, ma; auto; simpl in *; congruence.
Qed.


Load "imp1_phase1".

(** * FOO Voting Prtocol *)
(** ** Phase-II *)
(**
<<
- [V1 -> M : { (commit(v0, r), unblind(pkat, c, r'', pi1(x3))) }_pk5 ]

- [V2 -> M : { (commit(v1, r'), unblind(pkat, c, r''', pi2(x3))) }_pk5 ]
>> *)

(** - We assume that Administrator and collector are dishonest
    - Hence, random numbers used to generate keys (pair) of them included in attacker's knowledge *)

(** * ProtocolPhase-I: There are two voters [V1] and [V2], and their corresponding votes [va] and [vb] respectively. *)

(** * Frame [phi0]: initial knowledge *)
(** 1-Voter1, 2-Voter2, 3-Administrator, 4-Collector, 5-Mixnet
[(pk 0)] represents a public key used in commitment shceme. *)

(** Even though different symbols used to distinuish the key generations, for convenience, we use same symbols here in Coq *)

(** * Frame [phi3] *)

Definition x3tt n1 n2 :=  f (conv_mylist_listm (phi0 ++ [bol (eqm (to x1) (V 1));msg (pk 1, (e (b n1 7) 8, sign (sk 1) (e (b n1 7) 8) (r 30)))] ++  [bol (eqm (to (x2t n1)) (V 2)); msg (pk 2, (e (b n2 11) 12, sign (sk 2) (e (b n2 11) 12) (r 32)))])).

Definition x3tf n1 := f (conv_mylist_listm (phi0 ++ [bol (eqm (to x1) (V 1)); msg (pk 1, (e (b n1 7) 8, sign (sk 1) (e (b n1 7) 8) (r 30)))] ++  [bol (eqm (to (x2t n1)) (V 2)); msg O])).

Definition x3ftt n1 n2 := f (conv_mylist_listm (phi0 ++ [bol (eqm (to x1) (V 1)); bol (eqm (to x1) (V 2)); msg (pk 2, (e (b n2 9) 10, sign (sk 2) (e (b n2 9) 10) (r 31)))] ++ [bol (eqm (to (x2ft n2)) (V 1)); msg (pk 1, (e (b n1 13) 14, sign (sk 1) (e (b n1 13) 14) (r 33)))])).
Definition x3ftf n2 := f (conv_mylist_listm (phi0 ++ [bol (eqm (to x1) (V 1)); bol (eqm (to x1) (V 2)); msg (pk 2, (e (b n2 9) 10, sign (sk 2) (e (b n2 9) 10) (r 31)))] ++  [bol (eqm (to (x2ft n2)) (V 1)); msg O])).
Definition x3fff  := f (conv_mylist_listm (phi0 ++ [bol (eqm (to x1) (V 1))] ++ [bol (eqm (to x1) (V 2)); msg O; msg O])).

(** * Step 1 of Voting phase *)

Definition q11 n1 n2 := (ifm (eqm (to (x3tt n1 n2)) (V 1)) & (bacc pkat (b n1 7) (r 8) (pi1 (x3tt n1 n2))) (enc ((b n1 7), (unblind pkat (b n1 7) (r 8) (pi1 (x3tt n1 n2)))) (pk 5) (r 50)) (ifm (eqm (to (x3tt n1 n2)) (V 2)) & (bacc pkat (b n2 11) (r 12) (pi2 (x3tt n1 n2))) (enc ((b n2 11), (unblind pkat (b n2 11) (r 12) (pi2 (x3tt n1 n2)))) (pk 5) (r 51)) O)).

Definition q21 n1 n2 := (ifm (eqm (to (x3ftt n1 n2)) (V 2)) & (bacc pkat (b n2 9) (r 10) (pi2 (x3ftt n1 n2))) (enc ((b n2 9), (unblind pkat (b n2 9) (r 10) (pi2 (x3ftt n1 n2)))) (pk 5) (r 52)) (ifm (eqm (to (x3ftt n1 n2)) (V 1)) & (bacc pkat (b n1 13) (r 14) (pi1 (x3ftt n1 n2))) (enc ((b n1 13), (unblind pkat (b n1 13) (r 14) (pi1 (x3ftt n1 n2)))) (pk 5) (r 54)) O)).

(*****************************************************)

Definition q1_s n1 n2 := (* (ifm (eqm (to (x2t n1)) (V 1))& (bacc pkat (b n1 7) (r 8) (pi1 (x2t n1)))  ok *)
			            (ifm (eqm (to (x2t n1)) (V 2)) (q11 n1 n2) O).

Definition q2_s n1 n2 :=  (* (ifm (eqm (to (x2ft n2)) (V 2))& (bacc pkat (b n2 9) (r 10) (pi2 (x2ft n2)))  ok *)
			           (ifm (eqm (to (x2ft n2)) (V 1)) (q21 n1 n2)  O).

Definition t2 n1 n2 :=  (ifm (eqm (to x1) (V 1)) (q1_s n1 n2)
		    	     (ifm (eqm (to x1) (V 2))  (q2_s n1 n2) O)).

Definition phi3 n1 n2 := (phi2 n1 n2) ++ [msg (t2 n1 n2)].
 Theorem frame3ind: (phi3 0 1) ~ (phi3 1 0).
Proof. unfold phi3, phi2, phi1, phi0, q0, t1, t2, q1, q2, q1_s, q2_s, q11, q21.
simpl.
apply IFBRANCH_M3 with (ml1:= phi0) (ml2:= phi0).
simpl.
apply IFBRANCH_M2 with (ml1:= phi0 ++ [ bol (eqm (to x1) (V 1));
    msg (pk 1, (e (b 0 7) 8, sign (sk 1) (e (b 0 7) 8) (r 30)))])(ml2:= phi0 ++ [bol (eqm (to x1) (V 1));
   msg (pk 1, (e (b 1 7) 8, sign (sk 1) (e (b 1 7) 8) (r 30)))]).
simpl.
apply IFBRANCH_M1 with (ml1:= (phi0 ++[ bol (eqm (to x1) (V 1)); msg (pk 1, (e (b 0 7) 8, sign (sk 1) (e (b 0 7) 8) (r 30))); bol (eqm (to (x2t 0)) (V 2)); msg (pk 2, (e (b 1 11) 12, sign (sk 2) (e (b 1 11) 12) (r 32)))]))(ml2:= (phi0 ++ [ bol (eqm (to x1) (V 1)); msg (pk 1, (e (b 1 7) 8, sign (sk 1) (e (b 1 7) 8) (r 30))); bol (eqm (to (x2t 1)) (V 2)); msg (pk 2, (e (b 0 11) 12, sign (sk 2) (e (b 0 11) 12) (r 32)))])).
simpl.
(** TRACE-1 *)
 pose proof( commit_swap (vt 0) (vt 1) _ 7 7 11 11 [ msg (sk 1); msg (sk 2); msg (r 8); msg (r 12); msg (r 30); msg (r 32); msg (r 50); msg (pk 1); msg (pk 2); msg (pk 5); msg pkat; msg (vt 0); msg (vt 1); bol (eqm (to x1) (V 1))]); simpl in H;
assert( (L (vt 0)) # (L (vt 1)));
try eapply len_f1; try eapply len_nonce;
eapply EQmsg' in H0;
apply H in H0; clear H; rename H0 into H; auto;
appconst H;
x1checks x1 x1 H; simpl;
bterm 0 1 7 H.
(** Appending triple term of voter1 on both sides *)
funapptrmhyp (msg (sign (sk 1) (e (b 0 7) 8) (r 30))) (msg (sign (sk 1) (e (b 1 7) 8) (r 30))) H;
funapptrmhyp (msg ((e (b 0 7) 8), (sign (sk 1) (e (b 0 7) 8) (r 30)))) (msg ((e (b 1 7) 8),  (sign (sk 1) (e (b 1 7) 8) (r 30)))) H;
funapptrmhyp (msg (pk 1, (e (b 0 7) 8, sign (sk 1) (e (b 0 7) 8) (r 30)))) (msg (pk 1, (e (b 1 7) 8, sign (sk 1) (e (b 1 7) 8) (r 30)))) H;
x1checks (x2t 0) (x2t 1) H.
(** Appending second blind term *)
bterm 1 0 11 H;
(** Appending triple term of voter2 on both sides *)
funapptrmhyp (msg (sign (sk 2) (e (b 1 11) 12) (r 32))) (msg (sign (sk 2) (e (b 0 11) 12) (r 32))) H;
funapptrmhyp (msg ((e (b 1 11) 12), (sign (sk 2) (e (b 1 11) 12) (r 32)))) (msg ((e (b 0 11) 12),  (sign (sk 2) (e (b 0 11) 12) (r 32)))) H;
funapptrmhyp (msg (pk 2, (e (b 1 11) 12, sign (sk 2) (e (b 1 11) 12) (r 32)))) (msg (pk 2, (e (b 0 11) 12, sign (sk 2) (e (b 0 11) 12) (r 32)))) H.
x1checks (x3tt 0 1) (x3tt 1 0) H.
(** append bacc term on both sides *)
bacctac 0 1 7 8 (x3tt 0 1) (x3tt 1 0) H.
funapptrmhyp (bol (eqm (to (x3tt 0 1)) (V 1)) &
      (bacc pkat (b 0 7) (r 8) (pi1 (x3tt 0 1))))  (bol (eqm (to (x3tt 1 0)) (V 1)) & (bacc pkat (b 1 7) (r 8) (pi1 (x3tt 1 0)))) H.
(** append enc term on both sides *)
funapptrmhyp (msg (unblind pkat (b 0 7) (r 8) (pi1 (x3tt 0 1)))) (msg (unblind pkat (b 1 7) (r 8) (pi1 (x3tt 1 0)))) H.
funapptrmhyp (msg ((b 0 7), (unblind pkat (b 0 7) (r 8) (pi1 (x3tt 0 1))))) (msg ((b 1 7), (unblind pkat (b 1 7) (r 8) (pi1 (x3tt 1 0))))) H.
funapptrmhyp (msg (enc (b 0 7, unblind pkat (b 0 7) (r 8) (pi1 (x3tt 0 1))) (pk 5) (r 50))) (msg (enc (b 1 7, unblind pkat (b 1 7) (r 8) (pi1 (x3tt 1 0))) (pk 5) (r 50))) H.
restrsublis H.

(** TRACE-2 *)
simpl. Admitted.

(** Frame [phi 4] *)

Definition x4ttt n1 n2 := f (conv_mylist_listm (phi0 ++ [ bol (eqm (to x1) (V 1)); msg (pk 1, (e (b n1 7) 8, sign (sk 1) (e (b n1 7) 8) (r 30))); bol (eqm (to (x2t n1)) (V 2)); msg (pk 2, (e (b n2 11) 12, sign (sk 2) (e (b n2 11) 12) (r 32))); bol (eqm (to (x3tt n1 n2)) (V 1)) & (bacc pkat (b n1 7) (r 8) (pi1 (x3tt n1 n2))); msg (enc (b n1 7, unblind pkat (b n1 7) (r 8) (pi1 (x3tt n1 1))) (pk 5) (r 34))])).

Definition x4tftt n1 n2 := f (conv_mylist_listm (phi0 ++  [bol (eqm (to x1) (V 1)); msg (pk 1, (e (b n1 7) 8, sign (sk 1) (e (b n1 7) 8) (r 30))); bol (eqm (to (x2t n1)) (V 2)); msg (pk 2, (e (b n2 11) 12, sign (sk 2) (e (b n2 11) 12) (r 32))); bol (eqm (to (x3tt n1 n2)) (V 1)) & (bacc pkat (b n1 7) (r 8) (pi1 (x3tt n1 n2))); bol (eqm (to (x3tt n1 n2)) (V 2)) & (bacc pkat (b n2 11) (r 12) (pi2 (x3tt n1 n2)));
    msg (enc (b n2 11, unblind pkat (b n2 11) (r 12) (pi2 (x3tt n1 n2))) (pk 5) (r 35))])).

Definition x4ttff n1 n2 := f (conv_mylist_listm (phi0 ++ [bol (eqm (to x1) (V 1)); msg (pk 1, (e (b n1 7) 8, sign (sk 1) (e (b n1 7) 8) (r 30))); bol (eqm (to (x2t n1)) (V 2)); msg (pk 2, (e (b n2 11) 12, sign (sk 2) (e (b n2 11) 12) (r 32))); bol (eqm (to (x3tt n1 n2)) (V 1)) & (bacc pkat (b n1 7) (r 8) (pi1 (x3tt n1 n2))); bol (eqm (to (x3tt n1 n2)) (V 2)) & (bacc pkat (b n2 11) (r 12) (pi2 (x3tt n1 n2)));
    msg O])) .
Definition x4fttt n1 n2 := f (conv_mylist_listm (phi0 ++ [ bol (eqm (to x1) (V 1));  bol (eqm (to x1) (V 2)); msg (pk 2, (e (b n2 9) 10, sign (sk 2) (e (b n2 9) 10) (r 31))); bol (eqm (to (x2ft n2)) (V 1)); msg (pk 1, (e (b n1 13) 14, sign (sk 1) (e (b n1 13) 14) (r 33)));
    bol (eqm (to (x3ftt n1 n2)) (V 2)) & (bacc pkat (b n2 9) (r 10) (pi2 (x3ftt n1 n2))); msg (enc (b n2 9, unblind pkat (b n2 9) (r 10) (pi2 (x3ftt n1 n2))) (pk 5) (r 36))]) ).

Definition x4fttft n1 n2 := f (conv_mylist_listm (phi0 ++ [bol (eqm (to x1) (V 1));
    bol (eqm (to x1) (V 2));
    msg (pk 2, (e (b n2 9) 10, sign (sk 2) (e (b n2 9) 10) (r 31)));
    bol (eqm (to (x2ft n2)) (V 1));
    msg (pk 1, (e (b n1 13) 14, sign (sk 1) (e (b n1 13) 14) (r 33)));
    bol
      (eqm (to (x3ftt n1 n2)) (V 2)) &
      (bacc pkat (b n2 9) (r 10) (pi2 (x3ftt n1 n2)));
    bol
      (eqm (to (x3ftt n1 n2)) (V 1)) &
      (bacc pkat (b n1 13) (r 14) (pi1 (x3ftt n1 n2)));
    msg
      (enc (b n1 13, unblind pkat (b n1 13) (r 14) (pi1 (x3ftt n1 n2)))
         (pk 5) (r 37))])).

Definition x4fttff n1 n2 := f (conv_mylist_listm (phi0 ++ [bol (eqm (to x1) (V 1));
    bol (eqm (to x1) (V 2));
    msg (pk 2, (e (b n2 9) 10, sign (sk 2) (e (b n2 9) 10) (r 31)));
    bol (eqm (to (x2ft 1)) (V 1));
    msg (pk 1, (e (b n1 13) 14, sign (sk 1) (e (b n1 13) 14) (r 33)));
    bol
      (eqm (to (x3ftt n1 n2)) (V 2)) &
      (bacc pkat (b n2 9) (r 10) (pi2 (x3ftt n1 n2)));
    bol
      (eqm (to (x3ftt n1 n2)) (V 1)) &
      (bacc pkat (b n1 13) (r 14) (pi1 (x3ftt n1 n2)));
    msg O])) .

Definition x4ttft n1 n2 := f (conv_mylist_listm (phi0 ++ [ bol (eqm (to x1) (V 1));
  msg (pk 1, (e (b n1 7) 8, sign (sk 1) (e (b n1 7) 8) (r 30)));
  bol (eqm (to (x2t n1)) (V 2));
  msg (pk 2, (e (b n2 11) 12, sign (sk 2) (e (b n2 11) 12) (r 32)));
  bol
    (eqm (to (x3tt 0 1)) (V 1)) & (bacc pkat (b 0 7) (r 8) (pi1 (x3tt 0 1)));
  bol
    (eqm (to (x3tt 0 1)) (V 2)) &
    (bacc pkat (b n2 11) (r 12) (pi2 (x3tt n1 n2)));
  msg
    (enc (b n2 11, unblind pkat (b n2 11) (r 12) (pi2 (x3tt n1 n2)))
       (pk 5) (r 35))])).

(** * Step 2 of Voting phase *)
 Definition q111 n1 n2 :=  (ifm (eqm (to (x4ttt n1 n2)) (V 2)) & (bacc pkat (b n2 11) (r 12) (pi2 (x3tt n1 n2))) (enc ((b n2 11), (unblind pkat (b n2 11) (r 12) (pi2 (x4ttt n1 n2)))) (pk 5) (r 55)) O).

 Definition q112 n1 n2 :=    (ifm (eqm (to (x4ttft n1 n2)) (V 1)) & (bacc pkat (b n1 7) (r 8) (pi1 (x4ttft n1 n2))) (enc ((b n1 7), (unblind pkat (b n1 7) (r 8) (pi1 (x4ttft n1 n2)))) (pk 5) (r 56)) O).

 Definition q211 n1 n2 :=  (ifm (eqm (to (x4fttt n1 n2)) (V 1)) & (bacc pkat (b n1 13) (r 14) (pi1 (x4fttt n1 n2))) (enc ((b n1 13), (unblind pkat (b n1 13) (r 14) (pi1 (x4fttt n1 n2)))) (pk 5) (r 57)) O).

 Definition q212 n1 n2 :=  (ifm (eqm (to (x4fttft n1 n2)) (V 2)) & (bacc pkat (b n2 9) (r 10) (pi2 (x4fttft n1 n2))) (enc ((b n2 9), (unblind pkat (b n2 9) (r 10) (pi2 (x4fttft n1 n2)))) (pk 5) (r 58)) O).

 (** * Step 1 of Voting phase *)

Definition q11s n1 n2 := (ifm (eqm (to (x3tt n1 n2)) (V 1)) & (bacc pkat (b n1 7) (r 8) (pi1 (x3tt n1 n2))) (q111 n1 n2) (ifm (eqm (to (x3tt n1 n2)) (V 2)) & (bacc pkat (b n2 11) (r 12) (pi2 (x3tt n1 n2))) (q112 n1 n2) O)).

Definition q21s n1 n2 := (ifm (eqm (to (x3ftt n1 n2)) (V 2)) & (bacc pkat (b n2 9) (r 10) (pi2 (x3ftt n1 n2))) (q211 n1 n2) (ifm (eqm (to (x3ftt n1 n2)) (V 1)) & (bacc pkat (b n1 13) (r 14) (pi1 (x3ftt n1 n2))) (q212 n1 n2) O)).
(*****************************************************)
Definition q1s n1 n2 := (* (ifm (eqm (to (x2t n1)) (V 1))& (bacc pkat (b n1 7) (r 8) (pi1 (x2t n1)))  ok *)
			            (ifm (eqm (to (x2t n1)) (V 2)) (q11s n1 n2) O).

Definition q2s n1 n2 :=  (* (ifm (eqm (to (x2ft n2)) (V 2))& (bacc pkat (b n2 9) (r 10) (pi2 (x2ft n2)))  ok *)
			           (ifm (eqm (to (x2ft n2)) (V 1)) (q21s n1 n2)  O).

Definition t3 n1 n2 :=  (ifm (eqm (to x1) (V 1)) (q1s n1 n2)
		    	     (ifm (eqm (to x1) (V 2))  (q2s n1 n2) O)).

Definition phi4 n1 n2 := (phi3 n1 n2) ++ [msg (t3 n1 n2)].