Require Export Prelude.Basics.
Require Export Prelude.Functor.

Notation " // f " := (flip f) (at level 2).

Class CoMonad `{copointed : CoPointed μ} := {
  duplicate : Π {α}, μ α ---> μ (μ α) ;

  (* Natural transformations. *)
 
  duplicate_fmap : Π α β (f : α ---> β),
    duplicate ∙ fmap f == fmap (fmap f) ∙ duplicate ;

  (* Monad laws *)

  fmap_extract_duplicate : Π α, fmap extract ∙ duplicate == (@idS (μ α)) ;

  extract_duplicate : Π α,  (extract (α:=μ α)) ∙ duplicate == idS ;
  
  duplicate_duplicate : Π α, (@duplicate (μ α)) ∙ duplicate == fmap (@duplicate α) ∙ duplicate }.

Implicit Arguments CoMonad [[F] [copointed]].

Definition extend `{CoMonad μ} {α β} (f : μ α ---> β) : μ α ---> μ β :=
  fmap f ∙ duplicate.

Lemma extend_extract `{CoMonad μ} {α} : extend extract === (@idS (μ α)).
Proof.
  intros.
  unfold extend.
  apply fmap_extract_duplicate.
Qed.

Instance pointwise_is_eq A B (R : relation B) :
  subrelation (@eq A ==> R)%signature (pointwise_relation A R) | 5.
Proof. reduce. unfold pointwise_relation in *. subst ; apply H. reflexivity. Qed.

Instance eq_is_pointwise A B (R : relation B) :
  subrelation (pointwise_relation A R) (@eq A ==> R)%signature | 5.
Proof. reduce. unfold pointwise_relation in *. subst ; apply H. Qed.

Lemma extract_extend_right `{CoMonad μ} {α β} (f : μ α ---> β) : extract ∙ extend f == f.
Proof.
  intros.
  unfold extend.
  rewrite <- composeS_assoc.
  rewrite (extract_fmap f).
  rewrite composeS_assoc. 
  rewrite (extract_duplicate α).
  rewrite composeS_unit_right. reflexivity.
Qed.

Lemma extend_assoc `{CoMonad μ} {α β δ} (k : μ α ---> β) (l : μ β ---> δ) :
  extend (l ∙ extend k) === extend l ∙ extend k.
Proof.
  intros. unfold extend.
  setoid_rewrite fmap_compose at 1.
  setoid_rewrite fmap_compose at 1.
  do 2 rewrite composeS_assoc.
  setoid_rewrite <- (duplicate_duplicate α).
  setoid_rewrite <- composeS_assoc at 2.
  setoid_rewrite <- duplicate_fmap.
  setoid_rewrite composeS_assoc at 2.
  setoid_rewrite <- composeS_assoc at 3.
  reflexivity.
Qed.

Notation "T =>> U" := (extend T U) (at level 55) : comonad.
Notation "T --> X ; E" := (extend (fun X => E) T) (at level 30, X ident, right associativity) : comonad.

Hint Rewrite @extract_extend_right @extend_extract : comonad.
Hint Rewrite @extend_assoc : comonad.

Open Local Scope comonad.

Lemma comap `{CoMonad μ} {α β} (f : α ---> β) :
  extend (f ∙ extract) == fmap f.
Proof.
  intros. unfold extend.
  setoid_rewrite <- extract_fmap.
  setoid_rewrite fmap_compose.
  setoid_rewrite composeS_assoc.
  setoid_rewrite <- duplicate_fmap.
  setoid_rewrite <- composeS_assoc.
  setoid_rewrite fmap_extract_duplicate.
  rewrite composeS_unit_left; reflexivity.
Qed.