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 => EquationAtDepth.mk X lhs rhs

end Helpers

section DerivableEquations

inductive Derivable (T : Theory 𝔸) : {m : ℕ} → (eq : EquationAtDepth 𝔸 T.sig m) → Prop
  | refl {X : Type} [NominalSet 𝔸 X] [OrbitFinite 𝔸 X] {x : X} :
      Derivable T (X ⊢ .var x ≣ .var x)
  | symm {X : Type} [NominalSet 𝔸 X] [OrbitFinite 𝔸 X]
      {m : ℕ} {t u : TermAtDepth 𝔸 T.sig X m} :
      Derivable T (X ⊢ u ≣ t) →
      Derivable T (X ⊢ t ≣ u)
  | trans {X : Type} [NominalSet 𝔸 X] [OrbitFinite 𝔸 X]
      {m : ℕ} {t u v : TermAtDepth 𝔸 T.sig X m} :
      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} {m : ℕ} {ts us : T.sig.pureArgs f → TermAtDepth 𝔸 T.sig X m} :
      (∀ 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} {m : ℕ} {ts us : T.sig.freeArgs f → TermAtDepth 𝔸 T.sig X m} :
      (∀ 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} {m : ℕ} {ts us : T.sig.boundArgs f → TermAtDepth 𝔸 T.sig X m} :
      (∀ x, Derivable T (X ⊢ ts x ≣ us x)) →
      Derivable T (X ⊢ .bound a f ts ≣ .bound a f us)
  | ax' {X Y : Type} [NominalSet 𝔸 X] [OrbitFinite 𝔸 X] [NominalSet 𝔸 Y] [OrbitFinite 𝔸 Y]
      {m : ℕ} {r s : TermAtDepth 𝔸 T.sig Y m}
      {l : ℕ} {σ : UniformSubstitution 𝔸 T.sig Y X l} {τ : FPerm 𝔸} :
      ((Y ⊢ r ≣ s).toUniform ∈ T.axioms) →
      (∀ (y : Y) (π : FPerm 𝔸), Derivable T (X ⊢ σ (π • y) ≣ π • (σ y))) →
      Derivable T (X ⊢ σ.apply (τ • r) ≣ σ.apply (τ • s))
  | perm {X : Type} [NominalSet 𝔸 X] [OrbitFinite 𝔸 X]
      {a b : 𝔸} {f : T.sig.bound} {m : ℕ} {ts us : T.sig.boundArgs f → TermAtDepth 𝔸 T.sig X m} :
      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)

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

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

/-- Restated ax' rule using DerivablyEquivariant in the premiss. -/
lemma Derivable.ax
    {T : Theory 𝔸}
    {m : ℕ} {r s : TermAtDepth 𝔸 T.sig Y m}
    {l : ℕ} {σ : UniformSubstitution 𝔸 T.sig Y X l} {τ : FPerm 𝔸} :
    ((Y ⊢ r ≣ s).toUniform ∈ T.axioms) →
    σ.DerivablyEquivariant T →
    Derivable T (X ⊢ σ.apply (τ • r) ≣ σ.apply (τ • s)) :=
  .ax'

lemma Derivable.equivariant_substitution
    (T : Theory 𝔸)
    {l : ℕ} (σ : UniformSubstitution 𝔸 T.sig Y X l) (hequi : σ.DerivablyEquivariant T)
    {m : ℕ} (t : TermAtDepth 𝔸 T.sig Y m) {π : FPerm 𝔸} :
    Derivable T (X ⊢ σ.apply (π • t) ≣ π • (σ.apply t)) := by
  induction' t using TermAtDepth.induction <;> expose_names
  · sorry
  · sorry
  · sorry
  · sorry

end DerivableEquations

end Algebra