From Tealeaves Require Import
  Classes.Categorical.Applicative
  Classes.Categorical.Monad (Return, ret)
  Classes.Kleisli.DecoratedTraversableMonad
  Functors.Early.Batch
  Functors.Early.Writer.

Import Batch.Notations.
Import Monoid.Notations.
Import DecoratedTraversableMonad.Notations.
Import Product.Notations.
Import Applicative.Notations.

#[local] Generalizable Variables W T A B.

(** * Coalgebraic Decorated Traversable Monads *)
(**********************************************************************)

(** ** The <<toBatch7>> Operation *)
(**********************************************************************)
Class ToBatch7 (W: Type) (T U: Type -> Type) :=
  toBatch7: forall A B, U A -> Batch (W * A) (T B) (U B).

#[global] Arguments toBatch7 {W}%type_scope {T U}%function_scope
  {ToBatch7} {A B}%type_scope _.

(** ** The <<cojoin_Batch7>> Operation *)
(**********************************************************************)
Section cojoin.

  Context
    {W: Type}
    {T: Type -> Type}
    `{Monoid_op W}
    `{ToBatch7 W T T}.

  Context
    {A: Type} (* original leaves *)
    {A': Type} (* new leaves *)
    {A'': Type}. (* new type of new leaves *)

  Section auxiliary.

    Context {R: Type}.

    Definition cojoin_Type :=
      Batch (W * A) (T A') R ->
      Batch (W * A) (T A'') (Batch (W * A'') (T A') R).


    Definition key_function (w: W):
      Batch (W * A'') (T A') (T A' -> R) ->
      T A'' ->
      Batch (W * A'') (T A') R :=
      fun next_batch t =>
        next_batch <⋆> mapfst_Batch (incr w) (toBatch7 (B := A') t).

    Definition cojoin_Batch7_leaf_case:
      Batch (W * A) (T A'') (Batch (W * A'') (T A') (T A' -> R)) ->
      (* ^recursive call on cojoin of continuation *)
      W * A ->
      (* ^leaf in context *)
      Batch (W * A) (T A'') (T A'' -> Batch (W * A'') (T A') R) :=
      fun rec_continue '(w, a) =>
        map (F := Batch (W * A) (T A'')) (key_function w) rec_continue.

  End auxiliary.

  Fixpoint cojoin_Batch7 {R: Type}
    (b: Batch (W * A) (T A') R):
    Batch (W * A) (T A'') (Batch (W * A'') (T A') R) :=
    match b with
    | Done c => Done (Done c)
    | Step continuation (w, a) =>
        let new_continuation :=
          cojoin_Batch7_leaf_case
            (cojoin_Batch7 (R := T A' -> R) continuation) (w, a)
        in Step new_continuation (w, a)
    end.

End cojoin.

(** ** <<cojoin_Batch7>> as <<double_batch7>> *)
(**********************************************************************)
Section section.

  Context
    `{Monoid W}
    `{ToBatch7 W T T}.

  Definition double_batch7 {A A': Type} {R: Type}:
    W * A -> Batch (W * A) (T A') (Batch (W * A') (T R) (T R)) :=
    fun '(w, a) =>
      (map (F := Batch (W * A) (T A'))
         (mapfst_Batch (incr w) ∘ toBatch7) ∘
         batch (W * A) (T A')) (w, a).

  Lemma double_batch7_rw {A A': Type} {R: Type}:
    forall '(w, a),
      double_batch7 (A := A) (A' := A') (R := R) (w, a) =
        Done (mapfst_Batch (incr w) ∘ toBatch7) ⧆ (w, a).
W: Type
op: Monoid_op W
unit0: Monoid_unit W
H: Monoid W
T: Type -> Type
H0: ToBatch7 W T T
A, A', R: Type
forall '(w, a),
double_batch7 (w, a) =
Done (mapfst_Batch (incr w) ∘ toBatch7) ⧆ (w, a)
  Proof.
W: Type
op: Monoid_op W
unit0: Monoid_unit W
H: Monoid W
T: Type -> Type
H0: ToBatch7 W T T
A, A', R: Type
forall '(w, a),
double_batch7 (w, a) =
Done (mapfst_Batch (incr w) ∘ toBatch7) ⧆ (w, a)
    intros [w a].
W: Type
op: Monoid_op W
unit0: Monoid_unit W
H: Monoid W
T: Type -> Type
H0: ToBatch7 W T T
A, A', R: Type
w: W
a: A
double_batch7 (w, a) =
Done (mapfst_Batch (incr w) ∘ toBatch7) ⧆ (w, a)
    reflexivity.
  Qed.

  Lemma cojoin_Batch7_to_runBatch: forall (A A' A'': Type) (R: Type),
      cojoin_Batch7 (A := A) (A' := A') (A'' := A'') (R := R) =
        runBatch (G := Batch (W * A) (T A'') ∘ Batch (W * A'') (T A'))
          double_batch7.
W: Type
op: Monoid_op W
unit0: Monoid_unit W
H: Monoid W
T: Type -> Type
H0: ToBatch7 W T T
forall A A' A'' R : Type,
cojoin_Batch7 = runBatch double_batch7
  Proof.
W: Type
op: Monoid_op W
unit0: Monoid_unit W
H: Monoid W
T: Type -> Type
H0: ToBatch7 W T T
forall A A' A'' R : Type,
cojoin_Batch7 = runBatch double_batch7
    intros.
W: Type
op: Monoid_op W
unit0: Monoid_unit W
H: Monoid W
T: Type -> Type
H0: ToBatch7 W T T
A, A', A'', R: Type
cojoin_Batch7 = runBatch double_batch7 ext b.
W: Type
op: Monoid_op W
unit0: Monoid_unit W
H: Monoid W
T: Type -> Type
H0: ToBatch7 W T T
A, A', A'', R: Type
b: Batch (W * A) (T A') R
cojoin_Batch7 b = runBatch double_batch7 b
    induction b as [R r | R continuation IHcontinuation [w a]].
W: Type
op: Monoid_op W
unit0: Monoid_unit W
H: Monoid W
T: Type -> Type
H0: ToBatch7 W T T
A, A', A'', R: Type
r: R
cojoin_Batch7 (Done r) =
runBatch double_batch7 (Done r)
W: Type
op: Monoid_op W
unit0: Monoid_unit W
H: Monoid W
T: Type -> Type
H0: ToBatch7 W T T
A, A', A'', R: Type
continuation: Batch (W * A) (T A') (T A' -> R)
w: W
a: A
IHcontinuation: cojoin_Batch7 continuation =
runBatch double_batch7 continuation
cojoin_Batch7 (continuation ⧆ (w, a)) =
runBatch double_batch7 (continuation ⧆ (w, a))
    -
W: Type
op: Monoid_op W
unit0: Monoid_unit W
H: Monoid W
T: Type -> Type
H0: ToBatch7 W T T
A, A', A'', R: Type
r: R
cojoin_Batch7 (Done r) =
runBatch double_batch7 (Done r) cbn.
W: Type
op: Monoid_op W
unit0: Monoid_unit W
H: Monoid W
T: Type -> Type
H0: ToBatch7 W T T
A, A', A'', R: Type
r: R
Done (Done r) = pure r reflexivity.
    -
W: Type
op: Monoid_op W
unit0: Monoid_unit W
H: Monoid W
T: Type -> Type
H0: ToBatch7 W T T
A, A', A'', R: Type
continuation: Batch (W * A) (T A') (T A' -> R)
w: W
a: A
IHcontinuation: cojoin_Batch7 continuation =
runBatch double_batch7 continuation
cojoin_Batch7 (continuation ⧆ (w, a)) =
runBatch double_batch7 (continuation ⧆ (w, a)) cbn.
W: Type
op: Monoid_op W
unit0: Monoid_unit W
H: Monoid W
T: Type -> Type
H0: ToBatch7 W T T
A, A', A'', R: Type
continuation: Batch (W * A) (T A') (T A' -> R)
w: W
a: A
IHcontinuation: cojoin_Batch7 continuation =
runBatch double_batch7 continuation
map (key_function w) (cojoin_Batch7 continuation)
⧆ (w, a) =
map (compose (map (fun '(f, a) => f a)))
  (map (compose mult)
     (map strength_arrow
        (map
           (fun c : Batch (W * A'') (T A') (T A' -> R)
            =>
            (c, mapfst_Batch (incr w) ∘ toBatch7 ∘ id))
           (runBatch double_batch7 continuation))))
⧆ (w, a)
      do 3 compose near
        (runBatch
           (G := Batch (W * A) (T A'') ∘ Batch (W * A'') (T A'))
           double_batch7
           continuation).
W: Type
op: Monoid_op W
unit0: Monoid_unit W
H: Monoid W
T: Type -> Type
H0: ToBatch7 W T T
A, A', A'', R: Type
continuation: Batch (W * A) (T A') (T A' -> R)
w: W
a: A
IHcontinuation: cojoin_Batch7 continuation =
runBatch double_batch7 continuation
map (key_function w) (cojoin_Batch7 continuation)
⧆ (w, a) =
(map (compose (map (fun '(f, a) => f a)))
 ∘ (map (compose mult)
    ∘ (map strength_arrow
       ∘ map
           (fun c : Batch (W * A'') (T A') (T A' -> R)
            =>
            (c, mapfst_Batch (incr w) ∘ toBatch7 ∘ id)))))
  (runBatch double_batch7 continuation) ⧆ (w, a)
      do 3 rewrite (fun_map_map (F := Batch (W * A) (T A''))).
W: Type
op: Monoid_op W
unit0: Monoid_unit W
H: Monoid W
T: Type -> Type
H0: ToBatch7 W T T
A, A', A'', R: Type
continuation: Batch (W * A) (T A') (T A' -> R)
w: W
a: A
IHcontinuation: cojoin_Batch7 continuation =
runBatch double_batch7 continuation
map (key_function w) (cojoin_Batch7 continuation)
⧆ (w, a) =
map
  (compose (map (fun '(f, a) => f a))
   ∘ (compose mult
      ∘ (strength_arrow
         ∘ (fun c : Batch (W * A'') (T A') (T A' -> R)
            =>
            (c, mapfst_Batch (incr w) ∘ toBatch7 ∘ id)))))
  (runBatch double_batch7 continuation) ⧆ (w, a)
      rewrite IHcontinuation.
W: Type
op: Monoid_op W
unit0: Monoid_unit W
H: Monoid W
T: Type -> Type
H0: ToBatch7 W T T
A, A', A'', R: Type
continuation: Batch (W * A) (T A') (T A' -> R)
w: W
a: A
IHcontinuation: cojoin_Batch7 continuation =
runBatch double_batch7 continuation
map (key_function w)
  (runBatch double_batch7 continuation) ⧆ (w, a) =
map
  (compose (map (fun '(f, a) => f a))
   ∘ (compose mult
      ∘ (strength_arrow
         ∘ (fun c : Batch (W * A'') (T A') (T A' -> R)
            =>
            (c, mapfst_Batch (incr w) ∘ toBatch7 ∘ id)))))
  (runBatch double_batch7 continuation) ⧆ (w, a)
      reflexivity.
  Qed.

  Lemma cojoin_Batch7_batch: forall (A A' R: Type),
      cojoin_Batch7 (A'' := A') ∘ batch (W * A) (T R) = double_batch7.
W: Type
op: Monoid_op W
unit0: Monoid_unit W
H: Monoid W
T: Type -> Type
H0: ToBatch7 W T T
forall A A' R : Type,
cojoin_Batch7 ∘ batch (W * A) (T R) = double_batch7
  Proof.
W: Type
op: Monoid_op W
unit0: Monoid_unit W
H: Monoid W
T: Type -> Type
H0: ToBatch7 W T T
forall A A' R : Type,
cojoin_Batch7 ∘ batch (W * A) (T R) = double_batch7
    intros.
W: Type
op: Monoid_op W
unit0: Monoid_unit W
H: Monoid W
T: Type -> Type
H0: ToBatch7 W T T
A, A', R: Type
cojoin_Batch7 ∘ batch (W * A) (T R) = double_batch7
    rewrite cojoin_Batch7_to_runBatch.
W: Type
op: Monoid_op W
unit0: Monoid_unit W
H: Monoid W
T: Type -> Type
H0: ToBatch7 W T T
A, A', R: Type
runBatch double_batch7 ∘ batch (W * A) (T R) =
double_batch7
    rewrite (runBatch_batch
               (Batch (prod W A) (T A') ∘ Batch (prod W A') (T R))).
W: Type
op: Monoid_op W
unit0: Monoid_unit W
H: Monoid W
T: Type -> Type
H0: ToBatch7 W T T
A, A', R: Type
double_batch7 = double_batch7
    reflexivity.
  Qed.

  #[export] Instance AppMor_cojoin_Batch7: forall (A A' A'': Type),
      ApplicativeMorphism
        (Batch (W * A) (T A'))
        (Batch (W * A) (T A'') ∘ Batch (W * A'') (T A'))
        (@cojoin_Batch7 W T _ _ A A' A'').
W: Type
op: Monoid_op W
unit0: Monoid_unit W
H: Monoid W
T: Type -> Type
H0: ToBatch7 W T T
forall A A' A'' : Type,
ApplicativeMorphism (Batch (W * A) (T A'))
  (Batch (W * A) (T A'') ∘ Batch (W * A'') (T A'))
  (@cojoin_Batch7 W T op H0 A A' A'')
  Proof.
W: Type
op: Monoid_op W
unit0: Monoid_unit W
H: Monoid W
T: Type -> Type
H0: ToBatch7 W T T
forall A A' A'' : Type,
ApplicativeMorphism (Batch (W * A) (T A'))
  (Batch (W * A) (T A'') ∘ Batch (W * A'') (T A'))
  (@cojoin_Batch7 W T op H0 A A' A'')
    intros.
W: Type
op: Monoid_op W
unit0: Monoid_unit W
H: Monoid W
T: Type -> Type
H0: ToBatch7 W T T
A, A', A'': Type
ApplicativeMorphism (Batch (W * A) (T A'))
  (Batch (W * A) (T A'') ∘ Batch (W * A'') (T A'))
  (@cojoin_Batch7 W T op H0 A A' A'')
    assert (lemma:
             @cojoin_Batch7 W T op H0 A A' A'' =
               fun R => runBatch (G := Batch (W * A) (T A'') ∘
                                      Batch (W * A'') (T A'))
                       double_batch7).
W: Type
op: Monoid_op W
unit0: Monoid_unit W
H: Monoid W
T: Type -> Type
H0: ToBatch7 W T T
A, A', A'': Type
@cojoin_Batch7 W T op H0 A A' A'' =
(fun R : Type => runBatch double_batch7)
W: Type
op: Monoid_op W
unit0: Monoid_unit W
H: Monoid W
T: Type -> Type
H0: ToBatch7 W T T
A, A', A'': Type
lemma: @cojoin_Batch7 W T op H0 A A' A'' =
(fun R : Type => runBatch double_batch7)
ApplicativeMorphism (Batch (W * A) (T A'))
  (Batch (W * A) (T A'') ∘ Batch (W * A'') (T A'))
  (@cojoin_Batch7 W T op H0 A A' A'')
    {
W: Type
op: Monoid_op W
unit0: Monoid_unit W
H: Monoid W
T: Type -> Type
H0: ToBatch7 W T T
A, A', A'': Type
@cojoin_Batch7 W T op H0 A A' A'' =
(fun R : Type => runBatch double_batch7) ext R.
W: Type
op: Monoid_op W
unit0: Monoid_unit W
H: Monoid W
T: Type -> Type
H0: ToBatch7 W T T
A, A', A'', R: Type
cojoin_Batch7 = runBatch double_batch7 now rewrite cojoin_Batch7_to_runBatch. }
W: Type
op: Monoid_op W
unit0: Monoid_unit W
H: Monoid W
T: Type -> Type
H0: ToBatch7 W T T
A, A', A'': Type
lemma: @cojoin_Batch7 W T op H0 A A' A'' =
(fun R : Type => runBatch double_batch7)
ApplicativeMorphism (Batch (W * A) (T A'))
  (Batch (W * A) (T A'') ∘ Batch (W * A'') (T A'))
  (@cojoin_Batch7 W T op H0 A A' A'')
    rewrite lemma.
W: Type
op: Monoid_op W
unit0: Monoid_unit W
H: Monoid W
T: Type -> Type
H0: ToBatch7 W T T
A, A', A'': Type
lemma: @cojoin_Batch7 W T op H0 A A' A'' =
(fun R : Type => runBatch double_batch7)
ApplicativeMorphism (Batch (W * A) (T A'))
  (Batch (W * A) (T A'') ∘ Batch (W * A'') (T A'))
  (fun R : Type => runBatch double_batch7)
    apply ApplicativeMorphism_runBatch.
  Qed.

End section.

(** ** Typeclass *)
(**********************************************************************)
Class DecoratedTraversableMonad (W: Type) (T: Type -> Type)
  `{Monoid_op W} `{Monoid_unit W} `{Return T} `{ToBatch7 W T T} :=
  { dtm_monoid :> Monoid W;
    dtm_ret: forall (A B: Type),
      toBatch7 ∘ ret (T := T) (A := A) =
        Step (Done (@id (T B))) ∘ ret (T := (W ×));
    dtm_extract: forall (A: Type),
      extract_Batch ∘ mapfst_Batch (ret ∘ extract (W := (W ×))) ∘
        toBatch7 = @id (T A);
    dtm_duplicate: forall (A B C: Type),
      cojoin_Batch7 ∘ toBatch7 (A := A) (B := C) =
        map (F := Batch (W * A) (T B)) (toBatch7) ∘ toBatch7;
  }.