-- Natural transformations between functors
-- Part of Cumulus, the Grace17 mathematics library.
-- Copyright (c) 2026 Grant Passmore.
-- Released under the Apache License 2.0. See LICENSE in the repository root.
-- Authors: Grant Passmore

module Grace.Cumulus.Category.Natural
import Functor
import Set.Function.Space

-- η : F ⟹ G, for functors F G : C ⥤ D: a function η : obj(C) → mor(D) whose
-- component at a is a morphism η(a) : F(a) ⟶ G(a) and which commutes with
-- every morphism, η(b) ∘ F(f) = G(f) ∘ η(a). The categories are those of
-- the functors, so the predicate names only the three maps.
predicate NaturalTransformation (η F G : Set) :=
  ∀ C, ∀ D, F : C ⥤ D → G : C ⥤ D →
    (η : obj(C) → mor(D)) ∧ (∀ a ∈ obj(C), η(a) : F(a) ⟶ G(a)) ∧
    (∀ (a b ∈ obj(C)) (f : a ⟶ b), η(b) ∘ F(f) = G(f) ∘ η(a))

notation mixfix "_ : _ ⟹ _" => NaturalTransformation

-- The set of natural transformations from F to G.
def NatTrans (F G : Set) : Set :=
  { η ∈ functionSpace(Category.obj(src(F)), mor(tgt(F))) | η : F ⟹ G }

predicate NaturalIsomorphism (η F G : Set) :=
  η : F ⟹ G ∧ ∀ a ∈ Category.obj(src(F)), Isomorphism(tgt(F), η(a))

-- The constructor: a transformation is given by its components and the
-- naturality squares.
theorem NaturalTransformation.mk (η F G C D : Set)
  requires functorF : F : C ⥤ D
  requires functorG : G : C ⥤ D
  requires function : η : obj(C) → mor(D)
  requires components : ∀ a ∈ obj(C), η(a) : F(a) ⟶ G(a)
  requires naturality : ∀ (a b ∈ obj(C)) (f : a ⟶ b), η(b) ∘ F(f) = G(f) ∘ η(a)
  shows η : F ⟹ G
proof
  unfold NaturalTransformation
  intros C2 D2 functorF2 functorG2
  obtain ⟨srcF, tgtF, rest⟩ := functorF
  obtain ⟨srcF2, tgtF2, rest2⟩ := functorF2
  have sameSource : C2 = C by cc
  have sameTarget : D2 = D by cc
  rw under sameSource
  rw under sameTarget
  refine ⟨function, components, naturality⟩

section TransformationLaws
variable (η F G C D : Set)
hypothesis functorF : F : C ⥤ D
hypothesis functorG : G : C ⥤ D
hypothesis natural : η : F ⟹ G

theorem NaturalTransformation.function
  shows η : obj(C) → mor(D)
proof
  obtain ⟨function, components, naturality⟩ := natural with C, D, functorF, functorG
  exact function

instance NaturalTransformation.component_hom (a : Set)
  requires aInO : a ∈ obj(C)
  shows η(a) : F(a) ⟶ G(a)
proof
  obtain ⟨function, components, naturality⟩ := natural with C, D, functorF, functorG
  exact components with a, aInO

theorem NaturalTransformation.naturality (a b ∈ obj(C)) (f : a ⟶ b)
  shows η(b) ∘ F(f) = G(f) ∘ η(a)
proof
  obtain ⟨function, components, naturality⟩ := natural with C, D, functorF, functorG
  exact naturality with a, assumption, b, assumption, f, assumption

theorem NaturalTransformation.component (a : Set)
  requires aInO : a ∈ obj(C)
  shows η(a) ∈ hom(D, F(a), G(a))
proof
  have component := NaturalTransformation.component_hom with *
  exact Morphism.hom with component

theorem NaturalTransformation.component_closed (a : Set)
  requires aInO : a ∈ obj(C)
  shows η(a) ∈ mor(D)
proof
  have component := NaturalTransformation.component_hom with *
  exact Morphism.in_morphisms with component

theorem NaturalTransformation.component_dom (a : Set)
  requires aInO : a ∈ obj(C)
  shows dom(D)(η(a)) = F(a)
proof
  have component := NaturalTransformation.component_hom with *
  exact Morphism.dom with component

