From Tealeaves Require Export
  Categories.TypeFamily
  Classes.Monoid
  Classes.Functor
  Classes.Categorical.Applicative
  Classes.Multisorted.Multifunctor
  Functors.Writer.

Import Categorical.Monad (Return, ret).

Import TypeFamily.Notations.
Import Product.Notations.
Import Monoid.Notations.

#[local] Generalizable Variables A B C F G ϕ W T U.

Section MultisortedDTM_typeclasses.

  Context
    `{ix: Index}.

  Section operations.

    Context
      (W: Type)
      (T: K → Type → Type)
      (U: Type → Type).

    Class MReturn :=
      mret: ∀ (A: Type) (k: K), A → T k A.

    Class MBind :=
      mbinddt: ∀ (F: Type → Type) `{Map F} `{Pure F} `{Mult F} {A B: Type},
          (∀ (k: K), W × A → F (T k B)) → U A → F (U B).

  End operations.

  Definition compose_dtm
    {W: Type}
    {T: K → Type → Type}
    `{mn_op: Monoid_op W}
    `{mn_unit: Monoid_unit W}
    `{Map F} `{Mult F} `{Pure F}
    `{Map G} `{Mult G} `{Pure G}
    `{∀ k, MBind W T (T k)}
    {A B C: Type}
    (g: ∀ k, W × B → G (T k C))
    (f: ∀ k, W × A → F (T k B)): ∀ k, W × A → F (G (T k C)) :=
    fun (k: K) '(w, a) ⇒
      map (F := F)
        (mbinddt W T (T k) G (g ◻ allK (incr w))) (f k (w, a)).

  Infix "⋆dtm" := compose_dtm (at level 40): tealeaves_scope.

  Section PreModule.

    Context
      (W: Type)
      (T: K → Type → Type)
      (U: Type → Type)
      `{! MReturn T}
      `{! MBind W T U}
      `{! ∀ k, MBind W T (T k)}
      {mn_op: Monoid_op W}
      {mn_unit: Monoid_unit W}.

    Class MultiDecoratedTraversablePreModule :=
      { dtp_monoid :> Monoid W;
        dtp_mbinddt_mret: ∀ A,
          mbinddt W T U (fun a ⇒ a) (mret T A ◻ allK extract) = @id (U A);
        dtp_mbinddt_mbinddt: ∀
          (F: Type → Type)
          (G: Type → Type)
          `{Applicative F}
          `{Applicative G}
          `(g: ∀ k, W × B → G (T k C))
          `(f: ∀ k, W × A → F (T k B)),
          map (F := F) (mbinddt W T U G g) ∘ mbinddt W T U F f =
            mbinddt W T U (F ∘ G) (g ⋆dtm f);
        dtp_mbinddt_morphism: ∀
          (F: Type → Type)
          (G: Type → Type)
          `{ApplicativeMorphism F G ϕ}
          `(f: ∀ k, W × A → F (T k B)),
          ϕ (U B) ∘ mbinddt W T U F f =
            mbinddt W T U G ((fun k ⇒ ϕ (T k B)) ◻ f);
      }.

  End PreModule.

  Section DTM.

    Context
      (W: Type)
      (T: K → Type → Type)
      `{! MReturn T}
      `{! ∀ k, MBind W T (T k)}
      {mn_op: Monoid_op W}
      {mn_unit: Monoid_unit W}.

    Class MultiDecoratedTraversableMonad :=
      { dtm_pre :> ∀ k, MultiDecoratedTraversablePreModule W T (T k);
        dtm_mbinddt_comp_mret:
        ∀ k F `{Applicative F}
          `(f: ∀ k, W × A → F (T k B)),
          mbinddt W T (T k) F f ∘ mret T A k = f k ∘ pair Ƶ;
      }.

  End DTM.

End MultisortedDTM_typeclasses.

Arguments mret {ix} _%function_scope {MReturn} {A}%type_scope _ _.
Arguments mbinddt {ix} {W}%type_scope {T} (U)%function_scope {MBind} F%function_scope {H H0 H1} {A B}.

