import Mathlib
import Algebra.Terms
import Algebra.Derivations

open CategoryTheory

namespace NominalAlgebra

section AlgebraDefinition

variable (𝔸 : Type) [Infinite 𝔸] [Countable 𝔸] [DecidableEq 𝔸]
variable (sig : Signature)

structure Algebra where
  carrier : ℕ → Nom 𝔸
  pure (f : sig.pure) {m : ℕ} :
    (sig.pureArgs f → carrier m) →ₑ[𝔸] carrier (m + sig.pureDepth f)
  free (f : sig.free) {m : ℕ} :
    (𝔸 × (sig.freeArgs f → carrier m)) →ₑ[𝔸] carrier (m + sig.freeDepth f)
  bound (f : sig.bound) {m : ℕ} :
    NameAbstraction 𝔸 (sig.boundArgs f → carrier m) →ₑ[𝔸] carrier (m + sig.boundDepth f)

instance : CoeFun (Algebra 𝔸 sig) (fun _ => ℕ → Nom 𝔸) where
  coe A := A.carrier

@[ext]
structure Algebra.Hom (A B : Algebra 𝔸 sig) where
  toFun : (n : ℕ) → (A n →ₑ[𝔸] B n)
  preserve_pure : ∀ (f : sig.pure) {m : ℕ} (args : sig.pureArgs f → A m),
    toFun (m + sig.pureDepth f) (A.pure f args) = B.pure f (toFun m ∘ args)
  preserve_free : ∀ (a : 𝔸) (f : sig.free) {m : ℕ} (args : sig.freeArgs f → A m),
    toFun (m + sig.freeDepth f) (A.free f ⟨a, args⟩) = B.free f ⟨a, toFun m ∘ args⟩
  preserve_bound : ∀ (a : 𝔸) (f : sig.bound) {m : ℕ} (args : sig.boundArgs f → A m),
    toFun (m + sig.boundDepth f) (A.bound f (⟪a⟫ args)) = B.bound f (⟪a⟫ toFun m ∘ args)

instance {A B : Algebra 𝔸 sig} : CoeFun (Algebra.Hom 𝔸 sig A B) (fun _ => (n : ℕ) → (A n →ₑ[𝔸] B n)) where
  coe h := h.toFun

end AlgebraDefinition

section Algebra

variable {𝔸 : Type} [Infinite 𝔸] [Countable 𝔸] [DecidableEq 𝔸]
variable {sig : Signature}

namespace Algebra

def Hom.id (A : Algebra 𝔸 sig) : Hom 𝔸 sig A A where
  toFun _ := EquivariantMap.id 𝔸
  preserve_pure _ _ _ := rfl
  preserve_free _ _ _ _ := rfl
  preserve_bound _ _ _ _ := rfl

def Hom.comp {A B C : Algebra 𝔸 sig} (h : Hom 𝔸 sig B C) (g : Hom 𝔸 sig A B) : Hom 𝔸 sig A C where
  toFun n := EquivariantMap.comp 𝔸 (h n) (g n)
  preserve_pure f m args := by
    rw [EquivariantMap.comp_app, g.preserve_pure, h.preserve_pure]
    rfl
  preserve_free a f m args := by
    rw [EquivariantMap.comp_app, g.preserve_free, h.preserve_free]
    rfl
  preserve_bound a f m args := by
    rw [EquivariantMap.comp_app, g.preserve_bound, h.preserve_bound]
    rfl

instance : Category (Algebra 𝔸 sig) where
  Hom := Hom 𝔸 sig
  id := Hom.id
  comp h g := Hom.comp g h

def forgetfulFunctor : Algebra 𝔸 sig ⥤ Nom 𝔸 where
  obj X := X 0
  map f := f.toFun 0

end Algebra

end Algebra

end NominalAlgebra