theorem NaturalTransformation.component_cod (a : Set)
  requires aInO : a ∈ obj(C)
  shows cod(D)(η(a)) = G(a)
proof
  have component := NaturalTransformation.component_hom with *
  exact Morphism.cod with component

-- Naturality in hom-set form.
theorem NaturalTransformation.square (a b f : Set)
  requires fHom : f ∈ hom(C, a, b)
  shows η(b) ∘ F(f) = G(f) ∘ η(a)
proof
  have typed : f : a ⟶ b := Morphism.of_hom with fHom
  exact NaturalTransformation.naturality with *

theorem NatTrans.of_transformation
  shows η ∈ NatTrans(F, G)
proof
  have srcF := Functor.source with functorF
  have tgtF := Functor.target with functorF
  have function := NaturalTransformation.function with *
  have bounded := functionSpace.mem_iff.backward with function
  rw [← srcF, ← tgtF] at bounded
  simp [NatTrans]

end TransformationLaws

theorem NatTrans.natural (η F G : Set)
  requires member : η ∈ NatTrans(F, G)
  shows η : F ⟹ G
proof
  simp [NatTrans] at member
  exact member.right

-- The categorical rewrite set extended by the functor and transformation
-- laws: functoriality and naturality goals close by functor_simp when the
-- hom guards of the morphisms involved are in scope.
tactic functor_simp := simp [Category.id_left, Category.id_right, Category.assoc, Functor.map_id, Functor.map_comp, NaturalTransformation.naturality]

-- Two transformations with the same components are equal.
theorem NaturalTransformation.ext (η θ F G C D : Set)
  requires functorF : F : C ⥤ D
  requires functorG : G : C ⥤ D
  requires etaNatural : η : F ⟹ G
  requires thetaNatural : θ : F ⟹ G
  requires same : ∀ a ∈ obj(C), η(a) = θ(a)
  shows η = θ
proof
  have etaFunction := NaturalTransformation.function with functorF, functorG, etaNatural
  have thetaFunction := NaturalTransformation.function with functorF, functorG, thetaNatural
  exact FunctionBetween.ext with etaFunction, thetaFunction, same

-- The identity transformation of a functor.
defgraph identityTransformation (F : Set) (a ∈ Category.obj(src(F))) : mor(tgt(F)) :=
  id(tgt(F))(obj(F)(a))
requires functor : Functor(F, src(F), tgt(F))
proof
  have image := Functor.object_closed with functor, member
  have target := Functor.target_category with functor
  exact Category.id_closed with target, image

section IdentityTransformation
variable (F C D : Set)
hypothesis functor : F : C ⥤ D

theorem identityTransformation.component (a : Set)
  requires aInO : a ∈ obj(C)
  shows identityTransformation(F)(a) = id(D)(F(a))
proof
  have self := Functor.self with functor
  have srcF := Functor.source with functor
  have tgtF := Functor.target with functor
  have aInSrc : a ∈ obj(src(F)) by rwa srcF
  have value := identityTransformation.value with self, aInSrc
  rwa tgtF at value

theorem identityTransformation.function_hom
  shows identityTransformation(F) : obj(C) → mor(D)
proof
  have self := Functor.self with functor
  have srcF := Functor.source with functor
  have tgtF := Functor.target with functor
  have function := identityTransformation.function with self
  rwa [srcF, tgtF] at function

theorem identityTransformation.natural
  shows identityTransformation(F) : F ⟹ F
proof
  have function := identityTransformation.function_hom with functor
  apply NaturalTransformation.mk with functor, functor, function
  . intros a aIn
    rw identityTransformation.component
    synthesize
  . intros a aIn b bIn f fHom
    rw [identityTransformation.component, identityTransformation.component]
    category_simp

end IdentityTransformation

-- Vertical composition of transformations η : F ⟹ G and θ : G ⟹ H.
defgraph verticalComposition (η θ F G H : Set) (a ∈ Category.obj(src(F))) : mor(tgt(F)) :=
  comp(tgt(F))(θ(a), η(a))
