import Mathlib
import Nominal.Support
import Nominal.Equivariance
import Nominal.Category
import Nominal.Permutations
import Nominal.Instances.Basic
import Nominal.Instances.Abstraction
import Nominal.NameRestriction.Finsupp

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

-- M₀ algebra structure: Equivariant join semilattices with name restriction

class AlgebraStructure (X : Type) [NominalSet 𝔸 X] where
  bot : Unit →ₑ[𝔸] X
  join : (X × X) →ₑ[𝔸] X
  res : NameAbstraction 𝔸 X →ₑ[𝔸] X
  ax_neutral : ∀ x, join (x, bot ()) = x
  ax_idempotent : ∀ x, join (x, x) = x
  ax_commutative : ∀ x y, join (x, y) = join (y, x)
  ax_associative : ∀ x y z, join (join (x, y), z) = join (x, join (y, z))
  ax_fresh : ∀ a x, a # x → res (⟪a⟫ x) = x
  ax_swap : ∀ a b x, res (⟪a⟫ res (⟪b⟫ x)) = res (⟪b⟫ res (⟪a⟫ x))
  ax_distributive : ∀ a x y, res (⟪a⟫ join (x, y)) = join (res (⟪a⟫ x), res (⟪a⟫ y))
  ax_redundant : ∀ a x, x = join (x, res (⟪a⟫ x))

private abbrev AlgebraStructure.bot' {X : Type} [NominalSet 𝔸 X] [AlgebraStructure 𝔸 X] :=
  (AlgebraStructure.bot : Unit →ₑ[𝔸] X)
private abbrev AlgebraStructure.join' {X : Type} [NominalSet 𝔸 X] [AlgebraStructure 𝔸 X] :=
  (AlgebraStructure.join : (X × X) →ₑ[𝔸] X)

notation "⊥[" A "]" => AlgebraStructure.bot' A ()
notation:60 x " ∨[" A "] " y => AlgebraStructure.join' A (x, y)
notation:60 a " \\ " x => AlgebraStructure.res (⟪a⟫ x)

namespace AlgebraStructure

variable {X : Type} [NominalSet 𝔸 X] [AlgebraStructure 𝔸 X]

