From Tealeaves Require Export
  Classes.Functor.

#[local] Generalizable Variables F G A B.

Class Map2 (F: Type → Type → Type): Type :=
  map2: ∀ (B1 A1 B2 A2: Type) (g: B1 → B2) (f: A1 → A2),
      F B1 A1 → F B2 A2.

#[global] Arguments map2 {F}%function_scope {Map2}
  {B1 A1 B2 A2}%type_scope (g f)%function_scope.

Class Functor2 (F: Type → Type → Type) `{map2_F: Map2 F}: Type :=
  { fun2_map_id: ∀ (A1 A2: Type),
      map2 (@id A1) (@id A2) = @id (F A1 A2);
    fun2_map_map: ∀ (B1 A1 B2 A2 B3 A3: Type)
                    (g2: B2 → B3) (f2: A2 → A3)
                    (g1: B1 → B2) (f1: A1 → A2),
      map2 g2 f2 ∘ map2 g1 f1 = map2 (g2 ∘ g1) (f2 ∘ f1);
  }.

#[local] Notation "F ⇒2 G" :=
  (∀ B A: Type, F B A → G B A) (at level 60):
    tealeaves_scope.

Class Natural2 `{Map2 F} `{Map2 G} (ϕ: F ⇒2 G) :=
  { natural2_src: Functor2 F;
    natural2_tgt: Functor2 G;
    natural2: ∀ `(g: B1 → B2) `(f: A1 → A2),
      map2 (F := G) g f ∘ ϕ B1 A1 = ϕ B2 A2 ∘ map2 (F := F) g f
  }.

Section single_functor_instances.

  Context
    `{Functor2 F}.

  #[export] Instance Map2_1 {B}: Map (F B) :=
    fun A1 A2 f ⇒ map2 id f.

  #[export] Instance Map2_2 {A}: Map (fun B ⇒ F B A) :=
    fun B1 B2 g ⇒ map2 g id.

  #[export] Instance Functor_Map2_1 {B}:
    Functor (F B) (Map_F := Map2_1).
  Proof.
    constructor; intros; unfold_ops @Map2_1.
    - rewrite fun2_map_id.
      reflexivity.
    - rewrite fun2_map_map.
      reflexivity.
  Qed.

  #[export] Instance Functor_Map2_2 {A}:
    Functor (fun B ⇒ F B A) (Map_F := Map2_2).
  Proof.
    constructor; intros; unfold_ops @Map2_2.
    - rewrite fun2_map_id.
      reflexivity.
    - rewrite fun2_map_map.
      reflexivity.
  Qed.

  Context {B1 B2 A1 A2: Type}.
  Context {B0 A0 B3 A3: Type}
    {g: B1 → B2} {f: A1 → A2}
    {h: A0 → A1} {j: B0 → B1}
    {k: B2 → B3} {l: A2 → A3}.

  Lemma map2_to_21:
    map2 g (@id A1) = map (F := fun B ⇒ F B A1) g.
  Proof.
    reflexivity.
  Qed.

  Lemma map2_to_22:
    map2 (@id B1) f = map (F := F B1) f.
  Proof.
    reflexivity.
  Qed.

  Lemma fun2_map22_map21:
    map (Map := Map2_2) g ∘ map (Map := Map2_1) f =
      map2 (F := F) g f.
  Proof.
    unfold_ops @Map2_2 @Map2_1.
    rewrite fun2_map_map.
    reflexivity.
  Qed.

  Lemma fun2_map22_map21_commute:
    map (Map := Map2_2) g ∘ map (Map := Map2_1) f =
      map (Map := Map2_1) f ∘ map (Map := Map2_2) g.
  Proof.
    unfold_ops @Map2_2 @Map2_1.
    rewrite fun2_map_map.
    rewrite fun2_map_map.
    reflexivity.
  Qed.

  Lemma fun2_map2_map21:
    map2 (F := F) g f ∘ map (Map := Map2_1) h =
      map2 (F := F) g (f ∘ h).
  Proof.
    unfold_ops @Map2_1.
    rewrite fun2_map_map.
    reflexivity.
  Qed.

  Lemma fun2_map2_map22:
    map2 (F := F) g f ∘ map (Map := Map2_2) j =
      map2 (F := F) (g ∘ j) f.
  Proof.
    unfold_ops @Map2_2.
    rewrite fun2_map_map.
    reflexivity.
  Qed.

  Lemma fun2_map21_map2:
    map (Map := Map2_1) l ∘ map2 (F := F) g f =
      map2 g (l ∘ f).
  Proof.
    unfold_ops @Map2_1.
    rewrite fun2_map_map.
    reflexivity.
  Qed.

  Lemma fun2_map22_map2:
    map (Map := Map2_2) k ∘ map2 (F := F) g f =
      map2 (k ∘ g) f.
  Proof.
    unfold_ops @Map2_2.
    rewrite fun2_map_map.
    reflexivity.
  Qed.

  Definition vmap {B V1 V2:Type} `(f: V1 → V2) := map (F := F B) (Map := Map2_1) f.

  Definition bmap {V B1 B2:Type} `(f: B1 → B2) := map (F := fun B ⇒ F B V) (Map := Map2_2) f.

