import Mathlib
import Nominal.Permutations
import Nominal.Helpers.Finset

/-!
# Support

This file provides basic definitions for support of (FPerm 𝔸)-sets and nominal sets.
-/

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

/-!
## Support Definitions

We can define the notion of "supports" on elements of (FPerm 𝔸)-Sets as the atoms which need to be left unchanged to also fix the element.
-/
section Support

variable {X : Type} [PermSet 𝔸 X]

namespace PermSet

/-- The permutations that leave an element x unchanged. -/
def fix' (x : X) : Set (FPerm 𝔸) :=
  { π : FPerm 𝔸 | π • x = x }

/-- The permutations that leave all elements in A unchanged. -/
def fix (A : Set X) : Set (FPerm 𝔸) :=
  { π : FPerm 𝔸 | ∀ x ∈ A, π • x = x }

/-- A set S ⊆ 𝔸 is a support for x if every permutation that fixes S also leaves x unchanged. -/
def HasSupport (x : X) (S : Set 𝔸) : Prop :=
  fix 𝔸 S ⊆ fix' 𝔸 x

/-- fix is antitone: If S ⊆ T, then fix S ⊇ fix T. -/
lemma fix_antitone {X : Type} [PermSet 𝔸 X] {S T : Set X} : S ⊆ T → PermSet.fix 𝔸 T ⊆ PermSet.fix 𝔸 S := by
  intro hsub π hπ x hx
  exact hπ x (hsub hx)

/-- Restatement of `fix_antitone` that is sometimes more convenient for use with apply. -/
lemma fix_antitone_mem {X : Type} [PermSet 𝔸 X] {S T : Set X} {π : FPerm 𝔸} :
    π ∈ PermSet.fix 𝔸 T → S ⊆ T → π ∈ PermSet.fix 𝔸 S := by
  intro hπ hsub
  exact fix_antitone 𝔸 hsub hπ

/-- If π fixes S ∪ T, it fixes both S and T. -/
lemma fix_union {X : Type} [PermSet 𝔸 X] {S T : Set X} : PermSet.fix 𝔸 (S ∪ T) ⊆ PermSet.fix 𝔸 S ∩ PermSet.fix 𝔸 T := by
  apply Set.subset_inter <;> all_goals simp [fix_antitone]

