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

lemma Function.equivariant (f : X → Y) [Equivariant 𝔸 f] :
    ∀ (x : X) (π : FPerm 𝔸), f (π • x) = π • f x := by
  cases (inferInstance : Equivariant 𝔸 f); assumption

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

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

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

/-- For equivariant relations, we even have "if and only if". -/
lemma Rel.equivariant (R : Rel X Y) [Rel.Equivariant 𝔸 R] :
    ∀ (x : X) (y : Y) (π : FPerm 𝔸), R x y ↔ R (π • x) (π • y) := by
  cases (inferInstance : Equivariant 𝔸 R); rename_i h
  intro 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

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. -/
instance Function.Equivariant.id : Function.Equivariant 𝔸 (id : X → X) := by
  constructor
  simp

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

/-- If a function is equivariant, its inverse is also equivariant. -/
@[reducible]
def Function.Equivariant.inv
    {f : X → Y} [Function.Equivariant 𝔸 f] {inv : Y → X}
    (hleft : Function.LeftInverse inv f)
    (hright : Function.RightInverse inv f) :
    Function.Equivariant 𝔸 inv := by
  constructor; intro y π
  apply hleft.injective
  rw [Function.equivariant 𝔸 f, hright, hright]

/-- The subset relation is equivariant. -/
instance Set.Equivariant.subset : Rel.Equivariant 𝔸 (Set.Subset : Set X → Set X → Prop) := by
  constructor; 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} [Equivariant 𝔸 f] {S : Set 𝔸} :
    PermSet.HasSupport 𝔸 x S → PermSet.HasSupport 𝔸 (f x) S := by
  intro hsupp π hπ
  calc
    π • f x = f (π • x) := by rw [← equivariant 𝔸 f]
    _       = f x       := by rw [hsupp hπ]

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

end Function.Equivariant

/-- HasSupport is an equivariant relation. -/
instance PermSet.instHasSupportEquivariant : Rel.Equivariant 𝔸 (HasSupport 𝔸 : X → Set 𝔸 → Prop) := by
  constructor; 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
  apply NominalSet.support_subset_supp
  · exact preserve_support 𝔸 (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 hinj
  simp [NominalSet.supp, preserve_support_inj 𝔸 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 hx ha
  exact hx (supp_subset 𝔸 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 hinj
  rw [NominalSet.IsFresh, NominalSet.IsFresh, supp_eq_if_inj 𝔸 hinj]

end Function.Equivariant

/-- If there is an equivariant surjection from a nominal set into an (FPerm 𝔸)-Set, then that set is nominal. -/
@[reducible]
def NominalSet.ofEquivariantSurjective
    {X Y : Type} [NominalSet 𝔸 X] [PermSet 𝔸 Y]
    (f : X → Y) [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 𝔸 (supp_is_support 𝔸)
    · exact supp_finite 𝔸

/-- supp is an equivariant function. -/
instance NominalSet.instSuppEquivariant : Function.Equivariant 𝔸 (supp 𝔸 : X → Set 𝔸) := by
  constructor; intro x π
  ext a; constructor
  · apply support_subset_supp 𝔸
    · rw [← Rel.equivariant 𝔸 (PermSet.HasSupport 𝔸)]
      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]
      rw [← Rel.equivariant 𝔸 (PermSet.HasSupport 𝔸)]
      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. -/
instance NominalSet.instFreshEquivariant : Rel.Equivariant 𝔸 (NominalSet.IsFresh : 𝔸 → X → Prop) := by
  constructor; intro x y π h
  rwa [IsFresh, Function.equivariant 𝔸 (supp 𝔸), 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₁ : (И a, R a x) → (∃ a, a # x ∧ R a x) := by
  intro 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₂ [Rel.Equivariant 𝔸 R] : (∃ a, a # x ∧ R a x) → (∀ a, a # x → R a x) := by
  intro h
  obtain ⟨a, ⟨hfresha, ha⟩⟩ := h
  intro b hfreshb
  rw [Rel.equivariant 𝔸 R b x ⸨a, b⸩]
  rw [PermSet.atom_swap, Equiv.swap_apply_right, NominalSet.fresh_swap_fix 𝔸 hfresha hfreshb]
  exact ha

private lemma for_fresh_iff₃ : (∀ a, a # x → R a x) → (И a, R a x) := by
  intro 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 [Rel.Equivariant 𝔸 R] : ((И a, R a x) ↔ (∃ a, a # x ∧ R a x)) := by
  constructor
  · intro h
    apply for_fresh_iff₁
    exact h
  · intro h
    apply for_fresh_iff₃
    apply for_fresh_iff₂
    exact h

/-- The quantifier in `И a, a R x` is equivalent to "for all fresh `a`". -/
lemma for_fresh_iff_all [Rel.Equivariant 𝔸 R] : ((И a, R a x) ↔ (∀ a, a # x → R a x)) := by
  constructor
  · intro h
    apply for_fresh_iff₂
    apply for_fresh_iff₁
    exact h
  · intro h
    apply for_fresh_iff₃
    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 EquivariantMap.toFunEquivariant (f : X →ₑ[𝔸] Y) : Function.Equivariant 𝔸 f.toFun :=
  f.equivariant

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, inferInstance⟩

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

@[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

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.Equivariant.inv 𝔸 left_inv right_inv

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
  }

instance EquivariantEquiv.toFunEquivariant
    {X Y : Type} [PermSet 𝔸 X] [PermSet 𝔸 Y]
    (equiv : EquivariantEquiv 𝔸 X Y) : Function.Equivariant 𝔸 equiv.toFun :=
  equiv.toFun_equivariant

instance EquivariantEquiv.invFunEquivariant
    {X Y : Type} [PermSet 𝔸 X] [PermSet 𝔸 Y]
    (equiv : EquivariantEquiv 𝔸 X Y) : Function.Equivariant 𝔸 equiv.invFun :=
  equiv.invFun_equivariant

/-- If an (FPerm 𝔸)-Set is equivalent to a nominal set, it is also a nominal set. -/
@[reducible]
def NominalSet.ofEquivariantEquiv {X Y : Type} [PermSet 𝔸 X] [NominalSet 𝔸 Y] (equiv : EquivariantEquiv 𝔸 X Y) : NominalSet 𝔸 X :=
  NominalSet.ofEquivariantSurjective 𝔸 equiv.invFun 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.injective

end Equivalence