End single_functor_instances.

Section naturality_bmap_vmap.

  Context `{Functor2 F}
    {B B1 B2: Type}
    {V V1 V2: Type}.

  Lemma bmap_vmap_commute: ∀ (g: B1 → B2) (f: V1 → V2),
      bmap g ∘ vmap f = vmap f ∘ bmap g.
    Proof.
      intros.
      apply fun2_map22_map21_commute.
    Qed.

    #[export] Instance Natural_vmap: ∀ (f: V1 → V2),
        @Natural (fun B ⇒ F B V1) (@Map2_2 F _ V1) (fun B ⇒ F B V2) (@Map2_2 F _ V2) (fun B ⇒ vmap (F := F) f).
    Proof.
      intros. constructor.
      - typeclasses eauto.
      - typeclasses eauto.
      - intros.
        unfold vmap.
        now rewrite fun2_map22_map21_commute.
    Qed.

    #[export] Instance Natural_bmap: ∀ (f: B1 → B2),
        @Natural (F B1) (@Map2_1 F _ B1) (F B2) (@Map2_1 F _ B2) (fun B ⇒ bmap (F := F) f).
    Proof.
      intros. constructor.
      - typeclasses eauto.
      - typeclasses eauto.
      - intros.
        unfold bmap.
        now rewrite fun2_map22_map21_commute.
    Qed.

End naturality_bmap_vmap.

#[local] Notation "G ○12 F" :=
  (compose G ○ F) (at level 50):tealeaves_scope.
#[local] Notation "F ○21 G" :=
  (fun B A ⇒ F (G B) (G A)) (at level 50): tealeaves_scope.

Section composition_with_functor.

  Context
    `{Functor G}
    `{Functor2 F}.

  #[export] Instance Map21_compose: Map2 (F ○21 G) :=
    fun A1 B1 A2 B2 f1 f2 ⇒ map2 (F := F) (map f1) (map f2).

  #[export] Instance Functor21_compose: Functor2 (F ○21 G).
  Proof.
    constructor; intros; unfold_ops @Map21_compose.
    - rewrite fun_map_id.
      rewrite fun_map_id.
      rewrite fun2_map_id.
      reflexivity.
    - rewrite fun2_map_map.
      rewrite fun_map_map.
      rewrite fun_map_map.
      reflexivity.
  Qed.

  #[export] Instance Map12_compose: Map2 (G ○12 F) :=
    fun A1 B1 A2 B2 f1 f2 ⇒ map (F := G) (map2 f1 f2).

  #[export] Instance Functor12_compose: Functor2 (G ○12 F).
  Proof.
    constructor; intros; unfold_ops @Map12_compose.
    - rewrite fun2_map_id.
      rewrite fun_map_id.
      reflexivity.
    - rewrite (fun_map_map (F := G)
                 (F B1 A1) (F B2 A2) (F B3 A3)).
      rewrite fun2_map_map.
      reflexivity.
  Qed.

End composition_with_functor.

Definition strength2 {F: Type → Type → Type}
  `{Map2 F} {A B C}:
  ∀ (p : A × F B C), F (A × B) (A × C) :=
    fun '(a, t) ⇒ map2 (pair a) (pair a) t.

Module Notations.

  #[global] Notation "G ○12 F" :=
  (compose G ○ F) (at level 50):
    tealeaves_scope.

  #[global] Notation "F ○21 G" :=
    (fun B A ⇒ F (G B) (G A)) (at level 50):
      tealeaves_scope.

  #[global] Notation "F ⇒2 G" :=
    (∀ B A: Type, F B A → G B A) (at level 60):
      tealeaves_scope.

End Notations.

Section rewriting_maps.

  Import Notations.
  Context `{Functor2 F} `{Functor G}.

  Context {B1 B2 A1 A2: Type} (f: A1 → A2) (g: B1 → B2).

  Lemma map2_comp12_rw:
    map2 (F := (G ○12 F)) g f =
      map (map2 g f).
  Proof.
    reflexivity.
  Qed.

  Lemma map2_comp21_rw:
    map2 (F := (F ○21 G)) g f =
      map2 (map (F := G) g) (map f).
  Proof.
    reflexivity.
  Qed.

End rewriting_maps.