From Tealeaves Require Export
     Classes.RightModule
     Classes.DecoratedMonad.

Import Functor.Notations.
Import Monad.Notations.
Import Comonad.Notations.
Import Monoid.Notations.
Import DecoratedMonad.Notations.
#[local] Open Scope tealeaves_scope.

Section shift_module_lemmas.

  Context
    `{RightModule F T}
    `{Monoid W}.

  Lemma shift_right_action {A} (t : F (T (W * A))) (w : W) :
    shift F (w, right_action F t) =
    right_action F (fmap F (fun t => shift T (w, t)) t).
  Proof.
    rewrite (shift_spec F). compose near t on left.
    rewrite natural. unfold compose; cbn.
    fequal. unfold_ops @Fmap_compose.
    fequal. ext x. now rewrite (shift_spec T).
  Qed.

  Lemma shift_sub `(t : F (W * A)) (w : W) `(f : W * A -> T (W * B)) :
    shift F (w, sub F f t) =
    sub F (fun p => shift T (w, f p)) t.
  Proof.
    unfold_ops @Sub_RightAction. unfold compose.
    rewrite shift_right_action. fequal.
    compose near t on left.
    now rewrite (fun_fmap_fmap F).
  Qed.

End shift_module_lemmas.

Section DecoratedModule_class.

  Context
    (W : Type)
    (F T : Type -> Type)
    `{Fmap T} `{Return T} `{Join T} `{Decorate W T}
    `{Fmap F} `{RightAction F T} `{Decorate W F}
    `{Monoid_op W} `{Monoid_unit W}.

  Class DecoratedModule :=
    { drmod_monad :> DecoratedMonad W T;
      drmod_functor :> DecoratedFunctor W F;
      drmod_module :> RightModule F T;
      drmod_action :
        `(dec F ∘ right_action F (A:=A) =
          right_action F ∘ fmap F (shift T) ∘ dec F ∘ fmap F (dec T));
    }.

End DecoratedModule_class.

