Require Import UniMath.Foundations.UnivalenceAxiom.

Definition j_type
  := ∑ (j : ∏ (A : UU) (a : A) (D : ∏ (b : A), a = b -> UU),
            D a (idpath a) -> ∏ (b : A) (p : a = b), D b p),
     ∏ (A : UU) (a : A) (D : ∏ (b : A) (p : a = b), UU)
       (d : D a (idpath a)), j A a D d a (idpath a) = d.

Definition nat_trans {A : UU} (F G : A -> UU) := ∏ (a : A), F a -> G a.

Definition vert_comp {A : UU} {H G F : A -> UU}
           (α : nat_trans G F) (β : nat_trans H G) : nat_trans H F
  := λ a, (α a) ∘ (β a).

Notation "α # β" := (vert_comp α β).

Definition yoneda_type :=
  ∑ (y : ∏ (A : UU) (F : A -> UU) (a : A),
         F a -> (nat_trans (λ b, a = b) F)),
  (∏ (A : UU) (G F : A -> UU) (α : nat_trans G F) (a b : A) (p : a = b)
     (x : G a), (α # y A G a x) b p = y A F a (α a x) b p)
× ((∏ (A : UU) (a b : A) (p : a = b)
        , (y A (λ b, a = b) a (idpath a) b) p = p)
 × (∏ (A : UU) (F : A -> UU) (a : A) (x : F a),
        y A F a x a (idpath a) = x)).

Section j_to_yoneda.

  Context (j : j_type).

  Definition j_derive_transport {A : UU} (F : A -> UU):
    ∑ (transp : ∏ (a : A) (b : A) (p : a = b), F a -> F b),
    ∏ (a : A), transp a a (idpath a) = λ x, x.
  Proof.
    exists (λ a, (pr1 j) A a (λ b _, F a -> F b) (λ x, x)).
    exact (λ a, pr2 j A a (λ b _, F a -> F b) (λ x, x)).
  Defined.

  Definition j_derive_idtohomot {A B : UU} {f g : A -> B}:
    f = g -> ∏ (a : A), f a = g a
    := (pr1 j) (A -> B) f (λ g p, ∏ (a : A), f a = g a)
               (λ a, idpath _) g.

  Definition j_derive_path_concat {A : UU} {a b c : A}:
    a = b -> b = c -> a = c
    := (pr1 j) A a (λ b _, b = c -> a = c) (λ x, x) b.

  Definition j_derive_ap {A B : UU} (f : A -> B) {a b : A}:
    a = b -> f a = f b
    := (pr1 j) A a (λ b _, f a = f b) (idpath _) b.

  Definition j_derive_inv {A : UU} {a b : A}:
    a = b -> b = a
    := (pr1 j) A a (λ b _, b = a) (idpath _) b.

  Definition j_yoneda_fun (A : UU) (F : A -> UU) (a : A) (x : F a):
    nat_trans (λ b, a = b) F
    := λ b p, (pr1 (j_derive_transport F)) a b p x.

  Definition j_yoneda_naturality
             (A : UU) (G F : A -> UU) (α : nat_trans G F)
             (a b : A) (p : a = b) (x : G a) :
    (α # j_yoneda_fun A G a x) b p = j_yoneda_fun A F a (α a x) b p.
  Proof.
    apply (λ d, (pr1 j)
                  A a (λ b p, (α # j_yoneda_fun A G a x) b p =
                              j_yoneda_fun A F a (α a x) b p) d b).
    refine (j_derive_path_concat _ _).
    - use (j_derive_ap (α a)).
      + exact x.
      + exact (j_derive_idtohomot
                 ((pr2 j) A a (λ b _ , G a → G b) (λ x,x)) x).
    - apply j_derive_inv.
      exact (j_derive_idtohomot
               ((pr2 j) A a (λ b _, F a -> F b) (λ x,x))
               (α a x)).
  Defined.

  Definition j_yoneda_special (A : UU) (a b : A) (p : a = b):
    j_yoneda_fun A (λ b, a = b) a (idpath a) b p = p.
  Proof.
    apply ((pr1 j) A a
                   (λ b p,
                    j_yoneda_fun A (λ b0 : A, a = b0)
                                 a (idpath a) b p = p)).
    cbn.
    exact (j_derive_idtohomot ((pr2 j) A a (λ b _ , a = a → a = b)
                                       (λ x,x)) (idpath _)).
  Defined.

  Definition j_yoneda_from (A : UU) (F : A -> UU) (a : A)
    (α : nat_trans (λ b, a = b) F): F a
    := α a (idpath a).

  Definition j_yoneda_fromto (A : UU) (F : A -> UU) (a : A):
    (λ x, j_yoneda_from A F a (j_yoneda_fun A F a x)) = λ x, x
    := pr2 (j_derive_transport F) a.

  Definition j_derive_y : yoneda_type.
    exists j_yoneda_fun.
    apply dirprodpair.
    - exact j_yoneda_naturality.
    - apply dirprodpair.
      + exact j_yoneda_special.
      + exact (λ A F a x, j_derive_idtohomot (j_yoneda_fromto A F a) x).
  Defined.

End j_to_yoneda.

Section yoneda_to_j.

  Context (y_inhab : yoneda_type).

  Definition y := pr1 y_inhab.
  Definition y_naturality := pr1 (pr2 y_inhab).
  Definition y_special := pr1 (pr2 (pr2 y_inhab)).
  Definition y_inverse := pr2 (pr2 (pr2 y_inhab)).

  Definition y_derive_ap {A B : UU} (f : A -> B) {a b : A}
             (p : a = b) : f a = f b
    := y A (λ b, f a = f b) a (idpath _) b p.

  Definition y_derive_pathcomp {A : UU} {a b c : A}
             (p : a = b) (q : b = c) : a = c
    := y A (λ c, a = c) b p c q.

  Definition y_derive_transp {A : UU} {B : A -> UU}
             {a a' : A} : a = a' -> B a -> B a'
    := y A (λ a', B a -> B a') a (λ b, b) a'.

  Definition y_derive_reverse {A : UU} {a b : A}: a = b -> b = a
    := y A (λ x, x = a) a (idpath _) b.

  Definition y_total_paths {A : UU} {a b : A} (p : a = b): A -> UU
    := (λ x, ∑ (q : a = x), (b ,, p) = (x ,, q)).

  Definition y_underlying {A : UU} {a : A}:
    nat_trans (y_total_paths (idpath a)) (λ b, a = b) := λ _, pr1.

  Definition y_lift {A : UU} {a : A}:
    nat_trans (λ b, a = b) (y_total_paths (idpath a))
    := y A (y_total_paths (idpath a)) a (idpath _ ,, idpath _).

  Definition y_lift_fixes_paths {A : UU} {a b : A} (p : a = b):
    (y_underlying # y_lift) b p = p.
  Proof.
    refine (y_derive_pathcomp _ _) ; try apply y_naturality.
    apply y_special.
  Defined.

  Definition y_j_fun_bad (A : UU) (a : A) (D : ∏ (b : A), a = b -> UU)
             (d : D a (idpath a)) (b : A) (p : a = b) : D b p.
  Proof.
    apply (y_derive_transp (y_lift_fixes_paths p)).
    apply (λ d', y (∑ b0 : A, a = b0)
                   (λ bp, D (pr1 bp) (pr2 bp))
                   (a,,idpath _) d' (b ,, pr1 (y_lift b p))).
    - exact d.
    - exact (pr2 (y_lift b p)).
  Defined.

  Definition y_pathcomp_computes_right
             {A : UU} {a : A}: ∏ (b : A) (p : a = b),
    y_derive_pathcomp p (idpath b) = p.
  Proof.
    apply (y_j_fun_bad A a _).
    unfold y_derive_pathcomp.
    apply y_inverse.
  Defined.

  Definition y_pathcomp_reveres_computes
             {A : UU} {a : A}: ∏ (b : A) (p : a = b),
    y_derive_pathcomp (y_derive_reverse p) p = idpath b.
  Proof.
    apply (y_j_fun_bad A a _).
    refine (y_derive_pathcomp _ _).
    - apply y_pathcomp_computes_right.
    - exact (y_inverse A (λ x, x=a) a (idpath a)).
  Defined.

  Definition y_embed_paths1 {A : UU} {a b : A} (p : a=b):
    (a ,, idpath _) = (b ,, p).
  Proof.
    apply (y_derive_pathcomp (pr2 (y_lift b p))).
    exact (@y_derive_ap _ _ (λ q, (b,,q))
                               (pr1 (y_lift b p)) p
                               (y_lift_fixes_paths p)).
  Defined.

  Definition y_embed_paths2 {A : UU} {a b : A} (p : a=b):
    (a ,, idpath _) = (b ,, p).
  Proof.
    refine (y_derive_pathcomp _ _).
    - exact (y_derive_reverse (y_embed_paths1 (idpath a))).
    - exact (y_embed_paths1 p).
  Defined.

  Definition y_embed_paths2_computes {A : UU} {a : A}:
    y_embed_paths2 (idpath a) = idpath _.
  Proof.
    unfold y_embed_paths2.
    apply y_pathcomp_reveres_computes.
  Defined.

  Definition y_j_fun (A : UU) (a : A) (D : ∏ (b : A), a = b -> UU)
             (d : D a (idpath a)) (b : A) (p : a = b) : D b p.
  Proof.
    apply (λ d', y (∑ b0 : A, a = b0)
                   (λ bp, D (pr1 bp) (pr2 bp))
                   (a,,idpath _) d' (b ,, p)).
    - exact d.
    - apply y_embed_paths2.
  Defined.

  Definition y_derive_j : j_type.
  Proof.
    exists y_j_fun.
    intros.
    refine (y_derive_pathcomp _ _).
    unfold y_j_fun.
    - apply y_derive_ap.
      exact y_embed_paths2_computes.
    - exact (y_inverse (∑ b, a=b)
                       (λ bp, D (pr1 bp) (pr2 bp))
                       (a ,, idpath _) d).
  Defined.

End yoneda_to_j.

Print All Dependencies y_derive_j.
Print All Dependencies j_derive_y.


(** This proof makes use of function extensionality twice.

    Notice however that it is _not_ used to derive a term of type
    [yoneda_type]. With [j_derive_y] alone we may construct
    [y_derive_j], as done above, from which we may then construct
    [yoneda_tofrom].

    That is, in the presence of function extensionality, [yoneda_type]
    implies a Yoneda equivalence. *)
Definition yoneda_tofrom (j : j_type) (A : UU) (F : A -> UU) (a : A)
           (α : nat_trans (λ b, a = b) F) :
  j_yoneda_fun A F a (j_yoneda_from A F a α) = α.
Proof.
  apply isweqtoforallpaths ; intro b.
  apply funextfun ; intro p.
  apply ((pr1 j) A a (λ b p,
                      j_yoneda_fun A F a (j_yoneda_from A F a α) b p = α b p)).
  unfold j_yoneda_fun.
  rewrite (pr2 (j_derive_transport F) a).
  apply idpath.
Defined.