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 Equivariant (f : X → Y) : Prop :=
  ∀ (x : X) (π : FPerm 𝔸), f (π • x) = π • f x

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

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

end Definition

/-!
## Support

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

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

namespace 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 Equivariant

end Support

/-!
## Nominal Sets

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

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

namespace 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]

end Equivariant

end Nominal