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

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

namespace PermSet

/-- 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.eq
    exact MulAction.one_smul x.val
  mul_smul π τ x := by
    apply Subtype.eq
    exact MulAction.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.eq
      exact hS₁ hπ
    · exact hS₂

end Finsupp

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

instance {X Y : Type} [SMul (FPerm 𝔸) X] [SMul (FPerm 𝔸) Y] : SMul (FPerm 𝔸) (X → Y) where
  smul π f := fun x ↦ π • f (π⁻¹ • x)

@[simp] lemma smul_function {X Y : Type} [PermSet 𝔸 X] [PermSet 𝔸 Y] (π : FPerm 𝔸)
    (f : X → Y) (x : X) : (π • f) x = π • f (π⁻¹ • x) :=
  rfl

instance {X Y : Type} [PermSet 𝔸 X] [PermSet 𝔸 Y] : PermSet 𝔸 (X → Y) where
  one_smul f := by ext x; simp
  mul_smul π τ f := by ext x; simp [MulAction.mul_smul]

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

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.eq
  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
    -- Pitts
    sorry⟩

instance : Function.Equivariant 𝔸 (unflatten' 𝔸 X Y) := by
  constructor
  intro S π
  -- Pitts
  sorry

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.eq
  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

namespace FinsuppSet

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.eq
  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.eq
  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.eq
  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

abbrev FinSet (X : Type) [PermSet 𝔸 X] :=
  { S : Set X // S.Finite }

namespace FinSet

variable {X : Type} [PermSet 𝔸 X]

instance : PermSet 𝔸 (FinSet 𝔸 X) where
  smul π S := ⟨π • S.val, by
    rw [Set.finite_smul_set]
    exact S.prop⟩
  one_smul x := by
    apply Subtype.eq
    exact MulAction.one_smul x.val
  mul_smul π τ x := by
    apply Subtype.eq
    exact MulAction.mul_smul π τ x.val

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

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

variable {X : Type} [NominalSet 𝔸 X]

instance : NominalSet 𝔸 (FinSet 𝔸 X) where
  fin_supp := by sorry

@[simp] lemma supp {S : FinSet 𝔸 X} : NominalSet.supp 𝔸 S = ⋃₀ {NominalSet.supp 𝔸 x | x ∈ S.val} := by sorry

@[simp] lemma fresh {a : 𝔸} {S : FinSet 𝔸 X} : a # S ↔ ∀ x ∈ S.val, a # x := by
  simp [NominalSet.IsFresh]

def toFinsuppSet (S : FinSet 𝔸 X) : FinsuppSet 𝔸 X :=
  ⟨S.val, by
    use NominalSet.supp 𝔸 S; constructor
    · apply Function.Equivariant.preserve_support
      apply NominalSet.supp_is_support
    · apply NominalSet.supp_finite⟩

@[simp] lemma toFinsuppSet_val (S : FinSet 𝔸 X) : S.toFinsuppSet.val = S.val := rfl

noncomputable def toFinset (S : FinSet 𝔸 X) : Finset X :=
  S.prop.toFinset

def union : (FinSet 𝔸 X × FinSet 𝔸 X) → FinSet 𝔸 X :=
  fun (S₁, S₂) ↦ ⟨S₁.val ∪ S₂.val, Set.Finite.union S₁.prop S₂.prop⟩

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

instance : Function.Equivariant 𝔸 (union 𝔸 : (FinSet 𝔸 X × FinSet 𝔸 X → FinSet 𝔸 X)) := by
  constructor; intro (S₁, S₂) π
  apply Subtype.eq
  simp only [union, smul_val]
  exact Eq.symm Set.smul_set_union

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

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

def inter : (FinSet 𝔸 X × FinSet 𝔸 X) → FinSet 𝔸 X :=
  fun (S₁, S₂) ↦ ⟨S₁.val ∩ S₂.val, Set.Finite.inter_of_left S₁.prop S₂.val⟩

instance : Function.Equivariant 𝔸 (inter 𝔸 : (FinSet 𝔸 X × FinSet 𝔸 X → FinSet 𝔸 X)) := by
  constructor; intro (S₁, S₂) π
  apply Subtype.eq
  simp only [inter, smul_val]
  exact Eq.symm Set.smul_set_inter

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

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

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

def sdiff : (FinSet 𝔸 X × FinSet 𝔸 X) → FinSet 𝔸 X :=
  fun (S₁, S₂) ↦ ⟨S₁.val \ S₂.val, Set.Finite.diff S₁.prop⟩

instance : Function.Equivariant 𝔸 (sdiff 𝔸 : (FinSet 𝔸 X × FinSet 𝔸 X → FinSet 𝔸 X)) := by
  constructor; intro (S₁, S₂) π
  apply Subtype.eq
  simp only [sdiff, smul_val]
  exact Eq.symm Set.smul_set_sdiff

def empty : FinSet 𝔸 X :=
  ⟨∅, Set.finite_empty⟩

instance : EmptyCollection (FinSet 𝔸 X) := ⟨empty 𝔸⟩

@[simp] lemma empty_val : (∅ : FinSet 𝔸 X).val = ∅ := rfl

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

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

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

end FinSet

def Finset.toPermFinSet {S : Finset X} : PermSet.FinSet 𝔸 X :=
  ⟨S, Finset.finite_toSet S⟩

end PermSet