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 TermAtDepth.var (x : X) : TermAtDepth 𝔸 sig X 0 := ⟨.var x, by constructor⟩

def TermAtDepth.pure {n : ℕ} (f : sig.pure) (args : sig.pureArgs f → TermAtDepth 𝔸 sig X n) :
    TermAtDepth 𝔸 sig X (n + sig.pureDepth f) :=
  ⟨.pure f (Subtype.val ∘ args), by constructor; intro arg; exact (args arg).prop⟩

def TermAtDepth.free {n : ℕ} (a : 𝔸) (f : sig.free) (args : sig.freeArgs f → TermAtDepth 𝔸 sig X n) :
    TermAtDepth 𝔸 sig X (n + sig.freeDepth f) :=
  ⟨.free a f (Subtype.val ∘ args), by constructor; intro arg; exact (args arg).prop⟩

def TermAtDepth.bound {n : ℕ} (a : 𝔸) (f : sig.bound) (args : sig.boundArgs f → TermAtDepth 𝔸 sig X n) :
    TermAtDepth 𝔸 sig X (n + sig.boundDepth f) :=
  ⟨.bound a f (Subtype.val ∘ args), by constructor; intro arg; exact (args arg).prop⟩

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

def Substitution.apply (σ : Substitution 𝔸 sig X Y) : Term 𝔸 sig X → Term 𝔸 sig Y
  | .var x => σ x
  | .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))

@[simp] lemma Substitution.apply_var {x : X} {σ : Substitution 𝔸 sig X Y} :
    σ.apply 𝔸 (.var x) = (σ x) := rfl

@[simp] lemma Substitution.apply_pure
    {f : sig.pure} {args : sig.pureArgs f → Term 𝔸 sig X} {σ : Substitution 𝔸 sig X Y} :
    σ.apply 𝔸 (.pure f args) = .pure f (apply 𝔸 σ ∘ args) := rfl

@[simp] lemma Substitution.apply_free
    {a : 𝔸} {f : sig.free} {args : sig.freeArgs f → Term 𝔸 sig X} {σ : Substitution 𝔸 sig X Y} :
    σ.apply 𝔸 (.free a f args) = .free a f (apply 𝔸 σ ∘ args) := rfl

@[simp] lemma Substitution.apply_bound
    {a : 𝔸} {f : sig.bound} {args : sig.boundArgs f → Term 𝔸 sig X} {σ : Substitution 𝔸 sig X Y} :
    σ.apply 𝔸 (.bound a f args) = .bound a f (apply 𝔸 σ ∘ args) := rfl

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

def UniformSubstitution.toSubstitution {l : ℕ} (σ : UniformSubstitution 𝔸 sig X Y l) : Substitution 𝔸 sig X Y :=
  fun x => (σ x).val

def UniformSubstitution.apply {n l : ℕ} (σ : UniformSubstitution 𝔸 sig X Y l) (t : TermAtDepth 𝔸 sig X n) :
    TermAtDepth 𝔸 sig Y (n + l) :=
  ⟨σ.toSubstitution.apply 𝔸 t.val, by
    suffices ∀ t n, t.HasDepth 𝔸 n → (σ.toSubstitution.apply 𝔸 t).HasDepth 𝔸 (n + l) by
      exact this t.val n t.prop
    intro t; fun_induction Substitution.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⟩

@[simp] lemma UniformSubstitution.apply_var
    {x : X} {l : ℕ} {σ : UniformSubstitution 𝔸 sig X Y l} :
    (σ.apply 𝔸 (.var 𝔸 x)).val = (σ x).val := rfl

@[simp] lemma UniformSubstitution.apply_pure
    {n : ℕ} {f : sig.pure} {args : sig.pureArgs f → TermAtDepth 𝔸 sig X n}
    {l : ℕ} {σ : UniformSubstitution 𝔸 sig X Y l} :
    (σ.apply 𝔸 (.pure 𝔸 f args)).val = (TermAtDepth.pure 𝔸 f (apply 𝔸 σ ∘ args)).val := rfl

@[simp] lemma UniformSubstitution.apply_free
    {n : ℕ} {a : 𝔸} {f : sig.free} {args : sig.freeArgs f → TermAtDepth 𝔸 sig X n}
    {l : ℕ} {σ : UniformSubstitution 𝔸 sig X Y l} :
    (σ.apply 𝔸 (.free 𝔸 a f args)).val = (TermAtDepth.free 𝔸 a f (apply 𝔸 σ ∘ args)).val := rfl

@[simp] lemma UniformSubstitution.apply_bound
    {n : ℕ} {a : 𝔸} {f : sig.bound} {args : sig.boundArgs f → TermAtDepth 𝔸 sig X n}
    {l : ℕ} {σ : UniformSubstitution 𝔸 sig X Y l} :
    (σ.apply 𝔸 (.bound 𝔸 a f args)).val = (TermAtDepth.bound 𝔸 a f (apply 𝔸 σ ∘ args)).val := rfl

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 𝔸

