import Mathlib
import Nominal

namespace Algebra

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

/-- A graded signature. -/
structure Signature where
  pure : Type
  free : Type
  bound : Type
  /-- Type representing the arguments required for an operation. -/
  args : pure ⊕ free ⊕ bound → Type
  args_fintype : ∀ f, Fintype (args f) := by intro; infer_instance
  /-- Depth of an operation. -/
  depth : pure ⊕ free ⊕ bound → ℕ

def Signature.pureArgs (sig : Signature) (f : sig.pure) : Type := sig.args (.inl f)
def Signature.freeArgs (sig : Signature) (f : sig.free) : Type := sig.args (.inr (.inl f))
def Signature.boundArgs (sig : Signature) (f : sig.bound) : Type := sig.args (.inr (.inr f))

def Signature.pureArgs_fintype (sig : Signature) (f : sig.pure) : Fintype (sig.pureArgs f) := sig.args_fintype (.inl f)
def Signature.freeArgs_fintype (sig : Signature) (f : sig.free) : Fintype (sig.freeArgs f) := sig.args_fintype (.inr (.inl f))
def Signature.boundArgs_fintype (sig : Signature) (f : sig.bound) : Fintype (sig.boundArgs f) := sig.args_fintype (.inr (.inr f))

def Signature.arity (sig : Signature) (f : sig.pure ⊕ sig.free ⊕ sig.bound) : ℕ := (sig.args_fintype f).card
def Signature.pureArity (sig : Signature) (f : sig.pure) : ℕ := sig.arity (.inl f)
def Signature.freeArity (sig : Signature) (f : sig.free) : ℕ := sig.arity (.inr (.inl f))
def Signature.boundArity (sig : Signature) (f : sig.bound) : ℕ := sig.arity (.inr (.inr f))

def Signature.pureDepth (sig : Signature) (f : sig.pure) : ℕ := sig.depth (.inl f)
def Signature.freeDepth (sig : Signature) (f : sig.free) : ℕ := sig.depth (.inr (.inl f))
def Signature.boundDepth (sig : Signature) (f : sig.bound) : ℕ := sig.depth (.inr (.inr f))

variable {sig : Signature}
variable {X : Type} [NominalSet 𝔸 X]
variable {Y : Type} [NominalSet 𝔸 Y]

inductive Term (sig : Signature) (X : Type) [NominalSet 𝔸 X]
  | var (x : X) : Term sig X
  | pure (f : sig.pure) (args : sig.pureArgs f → Term sig X) : Term sig X
  | free (a : 𝔸) (f : sig.free) (args : sig.freeArgs f → Term sig X) : Term sig X
  | bound (a : 𝔸) (f : sig.bound) (args : sig.boundArgs f → Term sig X) : Term sig X

inductive Term.HasDepth : Term 𝔸 sig X → ℕ → Prop
  | var (x : X) : HasDepth (.var x) 0
  | pure (f : sig.pure) (args : sig.pureArgs f → Term 𝔸 sig X) (k : ℕ) :
      (∀ arg : sig.pureArgs f, HasDepth (args arg) k) → HasDepth (.pure f args) (k + sig.pureDepth f)
  | free (a : 𝔸) (f : sig.free) (args : sig.freeArgs f → Term 𝔸 sig X) (k : ℕ) :
      (∀ arg : sig.freeArgs f, HasDepth (args arg) k) → HasDepth (.free a f args) (k + sig.freeDepth f)
  | bound (a : 𝔸) (f : sig.bound) (args : sig.boundArgs f → Term 𝔸 sig X) (k : ℕ) :
      (∀ arg : sig.boundArgs f, HasDepth (args arg) k) → HasDepth (.bound a f args) (k + sig.boundDepth f)

def TermAtDepth (sig : Signature) (X : Type) [NominalSet 𝔸 X] (n : ℕ) :=
  { t : Term 𝔸 sig X // t.HasDepth 𝔸 n }

def Substitution (sig : Signature) (X : Type) [NominalSet 𝔸 X] (Y : Type) [NominalSet 𝔸 Y] (l : ℕ) :=
  X → TermAtDepth 𝔸 sig Y l

private def Substitution.apply' {l : ℕ} (σ : Substitution 𝔸 sig X Y l) : Term 𝔸 sig X → Term 𝔸 sig Y
  | .var x => (σ x).val
  | .pure f args => .pure f (fun x => σ.apply' (args x))
  | .free a f args => .free a f (fun x => σ.apply' (args x))
  | .bound a f args => .bound a f (fun x => σ.apply' (args x))

def Substitution.apply {n l : ℕ} (σ : Substitution 𝔸 sig X Y l) (t : TermAtDepth 𝔸 sig X n) :
    TermAtDepth 𝔸 sig Y (n + l) :=
  ⟨σ.apply' 𝔸 t.val, by
    suffices ∀ t n, t.HasDepth 𝔸 n → (σ.apply' 𝔸 t).HasDepth 𝔸 (n + l) by
      exact this t.val n t.prop
    intro t; fun_induction apply'
    · intro n h; cases h
      rw [Nat.zero_add]
      exact (σ _).prop
    all_goals
      intro n h; cases h; expose_names
      rw [Nat.add_right_comm]
      constructor; intro arg
      apply ih1; exact h arg⟩

private def rename' (π : FPerm 𝔸) : Term 𝔸 sig X → Term 𝔸 sig X
  | .var x => .var x
  | .pure f args => .pure f (fun x => rename' π (args x))
  | .free a f args => .free (π • a) f (fun x => rename' π (args x))
  | .bound a f args => .bound (π • a) f (fun x => rename' π (args x))

private lemma one_rename (t : Term 𝔸 sig X) : rename' 𝔸 1 t = t := by
  induction' t <;> all_goals simp [rename', *]

private lemma mul_rename (π τ : FPerm 𝔸) (t : Term 𝔸 sig X) :
    rename' 𝔸 (π * τ) t = rename' 𝔸 π (rename' 𝔸 τ t) := by
  induction' t <;> all_goals simp [rename', mul_smul, *]

instance : PermSet 𝔸 (Term 𝔸 sig X) where
  smul := rename' 𝔸
  one_smul := one_rename 𝔸
  mul_smul := mul_rename 𝔸

private def rename {n : ℕ} (π : FPerm 𝔸) (t : TermAtDepth 𝔸 sig X n) : TermAtDepth 𝔸 sig X n :=
  ⟨π • t.val, by
    suffices ∀ t n, t.HasDepth 𝔸 n → (rename' 𝔸 π t).HasDepth 𝔸 n by
      exact this t.val n t.prop
    intro t; fun_induction rename'
    · intro n h; cases h; constructor
    all_goals
        intro n h; cases h
        expose_names
        constructor; intro arg
        apply ih1; exact h arg⟩

instance {n : ℕ} : PermSet 𝔸 (TermAtDepth 𝔸 sig X n) where
  smul := rename 𝔸
  one_smul t := by apply Subtype.eq; exact one_rename 𝔸 t.val
  mul_smul π τ t := by apply Subtype.eq; exact mul_rename 𝔸 π τ t.val

end Algebra