requires functorF : Functor(F, src(F), tgt(F))
requires functorG : Functor(G, src(F), tgt(F))
requires functorH : Functor(H, src(F), tgt(F))
requires etaNatural : η : F ⟹ G
requires thetaNatural : θ : G ⟹ H
proof
  have etaHom := NaturalTransformation.component_hom with functorF, functorG, etaNatural, member
  have thetaHom := NaturalTransformation.component_hom with functorG, functorH, thetaNatural, member
  have composite : comp(tgt(F))(θ(a), η(a)) : F(a) ⟶ H(a) := Morphism.comp with *
  exact Morphism.in_morphisms with composite

section VerticalComposition
variable (η θ F G H C D : Set)
hypothesis functorF : F : C ⥤ D
hypothesis functorG : G : C ⥤ D
hypothesis functorH : H : C ⥤ D
hypothesis etaNatural : η : F ⟹ G
hypothesis thetaNatural : θ : G ⟹ H

theorem verticalComposition.component (a : Set)
  requires aInO : a ∈ obj(C)
  shows verticalComposition(η, θ, F, G, H)(a) = θ(a) ∘ η(a)
proof
  have srcF := Functor.source with functorF
  have tgtF := Functor.target with functorF
  have selfF := Functor.self with functorF
  have selfG : Functor(G, src(F), tgt(F)) by rwa [srcF, tgtF]
  have selfH : Functor(H, src(F), tgt(F)) by rwa [srcF, tgtF]
  have aInSrc : a ∈ obj(src(F)) by rwa srcF
  have value := verticalComposition.value with
    selfF, selfG, selfH, etaNatural, thetaNatural, aInSrc
  rwa tgtF at value

theorem verticalComposition.function_hom
  shows verticalComposition(η, θ, F, G, H) : obj(C) → mor(D)
proof
  have srcF := Functor.source with functorF
  have tgtF := Functor.target with functorF
  have selfF := Functor.self with functorF
  have selfG : Functor(G, src(F), tgt(F)) by rwa [srcF, tgtF]
  have selfH : Functor(H, src(F), tgt(F)) by rwa [srcF, tgtF]
  have function := verticalComposition.function with selfF, selfG, selfH, etaNatural, thetaNatural
  rwa [srcF, tgtF] at function

theorem verticalComposition.natural
  shows verticalComposition(η, θ, F, G, H) : F ⟹ H
proof
  have function : verticalComposition(η, θ, F, G, H) : obj(C) → mor(D) := verticalComposition.function_hom with *
  apply NaturalTransformation.mk with functorF, functorH, function
  . intros a aIn
    rw verticalComposition.component
    synthesize
  . intros a aIn b bIn f fHom
    rw [verticalComposition.component, verticalComposition.component]
    rw [Category.assoc, NaturalTransformation.naturality, ← Category.assoc, NaturalTransformation.naturality, Category.assoc]
    rfl

end VerticalComposition

-- Natural isomorphisms and their inverses.
section NaturalIsomorphismLaws
variable (η F G C D : Set)
hypothesis functorF : F : C ⥤ D
hypothesis functorG : G : C ⥤ D

theorem NaturalIsomorphism.natural
  requires iso : NaturalIsomorphism(η, F, G)
  shows η : F ⟹ G
proof
  obtain ⟨natural, components⟩ := iso
  exact natural

theorem NaturalIsomorphism.component (a : Set)
  requires iso : NaturalIsomorphism(η, F, G)
  requires aInO : a ∈ obj(C)
  shows Isomorphism(D, η(a))
proof
  obtain ⟨natural, components⟩ := iso
  have srcF := Functor.source with functorF
  have tgtF := Functor.target with functorF
  have aInSrc : a ∈ obj(src(F)) by rwa srcF
  have component := components with a, aInSrc
  rwa tgtF at component

theorem NaturalIsomorphism.isomorphic (a : Set)
  requires iso : NaturalIsomorphism(η, F, G)
  requires aInO : a ∈ obj(C)
  shows Isomorphic(D, F(a), G(a))
proof
  have natural := NaturalIsomorphism.natural with functorF, functorG, iso
  have component := NaturalTransformation.component with functorF, functorG, natural, aInO
  have componentIso := NaturalIsomorphism.component with functorF, functorG, iso, aInO
  exact Isomorphic.of_hom with component, componentIso

theorem NaturalIsomorphism.mk
  requires natural : η : F ⟹ G
  requires components : ∀ a ∈ obj(C), Isomorphism(D, η(a))
  shows NaturalIsomorphism(η, F, G)
