import Mathlib
import Nominal.Support
import Nominal.Equivariance
import Nominal.Category
import Nominal.Permutations
import Nominal.Instances.Basic

/-!
# Constructing Finitely Supported Sets

We can restrict arbitrary `(FPerm 𝔸)`-sets to finitely supported subsets, which are nominal sets.

Examples of this include finitely supported sets (`FinsuppSet`) and finitely supported functions (`FinsuppFunction`).
-/

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

namespace PermSet

/-!
## Finitely Supported Subsets

We first define how every `PermSet` can be restricted to its finitely supported elements (`Finsupp`), which form a
nominal set.
-/
section FSSubsets

/-- Restrict a type to its finitely supported elements -/
abbrev Finsupp (X : Type) [PermSet 𝔸 X] : Type :=
  { x : X // ∃ S, PermSet.HasSupport 𝔸 x S ∧ S.Finite }

namespace Finsupp

variable {X : Type} [PermSet 𝔸 X]

instance : PermSet 𝔸 (Finsupp 𝔸 X) where
  smul π x := ⟨π • x.val, by
    obtain ⟨S, ⟨hS₁, hS₂⟩⟩ := x.prop
    use π • S; constructor
    · rw [← Rel.equivariant 𝔸 (HasSupport 𝔸)]
      exact hS₁
    · exact Set.Finite.smul_set hS₂⟩
  one_smul x := by
    apply Subtype.ext
    exact MulAction.one_smul x.val
  mul_smul π τ x := by
    apply Subtype.ext
    exact mul_smul π τ x.val

@[simp] lemma smul_val (π : FPerm 𝔸) (x : Finsupp 𝔸 X) : (π • x).val = π • x.val := rfl

instance instValEquivariant : Function.Equivariant 𝔸 (Subtype.val : Finsupp 𝔸 X → X) := by
  constructor; simp

instance : NominalSet 𝔸 (Finsupp 𝔸 X) where
  fin_supp x := by
    obtain ⟨S, ⟨hS₁, hS₂⟩⟩ := x.prop
    use S; constructor
    · intro π hπ
      apply Subtype.ext
      exact hS₁ hπ
    · exact hS₂

end Finsupp

end FSSubsets

/-!
## Finitely Supported Sets

Since the powerset of an (FPerm 𝔸)-set is an (FPerm 𝔸)-set, we can restrict ourselves to finitely supported sets to
construct a nominal set `FinsuppSet`.

Furthermore, common binary operations such as union, intersection, and difference are equivariant functions.
-/
section FSSets

abbrev FinsuppSet (X : Type) [PermSet 𝔸 X] : Type :=
  Finsupp 𝔸 (Set X)

namespace FinsuppSet

variable (X Y : Type) [NominalSet 𝔸 X] [NominalSet 𝔸 Y]

private def union' : (FinsuppSet 𝔸 X × FinsuppSet 𝔸 X) → Set X :=
  fun (S₁, S₂) ↦ S₁.val ∪ S₂.val

instance : Function.Equivariant 𝔸 (union' 𝔸 X) := by
  constructor; intro (S₁, S₂) π
  simp only [union', PermSet.Finsupp.smul_val]
  exact Eq.symm Set.smul_set_union

def union : (FinsuppSet 𝔸 X × FinsuppSet 𝔸 X) → FinsuppSet 𝔸 X :=
  fun p ↦ ⟨union' 𝔸 X p, by
    use NominalSet.supp 𝔸 p; constructor
    · apply Function.Equivariant.preserve_support
      apply NominalSet.supp_is_support
    · apply NominalSet.supp_finite⟩

instance : Function.Equivariant 𝔸 (union 𝔸 X) := by
  constructor; intro p π
  apply Subtype.ext
  simp only [union, Finsupp.smul_val]
  apply Function.equivariant

instance : Union (FinsuppSet 𝔸 X) := ⟨Function.curry (union 𝔸 X)⟩

@[simp] lemma union_val {S₁ S₂ : FinsuppSet 𝔸 X} :
    (S₁ ∪ S₂).val = S₁.val ∪ S₂.val := rfl

@[simp] lemma union_equivariant {S₁ S₂ : FinsuppSet 𝔸 X} {π : FPerm 𝔸} :
    π • (S₁ ∪ S₂) = π • S₁ ∪ π • S₂ := by
  show π • (union _ _ (S₁, S₂)) = union _ _ (π • (S₁, S₂))
  rw [← Function.equivariant]

private def inter' : (FinsuppSet 𝔸 X × FinsuppSet 𝔸 X) → Set X :=
  fun (S₁, S₂) ↦ S₁.val ∩ S₂.val

instance : Function.Equivariant 𝔸 (inter' 𝔸 X) := by
  constructor; intro (S₁, S₂) π
  simp only [inter', PermSet.Finsupp.smul_val]
  exact Eq.symm Set.smul_set_inter

def inter : (FinsuppSet 𝔸 X × FinsuppSet 𝔸 X) → FinsuppSet 𝔸 X :=
  fun p ↦ ⟨inter' 𝔸 X p, by
    use NominalSet.supp 𝔸 p; constructor
    · apply Function.Equivariant.preserve_support
      apply NominalSet.supp_is_support
    · apply NominalSet.supp_finite⟩

instance : Function.Equivariant 𝔸 (inter 𝔸 X) := by
  constructor; intro p π
  apply Subtype.ext
  simp only [inter, Finsupp.smul_val]
  apply Function.equivariant

instance : Inter (FinsuppSet 𝔸 X) := ⟨Function.curry (inter 𝔸 X)⟩

@[simp] lemma inter_val {S₁ S₂ : FinsuppSet 𝔸 X} :
    (S₁ ∩ S₂).val = S₁.val ∩ S₂.val := rfl

@[simp] lemma inter_equivariant {S₁ S₂ : FinsuppSet 𝔸 X} {π : FPerm 𝔸} :
    π • (S₁ ∩ S₂) = π • S₁ ∩ π • S₂ := by
  show π • (inter _ _ (S₁, S₂)) = inter _ _ (π • (S₁, S₂))
  rw [← Function.equivariant]

private def sdiff' : (FinsuppSet 𝔸 X × FinsuppSet 𝔸 X) → Set X :=
  fun (S₁, S₂) ↦ S₁.val \ S₂.val

instance : Function.Equivariant 𝔸 (sdiff' 𝔸 X) := by
  constructor; intro (S₁, S₂) π
  simp only [sdiff', PermSet.Finsupp.smul_val]
  exact Eq.symm Set.smul_set_sdiff

def sdiff : (FinsuppSet 𝔸 X × FinsuppSet 𝔸 X) → FinsuppSet 𝔸 X :=
  fun p ↦ ⟨sdiff' 𝔸 X p, by
    use NominalSet.supp 𝔸 p; constructor
    · apply Function.Equivariant.preserve_support
      apply NominalSet.supp_is_support
    · apply NominalSet.supp_finite⟩

instance : Function.Equivariant 𝔸 (sdiff 𝔸 X) := by
  constructor; intro p π
  apply Subtype.ext
  simp only [sdiff, Finsupp.smul_val]
  apply Function.equivariant

instance : SDiff (FinsuppSet 𝔸 X) := ⟨Function.curry (sdiff 𝔸 X)⟩

@[simp] lemma sdiff_val {S₁ S₂ : FinsuppSet 𝔸 X} :
    (S₁ \ S₂).val = S₁.val \ S₂.val := rfl

@[simp] lemma sdiff_equivariant {S₁ S₂ : FinsuppSet 𝔸 X} {π : FPerm 𝔸} :
    π • (S₁ \ S₂) = π • S₁ \ π • S₂ := by
  show π • (sdiff _ _ (S₁, S₂)) = sdiff _ _ (π • (S₁, S₂))
  rw [← Function.equivariant]

end FinsuppSet

end FSSets

/-!
## Finitely Supported Functions

Since functions between (FPerm 𝔸)-sets are an (FPerm 𝔸)-set, we can restrict ourselves to finitely supported functions
to construct a nominal set `FinsuppFunction`.

Furthermore, there is an isomorphism between finietly supported functions (X →fs Y) and sets 𝒫fs (X × Y).
-/
section FSFunctions

abbrev FinsuppFunction (X Y : Type) [PermSet 𝔸 X] [PermSet 𝔸 Y] : Type :=
  Finsupp 𝔸 (X → Y)

lemma FinsuppFunction.supp_image {X Y : Type} [NominalSet 𝔸 X] [NominalSet 𝔸 Y] (f : PermSet.FinsuppFunction 𝔸 X Y) :
    ∀ x, PermSet.HasSupport 𝔸 (f.val x) (NominalSet.supp 𝔸 f ∪ NominalSet.supp 𝔸 x) := by
  intro x π hπ
  have hπ₁ : π ∈ PermSet.fix 𝔸 (NominalSet.supp 𝔸 f) := by
    apply PermSet.fix_antitone_mem 𝔸 hπ; exact Set.subset_union_left
  have hπ₂ : π ∈ PermSet.fix 𝔸 (NominalSet.supp 𝔸 x) := by
    apply PermSet.fix_antitone_mem 𝔸 hπ; exact Set.subset_union_right
  suffices (π • f).val (π • x) = f.val x by
    simp at this
    exact this
  rw [show π • f = f from NominalSet.supp_is_support 𝔸 hπ₁]
  rw [show π • x = x from NominalSet.supp_is_support 𝔸 hπ₂]

variable (X Y : Type) [NominalSet 𝔸 X] [NominalSet 𝔸 Y]

namespace Finsupp

private def flatten' (f : X → PermSet.FinsuppSet 𝔸 Y) : Set (X × Y) :=
  {p : X × Y | p.snd ∈ (f p.fst).val}

instance : Function.Equivariant 𝔸 (flatten' 𝔸 X Y) := by
  constructor
  intro f π; ext ⟨x, y⟩; constructor
  · simp only [flatten', PermSet.smul_function, PermSet.Finsupp.smul_val, Set.mem_setOf_eq]
    intro hy
    have hy' : π⁻¹ • y ∈ π⁻¹ • π • (f (π⁻¹ • x)).val := by
      rw [Set.smul_mem_smul_set_iff]; exact hy
    rw [inv_smul_smul] at hy'
    rw [show (x, y) = π • (π⁻¹ • x, π⁻¹ • y) by simp]
    rw [Set.smul_mem_smul_set_iff]
    exact hy'
  · intro hxy
    have hxy' : π⁻¹ • (x, y) ∈ π⁻¹ • π • flatten' 𝔸 X Y f := by
      rw [Set.smul_mem_smul_set_iff]; exact hxy
    simp only [inv_smul_smul, Prod.smul_mk, flatten', Set.mem_setOf_eq] at hxy'
    simp only [flatten', PermSet.smul_function, PermSet.Finsupp.smul_val, Set.mem_setOf_eq]
    have hy : π • π⁻¹ • y ∈ π • (f (π⁻¹ • x)).val := by
      rw [Set.smul_mem_smul_set_iff]; exact hxy'
    simp only [smul_inv_smul] at hy
    exact hy

def flatten (f : PermSet.FinsuppFunction 𝔸 X (PermSet.FinsuppSet 𝔸 Y)) :
    PermSet.FinsuppSet 𝔸 (X × Y) :=
  ⟨flatten' 𝔸 X Y f.val, by
    obtain ⟨S, ⟨hS₁, hS₂⟩⟩ := f.prop
    use S; constructor
    · apply Function.Equivariant.preserve_support
      exact hS₁
    · exact hS₂⟩

instance : Function.Equivariant 𝔸 (flatten 𝔸 X Y) := by
  constructor; intro p π
  apply Subtype.ext
  simp only [flatten, Finsupp.smul_val]
  apply Function.equivariant

private def unflatten' (S : PermSet.FinsuppSet 𝔸 (X × Y)) : X → (PermSet.FinsuppSet 𝔸 Y) :=
  fun x ↦ ⟨{y : Y | (x, y) ∈ S.val}, by
    use NominalSet.supp 𝔸 S ∪ NominalSet.supp 𝔸 x; constructor
    · intro π hπ
      have hfixx : π⁻¹ • x = x := by
        apply NominalSet.supp_is_support
        rw [← PermSet.fix_inv]
        exact PermSet.fix_antitone_mem 𝔸 hπ Set.subset_union_right
      have hfixS : π⁻¹ • S = S := by
        apply NominalSet.supp_is_support
        rw [← PermSet.fix_inv]
        exact PermSet.fix_antitone_mem 𝔸 hπ Set.subset_union_left
      ext y; constructor
      · intro hy
        simp only [Set.mem_smul_set_iff_inv_smul_mem, Set.mem_setOf_eq] at hy
        rw [← hfixx, ← hfixS, smul_val, Set.mem_inv_smul_set_iff, Prod.smul_mk, smul_inv_smul, smul_inv_smul] at hy
        exact hy
      · intro hy
        simp only [Set.mem_smul_set_iff_inv_smul_mem, Set.mem_setOf_eq]
        simp only [Set.mem_setOf_eq] at hy
        have hy' : (π⁻¹ • x, π⁻¹ • y) ∈ π⁻¹ • S.val := Set.smul_mem_smul_set hy
        rw [hfixx, ← smul_val, hfixS] at hy'
        exact hy'
    · simp only [Set.finite_union]; constructor
      · exact NominalSet.supp_finite 𝔸
      · exact NominalSet.supp_finite 𝔸⟩

instance : Function.Equivariant 𝔸 (unflatten' 𝔸 X Y) := by
  constructor
  intro S π
  ext x y
  simp only [unflatten', smul_val, Set.mem_setOf_eq, smul_function]
  simp only [Set.mem_smul_set_iff_inv_smul_mem, Set.mem_setOf_eq, Prod.smul_mk]

def unflatten (S : PermSet.FinsuppSet 𝔸 (X × Y)) : PermSet.FinsuppFunction 𝔸 X (PermSet.FinsuppSet 𝔸 Y) :=
  ⟨unflatten' 𝔸 X Y S, by
    obtain ⟨A, ⟨hS₁, hS₂⟩⟩ := S.prop
    use A; constructor
    · apply Function.Equivariant.preserve_support
      rw [← Function.Equivariant.preserve_support_inj 𝔸 Subtype.val_injective] at hS₁
      exact hS₁
    · exact hS₂⟩

instance : Function.Equivariant 𝔸 (unflatten 𝔸 X Y) := by
  constructor; intro p π
  apply Subtype.ext
  simp only [unflatten, Finsupp.smul_val]
  apply Function.equivariant

def flattenEquiv  :
    EquivariantEquiv 𝔸
      (PermSet.FinsuppFunction 𝔸 X (PermSet.FinsuppSet 𝔸 Y))
      (PermSet.FinsuppSet 𝔸 (X × Y)) where
  toFun := flatten 𝔸 X Y
  invFun := unflatten 𝔸 X Y
  toFun_equivariant := inferInstance

end Finsupp

end FSFunctions

end PermSet