-- Examples of small categories
-- 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.Examples
import Product
import Set

theorem IsPreorder.relation (A r : Set)
  requires preorder : IsPreorder(A, r)
  shows RelationBetween(r, A, A)
proof
  obtain ⟨relation, reflexive, transitive⟩ := preorder
  exact relation

theorem IsPreorder.reflexive (A r : Set)
  requires preorder : IsPreorder(A, r)
  shows ReflexiveOn(A, r)
proof
  obtain ⟨relation, reflexive, transitive⟩ := preorder
  exact reflexive

theorem IsPreorder.transitive (A r : Set)
  requires preorder : IsPreorder(A, r)
  shows TransitiveOn(A, r)
proof
  obtain ⟨relation, reflexive, transitive⟩ := preorder
  exact transitive

-- Every member of a relation on A is a related pair of members of A.
theorem RelationBetween.shape (A B r f : Set)
  requires relation : RelationBetween(r, A, B)
  requires fInR : f ∈ r
  shows ∃ x ∈ A, ∃ y ∈ B, f = (x, y) ∧ relates(r, x, y)
proof
  unfold RelationBetween at relation
  have bounded := relation with f, fInR
  obtain ⟨bound, x, hx, y, hy, shape⟩ := product.membership.forward with bounded
  refine ⟨x, hx, y, hy, shape, ?_⟩
  unfold relates
  rwa ← shape

-- A preorder (A, r) as a category: objects A, morphisms the related
-- pairs, composition by transitivity.
defgraph orderDom (A r : Set) (f ∈ r) : A := fst(f)
requires relation : RelationBetween(r, A, A)
proof
  obtain ⟨x, xInA, y, yInA, shape, related⟩ :=
    RelationBetween.shape with relation, member
  rw shape
  rw fst.pair
  exact xInA

defgraph orderCod (A r : Set) (f ∈ r) : A := snd(f)
requires relation : RelationBetween(r, A, A)
proof
  obtain ⟨x, xInA, y, yInA, shape, related⟩ :=
    RelationBetween.shape with relation, member
  rw shape
  rw snd.pair
  exact yInA

defgraph orderIdentity (A r : Set) (a ∈ A) : r := (a, a)
requires reflexive : ReflexiveOn(A, r)
proof
  have related := reflexive with a, member
  unfold relates at related
  exact related

theorem orderDom.pair (A r x y : Set)
  requires relation : RelationBetween(r, A, A)
  requires related : relates(r, x, y)
  shows orderDom(A, r)((x, y)) = x
proof
  unfold relates at related
  have value := orderDom.value with relation, related
  rw value
  exact fst.pair at x, y

theorem orderCod.pair (A r x y : Set)
  requires relation : RelationBetween(r, A, A)
  requires related : relates(r, x, y)
  shows orderCod(A, r)((x, y)) = y
proof
  unfold relates at related
  have value := orderCod.value with relation, related
  rw value
  exact snd.pair at x, y

theorem orderIdentity.pair (A r a : Set)
  requires reflexive : ReflexiveOn(A, r)
  requires aInA : a ∈ A
  shows orderIdentity(A, r)(a) = (a, a)
proof
  exact orderIdentity.value with reflexive, aInA

defgraph orderComposition (A r : Set)
    (p ∈ composablePairsOf(r, orderDom(A, r), orderCod(A, r))) : r :=
  (orderDom(A, r)(snd(p)), orderCod(A, r)(fst(p)))
