path: root/Nominal/Equivariance.lean

import Mathlib
import Nominal.Permutations
import Nominal.Support

/-!
# Definitions and Properties for Equivariance

We define equivariance for functions, sets, and binary relations.
-/

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

/-!
## Definitions

Basic definitions for equivariance.
-/
section Definition

variable {X Y : Type} [PermSet 𝔸 X] [PermSet 𝔸 Y]

/-- A function is equivariant if it is invariant to applying a permutation before or after the function. -/
def Function.Equivariant (f : X → Y) : Prop :=
  ∀ (x : X) (π : FPerm 𝔸), f (π • x) = π • f x

/-- A set is equivariant if it is closed under permutations. -/
def Set.Equivariant (S : Set X) : Prop :=
  ∀ (x : X) (π : FPerm 𝔸), x ∈ S → π • x ∈ S

/-- A binary relation is equivariant if it is closed under permutations. -/
def Rel.Equivariant (R : Rel X Y) : Prop :=
  ∀ (x : X) (y : Y) (π : FPerm 𝔸), R x y → R (π • x) (π • y)

/-- For equivariant sets, we even have "if and only if". -/
lemma Set.Equivariant.iff {S : Set X} : Set.Equivariant 𝔸 S ↔ ∀ (x : X) (π : FPerm 𝔸), x ∈ S ↔ π • x ∈ S := by
  constructor
  · intro h x π; constructor
    · intro hx; exact h x π hx
    · intro hx
      apply (h (π • x) π⁻¹) at hx
      simp only [inv_smul_smul] at hx
      exact hx
  · intro h x π hx
    exact (h x π).mp hx

/-- For equivariant relations, we even have "if and only if". -/
lemma Rel.Equivariant.iff {R : Rel X Y} : Rel.Equivariant 𝔸 R ↔ ∀ (x : X) (y : Y) (π : FPerm 𝔸), R x y ↔ R (π • x) (π • y) := by
  constructor
  · intro h x y π; constructor
    · intro hxy; exact h x y π hxy
    · intro hxy
      apply (h (π • x) (π • y) π⁻¹) at hxy
      simp only [inv_smul_smul] at hxy
      exact hxy
  · intro h x y π hxy
    exact (h x y π).mp hxy

end Definition

/-!
## Trivial Instances

Some trivial instances of equivariant functions, sets, and relations.
-/
section Trivial

variable {X Y Z : Type} [PermSet 𝔸 X] [PermSet 𝔸 Y] [PermSet 𝔸 Z]

/-- The identity function is equivariant. -/
lemma id_equivariant : Function.Equivariant 𝔸 (id : X → X) := by
  simp [Function.Equivariant]

/-- The composition of two equivariant functions is equivariant. -/
lemma comp_equivariant {f : Y → Z} {g : X → Y} : Function.Equivariant 𝔸 f → Function.Equivariant 𝔸 g → Function.Equivariant 𝔸 (f ∘ g) := by
  intro hf hg x π
  calc
    f (g (π • x)) = f (π • g x)   := by rw [hg]
    _             = π • (f (g x)) := by rw [hf]

/-- If a function is equivariant, its inverse is also equivariant. -/
lemma Function.inv_equivariant {f : X → Y} {inv : Y → X} : Function.LeftInverse inv f → Function.RightInverse inv f → Function.Equivariant 𝔸 f → Function.Equivariant 𝔸 inv := by
  intro hleft hright hequi y π
  apply hleft.injective
  rw [hequi, hright, hright]

/-- The subset relation is equivariant. -/
lemma Set.subset_equivariant : Rel.Equivariant 𝔸 (Set.Subset : Set X → Set X → Prop) := by
  intro S T π hST x hx
  rw [Set.mem_smul_set_iff_inv_smul_mem] at hx
  rw [Set.mem_smul_set_iff_inv_smul_mem]
  exact hST hx

end Trivial

/-!
## Support

Propositions relating equivariance to support in (FPerm 𝔸)-sets.
-/
section Support

variable {X Y : Type} [PermSet 𝔸 X] [PermSet 𝔸 Y]

namespace Function.Equivariant

/-- Equivariant functions preserve supports. -/
lemma preserve_support {x : X} {f : X → Y} {S : Set 𝔸} : Equivariant 𝔸 f → PermSet.HasSupport 𝔸 x S → PermSet.HasSupport 𝔸 (f x) S := by
  intro hequi hsupp π hπ
  calc
    π • f x = f (π • x) := by rw [← hequi]
    _       = f x       := by rw [hsupp hπ]

