import Mathlib
import Nominal.LambdaCalculus.Preterms
import Nominal.Instances.Basic
import Nominal.Instances.Abstraction
import Nominal.Permutations

/-!
# λ-Terms

The set of λ-terms is the quotient of λ-preterms modulo α-equivalence.

We will show that this set also forms a nominal set and that, for any term t, supp t is exactly the set of free variables in t.
-/

namespace Lambda

/-!
## α-Equivalence

We will define the notion of α-equivalence between λ-preterms in the usual way.
-/
section AlphaEquivalence

/--
α-equivalence is the congruence generated by
  λ x, t ≈ λ y, (x y) t if y ∉ fv(t)
-/
inductive AlphaEq : Preterm → Preterm → Prop where
  | refl {t : Preterm} : AlphaEq t t
  | symm {t u : Preterm} : AlphaEq t u → AlphaEq u t
  | trans {t u s : Preterm} : AlphaEq t u → AlphaEq u s → AlphaEq t s
  | cong_abs {x : Var} {t u : Preterm} : AlphaEq t u → AlphaEq (.abs x t) (.abs x u)
  | cong_app {t₁ t₂ u₁ u₂ : Preterm} : AlphaEq t₁ u₁ → AlphaEq t₂ u₂ → AlphaEq (.app t₁ t₂) (.app u₁ u₂)
  | rename {x y : Var} {t : Preterm} : y ∉ t.fv → AlphaEq (.abs x t) (.abs y (⸨x, y⸩ • t))

/-- By definition, AlphaEq is an equivariant equivalence relation. -/
instance Preterm.setoid : EquivariantSetoid Var Preterm where
  r := AlphaEq
  iseqv := ⟨fun x => AlphaEq.refl, AlphaEq.symm, AlphaEq.trans⟩
  equivariant := by
    intro t u π htu
    induction' htu <;> all_goals expose_names
    · exact AlphaEq.refl
    · exact AlphaEq.symm a_ih
    · exact AlphaEq.trans a_ih a_ih_1
    · rw [Preterm.smul_abs]; exact AlphaEq.cong_abs a_ih
    · rw [Preterm.smul_app]; exact AlphaEq.cong_app a_ih a_ih_1
    · simp only [Preterm.smul_abs]
      rw [← mul_smul, FPerm.mul_swap_eq_swap_mul, mul_smul]
      apply AlphaEq.rename
      rw [Preterm.fv_equivariant, Set.smul_mem_smul_set_iff]
      exact h

/-- Define λ-terms as the quotient Preterm/≈. -/
def Term := Quotient Preterm.setoid.toSetoid

namespace Term

/-- The equivalence class of a λ-preterm. -/
def mk : Preterm → Term := Quotient.mk Preterm.setoid.toSetoid

instance : NominalSet Var Term := Quotient.instNominalSet Var

@[simp]
lemma smul {π : FPerm Var} {t : Preterm} : π • Term.mk t = Term.mk (π • t) := by congr

private lemma alpha_eq {t u : Preterm} : t ≈ u ↔ AlphaEq t u := by rfl

lemma eq_iff_alpha_eq {t u : Preterm} : mk t = mk u ↔ t ≈ u := by
  rw [mk, Quotient.eq_iff_equiv]

def var : Var → Term := fun x => mk (.var x)