requires preorder : IsPreorder(A, r)
proof
  have relation := IsPreorder.relation with preorder
  have transitive := IsPreorder.transitive with preorder
  obtain ⟨bounded, g, gInR, f, fInR, shape, composable⟩ :=
    composablePairsOf.membership.forward with member
  obtain ⟨y, yInA, z, zInA, gShape, gRelated⟩ :=
    RelationBetween.shape with relation, gInR
  obtain ⟨x, xInA, w, wInA, fShape, fRelated⟩ :=
    RelationBetween.shape with relation, fInR
  have domG := orderDom.pair with relation, gRelated
  have codF := orderCod.pair with relation, fRelated
  have domF := orderDom.pair with relation, fRelated
  have codG := orderCod.pair with relation, gRelated
  have middle : y = w by cc
  have related : relates(r, x, z)
    rw ← middle at fRelated
    exact transitive with x, xInA, y, yInA, z, zInA, fRelated, gRelated
  unfold relates at related
  rw shape
  rw [fst.pair, snd.pair, gShape, fShape, domF, codG]
  exact related

theorem orderComposition.pair (A r x y z : Set)
  requires preorder : IsPreorder(A, r)
  requires first : relates(r, x, y)
  requires second : relates(r, y, z)
  shows orderComposition(A, r)((y, z), (x, y)) = (x, z)
proof
  have relation := IsPreorder.relation with preorder
  have gInR : (y, z) ∈ r := second
  have fInR : (x, y) ∈ r := first
  have domG := orderDom.pair with relation, second
  have codF := orderCod.pair with relation, first
  have domF := orderDom.pair with relation, first
  have codG := orderCod.pair with relation, second
  have composable : orderDom(A, r)((y, z)) = orderCod(A, r)((x, y)) by cc
  have pairIn := composablePairsOf.pair at r, orderDom(A, r), orderCod(A, r),
    (y, z), (x, y) with gInR, fInR, composable
  have value := orderComposition.value with preorder, pairIn
  have firstValue := fst.pair at (y, z), (x, y)
  have secondValue := snd.pair at (y, z), (x, y)
  unfold call2
  cc

def orderCategory (A r : Set) : Set :=
  Category.mk(A, r, orderDom(A, r), orderCod(A, r), orderIdentity(A, r),
    orderComposition(A, r))

theorem orderCategory.category (A r : Set)
  requires preorder : IsPreorder(A, r)
  shows Category(orderCategory(A, r))
proof
  have relation := IsPreorder.relation with preorder
  have reflexive := IsPreorder.reflexive with preorder
  have domFunction := orderDom.function with relation
  have codFunction := orderCod.function with relation
  have idFunction := orderIdentity.function with reflexive
  have compFunction := orderComposition.function with preorder
  unfold orderCategory
  apply Category.intro with domFunction, codFunction, idFunction, compFunction
  . intros a aInA
    have idValue := orderIdentity.pair with reflexive, aInA
    have related := reflexive with a, aInA
    have domValue := orderDom.pair with relation, related
    have codValue := orderCod.pair with relation, related
    rw idValue
    refine ⟨domValue, codValue⟩
  . intros g gInR f fInR composable
    obtain ⟨y, yInA, z, zInA, gShape, gRelated⟩ :=
      RelationBetween.shape with relation, gInR
    obtain ⟨x, xInA, w, wInA, fShape, fRelated⟩ :=
      RelationBetween.shape with relation, fInR
    have domG := orderDom.pair with relation, gRelated
    have codF := orderCod.pair with relation, fRelated
    have domF := orderDom.pair with relation, fRelated
    have codG := orderCod.pair with relation, gRelated
    have middle : y = w by cc
    rw ← middle at fRelated
    rw ← middle at fShape
    have composite := orderComposition.pair with preorder, fRelated, gRelated
    have transitive := IsPreorder.transitive with preorder
    have related := transitive with
      x, xInA, y, yInA, z, zInA, fRelated, gRelated
    have domComposite := orderDom.pair with relation, related
    have codComposite := orderCod.pair with relation, related
    rw [gShape, fShape, composite]
    refine ⟨?_, ?_⟩
    . cc
    . cc
  . intros f fInR
    obtain ⟨x, xInA, y, yInA, fShape, fRelated⟩ :=
      RelationBetween.shape with relation, fInR
    have domF := orderDom.pair with relation, fRelated
    have codF := orderCod.pair with relation, fRelated
    have xRelated := reflexive with x, xInA
    have yRelated := reflexive with y, yInA
    have xIdentity := orderIdentity.pair with reflexive, xInA
    have yIdentity := orderIdentity.pair with reflexive, yInA
    have leftComposite := orderComposition.pair with
      preorder, fRelated, yRelated
    have rightComposite := orderComposition.pair with
      preorder, xRelated, fRelated
    rw fShape
    rw [domF, codF, xIdentity, yIdentity]
    refine ⟨leftComposite, rightComposite⟩
  . intros h hInR g gInR f fInR outer inner
    obtain ⟨z, zInA, w, wInA, hShape, hRelated⟩ :=
      RelationBetween.shape with relation, hInR
    obtain ⟨y, yInA, v, vInA, gShape, gRelated⟩ :=
      RelationBetween.shape with relation, gInR
    obtain ⟨x, xInA, u, uInA, fShape, fRelated⟩ :=
      RelationBetween.shape with relation, fInR
    have domH := orderDom.pair with relation, hRelated
    have codG := orderCod.pair with relation, gRelated
    have domG := orderDom.pair with relation, gRelated
    have codF := orderCod.pair with relation, fRelated
    have outerMiddle : z = v by cc
    have innerMiddle : y = u by cc
    rw ← outerMiddle at gRelated
    rw ← outerMiddle at gShape
    rw ← innerMiddle at fRelated
    rw ← innerMiddle at fShape
    have transitive := IsPreorder.transitive with preorder
    have hgRelated := transitive with
      y, yInA, z, zInA, w, wInA, gRelated, hRelated
    have gfRelated := transitive with
      x, xInA, y, yInA, z, zInA, fRelated, gRelated
    have hg := orderComposition.pair with preorder, gRelated, hRelated
    have gf := orderComposition.pair with preorder, fRelated, gRelated
    have leftComposite := orderComposition.pair with
      preorder, fRelated, hgRelated
    have rightComposite := orderComposition.pair with
      preorder, gfRelated, hRelated
    rw [hShape, gShape, fShape, hg, gf, leftComposite, rightComposite]
    rfl

