import Mathlib
import Nominal

namespace NominalAlgebra

section Definitions

/-- 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))

instance Signature.instPureArgsFintype (sig : Signature) (f : sig.pure) : Fintype (sig.pureArgs f) := sig.args_fintype (.inl f)
instance Signature.instFreeArgsFintype (sig : Signature) (f : sig.free) : Fintype (sig.freeArgs f) := sig.args_fintype (.inr (.inl f))
instance Signature.instBoundArgsFintype (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 (𝔸 : Type) [Infinite 𝔸] [Countable 𝔸] [DecidableEq 𝔸]
variable (sig : Signature)

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

inductive Term.HasDepth
    {𝔸 : Type} [Infinite 𝔸] [Countable 𝔸] [DecidableEq 𝔸]
    {sig : Signature}
    {X : Type} [NominalSet 𝔸 X]
    : 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 (X : Type) [NominalSet 𝔸 X] (n : ℕ) :=
  { t : Term 𝔸 sig X // t.HasDepth n }

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

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

end Definitions

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

section ConvenienceConstructors

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⟩

theorem TermAtDepth.induction
  {motive : {n : ℕ} → TermAtDepth 𝔸 sig X n → Prop}
  (var : ∀ (x : X), motive (.var x))
  (pure : ∀ {n : ℕ} (f : sig.pure) (args : sig.pureArgs f → TermAtDepth 𝔸 sig X n),
    (∀ x, motive (args x)) → motive (.pure f args))
  (free : ∀ {n : ℕ} (a : 𝔸) (f : sig.free) (args : sig.freeArgs f → TermAtDepth 𝔸 sig X n),
    (∀ x, motive (args x)) → motive (.free a f args))
  (bound : ∀ {n : ℕ} (a : 𝔸) (f : sig.bound) (args : sig.boundArgs f → TermAtDepth 𝔸 sig X n),
    (∀ x, motive (args x)) → motive (.bound a f args))
  : ∀ {n : ℕ} (t : TermAtDepth 𝔸 sig X n), motive t := by
  intro n ⟨t, ht⟩
  induction' t generalizing n <;> cases ht <;> expose_names
  · exact var x
  · let args' : sig.pureArgs f → TermAtDepth 𝔸 sig X k := fun x => ⟨args x, h x⟩
    have args_ih' : ∀ x, motive (args' x) := fun x => args_ih x (h x)
    exact pure f args' args_ih'
  · let args' : sig.freeArgs f → TermAtDepth 𝔸 sig X k := fun x => ⟨args x, h x⟩
    have args_ih' : ∀ x, motive (args' x) := fun x => args_ih x (h x)
    exact free a f args' args_ih'
  · let args' : sig.boundArgs f → TermAtDepth 𝔸 sig X k := fun x => ⟨args x, h x⟩
    have args_ih' : ∀ x, motive (args' x) := fun x => args_ih x (h x)
    exact bound a f args' args_ih'

end ConvenienceConstructors

section Substitution

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.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

end Substitution

section NominalSet

namespace Term

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) : t.rename 1 = t := by
  induction' t <;> all_goals simp [rename, *]

private lemma mul_rename (π τ : FPerm 𝔸) (t : Term 𝔸 sig X) :
    t.rename (π * τ) = (t.rename τ).rename π := 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} :
    π • (.var x : Term 𝔸 sig X) = .var (π • x) := rfl

@[simp] lemma smul_pure {π : FPerm 𝔸} {f : sig.pure} {args : sig.pureArgs f → Term 𝔸 sig X} :
    π • Term.pure f args = .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 = .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 = .bound (π • a) f ((π • ·) ∘ args) := rfl

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

private lemma atoms_finite (t : Term 𝔸 sig X) : t.atoms.Finite := by
  fun_induction atoms <;> all_goals expose_names
  · exact NominalSet.supp_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
  all_goals
    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 t.atoms := 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 t.atoms
    exact ⟨t.atoms_support, t.atoms_finite⟩

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

end Term

namespace TermAtDepth

instance {n : ℕ} : NominalSet 𝔸 (TermAtDepth 𝔸 sig X n) :=
  (Subtype.instNominalSetElem 𝔸 : NominalSet 𝔸 {t : Term 𝔸 sig X | t.HasDepth n})

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

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

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

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

end TermAtDepth

end NominalSet

end NominalAlgebra