Section DecoratedModule_Monad.

  Context
    (T : Type -> Type)
    `{DecoratedMonad W T}.

  Existing Instance RightAction_Join.

  Instance DecoratedModule_Monad : DecoratedModule W T T :=
    {| drmod_action := dmon_join W T;
       drmod_functor := dmon_functor W T;
       drmod_module := RightModule_Monad T;
    |}.

End DecoratedModule_Monad.

Section DecoratedModule_compose.

  Context
    `{Monoid W}
    `{Fmap F} `{Fmap G} `{Fmap T}
    `{Return T} `{Join T} `{RightAction F T}
    `{Decorate W F} `{Decorate W G} `{Decorate W T}
    `{! DecoratedFunctor W G}
    `{! DecoratedModule W F T}.

  Instance: RightModule (G ∘ F) T := RightModule_compose.

  Theorem drmod_action_compose :
    `(dec (G ∘ F) ∘ right_action (G ∘ F) (A:=A) =
      right_action (G ∘ F) ∘ fmap (G ∘ F) (shift T) ∘ dec (G ∘ F) ∘ fmap (G ∘ F) (dec T)).
  Proof.
    intros. unfold_ops @Decorate_compose @RightAction_compose.
    (* Use <<drmod_action>> to write the LHS with a butterfly *)
    change (?x ∘ fmap G (dec F) ∘ fmap G (right_action F))
      with (x ∘ (fmap G (dec F) ∘ fmap G (right_action F))).
    rewrite (fun_fmap_fmap G). rewrite (drmod_action W F T).
    (* Use naturality to cancel out equal terms (decorations on F and T) *)
    rewrite <- (fun_fmap_fmap G).
    change (fmap G (fmap F (dec T))) with (fmap (G ∘ F) (dec T (A := A))).
    repeat reassociate <-. fequal.
    rewrite <- (fun_fmap_fmap G).
    repeat reassociate <-. fequal.
    (* Bring left <<join (prod W)>> across the distribution, then re-associate *)
    reassociate -> on right.
    unfold shift at 3.
    reassociate <- on right.
    rewrite <- (fun_fmap_fmap G).
    reassociate <- on left.
    unfold shift at 3.
    rewrite <- (fun_fmap_fmap G).
    reassociate <- on right.
    change (fmap G (fmap F (join (A := ?A) (prod W))))
      with (fmap (G ∘ F) (join (A := A) (prod W))).
    reassociate -> near (fmap (G ∘ F) (join (prod W))).
    rewrite (fun_fmap_fmap (G ∘ F)).
    reassociate -> near (join (prod W)).
    rewrite (strength_join_l).
    repeat reassociate <-.
    rewrite (fun_fmap_fmap T).
    rewrite (mon_join_join (prod W)).
    (* Bring the right action of <<T>> on <<F>> under the distribution *)
    unfold_ops @Fmap_compose.
    unfold_compose_in_compose.
    rewrite (fun_fmap_fmap G (g := (@right_action F T _ (W * A)))).
    repeat reassociate -> on right.
    rewrite <- (fun_fmap_fmap F).
    rewrite <- (fun_fmap_fmap F).
    repeat reassociate <- on right.
    change (fmap F (fmap T ?f)) with (fmap (F ∘ T) f).
    rewrite <- (fun_fmap_fmap (F ∘ T)).
    reassociate <- on right.
    #[local] Set Keyed Unification.
    rewrite <- (natural (ϕ := @right_action F T _)).
    #[local] Unset Keyed Unification.
    change (fmap (F ∘ T) (fmap (prod W) ?f)) with (fmap F (fmap (T ○ prod W) f)).
    reassociate -> near (fmap F (strength T)).
    rewrite (fun_fmap_fmap F).
    rewrite (natural (ϕ := @strength T _ _)).
    (* Bring up the <<fmap G (strength F) ∘ dec G>> *)
    rewrite (fun_fmap_fmap G).
    #[local] Set Keyed Unification.
    reassociate -> near (strength F).
    change (fmap F (fmap (prod W) ?f)) with (fmap (F ○ prod W) f).
    rewrite (natural (ϕ := @strength F _ _)).
    reassociate <-.
    rewrite <- (fun_fmap_fmap G).
    reassociate -> near (dec G).
    change (fmap G (fmap (prod W ○ F) (strength (B:=?B) (A:=?A) T)) ∘ dec G)
      with (fmap (G ∘ prod W) (fmap F (strength (B:=B) (A:=A) T)) ∘ dec G).
    rewrite (natural (ϕ := @dec W G _)).
    reassociate <-.
    rewrite <- (fun_fmap_fmap F).
    change (fmap F (fmap (prod W ○ T) ?f)) with (fmap (F ○ prod W) (fmap T f)).
    reassociate <-. reassociate -> near (strength F).
    rewrite (natural (ϕ := @strength F _ _)).
    reassociate <-.
    change (?x ∘ right_action F ∘ fmap F (strength T) ∘ strength F ∘ ?y)
      with (x ∘ (right_action F ∘ fmap F (strength T) ∘ strength F) ∘ y).
    rewrite <- (strength_right_action T F).
    reassociate -> on right.
    rewrite <- (fun_fmap_fmap G).
    rewrite <- (fun_fmap_fmap G).
    rewrite <- (fun_fmap_fmap G).
    change (fmap G (fmap (prod W ○ F) (fmap T (join (A:=?A) (prod W)))))
      with (fmap (G ∘ prod W) (fmap F (fmap T (join (A:= A) (prod W))))).
    reassociate <-. reassociate -> near (dec G).
    rewrite (natural (ϕ := @dec W G _)).
    do 2 reassociate <-.
    change (fmap G (fmap (prod W) (right_action (A:=?A) F)))
      with (fmap (G ∘ prod W) (right_action (A:=A) F)).
    reassociate -> near (dec G). rewrite (natural (ϕ := @dec W G _)).
    reassociate <-.
    rewrite (fun_fmap_fmap G).
    reassociate -> on right.
    change (fmap G (fmap F ?f)) with (fmap (G ∘ F) f).
    rewrite (fun_fmap_fmap (G ∘ F)).
    reflexivity.
  Qed.

  Instance DecoratedModule_compose : DecoratedModule W (G ∘ F) T :=
    {| drmod_action := drmod_action_compose; |}.

End DecoratedModule_compose.

Definition subd (F : Type -> Type) `{Fmap F} `{RightAction F T} `{Decorate W F} {A B}
  : (W * A -> T B) -> F A -> F B := fun f => sub F f ∘ dec F.

Section decoratedmodule_suboperations.

  Context
    (F T : Type -> Type)
    `{DecoratedModule W F T}.

  Lemma fmapd_to_subd : forall `(f : W * A -> B),
      fmapd F f = subd F (ret T ∘ f).
  Proof.
    introv. unfold fmapd, subd.
    unfold_compose_in_compose.
    rewrite <- (RightModule.sub_fmap F T).
    now rewrite (RightModule.sub_id F T).
  Qed.

  Lemma sub_to_subd : forall `(f : A -> T B),
      sub F f = subd F (f ∘ extract (prod W)).
  Proof.
    introv. unfold subd.
    rewrite <- (RightModule.sub_fmap F T).
    reassociate -> on right.
    now rewrite (dfun_dec_extract W F).
  Qed.

  Lemma fmap_to_subd : forall `(f : A -> B),
      fmap F f = subd F (ret T ∘ f ∘ extract (prod W)).
  Proof.
    introv. now rewrite (fmap_to_fmapd F), fmapd_to_subd.
  Qed.

End decoratedmodule_suboperations.

Section dec_subd.

  #[local] Set Keyed Unification.

  Context
    (F : Type -> Type)
    `{DecoratedModule W F T}.

  Theorem dec_subd : forall A B (f : W * A -> T B),
      dec F ∘ subd F f =
      subd F (shift T ∘ cobind (prod W) (dec T ∘ f)).
  Proof.
    introv. unfold subd. unfold_ops @Sub_RightAction.
    do 2 reassociate <- on left.
    rewrite (drmod_action W F T).
    reassociate -> near (fmap F f).
    rewrite (fun_fmap_fmap F).
    reassociate -> near (fmap F (dec T ∘ f)).
    rewrite <- (natural (G := F ∘ prod W) (F := F) (ϕ := @dec W F _)).
    reassociate <- on left.
    unfold_ops @Fmap_compose.
    change_left (right_action F ∘ (fmap F (shift T) ∘ fmap F (fmap (prod W) (dec T ∘ f))) ∘ dec F ∘ dec F).
    rewrite (fun_fmap_fmap F).
    reassociate -> on left.
    rewrite (dfun_dec_dec W F).
    reassociate <- on left.
    reassociate -> near (fmap F (cojoin (A:=A) (prod W))).
    now rewrite (fun_fmap_fmap F).
  Qed.

  Corollary dec_sub : forall A B (f : W * A -> T B),
      dec F ∘ sub F f =
      subd F (shift T ∘ fmap (prod W) (dec T ∘ f)).
  Proof.
    introv. rewrite (sub_to_subd F T).
    rewrite dec_subd.
    fequal. now ext [w a].
  Qed.