-- The discrete category on a set: the identity relation is a preorder.
theorem identity.preorder (X : Set)
  shows IsPreorder(X, identity(X))
proof
  unfold IsPreorder
  refine ⟨identity.relation at X, ?_, ?_⟩
  . intros x xInX
    exact identity.relates with xInX
  . intros x xInX y yInX z zInX first second
    unfold relates at first
    obtain ⟨bound, a, aInX, shape⟩ := identity.membership.forward with first
    obtain ⟨xa, ya⟩ := orderedPair.injective with shape
    have same : x = y by cc
    rw same
    exact second

def discreteCategory (X : Set) : Set := orderCategory(X, identity(X))

theorem discreteCategory.category (X : Set)
  shows Category(discreteCategory(X))
proof
  have preorder := identity.preorder at X
  unfold discreteCategory
  exact orderCategory.category with preorder

theorem discreteCategory.obj (X : Set)
  shows obj(discreteCategory(X)) = X
proof
  simp [discreteCategory, orderCategory, Category.mk.obj]

theorem discreteCategory.mor (X : Set)
  shows mor(discreteCategory(X)) = identity(X)
proof
  simp [discreteCategory, orderCategory, Category.mk.mor]

theorem discreteCategory.id (X : Set)
  shows id(discreteCategory(X)) = orderIdentity(X, identity(X))
proof
  simp [discreteCategory, orderCategory, Category.mk.id]

theorem discreteCategory.morphism (X f : Set)
  requires fInM : f ∈ mor(discreteCategory(X))
  shows ∃ x ∈ X, f = id(discreteCategory(X))(x)
proof
  have preorder := identity.preorder at X
  have reflexive := IsPreorder.reflexive with preorder
  rw discreteCategory.mor at fInM
  obtain ⟨bound, x, xInX, shape⟩ := identity.membership.forward with fInM
  have idValue := orderIdentity.pair with reflexive, xInX
  refine ⟨x, xInX, ?_⟩
  rw [discreteCategory.id, idValue]
  exact shape