lemma bot_supp : NominalSet.supp 𝔸 (⊥[𝔸] : X) = ∅ := by
  rw [← Set.subset_empty_iff]
  show NominalSet.supp 𝔸 ((bot' 𝔸).toFun ()) ⊆ ∅
  grw [Function.Equivariant.supp_subset]
  simp

lemma bot_fresh {a : 𝔸} : a # (⊥[𝔸] : X) := by
  simp [NominalSet.IsFresh, bot_supp]

lemma bot_equivariant {π : FPerm 𝔸} : π • (⊥[𝔸] : X) = ⊥[𝔸] := by
  show π ∈ PermSet.fix' 𝔸 (bot' 𝔸 ())
  apply NominalSet.supp_is_support
  rw [bot_supp]
  intro a ha; contradiction

lemma join_equivariant {x y : X} {π : FPerm 𝔸} : π • (x ∨[𝔸] y) = (π • x) ∨[𝔸] (π • y) := by
  show π • (join' 𝔸).toFun (x, y) = (join' 𝔸).toFun (π • (x, y))
  rw [← Function.equivariant]

lemma res_equivariant {a : 𝔸} {x : X} {π : FPerm 𝔸} : π • (a \ x) = (π • a) \ (π • x) := by
  show π • res.toFun (⟪a⟫ x) = res.toFun (π • (⟪a⟫ x))
  rw [← Function.equivariant]

end AlgebraStructure

-- Base Model: M₀ X

def BaseModel (X : Type) [NominalSet 𝔸 X] : Type :=
  PermSet.FinSet 𝔸 X

namespace BaseModel

variable {X : Type} [NominalSet 𝔸 X]

instance : NominalSet 𝔸 (BaseModel 𝔸 X) := PermSet.FinSet.instNominalSet 𝔸
instance : EmptyCollection (BaseModel 𝔸 X) := PermSet.FinSet.instEmptyCollection 𝔸
instance : Union (BaseModel 𝔸 X) := PermSet.FinSet.instUnion 𝔸
instance : Inter (BaseModel 𝔸 X) := PermSet.FinSet.instInter 𝔸
instance : SDiff (BaseModel 𝔸 X) := PermSet.FinSet.instSDiff 𝔸

@[simp] lemma empty_val : (∅ : BaseModel 𝔸 X).val = ∅ := rfl
@[simp] lemma union_val {S₁ S₂ : BaseModel 𝔸 X} : (S₁ ∪ S₂).val = S₁.val ∪ S₂.val := rfl
@[simp] lemma inter_val {S₁ S₂ : BaseModel 𝔸 X} : (S₁ ∩ S₂).val = S₁.val ∩ S₂.val := rfl
@[simp] lemma sdiff_val {S₁ S₂ : BaseModel 𝔸 X} : (S₁ \ S₂).val = S₁.val \ S₂.val := rfl
@[simp] lemma smul_val {L : BaseModel 𝔸 X} {π : FPerm 𝔸} : (π • L).val = π • L.val := rfl

def bot : Unit → BaseModel 𝔸 X :=
  fun _ ↦ ∅

instance : Function.Equivariant 𝔸 (bot 𝔸 : Unit → BaseModel 𝔸 X) := by
  constructor; intro _ π
  apply Subtype.eq; simp [bot]

def union : (BaseModel 𝔸 X × BaseModel 𝔸 X) → BaseModel 𝔸 X := PermSet.FinSet.union 𝔸

instance : Function.Equivariant 𝔸 (union 𝔸 : (BaseModel 𝔸 X × BaseModel 𝔸 X) → BaseModel 𝔸 X) :=
  PermSet.FinSet.instEquivariantProdUnion 𝔸

def restrict : NameAbstraction 𝔸 (BaseModel 𝔸 X) → BaseModel 𝔸 X := Quotient.lift
  (fun ⟨a, S⟩ ↦ ⟨{x ∈ S.val | a # x}, by
    apply Set.Finite.subset S.prop; simp⟩)
  (by
    intro ⟨a₁, S₁⟩ ⟨a₂, S₂⟩ haeq
    apply Subtype.eq
    simp only
    rw [← NameAbstraction.eq_iff_alpha_eq 𝔸 (BaseModel 𝔸 X), NameAbstraction.eq_iff_cases] at haeq
    cases haeq with
    | inl h => rw [h.left, h.right]
    | inr h =>
      ext x; constructor
      · intro ⟨hx, ha₁⟩
        have hx' : ⸨a₁, a₂⸩ • x ∈ S₂.val := by
          rw [← Set.mem_inv_smul_set_iff, FPerm.swap_inv]
          show x ∈ (⸨a₁, a₂⸩ • S₂).val
          rw [← h.right.right]; exact hx
        have ha₂ := h.right.left
        rw [PermSet.FinSet.fresh] at ha₂
        specialize ha₂ (⸨a₁, a₂⸩ • x) hx'
        nth_rw 1 [show a₁ = ⸨a₁, a₂⸩ • a₂ by simp] at ha₂
        rw [← Rel.equivariant 𝔸 NominalSet.IsFresh] at ha₂
        constructor
        · rw [← show ⸨a₁, a₂⸩ • x = x from NominalSet.fresh_swap_fix 𝔸 ha₁ ha₂]
          exact hx'
        · exact ha₂
      · intro ⟨hx, ha₂⟩
        have hx' : ⸨a₁, a₂⸩ • x ∈ S₁.val := by
          rw [← Set.mem_inv_smul_set_iff, FPerm.swap_inv]
          show x ∈ (⸨a₁, a₂⸩ • S₁).val
          rw [h.right.right, ← MulAction.mul_smul, FPerm.swap_mul_self, MulAction.one_smul]
          exact hx
        have ha₁ := h.right.left
        rw [PermSet.FinSet.fresh] at ha₁
        specialize ha₁ x hx
        rw [← show ⸨a₁, a₂⸩ • x = x from NominalSet.fresh_swap_fix 𝔸 ha₁ ha₂] at hx'
        rw [← MulAction.mul_smul, FPerm.swap_mul_self, MulAction.one_smul] at hx'
        exact ⟨hx', ha₁⟩)

instance : Function.Equivariant 𝔸 (restrict 𝔸 : NameAbstraction 𝔸 (BaseModel 𝔸 X) → BaseModel 𝔸 X) := by
  constructor
  apply Quotient.ind
  intro ⟨a, S⟩ π
  apply Subtype.eq
  simp only [restrict, NameAbstraction.smul, Quotient.lift_mk, smul_val]
  ext x; constructor
  · intro ⟨hx, hfresh⟩
    rw [Set.mem_smul_set_iff_inv_smul_mem]; constructor
    · rw [← Set.mem_smul_set_iff_inv_smul_mem]; exact hx
    · rw [show a = π⁻¹ • π • a by simp, ← Rel.equivariant 𝔸 NominalSet.IsFresh]
      exact hfresh
  · rw [Set.mem_smul_set_iff_inv_smul_mem]
    intro ⟨hx, hfresh⟩; constructor
    · rw [Set.mem_smul_set_iff_inv_smul_mem]; exact hx
    · rw [show x = π • π⁻¹ • x by simp, ← Rel.equivariant 𝔸 NominalSet.IsFresh]
      exact hfresh

lemma restrict_eq {a : 𝔸} {S : BaseModel 𝔸 X} :
    (restrict 𝔸 (⟪a⟫ S)).val = {x ∈ S.val | a # x} := rfl

instance : AlgebraStructure 𝔸 (BaseModel 𝔸 X) where
  bot := ⟨bot 𝔸, inferInstance⟩
  join := ⟨union 𝔸, inferInstance⟩
  res := ⟨restrict 𝔸, inferInstance⟩
  ax_neutral x := by
    apply Subtype.eq
    show x.val ∪ ∅ = x.val
    simp
  ax_idempotent x := by
    apply Subtype.eq
    show x.val ∪ x.val = x.val
    simp
  ax_commutative x y := by
    apply Subtype.eq
    show x.val ∪ y.val = y.val ∪ x.val
    apply Set.union_comm
  ax_associative x y z := by
    apply Subtype.eq
    show (x.val ∪ y.val) ∪ z.val = x.val ∪ (y.val ∪ z.val)
    apply Set.union_assoc
  ax_fresh a S := by
    intro hfresh
    show restrict 𝔸 (⟪a⟫ S) = S
    apply Subtype.eq
    rw [restrict_eq]
    rw [PermSet.FinSet.fresh] at hfresh
    ext x; constructor
    · intro ⟨hS, _⟩; exact hS
    · intro hS; exact ⟨hS, hfresh x hS⟩
  ax_swap a b S := by
    show restrict 𝔸 (⟪a⟫ restrict 𝔸 (⟪b⟫ S)) = restrict 𝔸 (⟪b⟫ restrict 𝔸 (⟪a⟫ S))
    apply Subtype.eq
    ext x
    simp only [restrict_eq, Set.mem_setOf_eq]
    ac_nf
  ax_distributive a S₁ S₂ := by
    show restrict 𝔸 (⟪a⟫ S₁ ∪ S₂) = restrict 𝔸 (⟪a⟫ S₁) ∪ restrict 𝔸 (⟪a⟫ S₂)
    apply Subtype.eq
    simp [restrict_eq]
  ax_redundant a S := by
    show S = S ∪ restrict 𝔸 (⟪a⟫ S)
    apply Subtype.eq; ext x
    simp only [union_val, restrict_eq, Set.mem_union, Set.mem_setOf_eq, iff_self_or, and_imp]
    intro h _; exact h

@[simp] lemma bot_eq : (⊥[𝔸] : BaseModel 𝔸 X) = ∅ := rfl
@[simp] lemma join_eq {S₁ S₂ : BaseModel 𝔸 X} : (S₁ ∨[𝔸] S₂) = S₁ ∪ S₂ := rfl
@[simp] lemma res_eq {a : 𝔸} {S : BaseModel 𝔸 X} : (a \ S) = restrict 𝔸 (⟪a⟫ S) := rfl

end BaseModel

-- Step Model: \bar{M₁} X

def Req1 {X : Type} [NominalSet 𝔸 X] [AlgebraStructure 𝔸 X] (f : PermSet.FinsuppFunction 𝔸 𝔸 X) :=
  -- Support for individual atoms may only additionally add that atom
  ∀ a, PermSet.HasSupport 𝔸 (f.val a) (NominalSet.supp 𝔸 f ∪ {a})

lemma Req1.choose_restriction
    {X : Type} [NominalSet 𝔸 X] [AlgebraStructure 𝔸 X] (f : PermSet.FinsuppFunction 𝔸 𝔸 X) :
    Req1 𝔸 f → ∀ (a c d : 𝔸), c # f → d # f → c ≠ a → d ≠ a →
      (c \ (⸨a, c⸩ • f).val a) = (d \ (⸨a, d⸩ • f).val a) := by
  intro hreq a c d hcfresh hdfresh hca hda
  nth_rw 1 [← show ⸨c, d⸩ • f = f from NominalSet.fresh_swap_fix 𝔸 hcfresh hdfresh]
  simp only [PermSet.Finsupp.smul_val, PermSet.smul_function, FPerm.swap_inv, PermSet.atom_swap,
    Equiv.swap_apply_left]
  rw [show ⟪d⟫ ⸨a,d⸩ • f.val d = ⟪⸨a, d⸩ • a⟫ ⸨a,d⸩ • f.val d by simp]
  rw [show ⟪⸨a, d⸩ • a⟫ ⸨a,d⸩ • f.val d = ⸨a, d⸩ • ⟪a⟫ f.val d by simp]
  rw [show ⟪c⟫ ⸨a,c⸩ • ⸨c, d⸩ • f.val d = ⟪⸨a, c⸩ • a⟫ ⸨a,c⸩ • ⸨c, d⸩ • f.val d by simp]
  rw [show ⟪⸨a, c⸩ • a⟫ ⸨a,c⸩ • ⸨c, d⸩ • f.val d = ⸨a, c⸩ • ⟪a⟫ ⸨c, d⸩ • f.val d by simp]
  nth_rw 2 [show a = ⸨c, d⸩ • a by
    symm; apply FPerm.swap_apply_of_ne_of_ne
    · symm; exact hca
    · symm; exact hda]
  rw [show ⟪⸨c, d⸩ • a⟫ ⸨c,d⸩ • f.val d = ⸨c, d⸩ • ⟪a⟫ f.val d by simp]
  suffices ⸨a,c⸩ • ⸨c,d⸩ • ⟪a⟫ f.val d = ⸨a,d⸩ • ⟪a⟫ f.val d by rw [this]
  rw [← MulAction.mul_smul]
  apply PermSet.support_permutations_eq 𝔸 (NominalSet.supp_is_support 𝔸)
  intro b hb
  by_cases hba : b = a <;> by_cases hbc : b = c <;> by_cases hbd : b = d
  · simp [*] at *
  · simp [*] at *
  · simp [*] at *
  · simp [*] at *
  · rw [← hbc, ← hbd]; simp
  · rw [hbc] at hb
    suffices c ∉ NominalSet.supp 𝔸 (⟪a⟫ f.val d) by contradiction
    suffices c ∉ NominalSet.supp 𝔸 (f.val d) by simp [this]
    intro hc
    have h := NominalSet.support_subset_supp 𝔸 (hreq d) (by simp [NominalSet.supp_finite])
    specialize h hc
    simp only [Set.union_singleton, Set.mem_insert_iff] at h
    cases h
    · simp [*] at *
    · contradiction
  · simp [*] at *
  · simp only [Subgroup.coe_mul, Equiv.Perm.coe_mul, Function.comp_apply]
    rw [FPerm.swap_apply_of_ne_of_ne hbc hbd]
    rw [FPerm.swap_apply_of_ne_of_ne hba hbc]
    rw [FPerm.swap_apply_of_ne_of_ne hba hbd]

def Req2 {X : Type} [NominalSet 𝔸 X] [AlgebraStructure 𝔸 X] (f : PermSet.FinsuppFunction 𝔸 𝔸 X) :=
  -- Traces for all atoms include those from bound traces
  ∀ a c, c # f → c ≠ a → f.val a = f.val a ∨[𝔸] (c \ (⸨a, c⸩ • f).val a)

def permissibleFunctions (X : Type) [NominalSet 𝔸 X] [AlgebraStructure 𝔸 X] :=
  { f : PermSet.FinsuppFunction 𝔸 𝔸 X | Req1 𝔸 f ∧ Req2 𝔸 f }

def StepModel (X : Type) [NominalSet 𝔸 X] [AlgebraStructure 𝔸 X] : Type :=
  permissibleFunctions 𝔸 X

namespace StepModel

variable {X : Type} [NominalSet 𝔸 X] [AlgebraStructure 𝔸 X]

private lemma permissibleFunctions_req1
    (f : PermSet.FinsuppFunction 𝔸 𝔸 X) (π : FPerm 𝔸) : Req1 𝔸 f → Req1 𝔸 (π • f) := by
  intro hf a
  specialize hf (π⁻¹ • a)
  rw [show NominalSet.supp 𝔸 (π • f) = π • NominalSet.supp 𝔸 f from Function.equivariant _ _ _ _]
  rw [show ({a} : Set 𝔸) = π • {π⁻¹ • a} by simp]
  rw [← Set.smul_set_union]
  rw [PermSet.Finsupp.smul_val, PermSet.smul_function]
  rw [← Rel.equivariant 𝔸 (PermSet.HasSupport 𝔸)]
  exact hf

private lemma permissibleFunctions_req2
    (f : PermSet.FinsuppFunction 𝔸 𝔸 X) (π : FPerm 𝔸) : Req2 𝔸 f → Req2 𝔸 (π • f) := by
  intro hf a c hcf hca
  specialize hf (π⁻¹ • a) (π⁻¹ • c)
  specialize hf (by
    rw [show f = π⁻¹ • π • f by simp]
    rw [← Rel.equivariant 𝔸 NominalSet.IsFresh]
    exact hcf)
  specialize hf (by
    intro hcontra
    have : π • π⁻¹ • c = π • π⁻¹ • a := by rw [hcontra]
    simp only [smul_inv_smul] at this
    exact hca this)
  simp only [PermSet.Finsupp.smul_val, PermSet.smul_function, FPerm.swap_inv, PermSet.atom_swap,
    Equiv.swap_apply_left]
  nth_rw 1 [hf]
  rw [AlgebraStructure.join_equivariant, AlgebraStructure.res_equivariant]
  simp only [smul_inv_smul, PermSet.Finsupp.smul_val, PermSet.smul_function, FPerm.swap_inv,
    PermSet.atom_swap, Equiv.swap_apply_left]
  suffices π • ⸨π⁻¹ • a,π⁻¹ • c⸩ • f.val (π⁻¹ • c) = ⸨a, c⸩ • π • f.val (π⁻¹ • c) by rw [this]
  rw [← MulAction.mul_smul, ← MulAction.mul_smul]
  suffices π * ⸨π⁻¹ • a,π⁻¹ • c⸩ = ⸨a,c⸩ * π by rw [this]
  rw [FPerm.mul_swap_eq_swap_mul, show π⁻¹ • a = π⁻¹ a from rfl, show π⁻¹ • c = π⁻¹ c from rfl]
  simp

instance : Set.Equivariant 𝔸 (permissibleFunctions 𝔸 X) := by
  constructor; intro f π hf
  constructor
  · exact permissibleFunctions_req1 𝔸 f π hf.left
  · exact permissibleFunctions_req2 𝔸 f π hf.right

instance : NominalSet 𝔸 (StepModel 𝔸 X) := Subtype.instNominalSetElem 𝔸

@[simp] lemma smul_val (π : FPerm 𝔸) (f : StepModel 𝔸 X) : (π • f).val = π • f.val :=
  Subtype.smul_val 𝔸 π f

instance instValEquivariant : Function.Equivariant 𝔸 (Subtype.val : StepModel 𝔸 X → PermSet.FinsuppFunction 𝔸 𝔸 X) :=
  Subtype.instValEquivariant 𝔸

@[simp] lemma supp {f : StepModel 𝔸 X} : NominalSet.supp 𝔸 f = NominalSet.supp 𝔸 f.val :=
  Subtype.supp 𝔸

@[simp] lemma fresh {a : 𝔸} {f : StepModel 𝔸 X} : a # f ↔ a # f.val :=
  Subtype.fresh 𝔸

private def bot' : PermSet.FinsuppFunction 𝔸 𝔸 X :=
  ⟨fun _ ↦ ⊥[𝔸], by
    use ∅; constructor
    · intro π hπ; ext a
      simp [AlgebraStructure.bot_equivariant]
    · exact Set.finite_empty⟩

private lemma bot_req1 : Req1 𝔸 (bot' 𝔸 : PermSet.FinsuppFunction 𝔸 𝔸 X) := by
  intro a π _
  apply AlgebraStructure.bot_equivariant

private lemma bot_req2 : Req2 𝔸 (bot' 𝔸 : PermSet.FinsuppFunction 𝔸 𝔸 X) := by
  intro a c _ _
  simp only [PermSet.Finsupp.smul_val, PermSet.smul_function, bot']
  rw [AlgebraStructure.bot_equivariant]
  rw [AlgebraStructure.ax_fresh c (⊥[𝔸]) (AlgebraStructure.bot_fresh 𝔸)]
  rw [AlgebraStructure.ax_idempotent]

def bot : Unit → StepModel 𝔸 X :=
  fun _ ↦ ⟨bot' 𝔸, ⟨bot_req1 𝔸, bot_req2 𝔸⟩⟩

instance : Function.Equivariant 𝔸 (bot 𝔸 : Unit → StepModel 𝔸 X) := by
  constructor; intro _ π
  apply Subtype.eq; ext a
  simp [bot, bot', AlgebraStructure.bot_equivariant]

private def union'' : (StepModel 𝔸 X × StepModel 𝔸 X) → 𝔸 → X :=
  fun (f, g) ↦ fun a ↦ f.val.val a ∨[𝔸] g.val.val a

instance : Function.Equivariant 𝔸 (union'' 𝔸 : (StepModel 𝔸 X × StepModel 𝔸 X) → 𝔸 → X) := by
  constructor; intro (f, g) π
  ext a
  simp only [union'', smul_val, PermSet.Finsupp.smul_val, PermSet.smul_function]
  rw [AlgebraStructure.join_equivariant]

private def union' : (StepModel 𝔸 X × StepModel 𝔸 X) → PermSet.FinsuppFunction 𝔸 𝔸 X :=
  fun p ↦ ⟨union'' 𝔸 p, by
    use NominalSet.supp 𝔸 p; constructor
    · apply Function.Equivariant.preserve_support
      apply NominalSet.supp_is_support
    · apply NominalSet.supp_finite⟩

private lemma union_req1 {f g : StepModel 𝔸 X} : Req1 𝔸 (union' 𝔸 (f, g)) := by
  intro a π hπ
  have hπ₁ : π ∈ PermSet.fix 𝔸 (NominalSet.supp 𝔸 (union' 𝔸 (f, g))) := by
    apply PermSet.fix_antitone_mem 𝔸 hπ; exact Set.subset_union_left
  have hπ₂ : π • a = a := by
    apply hπ; apply Set.mem_union_right; rfl
  show π • (f.val.val a ∨[𝔸] g.val.val a) = (union' 𝔸 (f, g)).val a
  rw [AlgebraStructure.join_equivariant]
  rw [show π • f.val.val a = (π • f).val.val (π • a) by simp]
  rw [show π • g.val.val a = (π • g).val.val (π • a) by simp]
  rw [show ((π • f).val.val (π • a) ∨[𝔸] (π • g).val.val (π • a)) = (union' 𝔸 (π • f, π • g)).val (π • a) from rfl]
  suffices union' 𝔸 (π • f, π • g) = π • union' 𝔸 (f, g) by
    rw [this]
    rw [show π • union' 𝔸 (f, g) = union' 𝔸 (f, g) from NominalSet.supp_is_support 𝔸 hπ₁]
    rw [hπ₂]
  ext b
  simp [union', union'', AlgebraStructure.join_equivariant]

private lemma union_req2 {f g : StepModel 𝔸 X} : Req2 𝔸 (union' 𝔸 (f, g)) := by
  intro a c hcfresh hca
  have hinf : (NominalSet.supp 𝔸 (union' 𝔸 (f, g)) ∪ NominalSet.supp 𝔸 f ∪ NominalSet.supp 𝔸 g ∪ {a})ᶜ.Infinite := by
    apply Set.Finite.infinite_compl
    rw [Set.finite_union, Set.finite_union, Set.finite_union]
    constructor
    · constructor
      · constructor
        · apply NominalSet.supp_finite
        · apply NominalSet.supp_finite
      · apply NominalSet.supp_finite
    · apply Set.finite_singleton
  obtain ⟨d, hd⟩ := hinf.nonempty
  simp only [supp, Set.compl_union, Set.mem_inter_iff, Set.mem_compl_iff] at hd
  have hdfresh : d # union' 𝔸 (f, g) := hd.left.left.left
  have hdf : d # f.val := hd.left.left.right
  have hdg : d # g.val := hd.left.right
  have hda : d ≠ a := hd.right
  simp only [PermSet.Finsupp.smul_val]
  have h := Req1.choose_restriction 𝔸 (union' 𝔸 (f, g)) (union_req1 𝔸) a c d hcfresh hdfresh hca hda
  simp only [PermSet.Finsupp.smul_val] at h
  rw [h]
  simp only [union', union'', PermSet.smul_function, FPerm.swap_inv, PermSet.atom_swap,
    Equiv.swap_apply_left]
  nth_rw 1 [show d = ⸨a, d⸩ • a by simp]
  rw [← AlgebraStructure.res_equivariant]
  rw [AlgebraStructure.ax_distributive]
  rw [AlgebraStructure.join_equivariant]
  rw [AlgebraStructure.res_equivariant]
  rw [AlgebraStructure.res_equivariant]
  rw [AlgebraStructure.ax_associative]
  nth_rw 3 [AlgebraStructure.ax_commutative]
  rw [AlgebraStructure.ax_associative]
  rw [← AlgebraStructure.ax_associative]
  nth_rw 4 [AlgebraStructure.ax_commutative]
  rw [show ⸨a, d⸩ • a = d by simp]
  have hf := f.prop.right a d hdf hda
  have hg := g.prop.right a d hdg hda
  simp only [PermSet.Finsupp.smul_val, PermSet.smul_function, FPerm.swap_inv, PermSet.atom_swap,
    Equiv.swap_apply_left] at hf hg
  rw [← hf, ← hg]

def union : (StepModel 𝔸 X × StepModel 𝔸 X) → StepModel 𝔸 X :=
  fun (f, g) ↦ ⟨union' 𝔸 (f, g), ⟨union_req1 𝔸, union_req2 𝔸⟩⟩

instance : Function.Equivariant 𝔸 (union 𝔸 : (StepModel 𝔸 X × StepModel 𝔸 X) → StepModel 𝔸 X) := by
  constructor; intro x π
  apply Subtype.eq; ext a
  simp only [union, union', Prod.mk.eta, smul_val, PermSet.Finsupp.smul_val]
  rw [show (π • x.fst, π • x.snd) = π • x by ext <;> simp]
  rw [← Function.equivariant]

private noncomputable def chooseFresh (f : StepModel 𝔸 X) (a : 𝔸) : 𝔸 :=
  have hinf : (NominalSet.supp 𝔸 f ∪ {a})ᶜ.Infinite := by
    apply Set.Finite.infinite_compl
    simp only [supp, Set.union_singleton, Set.finite_insert]
    apply NominalSet.supp_finite
  Classical.choose hinf.nonempty

private lemma chooseFresh_fresh (f : StepModel 𝔸 X) (a : 𝔸) :
    chooseFresh 𝔸 f a # f.val ∧ chooseFresh 𝔸 f a ≠ a := by
  have h : chooseFresh 𝔸 f a ∈ (NominalSet.supp 𝔸 f ∪ {a})ᶜ := Classical.choose_spec _
  simp only [supp, Set.union_singleton, Set.mem_compl_iff, Set.mem_insert_iff, not_or] at h
  exact ⟨h.right, h.left⟩

private noncomputable def restrict'' : NameAbstraction 𝔸 (StepModel 𝔸 X) → 𝔸 → X := Quotient.lift
  (fun ⟨a, f⟩ ↦
    let c := chooseFresh 𝔸 f a
    fun b ↦ if a = b then (c \ (⸨a, c⸩ • f.val).val a) else f.val.val b)
  (by
    intro ⟨a₁, f₁⟩ ⟨a₂, f₂⟩ haeq
    ext b
    rw [← NameAbstraction.eq_iff_alpha_eq 𝔸 (StepModel 𝔸 X), NameAbstraction.eq_iff_cases] at haeq
    cases haeq with
    | inl h => rw [h.left, h.right]
    | inr h =>
      simp only [PermSet.Finsupp.smul_val]
      by_cases h₁ : a₁ = b <;> by_cases h₂ : a₂ = b <;> simp only [h₁, h₂, ↓reduceIte]
      · rw [← h₂] at h₁
        have hcontra := h.left
        contradiction
      ·
        sorry
      · rw [h.right.right, ← h₂]
        simp only [smul_val, PermSet.Finsupp.smul_val]
        have heq := Req1.choose_restriction 𝔸 f₂.val f₂.prop.left
        sorry
      · sorry)

end StepModel