/-- Alternative definition with fix 𝔸 S unfolded. -/
lemma support_iff {x : X} {S : Set 𝔸} : HasSupport 𝔸 x S ↔ ∀ (π : FPerm 𝔸), (∀ x ∈ S, π x = x) → π ∈ fix' 𝔸 x := by
  simp [HasSupport, fix, fix']

/-- If a support is a subset of another set, that set is also a support.  -/
lemma support_superset {x : X} {S T : Set 𝔸} : HasSupport 𝔸 x S → S ⊆ T → HasSupport 𝔸 x T := by
  intro hS hsub π hπ
  apply hS
  intro a ha
  exact hπ a (hsub ha)

/-- For two finite supports, the intersection is also a support. -/
lemma support_inter {x : X} {S T : Set 𝔸} : HasSupport 𝔸 x S → S.Finite → HasSupport 𝔸 x T → T.Finite → HasSupport 𝔸 x (S ∩ T) := by
  intro h_S_support h_S_finite h_T_support h_T_finite π h_π_fix
  let π' := π.val
  let h_π_fin := π.prop

  -- First we obtain an injection m : S \ T → (π'.supp ∪ S ∪ T)ᶜ.
  have h_cod_inf : ((π'.supp ∪ S ∪ T)ᶜ).Infinite := by
    apply Set.Finite.infinite_compl
    simp only [Set.finite_union]
    exact ⟨⟨h_π_fin, h_S_finite⟩, h_T_finite⟩
  have h_dom_fin : (S \ T).Finite :=
    Set.Finite.diff h_S_finite

  obtain ⟨m, ⟨h_m_preimage, h_m_injective⟩⟩ :=
    Set.Finite.exists_injOn_of_encard_le h_dom_fin (by
      rw [Set.Infinite.encard_eq h_cod_inf, Set.Finite.encard_eq_coe h_dom_fin]
      simp)
  let m' := (S \ T).restrict m
  rw [Set.injOn_iff_injective, show (S \ T).restrict m = m' by rfl] at h_m_injective
  obtain h_m_inv := h_m_injective.isPartialInv
  have h_m_image : m '' (S \ T) ⊆ (π'.supp ∪ S ∪ T)ᶜ :=
    Set.image_subset_iff.mpr h_m_preimage

  -- Show some trivial statements about subsets of (π'.supp ∪ S ∪ T), i.e. the atoms excluded from the image of m.
  have h_sub_π   : π'.supp ⊆ (π'.supp ∪ S ∪ T) := Set.subset_union_of_subset_left Set.subset_union_left T
  have h_sub_S   : S       ⊆ (π'.supp ∪ S ∪ T) := Set.subset_union_of_subset_left Set.subset_union_right T
  have h_sub_T   : T       ⊆ (π'.supp ∪ S ∪ T) := Set.subset_union_right
  have h_sub_dom : (S \ T) ⊆ (π'.supp ∪ S ∪ T) := by trans; exact Set.diff_subset; exact h_sub_S
  have h_m_image_notIn_sub {U : Set 𝔸} (h_sub_U : U ⊆ π'.supp ∪ S ∪ T) {a : 𝔸} (h_a_dom : a ∈ S \ T) : m a ∉ U := by
    suffices m a ∈ (π'.supp ∪ S ∪ T)ᶜ by
      exact Set.notMem_subset h_sub_U (h_m_preimage h_a_dom)
    exact h_m_image (Set.mem_image_of_mem m h_a_dom)

  -- Next, extend m to a new permutation σ, about which we will then prove some properties.
  have {a : 𝔸} : Decidable (a ∈ S \ T) := Classical.propDecidable (a ∈ S \ T)
  let σ' : 𝔸 → 𝔸 := fun a =>
    if h : a ∈ S \ T then m' ⟨a, h⟩
    else if let .some b := m'.partialInv a then b
    else a

  -- First, σ is self-inverse and finitely supported, making it an element of FPerm 𝔸.
  have h_σ_involutive : Function.Involutive σ' := by
    intro a
    by_cases h₁ : a ∈ S \ T
    · have h := Function.partialInv_left h_m_injective ⟨a, h₁⟩
      simp [σ', m', h₁, h_m_image_notIn_sub h_sub_dom h₁, h] at *
    · cases h₂ : m'.partialInv a with
      | some b => simp [σ', h₂, h₁]; rwa [← h_m_inv]
      | none   => simp [σ', h₂, h₁]
  have h_σ_simple_case (a : 𝔸) (h₁ : a ∉ S \ T) (h₂ : a ∉ m '' (S \ T)) : σ' a = a := by
    cases h : m'.partialInv a with
    | some b => rw [h_m_inv] at h
                have : a ∈ (m '' (S \ T)) := by rw [← h, Set.mem_image]; use b; simp [m']
                contradiction
    | none   => simp [σ', h₁, h]
  have h_σ_fin : {a | σ' a ≠ a}.Finite := by
    apply Set.Finite.subset
    · exact Set.finite_union.mpr ⟨h_dom_fin, Set.Finite.image m h_dom_fin⟩
    · intro a ha₁; by_contra ha₂
      rw [← Set.mem_compl_iff, Set.compl_union, Set.mem_inter_iff] at ha₂
      suffices σ' a = a by contradiction
      apply h_σ_simple_case
      · rw [← Set.mem_compl_iff]; exact ha₂.left
      · exact Set.notMem_of_mem_diff ha₂
  let σ : FPerm 𝔸 := ⟨⟨σ', σ', h_σ_involutive, h_σ_involutive⟩, h_σ_fin⟩

  -- Second, σ fixes T and σπσ fixes S.
  have h_σ_fixT : σ ∈ fix 𝔸 T := by
    intro a haT; show σ a = a
    apply h_σ_simple_case
    · exact Set.notMem_diff_of_mem haT
    · exact fun h => h_m_image h (h_sub_T haT)
  have h_σ_fixS : (σ * π * σ) ∈ fix 𝔸 S := by
    intro a haS; show σ (π (σ a)) = a
    by_cases haT : a ∈ T
    · calc
        σ (π (σ a)) = σ (π (σ • a)) := rfl
        _           = σ (π a)       := by rw [h_σ_fixT a haT]
        _           = σ (π • a)     := rfl
        _           = σ a           := by rw [h_π_fix a (Set.mem_inter haS haT)]
        _           = σ • a         := rfl
        _           = a             := by rw [h_σ_fixT a haT]
    · suffices π (σ a) = σ a by rw [this]; exact h_σ_involutive a
      apply Equiv.Perm.supp_notin
      rw [show σ a = m a by simp [σ, σ', (Set.mem_diff_of_mem haS haT), m']]
      exact h_m_image_notIn_sub h_sub_π (Set.mem_diff_of_mem haS haT)

  -- Using this carefully constructed σ, we can finally prove the statement.
  calc
    π • x = (1 * π * 1) • x                 := rfl
    _     = ((σ * σ⁻¹) * π * (σ * σ⁻¹)) • x := by rw [@mul_inv_cancel]
    _     = ((σ * σ * π * σ) * σ) • x       := rfl
    _     = (σ * σ * π * σ) • x             := by rw [@mul_smul, h_T_support h_σ_fixT]
    _     = (σ * (σ * π * σ)) • x           := rfl
    _     = σ • x                           := by rw [@mul_smul, h_S_support h_σ_fixS]
    _     = x                               := by rw [h_T_support h_σ_fixT]

/-- For a finite number of finite supports, the intersection is also a support. -/
lemma support_sInter {x : X} {C : Set (Set 𝔸)}: C.Finite → C.Nonempty → (∀ S ∈ C, HasSupport 𝔸 x S ∧ S.Finite) → HasSupport 𝔸 x (⋂₀ C) := by
  suffices (C : Finset (Set 𝔸)) : C.Nonempty → (∀ S ∈ C, HasSupport 𝔸 x S ∧ S.Finite) → HasSupport 𝔸 x (⋂₀ C) ∧ (⋂₀ (↑C : Set (Set 𝔸))).Finite
  · intro hfin
    rw [show C = hfin.toFinset by simp]
    intro hne h
    exact (this hfin.toFinset hne h).left
  · have : DecidableEq (Set 𝔸) := Classical.typeDecidableEq (Set 𝔸)
    apply Finset.Nonempty.induction
    · intro S hS
      rw [Finset.coe_singleton S, Set.sInter_singleton S]
      apply hS S (Finset.mem_singleton.mpr rfl)
    · intro S C' hS hC' hsing
      rw [Finset.coe_insert S C', Set.sInter_insert S ↑C']
      specialize hC' (fun S' hS' => hsing S' (Finset.mem_insert_of_mem hS'))
      specialize hsing S (Finset.mem_insert_self S C')
      constructor
      · exact support_inter 𝔸 hsing.left hsing.right hC'.left hC'.right
      · apply Set.Finite.inter_of_left hsing.right

/-- If two permutations π and σ are equal for all atoms in a support of x, then πx = σx. -/
lemma support_permutations_eq {π σ : FPerm 𝔸} {x : X} {S : Set 𝔸} : PermSet.HasSupport 𝔸 x S → (∀ a ∈ S, π a = σ a) → π • x = σ • x := by
  intro hsupp heq
  have hfix : σ⁻¹ * π  ∈ PermSet.fix 𝔸 S := by
    intro a ha
    calc
      (σ⁻¹ * π) • a = σ⁻¹ (π a)        := rfl
      _             = σ⁻¹ (σ a)        := by rw [heq a ha]
      _             = (σ⁻¹ * σ) a      := rfl
      _             = (1 : FPerm 𝔸) a  := by rw [inv_mul_cancel]
      _             = a                := rfl
  calc
    π • x = (1 * π) • x         := rfl
    _     = (σ * σ⁻¹ * π) • x   := by rw [mul_inv_cancel]
    _     = (σ * (σ⁻¹ * π)) • x := rfl
    _     = σ • ((σ⁻¹ * π) • x) := by rw [smul_smul]
    _     = σ • x               := by rw [hsupp hfix]

end PermSet

end Support

/-!
## Nominal Sets

Nominal sets are then just (FPerm 𝔸)-Sets where all elements are finitely supported.
-/
section NominalSet

/-- An (FPerm 𝔸)-Set X is nominal if every x ∈ X is finitely supported. -/
class NominalSet (X : Type) extends PermSet 𝔸 X where
  fin_supp : ∀ (x : X), ∃ (S : Set 𝔸), PermSet.HasSupport 𝔸 x S ∧ S.Finite

/-- A discrete nominal set where permutations do not change anything. -/
class NominalSet.Discrete (X : Type) extends NominalSet 𝔸 X where
  is_discrete : ∀ (x : X) (π : FPerm 𝔸), π • x = x

namespace NominalSet

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

/-- For every element x of a nominal set, there exists a least finite support, denoted "supp x". --/
def supp (x : X) : Set 𝔸 :=
  ⋂₀ { S : Set 𝔸 | PermSet.HasSupport 𝔸 x S ∧ S.Finite }

/-- The least finite support in a nominal set is always finite. -/
lemma supp_finite {x : X} : (supp 𝔸 x).Finite := by
  unfold supp
  obtain ⟨S, ⟨hsuppS, hfinS⟩⟩ := nom.fin_supp x
  apply Set.Finite.subset
  · exact hfinS
  · apply Set.sInter_subset_of_mem
    exact Set.mem_sep hsuppS hfinS

/-- We can restrict the supp definition to subsets of a single finite support. -/
private lemma supp_equiv_finite {x : X} {S : Set 𝔸} : PermSet.HasSupport 𝔸 x S → S.Finite → supp 𝔸 x = ⋂₀ { T : Set 𝔸 | T ⊆ S ∧ PermSet.HasSupport 𝔸 x T ∧ T.Finite } := by
  intro hsuppS hfinS
  unfold supp; ext a
  constructor
  · intro h T hT
    exact h T hT.out.right
  · intro h T hT
    suffices a ∈ S ∩ T by apply Set.mem_of_mem_inter_right; assumption
    apply h
    constructor
    · exact Set.inter_subset_left
    · constructor
      · exact PermSet.support_inter 𝔸 hsuppS hfinS hT.out.left hT.out.right
      · exact Set.Finite.inter_of_left hfinS T

/-- The least finite support is itself a support. -/
lemma supp_is_support {x : X} : PermSet.HasSupport 𝔸 x (supp 𝔸 x) := by
  obtain ⟨S, ⟨hsuppS, hfinS⟩⟩ := nom.fin_supp x
  rw [supp_equiv_finite 𝔸 hsuppS hfinS]
  apply PermSet.support_sInter
  · apply Set.Finite.subset
    · exact Set.Finite.powerset hfinS
    · intro T hT
      exact hT.out.left
  · use S
    exact ⟨by rfl, ⟨hsuppS, hfinS⟩⟩
  · intro T hT
    exact hT.right

/-- If S is a finite support for x, then supp x is a subset of it. -/
lemma support_subset_supp {x : X} {S : Set 𝔸} : PermSet.HasSupport 𝔸 x S → S.Finite → supp 𝔸 x ⊆ S := by
  intro hsupp hfin
  apply Set.sInter_subset_of_mem
  exact ⟨hsupp, hfin⟩

end NominalSet

end NominalSet

section Freshness

/-- A predicate quantifying a statement for "suitably fresh" atoms.

We also write И a, P for for_fresh (fun a => P), using the cyrillic letter И.

A proposition holds for "suitably fresh" atoms iff the set of atoms that are not "suitably fresh" is finite. -/
def NominalSet.ForFresh {𝔸 : Type} [Infinite 𝔸] [Countable 𝔸] [DecidableEq 𝔸] (P : 𝔸 → Prop) : Prop :=
  {a : 𝔸 | P a}ᶜ.Finite

macro "И" x:ident "," t:term : term =>
  `(NominalSet.ForFresh (fun $x => $t))

/-- We write a # x for a ∉ supp x.

Note: Other conventions extend this meaning to x # y ↔ supp x ∩ supp y = ∅. We will restrict ourselves to x ∈ 𝔸. -/
def NominalSet.IsFresh {𝔸 : Type} [Infinite 𝔸] [Countable 𝔸] [DecidableEq 𝔸] {X : Type} [NominalSet 𝔸 X] (a : 𝔸) (x : X) : Prop :=
  a ∉ NominalSet.supp 𝔸 x

infix:50 " # " => NominalSet.IsFresh

namespace NominalSet

variable {X : Type} [NominalSet 𝔸 X]

/-- To show that an atom is fresh, it is enough to show that it doesn't occur in any finite support. -/
lemma fresh_support {a : 𝔸} {x : X} {S : Set 𝔸} : PermSet.HasSupport 𝔸 x S → S.Finite → a ∉ S → a # x := by
  intro hsuppS hfinS ha
  apply Set.notMem_subset
  · exact support_subset_supp 𝔸 hsuppS hfinS
  · exact ha

/-- One can always pick a fresh atom. -/
lemma fresh_exists {x : X} : ∃ (a : 𝔸), a # x := by
  unfold IsFresh; expose_names
  have h_inf : Set.univ.Infinite := Set.infinite_univ_iff.mpr inst
  have h_fin : (supp 𝔸 x).Finite := supp_finite 𝔸
  obtain ⟨a, ⟨_, ha⟩⟩ := Set.Infinite.exists_notMem_finite h_inf h_fin
  use a

/-- If a and b are fresh for x, then (a b) fixes x. -/
lemma fresh_swap_fix {a b : 𝔸} {x : X} : a # x → b # x → ⸨a, b⸩ ∈ PermSet.fix' 𝔸 x := by
  intro ha hb
  suffices ⸨a, b⸩ • x = (1 : FPerm 𝔸) • x by
    simp [one_smul, PermSet.fix'] at *; assumption
  apply PermSet.support_permutations_eq
  · exact supp_is_support 𝔸
  · intro c hc
    apply FPerm.swap_apply_of_ne_of_ne (ne_of_mem_of_not_mem hc ha) (ne_of_mem_of_not_mem hc hb)

/-- If a proposition holds for suitably fresh a, then one can always pick such an a. -/
lemma for_fresh_exists {P : 𝔸 → Prop} : (И a, P a) → (∃ a, P a) := by
  simp only [ForFresh]; intro h
  apply Set.Finite.infinite_compl at h
  rw [compl_compl] at h
  exact h.nonempty

/-- Extensionality lemma for the freshness quantifier. -/
lemma for_fresh_ext {P Q : 𝔸 → Prop} : (∀ a, P a ↔ Q a) → (И a, P a) = (И a, Q a) := by
  intro h
  simp only [ForFresh, eq_iff_iff]
  suffices {a | P a} = {a | Q a} by rw [this]
  exact Set.ext h

/-- The freshness quantifier preserves implications. -/
lemma for_fresh_imp {P Q : 𝔸 → Prop} : (И a, P a) → (∀ a, P a → Q a) → (И a, Q a) := by
  intro hfin himp
  apply Set.Finite.subset
  · exact hfin
  · simp only [Set.compl_subset_compl, Set.setOf_subset_setOf]
    exact himp

/-- The freshness quantifier commutes with conjunctions. -/
lemma for_fresh_and {P Q : 𝔸 → Prop} : (И a, P a) ∧ (И a, Q a) ↔ (И a, P a ∧ Q a) := by
  simp [ForFresh, Set.setOf_and, Set.compl_inter]

end NominalSet

end Freshness