theorem discreteCategory.isomorphism (X f : Set)
  requires fInM : f ∈ mor(discreteCategory(X))
  shows Isomorphism(discreteCategory(X), f)
proof
  have category := discreteCategory.category at X
  have xInObjects : ∀ x ∈ X, x ∈ obj(discreteCategory(X))
    intros x xInX
    rw discreteCategory.obj
    exact xInX
  obtain ⟨x, xInX, shape⟩ := discreteCategory.morphism with fInM
  have xInO := xInObjects with x, xInX
  rw shape
  exact Isomorphism.id with category, xInO

-- A monoid as a category with one object, the carrier itself.
theorem composablePairsOf.constant (A c : Set)
  shows composablePairsOf(A, constantFunction(A, {A}, A),
    constantFunction(A, {A}, A)) = A × A
proof
  have selfIn := singleton.self at A
  ext p
  constructor
  . intro member
    obtain ⟨bounded, rest⟩ := composablePairsOf.membership.forward with member
    exact bounded
  . intro member
    obtain ⟨bound, g, gInA, f, fInA, shape⟩ := product.membership.forward with
      member
    have gValue := constantFunction.apply with gInA, selfIn
    have fValue := constantFunction.apply with fInA, selfIn
    have composable : constantFunction(A, {A}, A)(g) =
        constantFunction(A, {A}, A)(f) by cc
    rw shape
    exact composablePairsOf.pair with gInA, fInA, composable

def monoidCategory (A op unit : Set) : Set :=
  Category.mk({A}, A, constantFunction(A, {A}, A), constantFunction(A, {A}, A),
    constantFunction({A}, A, unit), op)

theorem monoidCategory.category (A op unit : Set)
  requires monoid : IsMonoid(A, op, unit)
  shows Category(monoidCategory(A, op, unit))
proof
  obtain ⟨semigroup, identity⟩ := monoid
  obtain ⟨magma, associative⟩ := semigroup
  obtain ⟨unitInA, identityLaws⟩ := identity
  unfold IsMagma at magma
  have selfIn := singleton.self at A
  have constant := constantFunction.function at A, {A}, A with selfIn
  have identities := constantFunction.function at {A}, A, unit with unitInA
  have pairs := composablePairsOf.constant at A, A
  have operation : op : composablePairsOf(A, constantFunction(A, {A}, A),
      constantFunction(A, {A}, A)) → A
    rw pairs
    exact magma
  have closed : ∀ g ∈ A, ∀ f ∈ A, op(g, f) ∈ A
    intros g gInA f fInA
    have pairIn := product.pair with gInA, fInA
    unfold call2
    exact FunctionBetween.apply_closed with magma, pairIn
  unfold monoidCategory
  apply Category.intro with constant, constant, identities, operation
  . intros a aInO
    have aValue := singleton.membership.forward with aInO
    have idValue := constantFunction.apply with aInO, unitInA
    have domValue := constantFunction.apply with unitInA, selfIn
    rw idValue
    rw domValue
    refine ⟨?_, ?_⟩
    . symm
      exact aValue
    . symm
      exact aValue
  . intros g gInA f fInA composable
    have gfInA := closed with g, gInA, f, fInA
    have gValue := constantFunction.apply with gInA, selfIn
    have fValue := constantFunction.apply with fInA, selfIn
    have gfValue := constantFunction.apply with gfInA, selfIn
    refine ⟨?_, ?_⟩
    . cc
    . cc
  . intros f fInA
    have fValue := constantFunction.apply with fInA, selfIn
    have idValue := constantFunction.apply with selfIn, unitInA
    obtain ⟨leftUnit, rightUnit⟩ := identityLaws with f, fInA
    rw fValue
    rw idValue
    refine ⟨leftUnit, rightUnit⟩
  . intros h hInA g gInA f fInA outer inner
    exact associative with h, hInA, g, gInA, f, fInA