/-- Equivariant injective functions exactly preserve supports. -/
lemma preserve_support_inj {x : X} {f : X → Y} {S : Set 𝔸} : Equivariant 𝔸 f → f.Injective → (PermSet.HasSupport 𝔸 x S ↔ PermSet.HasSupport 𝔸 (f x) S) := by
  intro hequi hinj; constructor
  · exact preserve_support 𝔸 hequi
  · intro hsupp π hπ
    by_contra h
    have : π • f x ≠ f x := by rw [← hequi]; exact fun a => h (hinj a)
    have : π • f x = f x := hsupp hπ
    contradiction

end Function.Equivariant

/-- HasSupport is an equivariant relation. -/
lemma PermSet.support_equivariant : Rel.Equivariant 𝔸 (HasSupport 𝔸 : X → Set 𝔸 → Prop) := by
  intro x S π hS σ hσ
  have hfix : (π⁻¹ * σ * π) ∈ fix 𝔸 S := by
    intro a ha
    calc
      (π⁻¹ * σ * π) • a = π⁻¹ • (σ • (π • a)) := rfl
      _                 = π⁻¹ • (π • a)       := by rw [hσ (π • a) (Set.smul_mem_smul_set ha)]
      _                 = (π⁻¹ * π) • a       := rfl
      _                 = (1 : FPerm 𝔸) • a   := by rw [@Group.inv_mul_cancel]
      _                 = a                   := rfl
  calc
    σ • (π • x) = (σ * π) • x             := by apply smul_smul
    _           = (1 * σ * π) • x         := rfl
    _           = (π * π⁻¹ * σ * π) • x   := by rw [@mul_inv_cancel]
    _           = π • ((π⁻¹ * σ * π) • x) := by symm; apply smul_smul
    _           = π • x                   := by rw [hS hfix]

end Support

/-!
## Nominal Sets

Properties about equivariance relating to nominal sets.
-/
section Nominal

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

namespace Function.Equivariant

/-- If f is equivariant, then supp (f x) ⊆ supp x. -/
lemma supp_subset {x : X} {f : X → Y} : Equivariant 𝔸 f → NominalSet.supp 𝔸 (f x) ⊆ NominalSet.supp 𝔸 x := by
  intro hequi
  apply NominalSet.support_subset_supp
  · exact preserve_support 𝔸 hequi (NominalSet.supp_is_support 𝔸)
  · exact NominalSet.supp_finite 𝔸

/-- If f is equivariant and injective, then supp (f x) = supp x. -/
lemma supp_eq_if_inj {x : X} {f : X → Y} : Equivariant 𝔸 f → Function.Injective f → NominalSet.supp 𝔸 (f x) = NominalSet.supp 𝔸 x := by
  intro hequi hinj
  simp [NominalSet.supp, preserve_support_inj 𝔸 hequi hinj]

/-- If f is equivariant and a # x, then a # f x. -/
lemma preserve_fresh {a : 𝔸} {x : X} {f : X → Y} : Equivariant 𝔸 f → a # x → a # f x := by
  intro hequi hx ha
  exact hx ((supp_subset 𝔸 hequi) ha)

