import Mathlib
import Nominal.Permutations
import Nominal.Support
import Nominal.Equivariance
import Nominal.Instances.Basic
import Nominal.Helpers.CountableString

/-!
# λ-Preterms

As a practical example for nominal sets, this file defines λ-preterms (that is, λ-terms without α-equivalence).

For any preterm t, supp t is exactly the set of variables occurring in t (either bound or free).
-/

namespace Lambda

/-!
## Basic Syntax Definition

Define the AST and a parser for it.
-/
section Syntax

/-- We will use strings as the set of atoms, representing variable names. -/
def Var := String
  deriving DecidableEq, Countable, Infinite

/-- Basic λ-preterm definition.

We furthermore allow defining λ-terms as ⟫ t ⟪ with
  t, u ::= x | λ x y ..., t | t u,
where x, y, ... are variables. -/
inductive Preterm
  | var : Var → Preterm
  | abs : Var → Preterm → Preterm
  | app : Preterm → Preterm → Preterm
  deriving DecidableEq

/-- The set of variables mentioned by a preterm. -/
def Preterm.vars : Preterm → Set Var
  | .var x => {x}
  | .abs x t => {x} ∪ vars t
  | .app t u => vars t ∪ vars u

/-- The set of free variables in a preterm. -/
def Preterm.fv : Preterm → Set Var
  | .var x => {x}
  | .abs x t => fv t \ {x}
  | .app t u => fv t ∪ fv u

open Lean Elab Meta

declare_syntax_cat lam

syntax ident : lam
syntax "(" lam ")" : lam
syntax:10 (name := stxabs) "λ " ident+ ", " lam : lam
syntax:45 lam:45 lam:46 : lam

partial def elabLam : Syntax → MetaM Expr
  | `(lam | $x:ident) => mkAppM ``Preterm.var #[mkStrLit x.getId.toString]
  | `(lam | ($t:lam)) => elabLam t
  | `(lam | $t:lam $u:lam) => do mkAppM ``Preterm.app #[← elabLam t, ← elabLam u]
  | `(lam | λ $xs:ident*, $t:lam) => do
    xs.foldrM (fun x t' => mkAppM ``Preterm.abs #[mkStrLit x.getId.toString, t']) (← elabLam t)
  | _ => throwUnsupportedSyntax

elab "⟫ " t:lam " ⟪" : term => elabLam t

example : ⟫ x y z ⟪ = ⟫ (x y) z ⟪ := by decide
example : ⟫ λ x, λ y, x ⟪ = ⟫ λ x y, x ⟪ := by decide

end Syntax

/-!
## Nominal Set Instance

The set of preterms forms a nominal set, where the permutation action just renames all occurrences of a variable.
-/
section NominalSet

namespace Preterm

private def rename (π : FPerm Var) : Preterm → Preterm
  | .var x => .var (π • x)
  | .abs x t => .abs (π • x) (rename π t)
  | .app t u => .app (rename π t) (rename π u)

private lemma one_rename (t : Preterm) : rename 1 t = t := by
  induction' t <;> all_goals simp [rename, *]

private lemma mul_rename (π τ : FPerm Var) (t : Preterm) : rename (π * τ) t = rename π (rename τ t) := by
  induction' t <;> all_goals simp [rename, mul_smul, *]

instance : PermSet Var Preterm where
  smul := rename
  one_smul := one_rename
  mul_smul := mul_rename

@[simp] lemma smul_var {π : FPerm Var} {x : Var} : π • Preterm.var x = Preterm.var (π • x) := rfl
@[simp] lemma smul_abs {π : FPerm Var} {x : Var} {t : Preterm} : π • Preterm.abs x t = Preterm.abs (π • x) (π • t) := rfl
@[simp] lemma smul_app {π : FPerm Var} {t u : Preterm} : π • Preterm.app t u = Preterm.app (π • t) (π • u) := rfl

instance : NominalSet Var Preterm where
  fin_supp t := by
    use vars t
    induction' t <;> all_goals expose_names
    · constructor
      · simp [PermSet.HasSupport, PermSet.fix, vars, PermSet.fix']
      · simp [vars]
    · constructor
      · intro π hπ
        simp only [PermSet.fix', smul_abs, Preterm.abs.injEq, Set.mem_setOf_eq]
        simp only [PermSet.fix, vars, Set.singleton_union, Set.mem_insert_iff, forall_eq_or_imp, Set.mem_setOf_eq] at hπ
        exact ⟨hπ.left, a_ih.left hπ.right⟩
      · simp [vars, a_ih]
    · constructor
      · intro π hπ
        simp only [PermSet.fix', smul_app, Preterm.app.injEq, Set.mem_setOf_eq]
        simp only [PermSet.fix, vars, Set.mem_union, Set.mem_setOf_eq] at hπ
        constructor
        · exact a_ih.left (fun x hx => hπ x (Or.inl hx))
        · exact a_ih_1.left (fun x hx => hπ x (Or.inr hx))
      · simp [vars, a_ih, a_ih_1]

/-- It's easiest to prove the least support via an equivalence between preterms and a well-known nominal type. -/
private def equiv : EquivariantEquiv Var Preterm (Var ⊕ (Var × Preterm) ⊕ (Preterm × Preterm)) :=
  { toFun
      | .var x => .inl x
      | .abs x t => .inr (.inl ⟨x, t⟩)
      | .app t u => .inr (.inr ⟨t, u⟩)
    invFun
      | .inl x => .var x
      | .inr (.inl x) => .abs x.fst x.snd
      | .inr (.inr x) => .app x.fst x.snd
    left_inv := by intro x; cases x <;> all_goals simp
    right_inv := by
      intro x
      cases x with
      | inl _ => simp
      | inr x' => cases x' <;> all_goals simp
    toFun_equivariant := by
      constructor; intro t π
      cases t <;> all_goals simp
  }

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

/-- The support of a preterm consists exactly of the variables mentioned by it. -/
lemma supp {t : Preterm} : NominalSet.supp Var t = vars t := by
  fun_induction vars
  all_goals rw [EquivariantEquiv.supp_eq Var equiv]; simp [*]

end Preterm

end NominalSet

/-!
## Variables

Helpers for reasoning about variables in a preterm.
-/
section Variables

namespace Preterm

/-- The free variables are a subset of all variables in a preterm. -/
lemma fv_subset_vars (t : Preterm) : t.fv ⊆ t.vars := by
  induction' t <;> all_goals grind [Preterm.vars, Preterm.fv]

/-- The set of variables in a preterm is finite. -/
lemma vars_finite (t : Preterm) : t.vars.Finite := by
  rw [← Preterm.supp]
  exact NominalSet.supp_finite Var

/-- The set of free variables in a preterm is finite. -/
lemma fv_finite (t : Preterm) : t.fv.Finite := by
  apply Set.Finite.subset (vars_finite t) (fv_subset_vars t)

/-- The vars map is equivariant. -/
instance instVarsEquivariant : Function.Equivariant Var vars := by
  constructor; intro t π
  fun_induction vars <;> all_goals simp [vars, *, Set.smul_set_insert, Set.smul_set_union]

/-- The fv map is equivariant. -/
instance instFvEquivariant : Function.Equivariant Var fv := by
  constructor; intro t π
  fun_induction fv <;> all_goals simp [fv, *, Set.smul_set_sdiff, Set.smul_set_union]

end Preterm

end Variables

end Lambda