import Mathlib
import Nominal.Support
import Nominal.Equivariance
import Nominal.Category
import Nominal.Permutations
import Nominal.Instances.Basic

/-!
# Name Abstractions

Given a nominal set X, we define the set `[𝔸]X` of abstractions as pairs in 𝔸 × X modulo alpha-equivalence.
-/

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

/-!
## Raw Structure

We first define the raw type of name abstractions as a pair (𝔸, X).

This structure, along with the obvious permutation action, forms a nominal set where supp (a, x) = supp x ∪ {a}.
-/
section RawStructure

/-- The structure forming a name abstraction. -/
structure NameAbstraction' (𝔸 : Type) [Infinite 𝔸] [Countable 𝔸] [DecidableEq 𝔸] (X : Type) [NominalSet 𝔸 X] where
  /-- The atom that is abstracted away. -/
  atom : 𝔸
  /-- The body from which the atom is abstracted. -/
  body : X

variable (X : Type) [NominalSet 𝔸 X]

attribute [ext] NameAbstraction'

namespace NameAbstraction'

instance : SMul (FPerm 𝔸) (NameAbstraction' 𝔸 X) where
  smul π x := .mk (π • x.atom) (π • x.body)

@[simp] lemma smul_atom (π : FPerm 𝔸) (x : NameAbstraction' 𝔸 X) : (π • x).atom = π • x.atom := rfl

@[simp] lemma smul_body (π : FPerm 𝔸) (x : NameAbstraction' 𝔸 X) : (π • x).body = π • x.body := rfl

instance : PermSet 𝔸 (NameAbstraction' 𝔸 X) where
  one_smul := by intro _; ext <;> simp [one_smul]
  mul_smul := by intro _ _ _; ext <;> simp [mul_smul]

private def ofPair (x : (𝔸 × X)) : NameAbstraction' 𝔸 X := ⟨x.fst, x.snd⟩

private lemma ofPair_injective : Function.Injective (ofPair 𝔸 X) := by
  intro x y hxy
  simp only [ofPair, mk.injEq] at hxy
  ext
  · exact hxy.left
  · exact hxy.right

private lemma ofPair_surjective : Function.Surjective (ofPair 𝔸 X) := by
  intro x
  use ⟨x.atom, x.body⟩
  simp [ofPair]

private lemma ofPair_equivariant : Equivariant 𝔸 (ofPair 𝔸 X) := by
  intro x π
  ext <;> simp [ofPair]

instance : NominalSet 𝔸 (NameAbstraction' 𝔸 X) :=
  NominalSet.fromEquivariantSurjective 𝔸 (ofPair 𝔸 X) (ofPair_equivariant 𝔸 X) (ofPair_surjective 𝔸 X)

@[simp] lemma supp {x : NameAbstraction' 𝔸 X} : NominalSet.supp 𝔸 x = NominalSet.supp 𝔸 x.body ∪ {x.atom} := by
  have h : NominalSet.supp 𝔸 (ofPair 𝔸 X (x.atom, x.body)) = NominalSet.supp 𝔸 (x.atom, x.body) :=
    Equivariant.supp_eq_if_inj 𝔸 (ofPair_equivariant 𝔸 X) (ofPair_injective 𝔸 X)
  simp only [Prod.supp, NominalSet.supp_atom, ofPair] at h
  rwa [Set.union_comm] at h

@[simp] lemma fresh {a : 𝔸} {x : NameAbstraction' 𝔸 X} : a # x ↔ a ≠ x.atom ∧ a # x.body := by
  simp [NominalSet.IsFresh]

end NameAbstraction'

end RawStructure

/-!
## Alpha-Equivalence

We continue to define alpha-equivalence on this structure.
-/
section AlphaEquivalence

namespace NameAbstraction

variable {𝔸 : Type} [Infinite 𝔸] [Countable 𝔸] [DecidableEq 𝔸] {X : Type} [NominalSet 𝔸 X]

/-- (a, x) ≈ (b, y) if (c a) x = (c b) y for some fresh c. -/
def AlphaEq (x y : NameAbstraction' 𝔸 X) : Prop :=
  И c, ⸨c, x.atom⸩ • x.body = ⸨c, y.atom⸩ • y.body

/-- Alternatively define alpha equivalence as an equivariant relation.

This is used to show some properties using the Equivariance.for_fresh_* lemmas about the freshness quantifier. -/
private def AlphaEqPred (c : 𝔸) (xy : NameAbstraction' 𝔸 X × NameAbstraction' 𝔸 X) : Prop :=
  ⸨c, xy.fst.atom⸩ • xy.fst.body = ⸨c, xy.snd.atom⸩ • xy.snd.body

private lemma alpha_eq_iff_pred {x y : NameAbstraction' 𝔸 X} : AlphaEq x y = (И c, AlphaEqPred c (x, y)) := by
  apply NominalSet.for_fresh_ext
  simp [AlphaEqPred]

private lemma alpha_eq_pred_equivariant : EquivariantRelation 𝔸 (AlphaEqPred : 𝔸 → NameAbstraction' 𝔸 X × NameAbstraction' 𝔸 X → Prop) := by
  intro a x π
  rw [AlphaEqPred, AlphaEqPred]
  intro h
  have h' : π • ⸨a, x.fst.atom⸩ • x.fst.body = π • ⸨a, x.snd.atom⸩ • x.snd.body := by congr
  simp only [smul_smul, FPerm.mul_swap_eq_swap_mul] at h'
  simpa [smul_smul]

/-- The freshness quantifier in the definition is equivalent to "there exists some c # (x, y)". -/
lemma alpha_eq_iff_exists {x y : NameAbstraction' 𝔸 X} : AlphaEq x y ↔ ∃ c, c # x ∧ c # y ∧ ⸨c, x.atom⸩ • x.body = ⸨c, y.atom⸩ • y.body := by
  suffices (∃ c, c # x ∧ c # y ∧ ⸨c, x.atom⸩ • x.body = ⸨c, y.atom⸩ • y.body) = (∃ c, c # (x, y) ∧ AlphaEqPred c (x, y)) by
    rw [alpha_eq_iff_pred, this]
    apply Equivariance.for_fresh_iff_exists
    exact alpha_eq_pred_equivariant
  simp [Prod.fresh, AlphaEqPred, @and_assoc]

/-- The freshness quantifier in the definition is equivalent to "for all c # (x, y)". -/
lemma alpha_eq_iff_all {x y : NameAbstraction' 𝔸 X} : AlphaEq x y ↔ ∀ c, c # x ∧ c # y → ⸨c, x.atom⸩ • x.body = ⸨c, y.atom⸩ • y.body := by
  suffices (∀ c, c # x ∧ c # y → ⸨c, x.atom⸩ • x.body = ⸨c, y.atom⸩ • y.body) = (∀ c, c # (x, y) → AlphaEqPred c (x, y)) by
    rw [alpha_eq_iff_pred, this]
    apply Equivariance.for_fresh_iff_all
    exact alpha_eq_pred_equivariant
  simp [Prod.fresh, AlphaEqPred]

/-- An alternative definition of α-equivalence:

For ⟪a⟫ x = ⟪b⟫ y, we either have (a, x) = (b, y) or x is just (a b) y, requiring that a is not in the support of y. -/
lemma alpha_eq_iff_cases {x y : NameAbstraction' 𝔸 X} : AlphaEq x y ↔ (x.atom = y.atom ∧ x.body = y.body) ∨ (x.atom ≠ y.atom ∧ x.atom # y.body ∧ x.body = ⸨x.atom, y.atom⸩ • y.body) := by
  by_cases h_eq : x.atom = y.atom
  · constructor
    · sorry
    · sorry
  · constructor
    · rw [alpha_eq_iff_exists]
      intro h; right
      obtain ⟨c, ⟨hcx, ⟨hcy, hc⟩⟩⟩ := h
      constructor
      · exact h_eq
      · simp at hcx
        simp at hcy
        have h_atom : (⸨c, y.atom⸩ * ⸨c, x.atom⸩) • c # (⸨c, y.atom⸩ * ⸨c, x.atom⸩) • x.body :=
          NominalSet.fresh_equivariant 𝔸 _ _ _ hcx.right
        -- TODO: Refactor into lemma
        have h_swap : ⸨c, x.atom⸩ * ⸨c, y.atom⸩ = ⸨x.atom, y.atom⸩ * ⸨c, x.atom⸩ := by
          ext a
          simp [FPerm.swap]
          by_cases h₁ : (a = c) <;> by_cases h₂ : (a = x.atom) <;> by_cases h₃ : (a = y.atom)
          · simp [h₁, Equiv.swap_apply_of_ne_of_ne (Ne.symm hcy.left) (Ne.symm h_eq)]
          · simp [h₁, Equiv.swap_apply_of_ne_of_ne (Ne.symm hcy.left) (Ne.symm h_eq)]
          · simp [h₁, Equiv.swap_apply_of_ne_of_ne (Ne.symm hcy.left) (Ne.symm h_eq)]
          · simp [h₁, Equiv.swap_apply_of_ne_of_ne (Ne.symm hcy.left) (Ne.symm h_eq)]
          · simp [h₂, Equiv.swap_apply_of_ne_of_ne (Ne.symm hcx.left) h_eq, Equiv.swap_apply_of_ne_of_ne hcx.left hcy.left]
          · simp [h₂, Equiv.swap_apply_of_ne_of_ne (Ne.symm hcx.left) h_eq, Equiv.swap_apply_of_ne_of_ne hcx.left hcy.left]
          · simp [h₃, Equiv.swap_apply_of_ne_of_ne (Ne.symm hcy.left) (Ne.symm h_eq)]
          · simp [Equiv.swap_apply_of_ne_of_ne h₁ h₃, Equiv.swap_apply_of_ne_of_ne h₁ h₂, Equiv.swap_apply_of_ne_of_ne h₂ h₃]
        rw [mul_smul, PermSet.atom_swap, PermSet.atom_swap, Equiv.swap_apply_left, Equiv.swap_apply_of_ne_of_ne (Ne.symm hcx.left) h_eq] at h_atom
        rw [mul_smul, hc, ← mul_smul, FPerm.swap_mul_self, one_smul] at h_atom
        constructor
        · exact h_atom
        · calc
            x.body = ⸨c, x.atom⸩ • ⸨c, x.atom⸩ • x.body      := by rw [← mul_smul, FPerm.swap_mul_self, one_smul]
            _      = (⸨c, x.atom⸩ * ⸨c, y.atom⸩) • y.body    := by rw [mul_smul, hc]
            _      = ⸨x.atom, y.atom⸩ • ⸨c, x.atom⸩ • y.body := by rw [← mul_smul, h_swap]
            _      = ⸨x.atom, y.atom⸩ • y.body               := by rw [NominalSet.fresh_swap_fix 𝔸 hcy.right h_atom]
    · rw [alpha_eq_iff_all]
      intro h c ⟨hcx, hcy⟩
      simp at hcy
      simp only [h_eq, false_and, ne_eq, not_false_eq_true, true_and, false_or] at h
      -- TODO: Use new lemma from before
      have h_swap : ⸨c, x.atom⸩ * ⸨x.atom, y.atom⸩ = ⸨c, y.atom⸩ * ⸨c, x.atom⸩ := sorry
      rw [h.right, ← mul_smul, h_swap, mul_smul, NominalSet.fresh_swap_fix 𝔸 hcy.right h.left]

private lemma alpha_eq_refl (x : NameAbstraction' 𝔸 X) : AlphaEq x x := by sorry

private lemma alpha_eq_symm {x y : NameAbstraction' 𝔸 X} : AlphaEq x y → AlphaEq y x := by sorry

private lemma alpha_eq_trans {x y z : NameAbstraction' 𝔸 X} : AlphaEq x y → AlphaEq y z → AlphaEq x z := by sorry

lemma alpha_eq_equivariant : EquivariantRelation 𝔸 (AlphaEq : NameAbstraction' 𝔸 X → NameAbstraction' 𝔸 X → Prop) := by sorry

end NameAbstraction

end AlphaEquivalence

/-!
## Quotient Type

We then define `NameAbstraction` as the quotient type of this structure modulo alpha-equivalence.
-/
section Quotient

variable (X : Type) [NominalSet 𝔸 X]

/-- Associate the name abstraction structure with alpha-equivalence. -/
instance NameAbstraction.setoid : Setoid (NameAbstraction' 𝔸 X) where
  r := AlphaEq
  iseqv := ⟨alpha_eq_refl, alpha_eq_symm, alpha_eq_trans⟩

/-- Define name abstractions as the quotient set (𝔸 × X)/≈. -/
def NameAbstraction := Quotient (NameAbstraction.setoid 𝔸 X)

/-- The equivalence class of a name abstraction (a, x), often written `⟪a⟫ x` -/
def NameAbstraction.mk {𝔸 : Type} [Infinite 𝔸] [Countable 𝔸] [DecidableEq 𝔸] {X : Type} [NominalSet 𝔸 X] (a : 𝔸) (x : X) : NameAbstraction 𝔸 X :=
  Quotient.mk (setoid 𝔸 X) (.mk a x)

/-- We write `⟪a⟫ x` instead of the usual `⟨a⟩ x` because `⟨⟩` is already heavily used in Lean for constructors. -/
notation "⟪" a "⟫ " t => NameAbstraction.mk a t

end Quotient

/-!
## Nominal Set Instance

The equivalence classes of name abstractions form a nominal set where supp (⟪a⟫ x) = supp x \ {a}.
-/
section NominalSet

end NominalSet

/-!
## Abstraction Functor

Name abstraction can also be considered as a functor `Nom → Nom`.
-/
section Functor

end Functor