Require Import UniMath.Foundations.PartD.
Require Import UniMath.Foundations.Propositions.
Require Import UniMath.Foundations.Sets.
Require Import UniMath.MoreFoundations.Tactics.
Require Import UniMath.CategoryTheory.Categories.
Require Import UniMath.CategoryTheory.functor_categories.
Require Import UniMath.CategoryTheory.whiskering.
Require Import UniMath.CategoryTheory.Monads.Monads.
Require Import UniMath.CategoryTheory.Monads.RelativeMonads.

Local Open Scope cat.

Ltac pathvia b := (apply (@pathscomp0 _ _ b _ )).


Section Kleisli_defn.

  Context {C : precategory_data}.


Definition Kleisli_Data : UU := ∑ T : C → C,
  (∏ a : C, a --> T a) × (∏ a b : C, (a --> T b) → (T a --> T b)).


Definition Kleisli_Data_ob (T: Kleisli_Data) (c : C) : C := pr1 T c.
Coercion Kleisli_Data_ob : Kleisli_Data >-> Funclass.

Definition η (T : Kleisli_Data) : ∏ a : C, a --> T a := pr1 (pr2 T).

Definition bind (T : Kleisli_Data) {a b : C} :
  C⟦a,T b⟧ → C⟦T a,T b⟧ := pr2 (pr2 T) a b.


Definition Kleisli_Laws (T : Kleisli_Data) :=
  (∏ a, bind T (η T a) = identity (T a)) ×
  (∏ a b (f : C⟦a,T b⟧), bind T f ∘ η T a = f) ×
  (∏ a b c (f : C⟦a,T b⟧) (g : C⟦b,T c⟧), bind T g ∘ bind T f = bind T (bind T g ∘ f)).

Lemma isaprop_Kleisli_Laws (hs : has_homsets C) (T : Kleisli_Data) :
  isaprop (Kleisli_Laws T).
Proof.
  repeat apply isapropdirprod; repeat (apply impred_isaprop; intros); apply hs.
Defined.

Definition bind_bind {T : Kleisli_Data} (H: Kleisli_Laws T) :
  ∏ a b c (f : C⟦a,T b⟧) (g : C⟦b,T c⟧), bind T g ∘ bind T f = bind T (bind T g ∘ f) :=
  pr2 (pr2 H).

Definition bind_η {T : Kleisli_Data} (H : Kleisli_Laws T) :
  ∏ a b (f : C⟦a,T b⟧), bind T f ∘ η T a = f := pr1 (pr2 H).

Definition η_bind {T : Kleisli_Data} (H : Kleisli_Laws T) :
  ∏ a, bind T (η T a) = identity (T a) := pr1 H.


Definition Kleisli: UU :=
  ∑ (T : Kleisli_Data), Kleisli_Laws T.
Coercion Kleisli_Data_from_Kleisli (T : Kleisli) : Kleisli_Data := pr1 T.

Lemma Kleisli_is_rel_monad: Kleisli = RelMonad (functor_identity C).
Proof.
  apply idpath.
Qed.

Coercion kleisli_laws (T : Kleisli) : Kleisli_Laws (pr1 T) := pr2 T.

End Kleisli_defn.

Arguments Kleisli: clear implicits.
Arguments Kleisli_Data: clear implicits.


Section Kleisli_precategory.


