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)

/-- For equivariant sets, we even have "if and only if". -/
lemma equivariant_set_iff {S : Set X} : EquivariantSet 𝔸 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 at hx
      exact hx
  · intro h x π hx
    exact (h x π).mp hx

/-- For equivariant relations, we even have "if and only if". -/
lemma equivariant_relation_iff {R : Rel X Y} : EquivariantRelation 𝔸 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 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 : Equivariant 𝔸 (id : X → X) := by
  simp [Equivariant]

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

/-- The subset relation is equivariant. -/
lemma subset_equivariant : EquivariantRelation 𝔸 (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 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

/-- HasSupport is an equivariant relation. -/
lemma PermSet.support_equivariant : EquivariantRelation 𝔸 (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 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

/-- supp is an equivariant function. -/
lemma NominalSet.supp_equivariant : 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

end Nominal