import Mathlib
import Algebra.Terms

namespace Algebra

section Theories

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

@[ext]
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 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 UniformEquation.mk
    (X : Type) [NominalSet 𝔸 X] [OrbitFinite 𝔸 X] {n : ℕ} (t u : TermAtDepth 𝔸 sig X n) :
    UniformEquation 𝔸 sig :=
  ⟨⟨X, t.val, u.val⟩, ⟨n, t.prop, u.prop⟩⟩

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

end Helpers

section DerivableEquations

inductive Derivable (T : Theory 𝔸) : (eq : UniformEquation 𝔸 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]
      {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) ∈ 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)

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 ≣ u)

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

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

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

lemma Derivable.of_eq
    {m : ℕ} {eq : UniformEquation 𝔸 T.sig} (t u : TermAtDepth 𝔸 T.sig eq.val.vars m)
    (ht : t.val = eq.val.left := by rfl)
    (hu : u.val = eq.val.right := by rfl) :
    Derivable T (eq.val.vars ⊢ t ≣ u) →
    Derivable T eq := by
  intro hderiv
  have heq : (eq.val.vars ⊢ t ≣ u) = eq := by
    apply Subtype.eq; ext
    · rfl
    · rfl
    · rw [heq_eq_eq]; exact ht
    · rw [heq_eq_eq]; exact hu
  rw [← heq]; exact hderiv

lemma Derivable.equivariant_substitution
    {l : ℕ} (σ : UniformSubstitution 𝔸 T.sig Y X l) :
    σ.DerivablyEquivariant T →
    ∀ {n : ℕ} (t : TermAtDepth 𝔸 T.sig Y n) (π : FPerm 𝔸),
      Derivable T (X ⊢ σ.apply (π • t) ≣ π • (σ.apply t)) := by
  intro hequi n t π
  induction' t using TermAtDepth.induction <;> expose_names
  · apply of_eq (σ (π • x)) (π • σ x)
    exact hequi x π
  · apply of_eq (.pure f (σ.apply ∘ (π • ·) ∘ args)) (.pure f ((π • ·) ∘ σ.apply ∘ args))
    apply Derivable.cong_pure
    exact h
  · apply of_eq (.free (π • a) f (σ.apply ∘ (π • ·) ∘ args)) (.free (π • a) f ((π • ·) ∘ σ.apply ∘ args))
    apply Derivable.cong_free
    exact h
  · apply of_eq (.bound (π • a) f (σ.apply ∘ (π • ·) ∘ args)) (.bound (π • a) f ((π • ·) ∘ σ.apply ∘ args))
    apply Derivable.cong_bound
    exact h

lemma Derivable.equivariant {m : ℕ} {t u : TermAtDepth 𝔸 T.sig X m} :
    ∀ (π : FPerm 𝔸), Derivable T (X ⊢ t ≣ u) → Derivable T (X ⊢ π • t ≣ π • u) := by
  intro π
  let rename (π : FPerm 𝔸) (eq : UniformEquation 𝔸 T.sig) : UniformEquation 𝔸 T.sig :=
    ⟨⟨eq.val.vars, π • eq.val.left, π • eq.val.right⟩, by
      obtain ⟨n, ⟨hn₁, hn₂⟩⟩ := eq.prop
      let left : TermAtDepth 𝔸 T.sig eq.val.vars n := ⟨eq.val.left, hn₁⟩
      let right : TermAtDepth 𝔸 T.sig eq.val.vars n := ⟨eq.val.right, hn₂⟩
      use n; exact ⟨(π • left).prop, (π • right).prop⟩⟩
  have rename_eq
      {Y : Type} [NominalSet 𝔸 Y] [OrbitFinite 𝔸 Y] {π : FPerm 𝔸} {k : ℕ}
      {r s : TermAtDepth 𝔸 T.sig Y k} : rename π (Y ⊢ r ≣ s) = (Y ⊢ π • r ≣ π • s) := rfl
  suffices ∀ eq, Derivable T eq → Derivable T (rename π eq) by
    exact this (X ⊢ t ≣ u)
  intro eq hderiv
  induction' hderiv <;> expose_names
  · exact .refl
  · rw [rename_eq]
    apply Derivable.symm
    exact a_ih
  · rw [rename_eq] at *
    apply Derivable.trans
    · exact a_ih
    · exact a_ih_1
  · apply of_eq (TermAtDepth.pure f ((π • ·) ∘ ts)) (TermAtDepth.pure f ((π • ·) ∘ us))
    apply Derivable.cong_pure
    exact a_ih
  · apply of_eq (TermAtDepth.free (π • a) f ((π • ·) ∘ ts)) (TermAtDepth.free (π • a) f ((π • ·) ∘ us))
    apply Derivable.cong_free
    exact a_ih
  · apply of_eq (TermAtDepth.bound (π • a) f ((π • ·) ∘ ts)) (TermAtDepth.bound (π • a) f ((π • ·) ∘ us))
    apply Derivable.cong_bound
    exact a_ih
  · rw [rename_eq]
    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 of_eq (TermAtDepth.bound (π • a) f ((π • ·) ∘ ts)) (TermAtDepth.bound (π • b) f ((π • ·) ∘ us))
    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

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

@[refl]
lemma DerivableEq.refl (t : TermAtDepth 𝔸 T.sig X n) : DerivableEq T X n t t := by
  induction' t using TermAtDepth.induction <;> 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