proof
  have srcF := Functor.source with functorF
  have tgtF := Functor.target with functorF
  unfold NaturalIsomorphism
  rw under srcF
  rw under tgtF
  refine ⟨natural, components⟩

end NaturalIsomorphismLaws

-- The inverse of a natural isomorphism, componentwise.
defgraph inverseTransformation (η F G : Set) (a ∈ Category.obj(src(F))) : mor(tgt(F)) :=
  inverseMorphism(tgt(F), η(a))
requires functorF : Functor(F, src(F), tgt(F))
requires iso : NaturalIsomorphism(η, F, G)
proof
  have target := Functor.target_category with functorF
  obtain ⟨natural, components⟩ := iso
  have componentIso := components with a, member
  exact inverseMorphism.morphism with target, componentIso

section InverseTransformation
variable (η F G C D : Set)
hypothesis functorF : F : C ⥤ D
hypothesis functorG : G : C ⥤ D
hypothesis iso : NaturalIsomorphism(η, F, G)

theorem inverseTransformation.component (a : Set)
  requires aInO : a ∈ obj(C)
  shows inverseTransformation(η, F, G)(a) = inverseMorphism(D, η(a))
proof
  have srcF := Functor.source with functorF
  have tgtF := Functor.target with functorF
  have selfF := Functor.self with functorF
  have aInSrc : a ∈ obj(src(F)) by rwa srcF
  have value := inverseTransformation.value with selfF, iso, aInSrc
  rwa tgtF at value

theorem inverseTransformation.function_hom
  shows inverseTransformation(η, F, G) : obj(C) → mor(D)
proof
  have srcF := Functor.source with functorF
  have tgtF := Functor.target with functorF
  have selfF := Functor.self with functorF
  have function := inverseTransformation.function with selfF, iso
  rwa [srcF, tgtF] at function

theorem inverseTransformation.component_hom (a : Set)
  requires aInO : a ∈ obj(C)
  shows inverseTransformation(η, F, G)(a) : G(a) ⟶ F(a)
proof
  have target := Functor.target_category with functorF
  have natural := NaturalIsomorphism.natural with functorF, functorG, iso
  have component : η(a) : F(a) ⟶ G(a) :=
    NaturalTransformation.component_hom with functorF, functorG, natural, aInO
  have componentIso := NaturalIsomorphism.component with functorF, functorG, iso, aInO
  have pair := inverseMorphism.pair with target, componentIso
  rw inverseTransformation.component
  exact InversePair.inverse_of_hom with *

theorem inverseTransformation.natural
  shows inverseTransformation(η, F, G) : G ⟹ F
proof
  have target := Functor.target_category with functorF
  have natural := NaturalIsomorphism.natural with functorF, functorG, iso
  have function := inverseTransformation.function_hom with functorF, functorG, iso
  apply NaturalTransformation.mk with functorG, functorF, function
  . intros a aIn
    exact inverseTransformation.component_hom with functorF, functorG, iso, aIn
  . intros a aIn b bIn f fHom
    have fHomSet : f ∈ hom(C, a, b) := Morphism.hom with fHom
    have etaA := NaturalTransformation.component with functorF, functorG, natural, aIn
    have etaB := NaturalTransformation.component with functorF, functorG, natural, bIn
    have isoA := NaturalIsomorphism.component with functorF, functorG, iso, aIn
    have isoB := NaturalIsomorphism.component with functorF, functorG, iso, bIn
    have pairA := inverseMorphism.pair with target, isoA
    have pairB := inverseMorphism.pair with target, isoB
    have imageF := Functor.hom with functorF, fHomSet
    have imageG := Functor.hom with functorG, fHomSet
    have square : η(b) ∘ F(f) = G(f) ∘ η(a) :=
      NaturalTransformation.naturality with functorF, functorG, natural, aIn, bIn, fHom
    rw [inverseTransformation.component, inverseTransformation.component]
    exact InversePair.transpose with target, pairA, pairB, etaA, etaB, imageF, imageG, square

theorem inverseTransformation.natural_iso
  shows NaturalIsomorphism(inverseTransformation(η, F, G), G, F)
proof
  have target := Functor.target_category with functorF
  have natural := inverseTransformation.natural with functorF, functorG, iso
  apply NaturalIsomorphism.mk with functorG, functorF, natural
  intros a aIn
  have componentIso := NaturalIsomorphism.component with functorF, functorG, iso, aIn
  rw inverseTransformation.component
  exact inverseMorphism.isomorphism with target, componentIso