def abs : Var → Term → Term := fun x =>
  Quotient.lift
    (fun t => mk (.abs x t))
    (by intro t t'; rw [eq_iff_alpha_eq]; exact AlphaEq.cong_abs)

def app : Term → Term → Term :=
  Quotient.lift₂
    (fun t u => mk (.app t u))
    (by intro t₁ t₂ u₁ u₂; rw [eq_iff_alpha_eq]; exact AlphaEq.cong_app)

@[simp]
lemma abs_preterm {x : Var} {t : Preterm} : abs x (mk t) = mk (.abs x t) := rfl

@[simp]
lemma app_preterm {t u : Preterm} : app (mk t) (mk u) = mk (.app t u) := rfl

/-- Alternative characterization of α-equivalence for each preterm constructor. -/
inductive AlphaEq' : Preterm → Preterm → Prop where
  | var {x : Var} : AlphaEq' (.var x) (.var x)
  | abs {x y : Var} {t u : Preterm} : ⟪x⟫ (⟦t⟧ : Term) = ⟪y⟫ (⟦u⟧ : Term) → AlphaEq' (.abs x t) (.abs y u)
  | app {t₁ t₂ u₁ u₂ : Preterm} : (⟦t₁⟧ : Term) = (⟦u₁⟧ : Term) → (⟦t₂⟧ : Term) = (⟦u₂⟧ : Term) → AlphaEq' (.app t₁ t₂) (.app u₁ u₂)

namespace AlphaEq'

lemma refl {t : Preterm} : AlphaEq' t t := by
  cases t
  · exact var
  · exact abs rfl
  · exact app rfl rfl

lemma symm {t u : Preterm} : AlphaEq' t u → AlphaEq' u t := by
  intro h
  cases h <;> all_goals expose_names
  · exact var
  · exact abs h.symm
  · exact app h.symm h_1.symm

lemma trans {t u s : Preterm} : AlphaEq' t u → AlphaEq' u s → AlphaEq' t s := by
  intro h₁ h₂
  cases h₁ <;> all_goals expose_names
  · cases h₂; exact var
  · cases h₂; expose_names; exact abs (h.trans h_1)
  · cases h₂; expose_names; exact app (h.trans h_2) (h_1.trans h_3)

end AlphaEq'

private lemma pick_fresh {t : Preterm} : ∀ x ∉ t.fv, ∃ u, t ≈ u ∧ x ∉ u.vars := by
  suffices ∀ (π : FPerm Var), ∀ x ∉ (π • t).fv, ∃ u, (π • t) ≈ u ∧ x ∉ u.vars by
    specialize this 1
    simp only [one_smul] at this
    exact this
  induction' t <;> all_goals expose_names
  · intro π x h
    use .var (π a)
    exact ⟨by rfl, h⟩
  · intro π x h
    have hinf : ({x} ∪ (π • a_1).vars)ᶜ.Infinite := by
      apply Set.Finite.infinite_compl
      exact Set.Finite.union (Set.finite_singleton x) (Preterm.vars_finite (π • a_1))
    obtain ⟨c, hc⟩ := hinf.nonempty
    rw [Preterm.smul_abs, Preterm.fv, Set.mem_diff, Set.mem_singleton_iff, not_and, not_not] at h
    rw [Set.singleton_union, Set.mem_compl_iff, Set.mem_insert_iff, not_or] at hc
    have hc' : c ∉ (π • a_1).fv := fun h => hc.right (Preterm.fv_subset_vars _ h)
    obtain ⟨u, hu⟩ := a_ih (⸨π • a, c⸩ * π) x (by
      rw [mul_smul, Preterm.fv_equivariant]
      intro hcontra
      by_cases hxa : x = π • a
      · rw [hxa, ← PermSet.atom_set_swap_right] at hcontra
        exact hc' hcontra
      · rw [← PermSet.atom_set_swap_ne_ne _ hxa (Ne.symm hc.left)] at hcontra
        exact hxa (h hcontra))
    use .abs c u
    constructor
    · symm; trans
      · exact AlphaEq.cong_abs hu.left.symm
      · rw [Preterm.smul_abs, mul_smul]; symm
        apply AlphaEq.rename hc'
    · rw [Preterm.vars, Set.singleton_union, Set.mem_insert_iff, not_or]
      exact ⟨Ne.symm hc.left, hu.right⟩
  · intro π x h
    rw [Preterm.smul_app, Preterm.fv, Set.mem_union, not_or] at h
    obtain ⟨t₁, ht₁⟩ := a_ih π x h.left
    obtain ⟨t₂, ht₂⟩ := a_ih_1 π x h.right
    use .app t₁ t₂
    constructor
    · exact AlphaEq.cong_app ht₁.left ht₂.left
    · rw [Preterm.vars, Set.mem_union, not_or]
      exact ⟨ht₁.right, ht₂.right⟩

lemma alpha_eq_iff : AlphaEq = AlphaEq' := by
  ext t u
  constructor
  · intro h; induction' h <;> all_goals expose_names
    · exact AlphaEq'.refl
    · exact AlphaEq'.symm a_ih
    · exact AlphaEq'.trans a_ih a_ih_1
    · apply AlphaEq'.abs
      rw [← alpha_eq, ← Quotient.eq_iff_equiv] at h
      rw [h]
    · apply AlphaEq'.app
      · rw [← alpha_eq, ← Quotient.eq_iff_equiv] at h
        rw [h]
      · rw [← alpha_eq, ← Quotient.eq_iff_equiv] at h_1
        rw [h_1]
    · obtain ⟨z, ⟨hz₁, hz₂⟩⟩ := pick_fresh y h
      apply AlphaEq'.abs
      rw [Quotient.mk_equivariant, Quotient.sound hz₁, NameAbstraction.eq_iff_cases]
      by_cases hxy : x = y
      · left
        simp [hxy, FPerm.swap_self, one_smul]
      · right
        constructor
        · exact hxy
        · constructor
          · rw [← Preterm.supp] at hz₂
            nth_rw 1 [show x = ⸨x, y⸩ • y by simp]
            apply NominalSet.fresh_equivariant
            exact Function.Equivariant.preserve_fresh Var (Quotient.mk_equivariant Var) hz₂
          · rw [← mul_smul, FPerm.swap_mul_self, one_smul]
  · intro h; cases h <;> all_goals expose_names
    · exact AlphaEq.refl
    · rw [NameAbstraction.eq_iff_cases] at h
      cases h with
      | inl h =>
        rw [h.left]
        apply AlphaEq.cong_abs
        rw [← alpha_eq, ← Quotient.eq_iff_equiv]
        exact h.right
      | inr h =>
        obtain ⟨s, ⟨hs₁, hs₂⟩⟩ := Quotient.supp_pick_fresh Var h.right.left
        rw [← Quotient.sound hs₁, ← Quotient.mk_equivariant] at h
        have hrename : Preterm.abs y s ≈ Preterm.abs x (⸨y, x⸩ • s) := by
          apply AlphaEq.rename
          rw [Preterm.supp] at hs₂
          exact fun h => hs₂ (Preterm.fv_subset_vars s h)
        calc
          Preterm.abs x t ≈ Preterm.abs x (⸨x, y⸩ • s) := AlphaEq.cong_abs (eq_iff_alpha_eq.mp h.right.right)
          _               ≈ Preterm.abs y s            := by rw [FPerm.swap_comm]; symm; exact hrename
          _               ≈ Preterm.abs y u            := AlphaEq.cong_abs hs₁
    · apply AlphaEq.cong_app
      · rw [Quotient.eq_iff_equiv] at h
        exact h
      · rw [Quotient.eq_iff_equiv] at h_1
        exact h_1


/-- Alternative characterization of α-equivalence for variable preterms. -/
lemma alpha_eq_var {x y : Var} : Preterm.var x ≈ Preterm.var y ↔ x = y := by
  rw [alpha_eq, alpha_eq_iff]
  constructor
  · intro h; cases h; rfl
  · intro h; rw [h]; constructor

@[inherit_doc alpha_eq_var]
lemma eq_var {x y : Var} : Term.var x = Term.var y ↔ x = y := by
  rw [var, var, mk, Quotient.eq_iff_equiv, alpha_eq_var]

/-- Alternative characterization of α-equivalence for abstraction preterms. -/
lemma alpha_eq_abs {x y : Var} {t u : Preterm} : Preterm.abs x t ≈ Preterm.abs y u ↔ ⟪x⟫ (⟦t⟧ : Term) = ⟪y⟫ (⟦u⟧ : Term) := by
  rw [alpha_eq, alpha_eq_iff]
  constructor
  · intro h; cases h; assumption
  · intro h; constructor; assumption

@[inherit_doc alpha_eq_abs]
lemma eq_abs {x y : Var} {t u : Term} : Term.abs x t = Term.abs y u ↔ ⟪x⟫ t = ⟪y⟫ u := by
  obtain ⟨t', ht'⟩ := t.exists_rep
  obtain ⟨u', hu'⟩ := u.exists_rep
  simp only [abs, mk, ← ht', Quotient.lift_mk, ← hu']
  rw [Quotient.eq_iff_equiv, alpha_eq_abs]

/-- Alternative characterization of α-equivalence for application preterms. -/
lemma alpha_eq_app {t₁ t₂ u₁ u₂ : Preterm} : Preterm.app t₁ t₂ ≈ Preterm.app u₁ u₂ ↔ t₁ ≈ u₁ ∧ t₂ ≈ u₂ := by
  rw [alpha_eq, alpha_eq_iff]
  constructor
  · intro h; cases h; rw [← Quotient.eq_iff_equiv, ← Quotient.eq_iff_equiv]; constructor <;> assumption
  · rw [← Quotient.eq_iff_equiv, ← Quotient.eq_iff_equiv]; intro ⟨h₁, h₂⟩; constructor <;> assumption

@[inherit_doc alpha_eq_app]
lemma eq_app {t₁ t₂ u₁ u₂ : Term} : Term.app t₁ t₂ = Term.app u₁ u₂ ↔ t₁ = u₁ ∧ t₂ = u₂ := by
  obtain ⟨t₁', ht₁'⟩ := t₁.exists_rep
  obtain ⟨t₂', ht₂'⟩ := t₂.exists_rep
  obtain ⟨u₁', hu₁'⟩ := u₁.exists_rep
  obtain ⟨u₂', hu₂'⟩ := u₂.exists_rep
  simp only [app, mk, ← ht₁', ← ht₂', Quotient.lift_mk, ← hu₁', ← hu₂']
  rw [Quotient.eq_iff_equiv, Quotient.eq_iff_equiv, Quotient.eq_iff_equiv, alpha_eq_app]

private def equiv : EquivariantEquiv Var Term (Var ⊕ (NameAbstraction Var Term) ⊕ (Term × Term)) :=
  { toFun t := Quotient.liftOn t
      (fun t => match t with
        | .var x => .inl x
        | .abs x t' => .inr (.inl (⟪x⟫ (⟦t'⟧ : Term)))
        | .app t' u' => .inr (.inr ⟨(⟦t'⟧ : Term), (⟦u'⟧ : Term)⟩))
      (by
        intro t u htu
        rw [alpha_eq, alpha_eq_iff] at htu
        cases htu <;> simp [*])
    invFun
      | .inl x => .var x
      | .inr (.inl x) => Quotient.liftOn x
        (fun x => .abs x.atom x.body)
        (by
          intro x y hxy
          simp only [eq_abs]
          exact Quotient.sound hxy)
      | .inr (.inr x) => .app x.fst x.snd
    left_inv := by
      apply Quotient.ind; intro t
      cases t <;> rfl
    right_inv := by
      intro x
      cases x with
      | inl x => rfl
      | inr x =>
        cases x with
        | inl x =>
          obtain ⟨x', hx'⟩ := x.exists_rep
          obtain ⟨body, hbody⟩ := x'.body.exists_rep
          simp only [← hx', abs, mk, Quotient.lift_mk, ← hbody, Sum.inr.injEq, Sum.inl.injEq]
          rw [hbody]; rfl
        | inr x =>
          obtain ⟨x', hx'⟩ := x.fst.exists_rep
          obtain ⟨y', hy'⟩ := x.snd.exists_rep
          simp only [app, mk, ← hx', ← hy', Quotient.lift_mk, Sum.inr.injEq]
          rw [hx', hy']
    toFun_equivariant := by
      intro t π
      obtain ⟨t', ht'⟩ := t.exists_rep
      rw [← ht']
      cases t' <;> all_goals rfl
  }

@[simp] private lemma equiv_var {x : Var} : equiv.toFun ⟦.var x⟧ = .inl x := rfl
@[simp] private lemma equiv_abs {x : Var} {t : Preterm} : equiv.toFun ⟦.abs x t⟧ = .inr (.inl (⟪x⟫ ⟦t⟧)) := rfl
@[simp] private lemma equiv_app {t u : Preterm} : equiv.toFun ⟦.app t u⟧ = .inr (.inr (⟦t⟧, ⟦u⟧)) := rfl

private lemma supp' (t : Preterm) : NominalSet.supp Var (⟦t⟧ : Term) = t.fv := by
  induction' t; all_goals
    rw [EquivariantEquiv.supp_eq Var equiv]
    simp [Preterm.fv, *]

def fv : Term → Set Var := Quotient.lift Preterm.fv (by
  intro t u htu
  rw [← supp', ← supp', Quotient.sound htu])

lemma supp : NominalSet.supp Var = fv := by
  apply funext
  apply Quotient.ind
  intro t
  rw [supp']; rfl

end Term

end AlphaEquivalence

end Lambda