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 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 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 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 : Equivariant 𝔸 (id : X → X) := by
  simp [Equivariant]

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

/-- 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

/-!
## 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 Equivariance

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₁ : EquivariantRelation 𝔸 R → (И a, R a x) → (∃ a, a # x ∧ R a x) := by
  intro hequi 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₂ : EquivariantRelation 𝔸 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₃ : EquivariantRelation 𝔸 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.infinite_of_finite_compl
  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 : EquivariantRelation 𝔸 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 : EquivariantRelation 𝔸 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 Equivariance

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 : 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 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

end BundledFunction