theorem inverseTransformation.vertical_left
  shows verticalComposition(η, inverseTransformation(η, F, G), F, G, F) =
    identityTransformation(F)
proof
  have target := Functor.target_category with functorF
  have natural := NaturalIsomorphism.natural with functorF, functorG, iso
  have inverseNatural := inverseTransformation.natural with functorF, functorG, iso
  have composite := verticalComposition.natural with
    functorF, functorG, functorF, natural, inverseNatural
  have identityNatural := identityTransformation.natural with functorF
  apply NaturalTransformation.ext with functorF, functorF, composite, identityNatural
  intros a aIn
  have component : η(a) : F(a) ⟶ G(a) :=
    NaturalTransformation.component_hom with functorF, functorG, natural, aIn
  have componentIso := NaturalIsomorphism.component with functorF, functorG, iso, aIn
  have pair := inverseMorphism.pair with target, componentIso
  have left : inverseMorphism(D, η(a)) ∘ η(a) = id(D)(F(a)) := InversePair.left_hom with *
  rw [verticalComposition.component, inverseTransformation.component, identityTransformation.component]
  exact left

theorem inverseTransformation.vertical_right
  shows verticalComposition(inverseTransformation(η, F, G), η, G, F, G) =
    identityTransformation(G)
proof
  have target := Functor.target_category with functorF
  have natural := NaturalIsomorphism.natural with functorF, functorG, iso
  have inverseNatural := inverseTransformation.natural with functorF, functorG, iso
  have composite := verticalComposition.natural with
    functorG, functorF, functorG, inverseNatural, natural
  have identityNatural := identityTransformation.natural with functorG
  apply NaturalTransformation.ext with functorG, functorG, composite, identityNatural
  intros a aIn
  have component : η(a) : F(a) ⟶ G(a) :=
    NaturalTransformation.component_hom with functorF, functorG, natural, aIn
  have componentIso := NaturalIsomorphism.component with functorF, functorG, iso, aIn
  have pair := inverseMorphism.pair with target, componentIso
  have right : η(a) ∘ inverseMorphism(D, η(a)) = id(D)(G(a)) := InversePair.right_hom with *
  rw [verticalComposition.component, inverseTransformation.component, identityTransformation.component]
  exact right

end InverseTransformation

-- The identity transformation is a natural isomorphism, and it is a unit
-- for vertical composition.
section IdentityLaws
variable (η F G C D : Set)
hypothesis functorF : F : C ⥤ D
hypothesis functorG : G : C ⥤ D
hypothesis natural : η : F ⟹ G

theorem identityTransformation.natural_iso
  shows NaturalIsomorphism(identityTransformation(F), F, F)
proof
  have target := Functor.target_category with functorF
  have identityNatural := identityTransformation.natural with functorF
  apply NaturalIsomorphism.mk with functorF, functorF, identityNatural
  intros a aIn
  have image := Functor.object_closed with functorF, aIn
  rw identityTransformation.component
  exact Isomorphism.id with target, image

theorem verticalComposition.identity_left
  shows verticalComposition(η, identityTransformation(G), F, G, G) = η
proof
  have identityNatural := identityTransformation.natural with functorG
  have composite := verticalComposition.natural with
    functorF, functorG, functorG, natural, identityNatural
  apply NaturalTransformation.ext with functorF, functorG, composite, natural
  intros a aIn
  have component : η(a) : F(a) ⟶ G(a) :=
    NaturalTransformation.component_hom with functorF, functorG, natural, aIn
  rw [verticalComposition.component, identityTransformation.component]
  category_simp

theorem verticalComposition.identity_right
  shows verticalComposition(identityTransformation(F), η, F, F, G) = η
proof
  have identityNatural := identityTransformation.natural with functorF
  have composite := verticalComposition.natural with
    functorF, functorF, functorG, identityNatural, natural
  apply NaturalTransformation.ext with functorF, functorG, composite, natural
  intros a aIn
  have component : η(a) : F(a) ⟶ G(a) :=
    NaturalTransformation.component_hom with functorF, functorG, natural, aIn
  rw [verticalComposition.component, identityTransformation.component]
  category_simp

end IdentityLaws