/-- If f is equivariant and injective, then a # x iff a # f x. -/
lemma preserve_fresh_inj {a : 𝔸} {x : X} {f : X → Y} : Equivariant 𝔸 f → Function.Injective f → (a # x ↔ a # f x) := by
  intro hequi hinj
  rw [NominalSet.IsFresh, NominalSet.IsFresh, supp_eq_if_inj 𝔸 hequi hinj]

end Function.Equivariant

/-- If there is an equivariant surjection from a nominal set into an (FPerm 𝔸)-Set, then that set is nominal. -/
def NominalSet.fromEquivariantSurjective {X Y : Type} [NominalSet 𝔸 X] [PermSet 𝔸 Y] (f : X → Y) (hequi : Function.Equivariant 𝔸 f) (hsurj : Function.Surjective f) : NominalSet 𝔸 Y where
  fin_supp := by
    intro y
    obtain ⟨x, hx⟩ := hsurj y
    use (NominalSet.supp 𝔸 x)
    constructor
    · rw [← hx]
      exact Function.Equivariant.preserve_support 𝔸 hequi (supp_is_support 𝔸)
    · exact supp_finite 𝔸

/-- supp is an equivariant function. -/
lemma NominalSet.supp_equivariant : Function.Equivariant 𝔸 (supp 𝔸 : X → Set 𝔸) := by
  intro x π; ext a; constructor
  · apply support_subset_supp 𝔸
    · apply PermSet.support_equivariant
      exact supp_is_support 𝔸
    · apply Set.Finite.smul_set
      exact supp_finite 𝔸
  · intro ha S hS
    have hsupp : PermSet.HasSupport 𝔸 x (π⁻¹ • S) := by
      rw [show x = π⁻¹ • π • x by exact eq_inv_smul_iff.mpr rfl]
      apply PermSet.support_equivariant
      exact hS.out.left
    rw [Set.mem_smul_set_iff_inv_smul_mem] at ha
    suffices π⁻¹ • a ∈ π⁻¹ • S by rw [Set.smul_mem_smul_set_iff] at this; assumption
    suffices supp 𝔸 x ⊆ π⁻¹ • S by exact this ha
    apply support_subset_supp 𝔸
    · exact hsupp
    · apply Set.Finite.smul_set
      exact hS.out.right

/-- Freshness is an equivariant relation. -/
lemma NominalSet.fresh_equivariant : Rel.Equivariant 𝔸 (NominalSet.IsFresh : 𝔸 → X → Prop) := by
  intro x y π h
  rwa [IsFresh, supp_equivariant, Set.smul_mem_smul_set_iff]

end Nominal

/-!
## Freshness

For propositions of the form `И a, a R x`, where `R` is an equivariant relation, we do not have to differentiate between "for any fresh a" and "for all fresh a".
-/
section Freshness

namespace Rel.Equivariant

variable {X : Type} [NominalSet 𝔸 X] {R: Rel 𝔸 X} {x : X}

-- We continue with a circular implication chain argument to show these equivalences.

private lemma for_fresh_iff₁ : Equivariant 𝔸 R → (И a, R a x) → (∃ a, a # x ∧ R a x) := by
  intro hequi h
  apply Set.Finite.infinite_compl at h
  simp at h
  have hinf : ({a | R a x} \ NominalSet.supp 𝔸 x).Infinite := Set.Infinite.diff h (NominalSet.supp_finite 𝔸)
  obtain ⟨a, ha⟩ := hinf.nonempty
  simp only [Set.mem_diff, Set.mem_setOf_eq] at ha
  use a; exact ha.symm

private lemma for_fresh_iff₂ : Equivariant 𝔸 R → (∃ a, a # x ∧ R a x) → (∀ a, a # x → R a x) := by
  intro hequi h
  obtain ⟨a, ⟨hfresha, ha⟩⟩ := h
  intro b hfreshb
  have hmem := hequi a x ⸨a, b⸩ ha
  rw [PermSet.atom_swap, Equiv.swap_apply_left] at hmem
  rw [NominalSet.fresh_swap_fix 𝔸 hfresha hfreshb] at hmem
  exact hmem

private lemma for_fresh_iff₃ : Equivariant 𝔸 R → (∀ a, a # x → R a x) → (И a, R a x) := by
  intro hequi h
  have hsub : {a : 𝔸 | R a x}ᶜ ⊆ NominalSet.supp 𝔸 x := by
    rw [Set.compl_subset_comm]
    intro a ha
    rw [Set.mem_compl_iff] at ha
    exact h a ha
  simp only [NominalSet.ForFresh]
  apply Set.Finite.subset (NominalSet.supp_finite 𝔸) hsub

/-- The quantifier in `И a, a R x` is equivalent to "for any fresh `a`". -/
lemma for_fresh_iff_exists : Equivariant 𝔸 R → ((И a, R a x) ↔ (∃ a, a # x ∧ R a x)) := by
  intro hequi; constructor
  · intro h
    apply for_fresh_iff₁ 𝔸 hequi
    exact h
  · intro h
    apply for_fresh_iff₃ 𝔸 hequi
    apply for_fresh_iff₂ 𝔸 hequi
    exact h

/-- The quantifier in `И a, a R x` is equivalent to "for all fresh `a`". -/
lemma for_fresh_iff_all : Equivariant 𝔸 R → ((И a, R a x) ↔ (∀ a, a # x → R a x)) := by
  intro hequi; constructor
  · intro h
    apply for_fresh_iff₂ 𝔸 hequi
    apply for_fresh_iff₁ 𝔸 hequi
    exact h
  · intro h
    apply for_fresh_iff₃ 𝔸 hequi
    exact h

end Rel.Equivariant

end Freshness

/-!
## Bundled Functions

In the following, we also define a bundled function type for equivariant functions.
-/
section BundledFunction

/-- A function carrying a proof of its equivariance. -/
structure EquivariantMap (X Y : Type) [PermSet 𝔸 X] [PermSet 𝔸 Y] where
  /-- The underlying function. -/
  toFun : X → Y
  /-- Proof of equivariance. -/
  equivariant : Function.Equivariant 𝔸 toFun

/-- We write `X →ₑ Y` or `X →ₑ[𝔸] Y` for an equivariant map X → Y. -/
infixr:25 " →ₑ " => EquivariantMap _

/-- We write `X →ₑ Y` or `X →ₑ[𝔸] Y` for an equivariant map X → Y. -/
notation:25 X " →ₑ[" A "] " Y => EquivariantMap A X Y

variable {X Y Z : Type} [PermSet 𝔸 X] [PermSet 𝔸 Y] [PermSet 𝔸 Z]

instance : FunLike (X →ₑ[𝔸] Y) X Y where
  coe := EquivariantMap.toFun
  coe_injective' f g h := by cases f; cases g; congr

namespace EquivariantMap

/-- The identity function. -/
def id : X →ₑ[𝔸] X := ⟨_root_.id, id_equivariant 𝔸⟩

/-- Function composition. -/
def comp (f : Y →ₑ[𝔸] Z) (g : X →ₑ[𝔸] Y) : X →ₑ[𝔸] Z := ⟨f ∘ g, comp_equivariant 𝔸 f.equivariant g.equivariant⟩

@[simp]
lemma id_app (x : X) : id 𝔸 x = x := rfl

@[simp]
lemma comp_app (f : Y →ₑ[𝔸] Z) (g : X →ₑ[𝔸] Y) (x : X) : comp 𝔸 f g x = f (g x) := rfl

@[simp]
lemma toFun_eq_coe {f : X →ₑ[𝔸] Y} : f.toFun = (f : X → Y) :=
  rfl

@[ext]
lemma ext {f g : X →ₑ[𝔸] Y} (h : ∀ a, f a = g a) : f = g :=
  DFunLike.ext _ _ h

end EquivariantMap

/-- Extended setoid definition that also requires the relation to be equivariant. -/
class EquivariantSetoid (X : Type) [PermSet 𝔸 X] extends Setoid X where
  equivariant : Rel.Equivariant 𝔸 r

end BundledFunction

/-!
## Equivalences

If two types are equivalent and the conversion function is also equivariant, we know some additional information about nominal set structures.
-/
section Equivalence

/-- An equivalence where the function is equivariant. -/
structure EquivariantEquiv (X Y : Type) [PermSet 𝔸 X] [PermSet 𝔸 Y] extends Equiv X Y where
  toFun_equivariant : Function.Equivariant 𝔸 toFun
  invFun_equivariant : Function.Equivariant 𝔸 invFun := Function.inv_equivariant 𝔸 left_inv right_inv toFun_equivariant

def Equivariant.Equiv.symm {𝔸 X Y : Type} [Infinite 𝔸] [Countable 𝔸] [DecidableEq 𝔸] [PermSet 𝔸 X] [PermSet 𝔸 Y] (equiv : EquivariantEquiv 𝔸 X Y) : EquivariantEquiv 𝔸 Y X :=
  { toFun := equiv.invFun
    invFun := equiv.toFun
    toFun_equivariant := equiv.invFun_equivariant
    left_inv := equiv.right_inv
    right_inv := equiv.left_inv
  }

/-- If an (FPerm 𝔸)-Set is equivalent to a nominal set, it is also a nominal set. -/
def NominalSet.fromEquivariantEquiv {X Y : Type} [PermSet 𝔸 X] [NominalSet 𝔸 Y] (equiv : EquivariantEquiv 𝔸 X Y) : NominalSet 𝔸 X :=
  NominalSet.fromEquivariantSurjective 𝔸 equiv.invFun equiv.invFun_equivariant equiv.symm.surjective

/-- If two nominal sets are equivalent, the corresponding least supports are equal. -/
lemma EquivariantEquiv.supp_eq {X Y : Type} [NominalSet 𝔸 X] [NominalSet 𝔸 Y] (equiv : EquivariantEquiv 𝔸 X Y) (x : X) : NominalSet.supp 𝔸 x = NominalSet.supp 𝔸 (equiv.toFun x) := by
  symm
  apply Function.Equivariant.supp_eq_if_inj
  · exact equiv.toFun_equivariant
  · exact equiv.injective

end Equivalence