#[local] Infix "⋆dtm" := compose_dtm (at level 40): tealeaves_scope.

Section mapMret.

  Context
  `{ix: Index}
  `{! MReturn T}.

  Definition mapMret
    `{Map F} {A:Type}:
    ∀ (k: K), F A → F (T k A) :=
    vec_apply (fun k ⇒ map (A := A) (B := T k A)) (mret T).

  Lemma vec_compose_mret {W A B}:
    ∀ (f: K → W × A → B) (k:K),
      (mret T ◻ f) k =
        (mret T k ∘ f k).
  Proof.
    reflexivity.
  Qed.

  Lemma vec_compose_mapMret {W A B} `{Functor F}:
    ∀ (f: K → W × A → F B) (k:K),
      (mapMret (F := F) ◻ f) k =
        (map (F := F) (mret T k) ∘ f k).
  Proof.
    reflexivity.
  Qed.

End mapMret.

Section multisorted_dtm_kleisli_composition.

  Context
    `{ix: Index}
    {W: Type}
    {T: K → Type → Type}
    {U: Type → Type}
    `{! MReturn T}
    `{! MBind W T U}
    `{! ∀ k, MBind W T (T k)}
    {mn_op: Monoid_op W}
    {mn_unit: Monoid_unit W}.

  Context
    `{Applicative G}
    `{Applicative F}
    `{! Monoid W}
    {A B C: Type}
    {g: ∀ k, W × B → G (T k C)}.

  Lemma compose_dtm_incr
    {f: ∀ k, W × A → F (T k B)}:
      ∀ (w: W),
        (fun (k: K) ⇒ (g ⋆dtm f) k ∘ incr w) =
          ((fun (k: K) ⇒ g k ∘ incr w) ⋆dtm (fun (k: K) ⇒ f k ∘ incr w)).
  Proof.
    intros. ext k [w' a].
    cbn. do 2 fequal.
    ext j [w'' b].
    unfold vec_compose, compose. cbn. fequal.
    now rewrite monoid_assoc.
  Qed.

  Lemma compose_dtm_incr_alt
    {f: ∀ k, W × A → F (T k B)}:
    ∀ (w: W),
      vec_compose
        (C := fun (k: K) ⇒ F (G (T k C)))
        (g ⋆dtm f) (allK (incr w)) =
        (g ◻ allK (incr w)) ⋆dtm (f ◻ allK (incr w)).
  Proof.
    intros.
    ext k [w' a].
    cbn. do 2 fequal.
    ext j [w'' b].
    unfold vec_compose, compose.
    cbn. fequal.
    now rewrite monoid_assoc.
  Qed.

  Context
    `{! MultiDecoratedTraversableMonad W T}.

  Lemma compose_dtm_lemma1
    {f: ∀ (k:K), W × A → B}:
      g ⋆dtm (mret T ◻ f) =
        (fun (k: K) '(w, a) ⇒ g k (w, f k (w, a))).
  Proof.
    intros.
    unfold compose_dtm.
    ext k [w a].
    unfold_ops @Map_I.
    rewrite vec_compose_mret.
    unfold compose.
    compose near (f k (w, a)) on left.
    rewrite (dtm_mbinddt_comp_mret W T k); auto.
    rewrite vec_compose_k.
    reassociate → on left.
    unfold allK, const.
    rewrite pair_incr_zero.
    reflexivity.
  Qed.

  Lemma compose_dtm_lemma2
    {f: ∀ (k:K), W × A → F B}:
      g ⋆dtm (mapMret (F := F) ◻ f) =
        (fun (k: K) '(w, a) ⇒
           map (F := F) (g k ∘ pair w) (f k (w, a))).
  Proof.
    intros.
    unfold compose_dtm.
    ext k [w a].
    rewrite vec_compose_mapMret.
    unfold compose.
    compose near (f k (w, a)) on left.
    rewrite fun_map_map.
    rewrite (dtm_mbinddt_comp_mret W T k); auto.
    rewrite vec_compose_k.
    reassociate → on left.
    unfold allK, const.
    rewrite pair_incr_zero.
    reflexivity.
  Qed.

End multisorted_dtm_kleisli_composition.

Section DTM_laws.

  Context
    (U: Type → Type)
    `{MultiDecoratedTraversablePreModule W T U}
    `{! MultiDecoratedTraversableMonad W T}.

  Lemma mbinddt_pure: ∀ (A: Type) `{Applicative F},
      (let p := ((fun k ⇒ (@pure F _ (T k A))): ∀ k, T k A → F (T k A)) in
       let r := (@mret ix T _ A) in
       let e := (allK extract): ∀ k, W × A → A
       in mbinddt U F (p ◻ r ◻ e)) = pure (A := U A).
  Proof.
    intros.
    cbn zeta.
    rewrite vec_compose_assoc.
    rewrite <- (dtp_mbinddt_morphism W T U (fun x ⇒ x) F (ϕ := @pure F _)).
    rewrite (dtp_mbinddt_mret W T U).
    reflexivity.
  Qed.

End DTM_laws.

Section derived_operations.

  Context
    (U: Type → Type)
    `{MultiDecoratedTraversablePreModule W T U}.

  Section definitions.

    Context
      {A B: Type}
      (F: Type → Type)
      `{Map F} `{Mult F} `{Pure F}.

    Definition mbindd (f: ∀ k, W × A → T k B): U A → U B :=
      mbinddt U (fun x ⇒ x) f.

    Definition mmapdt (f: ∀ (k: K), W × A → F B): U A → F (U B) :=
      mbinddt U F (mapMret (T := T) ◻ f).

    Definition mbindt (f: ∀ k, A → F (T k B)): U A → F (U B) :=
      mbinddt U F (f ◻ allK extract).

    Definition mbind (f: ∀ k, A → T k B): U A → U B :=
      mbindd (f ◻ allK extract).

    Definition mmapd (f: ∀ k, W × A → B): U A → U B :=
      mbindd (mret T ◻ f).

    Definition mmapt (f: ∀ k, A → F B): U A → F (U B) :=
      mmapdt (f ◻ allK extract).

    Definition mmap (f: ∀ k, A → B): U A → U B :=
      mmapd (f ◻ allK extract).

  End definitions.

  Section special_cases.

    Context
      {A B: Type}
      (F: Type → Type)
      `{Map F} `{Mult F} `{Pure F}.

    Lemma mbindt_to_mbinddt (f: ∀ k, A → F (T k B)):
      mbindt F f = mbinddt U F (f ◻ allK extract).
    Proof.
      reflexivity.
    Qed.

    Lemma mbindd_to_mbinddt (f: ∀ k, W × A → T k B):
      mbindd f = mbinddt U (fun A ⇒ A) f.
    Proof.
      reflexivity.
    Qed.

    Lemma mmapdt_to_mbinddt (f: K → W × A → F B):
      mmapdt F f = mbinddt U F (mapMret (T := T) ◻ f).
    Proof.
      reflexivity.
    Qed.

    Lemma mbind_to_mbinddt (f: ∀ k, A → T k B):
      mbind f = mbinddt U (fun A ⇒ A) (f ◻ allK extract).
    Proof.
      reflexivity.
    Qed.

    Lemma mmapd_to_mbinddt (f: K → W × A → B):
      mmapd f = mbinddt U (fun A ⇒ A) (mret T ◻ f).
    Proof.
      reflexivity.
    Qed.

    Lemma mmapt_to_mbinddt (f: K → A → F B):
      mmapt F f = mbinddt U F (mapMret (T := T)
                                 ◻ f ◻ allK extract).
    Proof.
      reflexivity.
    Qed.

    Lemma mmap_to_mbinddt (f: K → A → B):
      mmap f = mbinddt U (fun A ⇒ A) (mret T ◻ f ◻ allK extract).
    Proof.
      reflexivity.
    Qed.

    Lemma mbind_to_mbindt (f: ∀ k, A → T k B):
      mbind f = mbindt (fun A ⇒ A) f.
    Proof.
      reflexivity.
    Qed.

    Lemma mmapt_to_mbindt (f: K → A → F B):
      mmapt F f = mbindt F (mapMret (T := T) ◻ f).
    Proof.
      reflexivity.
    Qed.

    Lemma mmap_to_mbindt (f: K → A → B):
      mmap f = mbindt (fun A ⇒ A) (mret T ◻ f).
    Proof.
      reflexivity.
    Qed.

    Lemma mbind_to_mbindd (f: ∀ k, A → T k B):
      mbind f = mbindd (f ◻ allK extract).
    Proof.
      reflexivity.
    Qed.

    Lemma mmapd_to_mbindd (f: W × A -k→ B):
      mmapd f = mbindd (mret T ◻ f).
    Proof.
      reflexivity.
    Qed.

    Lemma mmap_to_mbindd (f: A -k→ B):
      mmap f = mbindd (mret T ◻ f ◻ allK extract).
    Proof.
      reflexivity.
    Qed.

    Lemma mmapd_to_mmapdt (f: K → W × A → B):
      mmapd f = mmapdt (fun A ⇒ A) f.
    Proof.
      reflexivity.
    Qed.

    Lemma mmap_to_mmapdt (f: K → A → B):
      mmap f = mmapdt (fun A ⇒ A) (f ◻ allK extract).
    Proof.
      reflexivity.
    Qed.

    Lemma mmapt_to_mmapdt (f: K → A → F B):
      mmapt F f = mmapdt F (f ◻ allK extract).
    Proof.
      reflexivity.
    Qed.

    Lemma mmap_to_mmapt (f: K → A → B):
      mmap f = mmapt (fun A ⇒ A) f.
    Proof.
      reflexivity.
    Qed.

    Lemma mmap_to_mmapd (f: K → A → B):
      mmap f = mmapd (f ◻ allK extract).
    Proof.
      reflexivity.
    Qed.

    Lemma mmap_to_mbind (f: K → A → B):
      mmap f = mbind (mret T ◻ f).
    Proof.
      reflexivity.
    Qed.

  End special_cases.

End derived_operations.

Section derived_operations_composition.

  Context
    (U: Type → Type)
    `{MultiDecoratedTraversablePreModule W T U}
    `{! MultiDecoratedTraversableMonad W T}
    `{Applicative F}
    `{Applicative G}
    {A B C: Type}.

  Lemma mbindd_mbinddt: ∀
      (g: ∀ k, W × B → T k C)
      (f: ∀ k, W × A → F (T k B)),
      map (F := F) (mbindd U g) ∘ mbinddt U F f =
        mbinddt U F
          (fun k '(w, a) ⇒
             map (F := F)
               (mbindd (T k) (g ◻ allK (incr w))) (f k (w, a))).
  Proof.
    intros. rewrite mbindd_to_mbinddt.
    rewrite (dtp_mbinddt_mbinddt W T U F (fun A ⇒ A)).
    fequal. now erewrite Mult_compose_identity1.
  Qed.

  Lemma mmapdt_mbinddt: ∀
      (g: K → W × B → G C)
      (f: ∀ k, W × A → F (T k B)),
      map (F := F) (mmapdt U G g) ∘ mbinddt U F f =
        mbinddt U (F ∘ G)
          (fun k '(w, a) ⇒
             map (F := F)
               (mmapdt (T k) G (g ◻ allK (incr w))) (f k (w, a))).
  Proof.
    intros. rewrite mmapdt_to_mbinddt.
    now rewrite (dtp_mbinddt_mbinddt W T U F G).
  Qed.

  Lemma mbindt_mbinddt: ∀
      (g: ∀ k, B → G (T k C))
      (f: ∀ k, W × A → F (T k B)),
      map (F := F) (mbindt U G g) ∘ mbinddt U F f =
        mbinddt U (F ∘ G)
          (fun k ⇒ map (F := F) (mbindt (T k) G g) ∘ f k).
  Proof.
    intros. rewrite mbindt_to_mbinddt.
    rewrite (dtp_mbinddt_mbinddt W T U F G).
    fequal. ext k [w a]. unfold compose; cbn.
    fequal. rewrite mbindt_to_mbinddt.
    fequal. now ext j [w2 b].
  Qed.

  Lemma mbind_mbinddt: ∀
      (g: ∀ k, B → T k C)
      (f: ∀ k, W × A → F (T k B)),
      map (F := F) (mbind U g) ∘ mbinddt U F f =
        mbinddt U F ((fun k ⇒ map (F := F) (mbind (T k) g)) ◻ f).
  Proof.
    intros. rewrite mbind_to_mbinddt.
    rewrite (dtp_mbinddt_mbinddt W T U F (fun A ⇒ A)).
    fequal.
    - now erewrite Mult_compose_identity1.
    - unfold vec_compose, compose, compose_dtm.
      ext k [w a].
      fequal. rewrite mbind_to_mbinddt. fequal.
      now ext j [w2 b].
  Qed.

  Lemma mmapd_mbinddt: ∀
      (g: K → W × B → C)
      (f: ∀ k, W × A → F (T k B)),
      map (F := F) (mmapd U g) ∘ mbinddt U F f =
        mbinddt U F
          (fun k '(w, a) ⇒
             map (F := F)
               (mmapd (T k) (g ◻ allK (incr w))) (f k (w, a))).
  Proof.
    intros. rewrite mmapd_to_mbinddt.
    rewrite (dtp_mbinddt_mbinddt W T U F (fun A ⇒ A)).
    fequal. now erewrite Mult_compose_identity1.
  Qed.

  Lemma mmapt_mbinddt: ∀
      (g: K → B → G C)
      (f: ∀ k, W × A → F (T k B)),
      map (F := F) (mmapt U G g) ∘ mbinddt U F f =
        mbinddt U (F ∘ G)
          (fun k ⇒ map (F := F) (mmapt (T k) G g) ∘ f k).
  Proof.
    intros. rewrite mmapt_to_mbinddt.
    rewrite (dtp_mbinddt_mbinddt W T U F G).
    fequal. ext k [w a]. unfold compose; cbn.
    fequal. rewrite mmapt_to_mbinddt.
    fequal. now ext j [w2 b].
  Qed.

  Lemma mmap_mbinddt: ∀
      (g: K → B → C)
      (f: ∀ k, W × A → F (T k B)),
      map (F := F) (mmap U g) ∘ mbinddt U F f =
        mbinddt U F (fun k ⇒ map (F := F) (mmap (T k) g) ∘ f k).
  Proof.
    intros. unfold mmap. rewrite mmapd_mbinddt.
    fequal. ext k [w a]. unfold compose; cbn.
    fequal. fequal. now ext j [w2 b].
  Qed.

  Lemma mbinddt_mbindd: ∀
      (g: ∀ k, W × B → G (T k C))
      (f: ∀ k, W × A → T k B),
      mbinddt U G g ∘ mbindd U f =
        mbinddt U G
          (fun k '(w, a) ⇒
             mbinddt (T k) G (g ◻ allK (incr w)) (f k (w, a))).
  Proof.
    intros. rewrite mbindd_to_mbinddt.
    change (mbinddt U G g) with
      (map (F := (fun A ⇒ A)) (mbinddt U G g)).
    rewrite (dtp_mbinddt_mbinddt W T U (fun A ⇒ A) G).
    fequal. now rewrite (Mult_compose_identity2 G).
  Qed.

  Lemma mbinddt_mmapdt: ∀
      (g: ∀ k, W × B → G (T k C))
      (f: K → W × A → F B),
      map (F := F) (mbinddt U G g) ∘ mmapdt U F f =
        mbinddt U (F ∘ G)
          (fun k '(w, a) ⇒ map (F := F) (fun b ⇒ g k (w, b)) (f k (w, a))).
  Proof.
    intros.
    rewrite mmapdt_to_mbinddt.
    rewrite (dtp_mbinddt_mbinddt W T U F G).
    fequal.
    ext k [w a]. unfold compose; cbn.
    rewrite vec_compose_mapMret.
    unfold compose at 1.
    compose near (f k (w, a)) on left.
    rewrite fun_map_map.
    rewrite (dtm_mbinddt_comp_mret W T k G).
    rewrite vec_compose_allK2.
    reassociate → on left.
    rewrite pair_incr_zero.
    reflexivity.
  Qed.

  Lemma mbinddt_mbindt: ∀
      (g: ∀ k, W × B → G (T k C))
      (f: ∀ k, A → F (T k B)),
      map (F := F) (mbinddt U G g) ∘ mbindt U F f =
        mbinddt U (F ∘ G)
          (fun k '(w, a) ⇒
             map (F := F)
               (mbinddt (T k) G (g ◻ allK (incr w))) (f k a)).
  Proof.
    intros.
    rewrite mbindt_to_mbinddt.
    rewrite (dtp_mbinddt_mbinddt W T U F G).
    reflexivity.
  Qed.

  Lemma mbinddt_mbind: ∀
      (g: ∀ k, W × B → G (T k C))
      (f: ∀ k, A → T k B),
      mbinddt U G g ∘ mbind U f =
        mbinddt U G
          (fun k '(w, a) ⇒
             mbinddt (T k) G (g ◻ allK (incr w)) (f k a)).
  Proof.
    intros. rewrite mbind_to_mbinddt.
    change (mbinddt U G g) with (map (F := fun A ⇒ A) (mbinddt U G g)).
    rewrite (dtp_mbinddt_mbinddt W T U (fun A ⇒ A) G).
    fequal. now rewrite (Mult_compose_identity2 G).
  Qed.

  Lemma mbinddt_mmapd: ∀
      (g: ∀ k, W × B → G (T k C))
      (f: ∀ k, W × A → B),
      mbinddt U G g ∘ mmapd U f =
        mbinddt U G (fun k '(w, a) ⇒ g k (w, f k (w, a))).
  Proof.
    intros.
    erewrite mmapd_to_mbinddt.
    change (mbinddt U G g) with (map (F := fun A ⇒ A) (mbinddt U G g)).
    rewrite (dtp_mbinddt_mbinddt W T U (fun A ⇒ A) G).
    fequal. now rewrite (Mult_compose_identity2 G).
    rewrite compose_dtm_lemma1.
    reflexivity.
  Qed.

  Lemma mbinddt_mmapt: ∀
      (g: ∀ k, W × B → G (T k C))
      (f: K → A → F B),
      map (F := F) (mbinddt U G g) ∘ mmapt U F f =
        mbinddt U (F ∘ G)
          (fun k '(w, a) ⇒ map (F := F) (fun b ⇒ g k (w, b)) (f k a)).
  Proof.
    intros.
    rewrite mmapt_to_mmapdt.
    rewrite mbinddt_mmapdt.
    reflexivity.
  Qed.

  Lemma mbinddt_mmap: ∀
      (g: ∀ k, W × B → G (T k C))
      (f: K → A → B),
      mbinddt U G g ∘ mmap U f =
        mbinddt U G (fun k '(w, a) ⇒ g k (w, f k a)).
  Proof.
    intros.
    rewrite mmap_to_mmapd.
    rewrite mbinddt_mmapd.
    reflexivity.
  Qed.

End derived_operations_composition.

Section mixed_composition_laws.

  Context
    (U: Type → Type)
    (F G: Type → Type)
    `{Applicative F}
    `{Applicative G}
    `{MultiDecoratedTraversablePreModule W T U}
    `{! MultiDecoratedTraversableMonad W T} {A B C: Type}.

  Lemma mbindd_mmapdt: ∀
      (g: ∀ k, W × B → T k C)
      (f: ∀ k, W × A → F B),
      map (F := F) (mbindd U g) ∘ mmapdt U F f =
        mbinddt U F (fun (k: K) '(w, a) ⇒
                       map (F := F) (fun b ⇒ g k (w, b)) (f k (w, a))).
  Proof.
    intros. rewrite mmapdt_to_mbinddt.
    rewrite mbindd_to_mbinddt.
    rewrite (dtp_mbinddt_mbinddt W T U F (fun x ⇒ x)).
    rewrite compose_dtm_lemma2.
    rewrite (Mult_compose_identity1 F).
    reflexivity.
  Qed.

  Lemma mbindd_mbindt: ∀
      (g: ∀ k, W × B → T k C)
      (f: ∀ k, A → F (T k B)),
      map (F := F) (mbindd U g) ∘ mbindt U F f =
        mbinddt U F
          (fun (k: K) '(w, a) ⇒
             map (F := F) (mbindd (T k) (g ◻ allK (incr w))) (f k a)).
  Proof.
    intros.
    rewrite mbindd_to_mbinddt.
    rewrite mbindt_to_mbinddt.
    rewrite (dtp_mbinddt_mbinddt W T U F (fun x ⇒ x)).
    rewrite (Mult_compose_identity1 F).
    reflexivity.
  Qed.

  Lemma mbindd_mmapt: ∀
      (g: ∀ k, W × B → T k C)
      (f: ∀ k, A → F B),
      map (F := F) (mbindd U g) ∘ mmapt U F f =
        mbinddt U F
          (fun (k: K) '(w, a) ⇒
             map (F := F) (fun b ⇒ g k (w, b)) (f k a)).
  Proof.
    intros.
    rewrite mbindd_to_mbinddt.
    rewrite mmapt_to_mbinddt.
    rewrite (dtp_mbinddt_mbinddt W T U F (fun x ⇒ x)).
    rewrite vec_compose_assoc.
    rewrite compose_dtm_lemma2.
    rewrite (Mult_compose_identity1 F).
    reflexivity.
  Qed.

  (* TODO Also need to put <<mmapt_mbindd>> somewhere. *)

  Lemma mmapdt_mbindd: ∀
      (g: ∀ k, W × B → G C)
      (f: ∀ k, W × A → T k B),
      mmapdt U G g ∘ mbindd U f =
        mbinddt U G
          (fun (k: K) '(w, a) ⇒
             mmapdt (T k) G (g ◻ allK (incr w)) (f k (w, a))).
  Proof.
    (* TODO *)
  Abort.

  Lemma mmapdt_mbindt: ∀
      (g: ∀ k, W × B → G C)
      (f: ∀ k, A → F (T k B)),
      map (F := F) (mmapdt U G g) ∘ mbindt U F f =
        mbinddt U (F ∘ G)
          (fun (k: K) '(w, a) ⇒
             map (F := F) (mmapdt (T k) G (g ◻ allK (incr w))) (f k a)).
  Proof.
    (* TODO *)
  Abort.

  Lemma mmapdt_mbind: ∀
      (g: K → W × B → G C) (f: ∀ k, A → T k B),
      mmapdt U G g ∘ mbind U f =
        mbinddt U G (fun k '(w, a) ⇒
                       mmapdt (T k) G (g ◻ allK (incr w)) (f k a)).
  Proof.
    (* TODO *)
  Abort.

  (* TODO Also need to put <<mbind_mmapdt>> somewhere. *)

  Lemma mbindt_mbindd: ∀
      (g: ∀ k, B → G (T k C))
      (f: ∀ k, W × A → T k B),
      mbindt U G g ∘ mbindd U f =
        mbinddt U G (fun (k: K) '(w, a) ⇒ mbindt (T k) G g (f k (w, a))).
  Proof.
    (* TODO *)
  Abort.

  Lemma mbindt_mmapdt: ∀
      (g: ∀ k, B → G (T k C))
      (f: ∀ k, W × A → F B),
      map (F := F) (mbindt U G g) ∘ mmapdt U F f =
        mbinddt U (F ∘ G)
          (fun (k: K) '(w, a) ⇒ map (F := F) (g k) (f k (w, a))).
  Proof.
    (* TODO *)
  Abort.

  (* TODO <<mbindt_mmapd>> *)

  (* TODO Also need to put <<mmapd_mbindt>> somewhere. *)

End mixed_composition_laws.

Definition compose_dm
  `{ix: Index}
  {W: Type}
  {T: K → Type → Type}
  `{mn_op: Monoid_op W}
  `{mn_unit: Monoid_unit W}
  `{∀ k, MBind W T (T k)}
  {A B C: Type}
  (g: ∀ k, W × B → T k C)
  (f: ∀ k, W × A → T k B): ∀ k, W × A → T k C :=
  fun k '(w, a) ⇒ mbindd (T k) (g ◻ allK (incr w)) (f k (w, a)).

Infix "⋆dm" := compose_dm (at level 40).

Section DecoratedMonad.

  Context
    (U: Type → Type)
    `{MultiDecoratedTraversablePreModule W T U}
    `{! MultiDecoratedTraversableMonad W T} {A B C: Type}.

  Theorem mbindd_id:
    mbindd U (mret T ◻ allK extract) = @id (U A).
  Proof.
    intros. rewrite mbindd_to_mbinddt.
    now rewrite <- (dtp_mbinddt_mret W T U).
  Qed.

  Theorem mbindd_mbindd: ∀
      (g: W × B ¬k~> T C)
      (f: W × A ¬k~> T B),
      mbindd U g ∘ mbindd U f =
        mbindd U (fun (k: K) '(w, a) ⇒
                    mbindd (T k) (g ◻ allK (incr w)) (f k (w, a))).
  Proof.
    intros.
    rewrite 3mbindd_to_mbinddt.
    change_left (map (F := fun x ⇒ x)
                   (mbinddt U (fun x: Type ⇒ x) g) ∘
                   (mbinddt U (fun x: Type ⇒ x) f)).
    rewrite (dtp_mbinddt_mbinddt W T U (fun x ⇒ x) (fun x ⇒ x)).
    rewrite Mult_compose_identity2.
    reflexivity.
  Qed.

  Theorem mbindd_comp_mret: ∀
      (k: K) (f: ∀ k, W × A → T k B),
      mbindd (T k) f ∘ mret T k = f k ∘ ret (T := (W ×)).
  Proof.
    intros. rewrite mbindd_to_mbinddt.
    now rewrite (dtm_mbinddt_comp_mret W T k (fun A ⇒ A)).
  Qed.

  Lemma mbindd_mbind: ∀
      (g: ∀ k, W × B → T k C)
      (f: A ¬k~> T B),
      mbindd U g ∘ mbind U f =
        mbindd U (fun (k: K) '(w, a) ⇒
                    mbindd (T k) (g ◻ allK (incr w)) (f k a)).
  Proof.
    intros. rewrite mbind_to_mbindd.
    now rewrite mbindd_mbindd.
  Qed.

  Lemma mbindd_mmapd: ∀
      (g: ∀ k, W × B → T k C)
      (f: ∀ k, W × A → B),
      mbindd U g ∘ mmapd U f =
        mbindd U (fun (k: K) '(w, a) ⇒ g k (w, f k (w, a))).
  Proof.
    intros. rewrite mmapd_to_mbindd.
    rewrite (mbindd_mbindd).
    fequal. ext k [w a].
    unfold vec_compose, compose; cbn.
    compose near (f k (w, a)).
    rewrite mbindd_to_mbinddt.
    rewrite (dtm_mbinddt_comp_mret W T k (fun A ⇒ A)).
    unfold compose; cbn. now simpl_monoid.
  Qed.

  Lemma mbindd_mmap: ∀
      (g: ∀ k, W × B → T k C)
      (f: ∀ k, A → B),
      mbindd U g ∘ mmap U f =
        mbindd U (fun (k: K) '(w, a) ⇒ g k (w, f k a)).
  Proof.
    intros. unfold mmap.
    now rewrite (mbindd_mmapd).
  Qed.