import Mathlib
import Algebra.Terms

namespace Algebra

section Theories

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

structure Equation where
  vars : Type
  [vars_nominal : NominalSet 𝔸 vars]
  [vars_orbitFinite : OrbitFinite 𝔸 vars]
  left : Term 𝔸 sig vars
  right : Term 𝔸 sig vars

namespace Equation

instance (eq : Equation 𝔸 sig) : NominalSet 𝔸 eq.vars :=
  eq.vars_nominal

instance (eq : Equation 𝔸 sig) : OrbitFinite 𝔸 eq.vars :=
  eq.vars_orbitFinite

def HasDepth
    {𝔸 : Type} [Infinite 𝔸] [Countable 𝔸] [DecidableEq 𝔸] {sig : Signature}
    (eq : Equation 𝔸 sig) (n : ℕ) : Prop :=
  eq.left.HasDepth n ∧ eq.right.HasDepth n

end Equation

def EquationAtDepth (n : ℕ) :=
  { eq : Equation 𝔸 sig // eq.HasDepth n }

def UniformEquation :=
  { eq : Equation 𝔸 sig // ∃ n, eq.HasDepth n }

structure Theory where
  sig : Signature
  axioms : Set (UniformEquation 𝔸 sig)

end Theories

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

section Helpers

def EquationAtDepth.left {n : ℕ} (eq : EquationAtDepth 𝔸 sig n) : TermAtDepth 𝔸 sig eq.val.vars n :=
  ⟨eq.val.left, eq.prop.left⟩

def EquationAtDepth.right {n : ℕ} (eq : EquationAtDepth 𝔸 sig n) : TermAtDepth 𝔸 sig eq.val.vars n :=
  ⟨eq.val.right, eq.prop.right⟩

def EquationAtDepth.mk
    (X : Type) [NominalSet 𝔸 X] [OrbitFinite 𝔸 X] {n : ℕ} (t u : TermAtDepth 𝔸 sig X n) :
    EquationAtDepth 𝔸 sig n :=
  ⟨⟨X, t.val, u.val⟩, ⟨t.prop, u.prop⟩⟩

def EquationAtDepth.toUniform {n : ℕ} (eq : EquationAtDepth 𝔸 sig n) : UniformEquation 𝔸 sig :=
  ⟨eq.val, ⟨n, eq.prop⟩⟩

notation:50 X " ⊢ " lhs " ≣ " rhs => Equation.mk X lhs rhs

end Helpers

section DerivableEquations

inductive Derivable (T : Theory 𝔸) : (eq : Equation 𝔸 T.sig) → Prop
  | refl {X : Type} [NominalSet 𝔸 X] [OrbitFinite 𝔸 X] {x : X} :
      Derivable T (X ⊢ .var x ≣ .var x)
  | symm {X : Type} [NominalSet 𝔸 X] [OrbitFinite 𝔸 X]
      {t u : Term 𝔸 T.sig X} :
      Derivable T (X ⊢ u ≣ t) →
      Derivable T (X ⊢ t ≣ u)
  | trans {X : Type} [NominalSet 𝔸 X] [OrbitFinite 𝔸 X]
      {t u v : Term 𝔸 T.sig X} :
      Derivable T (X ⊢ t ≣ v) →
      Derivable T (X ⊢ v ≣ u) →
      Derivable T (X ⊢ t ≣ u)
  | cong_pure {X : Type} [NominalSet 𝔸 X] [OrbitFinite 𝔸 X]
      {f : T.sig.pure} {ts us : T.sig.pureArgs f → Term 𝔸 T.sig X} :
      (∀ x, Derivable T (X ⊢ ts x ≣ us x)) →
      Derivable T (X ⊢ .pure f ts ≣ .pure f us)
  | cong_free {X : Type} [NominalSet 𝔸 X] [OrbitFinite 𝔸 X]
      {a : 𝔸} {f : T.sig.free} {ts us : T.sig.freeArgs f → Term 𝔸 T.sig X} :
      (∀ x, Derivable T (X ⊢ ts x ≣ us x)) →
      Derivable T (X ⊢ .free a f ts ≣ .free a f us)
  | cong_bound {X : Type} [NominalSet 𝔸 X] [OrbitFinite 𝔸 X]
      {a : 𝔸} {f : T.sig.bound} {ts us : T.sig.boundArgs f → Term 𝔸 T.sig X} :
      (∀ x, Derivable T (X ⊢ ts x ≣ us x)) →
      Derivable T (X ⊢ .bound a f ts ≣ .bound a f us)
  | ax' {X : Type} [NominalSet 𝔸 X] [OrbitFinite 𝔸 X]
      {eq : UniformEquation 𝔸 T.sig}
      {l : ℕ} {σ : UniformSubstitution 𝔸 T.sig eq.val.vars X l} {τ : FPerm 𝔸} :
      (eq ∈ T.axioms) →
      (∀ (y : eq.val.vars) (π : FPerm 𝔸), Derivable T (X ⊢ σ.toSubstitution (π • y) ≣ π • (σ.toSubstitution y))) →
      Derivable T (X ⊢ σ.toSubstitution.apply (τ • eq.val.left) ≣ σ.toSubstitution.apply (τ • eq.val.right))
  | perm {X : Type} [NominalSet 𝔸 X] [OrbitFinite 𝔸 X]
      {a b : 𝔸} {f : T.sig.bound} {ts us : T.sig.boundArgs f → Term 𝔸 T.sig X} :
      a ≠ b →
      (∀ x, a # us x) →
      (∀ x, Derivable T (X ⊢ ts x ≣ ⸨a, b⸩ • us x)) →
      Derivable T (X ⊢ .bound a f ts ≣ .bound b f us)

def DerivableEq (T : Theory 𝔸) (X : Type) [NominalSet 𝔸 X] [OrbitFinite 𝔸 X] (n : ℕ) :
    Rel (TermAtDepth 𝔸 T.sig X n) (TermAtDepth 𝔸 T.sig X n) :=
  fun t u => Derivable T (X ⊢ t.val ≣ u.val)

variable {X Y : Type} [NominalSet 𝔸 X] [OrbitFinite 𝔸 X] [NominalSet 𝔸 Y] [OrbitFinite 𝔸 Y]

def Substitution.DerivablyEquivariant
    (T : Theory 𝔸) (σ : Substitution 𝔸 T.sig X Y) : Prop :=
  (∀ (x : X) (π : FPerm 𝔸), Derivable T (Y ⊢ σ (π • x) ≣ π • σ x))

/-- Restated ax' rule using DerivablyEquivariant in the premiss. -/
lemma Derivable.ax
    {T : Theory 𝔸}
    {eq : UniformEquation 𝔸 T.sig}
    {l : ℕ} {σ : UniformSubstitution 𝔸 T.sig eq.val.vars X l} {τ : FPerm 𝔸} :
    (eq ∈ T.axioms) →
    σ.toSubstitution.DerivablyEquivariant T →
    Derivable T (X ⊢ σ.toSubstitution.apply (τ • eq.val.left) ≣ σ.toSubstitution.apply (τ • eq.val.right)) :=
  .ax'

lemma Derivable.equivariant_substitution
    (T : Theory 𝔸)
    {l : ℕ} (σ : UniformSubstitution 𝔸 T.sig Y X l) :
    σ.toSubstitution.DerivablyEquivariant T →
    ∀ (t : Term 𝔸 T.sig Y) (π : FPerm 𝔸),
      Derivable T (X ⊢ σ.toSubstitution.apply (π • t) ≣ π • (σ.toSubstitution.apply t)) := by
  intro hequi t π
  induction' t <;> expose_names
  · exact hequi x π
  · simp only [Term.smul_pure, Substitution.apply_pure]
    apply Derivable.cong_pure
    exact args_ih
  · simp only [Term.smul_free, Substitution.apply_free]
    apply Derivable.cong_free
    exact args_ih
  · simp only [Term.smul_bound, Substitution.apply_bound]
    apply Derivable.cong_bound
    exact args_ih

lemma Derivable.equivariant
    (T : Theory 𝔸) {t u : Term 𝔸 T.sig X} :
    ∀ (π : FPerm 𝔸), Derivable T (X ⊢ t ≣ u) → Derivable T (X ⊢ π • t ≣ π • u) := by
  intro π
  suffices ∀ eq, Derivable T eq → Derivable T (eq.vars ⊢ π • eq.left ≣ π • eq.right) by
    exact this (X ⊢ t ≣ u)
  intro eq hderiv
  induction' hderiv <;> expose_names
  · exact .refl
  · apply Derivable.symm
    exact a_ih
  · apply Derivable.trans
    · exact a_ih
    · exact a_ih_1
  · apply Derivable.cong_pure
    exact a_ih
  · apply Derivable.cong_free
    exact a_ih
  · apply Derivable.cong_bound
    exact a_ih
  · apply Derivable.trans (by apply Derivable.symm; apply equivariant_substitution; exact h_1) _
    apply Derivable.trans _ (by apply equivariant_substitution; exact h_1)
    rw [← MulAction.mul_smul, ← MulAction.mul_smul]
    apply Derivable.ax
    · exact h
    · exact h_1
  · apply Derivable.perm
    · intro hcontra
      apply π.val.injective at hcontra
      contradiction
    · intro x
      show π • a # π • (us x)
      rw [← Rel.equivariant 𝔸 NominalSet.IsFresh]
      exact h_1 x
    · intro x
      show Derivable T (X_1 ⊢ π • ts x ≣ ⸨π a, π b⸩ • π • us x)
      rw [← MulAction.mul_smul, ← FPerm.mul_swap_eq_swap_mul, MulAction.mul_smul]
      exact a_ih x

variable {T : Theory 𝔸} {n : ℕ}

instance DerivableEq.instEquivariant : Rel.Equivariant 𝔸 (DerivableEq T X n) := by
  constructor; intro t u
  exact Derivable.equivariant T

@[refl]
lemma DerivableEq.refl (t : TermAtDepth 𝔸 T.sig X n) : DerivableEq T X n t t := by
  suffices ∀ t, Derivable T (X ⊢ t ≣ t) by exact this t.val
  intro t
  induction' t <;> expose_names
  · apply Derivable.refl
  · apply Derivable.cong_pure; assumption
  · apply Derivable.cong_free; assumption
  · apply Derivable.cong_bound; assumption

@[symm]
lemma DerivableEq.symm {t u : TermAtDepth 𝔸 T.sig X n} :
    DerivableEq T X n u t → DerivableEq T X n t u := by
  intro h
  apply Derivable.symm
  exact h

@[trans]
lemma DerivableEq.trans {t u v : TermAtDepth 𝔸 T.sig X n} :
    DerivableEq T X n t v → DerivableEq T X n v u → DerivableEq T X n t u := by
  intro htv hvu
  apply Derivable.trans
  · exact htv
  · exact hvu

instance TermAtDepth.setoid : Setoid (TermAtDepth 𝔸 T.sig X n) where
  r := DerivableEq T X n
  iseqv :=
    { refl := DerivableEq.refl
      symm := DerivableEq.symm
      trans := DerivableEq.trans }

end DerivableEquations

end Algebra