End dec_subd.

Section decoratedmodule_kleisli_laws.

  Context
    (F T : Type -> Type)
    `{DecoratedModule W F T}.

  Lemma subd_id :
    `(subd F (ret T ∘ extract (prod W)) = @id (F A)).
  Proof.
    intros. unfold subd.
    rewrite <- (RightModule.sub_fmap F T).
    reassociate -> on left.
    rewrite (dfun_dec_extract W F).
    now rewrite (RightModule.sub_id F T).
  Qed.

  Lemma subd_subd {A B C} : forall (g : W * B -> T C) (f : W * A -> T B),
      subd F (g ⋆d f) = subd F g ∘ subd F f.
  Proof.
    intros. unfold subd at 2.
    reassociate -> on right.
    rewrite (dec_subd F).
    unfold subd at 2.
    reassociate <- on right.
    rewrite (sub_sub F T).
    unfold kcomposed; reassociate -> on left.
    reflexivity.
  Qed.

End decoratedmodule_kleisli_laws.

Section decoratedmonad_suboperation_composition.

  Context
    (F : Type -> Type)
    `{DecoratedModule W F T}.

  Lemma subd_fmapd {A B C} : forall (g : W * B -> T C) (f : W * A -> B),
      subd F g ∘ fmapd F f = subd F (g co⋆ f).
  Proof.
    introv. rewrite (fmapd_to_subd F T).
    rewrite <- (subd_subd F T).
    now rewrite (dm_kleisli_star1).
  Qed.

  Corollary sub_subd {A B C} : forall (g : B -> T C) (f : W * A -> T B),
      sub F g ∘ subd F f = subd F (g ⋆ f).
  Proof.
    intros. rewrite (sub_to_subd F T).
    rewrite <- (subd_subd F T).
    now rewrite (dm_kleisli_star4).
  Qed.

  Corollary fmapd_subd {A B C} : forall (g : W * B -> C) (f : W * A -> T B),
      fmapd F g ∘ subd F f = subd F (fmap T g ∘ shift T ∘ cobind (prod W) (dec T ∘ f)).
  Proof.
    intros. rewrite (fmapd_to_subd F T).
    rewrite <- (subd_subd F T).
    unfold kcomposed. rewrite <- (fmap_to_bind T).
    now rewrite <- (fmap_cobind (prod W)).
  Qed.

  Corollary subd_sub {A B C} : forall (g : W * B -> T C) (f : A -> T B),
      subd F g ∘ sub F f = subd F (bind T g ∘ shift T ∘ fmap (prod W) (dec T ∘ f)).
  Proof.
    introv. rewrite (sub_to_subd F T).
    rewrite <- (subd_subd F T).
    unfold kcomposed. reassociate <-.
    now rewrite <- (fmap_to_cobind (prod W)).
  Qed.

  Lemma subd_fmap {A B C} : forall (g : W * B -> T C) (f : A -> B),
      subd F g ∘ fmap F f = subd F (g ∘ fmap (prod W) f).
  Proof.
    intros. rewrite (fmap_to_subd F T).
    rewrite <- (subd_subd F T).
    reassociate -> on left. rewrite (dm_kleisli_star1).
    unfold cokcompose. now rewrite (fmap_to_cobind (prod W)).
  Qed.

  Corollary fmap_subd {A B C} : forall (g : B -> C) (f : W * A -> T B),
      fmap F g ∘ subd F f = subd F (fmap T g ∘ f).
  Proof.
    intros. rewrite (fmap_to_subd F T).
    rewrite <- (subd_subd F T).
    rewrite (dm_kleisli_star4). unfold kcompose.
    now rewrite <- (fmap_to_bind T).
  Qed.

End decoratedmonad_suboperation_composition.