Definition Kleisli_Mor_laws {C : precategory_data} (T T': Kleisli_Data C)
             (α : ∏ a : C, T a --> T' a) : UU :=
  (∏ a : C, α a ∘ η T a = η T' a) ×
  (∏ (a b : C) (f : C⟦a,T b⟧), bind T' (α b ∘ f) ∘ α a = α b ∘ (bind T f)).

Lemma isaprop_Kleisli_Mor_laws {C : precategory_data} (hs : has_homsets C) (T T' : Kleisli_Data C)
        (α : ∏ a : C, T a --> T' a) :
  isaprop (Kleisli_Mor_laws T T' α).
Proof.
  apply isapropdirprod; repeat (apply impred_isaprop; intros); apply hs.
Defined.

Definition Kleisli_Mor {C : precategory_data} (T T' : Kleisli_Data C) : UU :=
  ∑ (α : ∏ a : C, T a --> T' a), Kleisli_Mor_laws T T' α.

Definition nat_trans_from_kleisli_mor {C : precategory_data}
           {T T' : Kleisli_Data C} (s : Kleisli_Mor T T') :
  ∏ a : C, T a --> T' a := pr1 s.

Definition Kleisli_Mor_η {C : precategory_data} {T T' : Kleisli C} (α : Kleisli_Mor T T') :
  ∏ a : C, η T a · nat_trans_from_kleisli_mor α a = η T' a :=
  pr1 (pr2 α).

Definition Kleisli_Mor_bind {C : precategory_data} {T T' : Kleisli C} (α : Kleisli_Mor T T') :
  ∏ (a b : C) (f : C⟦a,T b⟧),
    bind T' (nat_trans_from_kleisli_mor α b ∘ f) ∘ nat_trans_from_kleisli_mor α a =
    nat_trans_from_kleisli_mor α b ∘ (bind T f) :=
  pr2 (pr2 α).

Definition Kleisli_Mor_equiv {C : precategory_data} (hs : has_homsets C) {T T' : Kleisli C}
           (α β : Kleisli_Mor T T') :
  α = β ≃ (nat_trans_from_kleisli_mor α = nat_trans_from_kleisli_mor β).
Proof.
  apply subtypeInjectivity.
  intro a.
  now apply isaprop_Kleisli_Mor_laws.
Defined.


Definition map {C : precategory_data} (T : Kleisli_Data C) {a b : C} (f : a --> b) :
  T a --> T b :=
  bind T (η T b ∘ f).

Lemma map_id {C : precategory} {T : Kleisli_Data C} (H : Kleisli_Laws T) :
  ∏ a : C, map T (identity a) = identity (T a).
Proof.
  intro. unfold map.
  rewrite id_left.
  now apply η_bind.
Defined.

Lemma map_map {C : precategory} {T : Kleisli_Data C} (H : Kleisli_Laws T) :
  ∏ (a b c : C) (f : a --> b) (g : b --> c), map T (g ∘ f) = map T g ∘ map T f.
Proof.
  intros. unfold map.
  rewrite (bind_bind H).
  do 2 rewrite <- assoc.
  now rewrite (bind_η H).
Defined.

Lemma map_bind {C : precategory} {T : Kleisli_Data C} (H : Kleisli_Laws T) :
  ∏ (a b c : C) (f : b --> c) (g : a --> T b), map T f ∘ bind T g = bind T (map T f ∘ g).
Proof.
  intros. unfold map.
  now rewrite (bind_bind H).
Defined.

Lemma bind_map {C : precategory} {T : Kleisli_Data C} (H : Kleisli_Laws T) :
  ∏ (a b c : C) (f : b --> T c) (g : a --> b),
  bind T f ∘ map T g = bind T (f ∘ g).
Proof.
  intros. unfold map.
  rewrite (bind_bind H).
  rewrite <- assoc.
  now rewrite (bind_η H).
Defined.

Lemma map_η {C : precategory} {T : Kleisli_Data C} (H : Kleisli_Laws T) :
  ∏ (a b : C) (f : a --> b),
  map T f ∘ η T a = η T b ∘ f.
Proof.
  intros. unfold map.
  now rewrite (bind_η H).
Defined.

Definition μ {C : precategory} (T : Kleisli_Data C) (a : C) :
  T (T a) --> T a :=
  bind T (identity (T a)).


Definition kleisli_functor_data {C : precategory} (T : Kleisli_Data C) :
  functor_data C C :=
  mk_functor_data T (@map C T).

Definition is_functor_kleisli {C : precategory}
           {T : Kleisli_Data C} (H : Kleisli_Laws T) :
  is_functor(kleisli_functor_data T) :=
  map_id H ,, map_map H.

Definition kleisli_functor {C : precategory} (T : Kleisli C) : functor C C :=
  mk_functor (kleisli_functor_data T) (is_functor_kleisli T).

Lemma is_nat_trans_kleisli_mor {C : precategory} {T T' : Kleisli C} (α : Kleisli_Mor T T') :
  is_nat_trans (kleisli_functor T) (kleisli_functor T') (nat_trans_from_kleisli_mor α).
Proof.
  unfold is_nat_trans. intros. simpl. unfold map.
  rewrite <- (Kleisli_Mor_bind α).
  rewrite <- assoc.
  now rewrite (Kleisli_Mor_η α).
Defined.

Definition nat_trans_kleisli_mor {C : precategory} {T T' : Kleisli C} (α : Kleisli_Mor T T') :
  nat_trans (kleisli_functor T) (kleisli_functor T') :=
  mk_nat_trans (kleisli_functor T) (kleisli_functor T')
               (nat_trans_from_kleisli_mor α) (is_nat_trans_kleisli_mor α).

Lemma Kleisli_Mor_eq {C : precategory} (hs : has_homsets C)
      {T T' : Kleisli C} (α α' : Kleisli_Mor T T') :
  nat_trans_from_kleisli_mor α = nat_trans_from_kleisli_mor α' → α = α'.
Proof.
  apply Kleisli_Mor_equiv; assumption.
Defined.

Lemma is_nat_trans_η {C : precategory} (T : Kleisli C) :
  is_nat_trans (functor_identity C) (kleisli_functor T) (η T).
Proof.
  unfold is_nat_trans. simpl. intros.
  now rewrite (map_η T).
Defined.

Definition nat_trans_η {C : precategory} (T : Kleisli C) :
  functor_identity C ⟹ kleisli_functor T :=
  (η T,, is_nat_trans_η T).


Lemma is_nat_trans_μ {C : precategory} (T : Kleisli C) :
  is_nat_trans (kleisli_functor T □ kleisli_functor T) (kleisli_functor T) (μ T).
Proof.
  unfold is_nat_trans, μ. simpl. intros.
  rewrite (map_bind T), (bind_map T).
  now rewrite id_left, id_right.
Defined.

Definition nat_trans_μ {C : precategory} (T : Kleisli C) :
  kleisli_functor T □ kleisli_functor T ⟹ kleisli_functor T :=
  μ T,, is_nat_trans_μ T.


Lemma Kleisli_identity_laws {C : precategory} (T : Kleisli C) :
  Kleisli_Mor_laws T T (λ a : C, identity (T a)).
Proof.
  split; simpl; intros a.
  - apply id_right.
  - intros. do 2 rewrite id_right; apply id_left.
Defined.

Definition Kleisli_identity {C : precategory} (T : Kleisli C) : Kleisli_Mor T T :=
  (λ a : C, identity (T a)),, Kleisli_identity_laws T.


Lemma Kleisli_composition_laws {C : precategory} {T T' T'' : Kleisli C}
        (α : Kleisli_Mor T T') (α' : Kleisli_Mor T' T'') :
  Kleisli_Mor_laws T T''
    (λ a : C, nat_trans_from_kleisli_mor α a · nat_trans_from_kleisli_mor α' a).
Proof.
  split; intros; simpl.
  - rewrite assoc.
    set (H := Kleisli_Mor_η α). rewrite H.
    apply Kleisli_Mor_η.
  - pathvia (nat_trans_from_kleisli_mor α a ·
             (nat_trans_from_kleisli_mor α' a ·
              bind T'' ((f · nat_trans_from_kleisli_mor α b) ·
                        nat_trans_from_kleisli_mor α' b))).
    × now repeat rewrite assoc.
    × rewrite (Kleisli_Mor_bind α').
      rewrite assoc. rewrite (Kleisli_Mor_bind α).
      apply pathsinv0. apply assoc.
Defined.

Definition Kleisli_composition {C : precategory} {T T' T'' : Kleisli C}
           (α : Kleisli_Mor T T') (α' : Kleisli_Mor T' T'') :
  Kleisli_Mor T T'' :=
  (λ a : C, nat_trans_from_kleisli_mor α a · nat_trans_from_kleisli_mor α' a),,
  Kleisli_composition_laws α α'.


Definition precategory_Kleisli_ob_mor (C : precategory) : precategory_ob_mor :=
  precategory_ob_mor_pair (Kleisli C) Kleisli_Mor.

Definition precategory_Kleisli_Data (C : precategory) : precategory_data :=
  precategory_data_pair (precategory_Kleisli_ob_mor C)
                        (@Kleisli_identity C)
                        (@Kleisli_composition C).

Lemma precategory_Kleisli_axioms (C : precategory) (hs : has_homsets C) :
  is_precategory (precategory_Kleisli_Data C).
Proof.
  repeat split; simpl; intros.
  - apply (invmap (Kleisli_Mor_equiv hs _ _ )).
    apply funextsec. intros x. apply id_left.
  - apply (invmap (Kleisli_Mor_equiv hs _ _ )).
    apply funextsec. intros x. apply id_right.
  - apply (invmap (Kleisli_Mor_equiv hs _ _ )).
    apply funextsec. intros x. apply assoc.
  - apply (invmap (Kleisli_Mor_equiv hs _ _ )).
    apply funextsec. intros x. apply assoc'.
Defined.

Definition precategory_Kleisli (C : precategory) (hs : has_homsets C) : precategory :=
  precategory_Kleisli_Data C,, precategory_Kleisli_axioms C hs.

Lemma has_homsets_Kleisli (C : category) :
  has_homsets (precategory_Kleisli C (homset_property C)).
Proof.
  intros F G. simpl. unfold Kleisli_Mor.
  apply isaset_total2 .
  - apply impred_isaset.
    intro. apply C.
  - intros.
    apply isasetaprop.
    apply isaprop_Kleisli_Mor_laws.
    apply C.
Defined.


Definition category_Kleisli (C : category) : category :=
  precategory_Kleisli C (homset_property C),, has_homsets_Kleisli C.

Definition forgetfunctor_Kleisli (C : category) :
  functor (category_Kleisli C) (functor_category C C).
Proof.
  use mk_functor.
  - simpl.
    use mk_functor_data.
    + simpl.
      exact (λ T : Kleisli C, kleisli_functor T).
    + simpl. intros T T' α.
      exact (nat_trans_kleisli_mor α).
  - split.
    + red. intros. simpl. apply nat_trans_eq.
      × apply C.
      × intros; apply idpath.
    + unfold functor_compax. simpl. intros. apply nat_trans_eq.
      × apply C.
      × intros. apply idpath.
Defined.

Lemma forgetKleisli_faithful (C : category) : faithful (forgetfunctor_Kleisli C).
Proof.
  intros T T'. simpl.
  apply isinclbetweensets.
  - apply isaset_total2.
    + apply impred_isaset.
      intros. apply C.
    + intros. apply isasetaprop. apply isaprop_Kleisli_Mor_laws. apply C.
  - apply isaset_nat_trans. apply C.
  - intros α α' p.
    apply Kleisli_Mor_eq.
    + apply C.
    + apply funextsec. intro c.
      change (pr1 (nat_trans_kleisli_mor α) c = pr1 (nat_trans_kleisli_mor α') c).
      now rewrite p.
Defined.

Lemma is_univalent_precategory_Kleisli {C : precategory}
      (H: is_univalent C) (R R': Kleisli C)
  : is_univalent (precategory_Kleisli C (pr2 H)).
Proof.
  exact (is_univalent_RelMonad H (functor_identity C) R R').
Qed.

End Kleisli_precategory.