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
  induction' t <;> all_goals expose_names
  · intro x h
    use .var a
    exact ⟨by rfl, h⟩
  · rw [Preterm.fv]
    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
    simp only [Set.singleton_union, Set.mem_compl_iff, Set.mem_insert_iff, not_or] at hc
    obtain ⟨u, hu⟩ := a_ih c (fun h => hc.right (Preterm.fv_subset_vars a_1 h))
    use Preterm.abs c (⸨a, c⸩ • u)
    simp [Preterm.vars]
    constructor
    · trans
      · exact AlphaEq.cong_abs hu.left
      · exact AlphaEq.rename (fun h => hu.right (Preterm.fv_subset_vars u h))
    · sorry
  · sorry

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
  sorry

@[inherit_doc alpha_eq_var]
lemma eq_var {x y : Var} : Term.var x = Term.var y ↔ x = y := by
  sorry

/-- 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⟫ Term.mk t = ⟪y⟫ Term.mk u := by
  sorry

@[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
  sorry

/-- 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
  sorry

@[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
  sorry

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⟫ Term.mk t'))
        | .app t' u' => .inr (.inr ⟨(Term.mk t'), (Term.mk u')⟩))
      (by
        intro t u htu
        cases htu
        · sorry
        · sorry
        · sorry
        · sorry
        · sorry
        · sorry)
    invFun
      | .inl x => .var x
      | .inr (.inl x) => Quotient.liftOn x
        (fun x => .abs x.atom x.body)
        (by sorry)
      | .inr (.inr x) => .app x.fst x.snd
    left_inv := by sorry
    right_inv := by sorry
    toFun_equivariant := by sorry
  }

@[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

def fv : Term → Set Var := Quotient.lift Preterm.fv
  (by sorry)

lemma supp : NominalSet.supp Var = fv := by
  apply funext
  apply Quotient.ind
  intro t
  induction' t; all_goals
    expose_names
    rw [EquivariantEquiv.supp_eq Var equiv]
    simp [fv, Preterm.fv, *]

end Term

end AlphaEquivalence

end Lambda