@[simp] lemma smul_var {π : FPerm 𝔸} {x : X} :
    π • (Term.var x : Term 𝔸 sig X) = Term.var (π • x) := rfl

@[simp] lemma smul_pure {π : FPerm 𝔸} {f : sig.pure} {args : sig.pureArgs f → Term 𝔸 sig X} :
    π • Term.pure f args = Term.pure f ((π • ·) ∘ args) := rfl

@[simp] lemma smul_free {π : FPerm 𝔸} {a : 𝔸} {f : sig.free} {args : sig.freeArgs f → Term 𝔸 sig X} :
    π • Term.free a f args = Term.free (π • a) f ((π • ·) ∘ args) := rfl

@[simp] lemma smul_bound {π : FPerm 𝔸} {a : 𝔸} {f : sig.bound} {args : sig.boundArgs f → Term 𝔸 sig X} :
    π • Term.bound a f args = Term.bound (π • a) f ((π • ·) ∘ args) := rfl

private def atoms : Term 𝔸 sig X → Set 𝔸
  | .var x => NominalSet.supp 𝔸 x
  | .pure f args => ⋃₀ {atoms (args x) | x : sig.pureArgs f}
  | .free a f args => {a} ∪ ⋃₀ {atoms (args x) | x : sig.freeArgs f}
  | .bound a f args => {a} ∪ ⋃₀ {atoms (args x) | x : sig.boundArgs f}

private lemma atoms_finite (t : Term 𝔸 sig X) : (atoms 𝔸 t).Finite := by
  fun_induction atoms <;> all_goals expose_names
  · exact NominalSet.supp_finite 𝔸
  · have hfin := (sig.pureArgs_fintype f).finite
    apply Set.Finite.sUnion
    · simp only [Set.setOf_exists, Set.setOf_eq_eq_singleton', Set.iUnion_singleton_eq_range]
      apply Set.finite_range
    · intro S ⟨x, hS⟩
      rw [← hS]
      exact ih1 x
  · have hfin := (sig.freeArgs_fintype f).finite
    apply Set.Finite.union
    · exact Set.finite_singleton a
    · apply Set.Finite.sUnion
      · simp only [Set.setOf_exists, Set.setOf_eq_eq_singleton', Set.iUnion_singleton_eq_range]
        apply Set.finite_range
      · intro S ⟨x, hS⟩
        rw [← hS]
        exact ih1 x
  · have hfin := (sig.boundArgs_fintype f).finite
    apply Set.Finite.union
    · exact Set.finite_singleton a
    · apply Set.Finite.sUnion
      · simp only [Set.setOf_exists, Set.setOf_eq_eq_singleton', Set.iUnion_singleton_eq_range]
        apply Set.finite_range
      · intro S ⟨x, hS⟩
        rw [← hS]
        exact ih1 x

private lemma atoms_support (t : Term 𝔸 sig X) : PermSet.HasSupport 𝔸 t (atoms 𝔸 t) := by
  intro π
  fun_induction atoms <;> all_goals expose_names
  · have hsupp : PermSet.HasSupport 𝔸 x (NominalSet.supp 𝔸 x) := NominalSet.supp_is_support 𝔸
    intro hπ
    simp only [PermSet.fix', smul_var, Term.var.injEq, Set.mem_setOf_eq]
    exact hsupp hπ
  · intro hπ
    simp only [PermSet.fix', smul_pure, Term.pure.injEq, heq_eq_eq, true_and, Set.mem_setOf_eq]
    ext x
    rw [Function.comp_apply]
    apply ih1
    apply PermSet.fix_antitone_mem 𝔸 hπ
    apply Set.subset_sUnion_of_mem
    use x
  all_goals
    intro hπ
    simp only [PermSet.fix', smul_free, smul_bound, Term.free.injEq, Term.bound.injEq, heq_eq_eq,
      true_and, Set.mem_setOf_eq]
    constructor
    · apply PermSet.fix_antitone_mem 𝔸 hπ
      apply Set.subset_union_left
      rfl
    · ext x
      rw [Function.comp_apply]
      apply ih1
      apply PermSet.fix_antitone_mem 𝔸 hπ
      apply Set.subset_union_of_subset_right
      apply Set.subset_sUnion_of_mem
      use x

instance : NominalSet 𝔸 (Term 𝔸 sig X) where
  fin_supp t := by
    use atoms 𝔸 t
    exact ⟨atoms_support 𝔸 t, atoms_finite 𝔸 t⟩

instance Term.instHasDepthEquivariant {n : ℕ} :
    Set.Equivariant 𝔸 {t : Term 𝔸 sig X | HasDepth 𝔸 t n} := by
  constructor; intro t π ht
  suffices ∀ t n, t.HasDepth 𝔸 n → (rename 𝔸 π t).HasDepth 𝔸 n by
    exact this t n ht
  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 : ℕ} : NominalSet 𝔸 (TermAtDepth 𝔸 sig X n) :=
  (Subtype.instNominalSetElem 𝔸 : NominalSet 𝔸 {t : Term 𝔸 sig X | Term.HasDepth 𝔸 t n})

end Algebra