import Mathlib
import Nominal.Support
import Nominal.Equivariance

/-!
# Basic Nominal Sets

This file contains common instances of nominal sets and their smallest supports.
-/

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

/-!
## Atoms

The set 𝔸 of atoms itself is a nominal set, where supp a = {a}.
-/
section Atoms

private lemma NominalSet.supp_atom' {a : 𝔸} : PermSet.HasSupport 𝔸 a {a} := by
  simp [PermSet.HasSupport, PermSet.fix, PermSet.fix']

instance : NominalSet 𝔸 𝔸 where
  fin_supp := by
    intro a; use {a}
    exact ⟨NominalSet.supp_atom' 𝔸, Set.finite_singleton a⟩

@[simp] lemma NominalSet.supp_atom {a : 𝔸} : supp 𝔸 a = {a} := by
  ext b; constructor
  · apply support_subset_supp 𝔸
    · exact NominalSet.supp_atom' 𝔸
    · exact Set.finite_singleton a
  · suffices a ∈ supp 𝔸 a by
      simp only [Set.mem_singleton_iff]
      intro h
      rwa [h]
    intro S ⟨hsuppS, hfinS⟩
    by_contra ha
    -- Show that we can find a fresh c with (a c) ∈ fix S
    obtain ⟨c, hc⟩ := Infinite.exists_notMem_finset ({a} ∪ hfinS.toFinset)
    simp only [Finset.singleton_union, Finset.mem_insert, Set.Finite.mem_toFinset, not_or] at hc
    have hfix : ⸨a, c⸩ ∈ PermSet.fix 𝔸 S := by
      intro d hd
      apply FPerm.swap_apply_of_ne_of_ne (ne_of_mem_of_not_mem hd ha) (ne_of_mem_of_not_mem hd hc.right)
    have h₁ : ⸨a, c⸩ • a = c := by rw [PermSet.atom_swap, Equiv.swap_apply_left]
    have h₂ : ⸨a, c⸩ • a = a := by apply hsuppS; assumption
    rw [h₁] at h₂
    exact hc.left h₂

@[simp] lemma NominalSet.fresh_atom {a b : 𝔸} : a # b ↔ a ≠ b := by
  simp [IsFresh]

end Atoms

/-!
## Discrete Nominal Sets

Every set can be made into a discrete nominal set, where supp x = ∅.
-/
section Discrete

def NominalSet.mkDiscrete (X : Type) : NominalSet.Discrete 𝔸 X where
  smul π x := x
  one_smul := by intro; rfl
  mul_smul := by intro _ _ _; rfl
  is_discrete := by intro _ _; rfl
  fin_supp := by
    intro x; use {}
    constructor
    · intro π hπ; simp only [PermSet.fix', Set.mem_setOf_eq]; rfl
    · exact Set.finite_empty

instance {X : Type} [NominalSet.Discrete 𝔸 X] [Finite X] : OrbitFinite 𝔸 X where
  fin_orbits := Finite.of_equiv X
    { toFun x := Quotient.mk _ x
      invFun x := x.out
      left_inv x := by
        suffices ∀ (x y : X), MulAction.orbitRel (FPerm 𝔸) X x y → x = y by
          apply this
          rw [← Quotient.eq]
          apply Quotient.out_eq
        intro x y ⟨π, hπ⟩
        simp only [NominalSet.Discrete.is_discrete] at hπ
        exact hπ.symm
      right_inv x := by simp
    }

@[simp] lemma NominalSet.Discrete.supp {X : Type} [NominalSet.Discrete 𝔸 X] {x : X} : supp 𝔸 x = ∅ := by
  rw [← Set.subset_empty_iff]
  apply Set.sInter_subset_of_mem
  constructor
  · intro π hπ; simp [PermSet.fix', is_discrete]
  · exact Set.finite_empty

@[simp] lemma NominalSet.Discrete.fresh {X : Type} [NominalSet.Discrete 𝔸 X] {a : 𝔸} {x : X} : a # x := by
  simp [IsFresh]

/-- We consider Unit a discrete nominal set. -/
instance : NominalSet.Discrete 𝔸 Unit := NominalSet.mkDiscrete 𝔸 Unit

/-- We consider Bool a discrete nominal set. -/
instance : NominalSet.Discrete 𝔸 Bool := NominalSet.mkDiscrete 𝔸 Bool

end Discrete

/-!
## Products

Pairs of two nominal sets also form a nominal set, where supp (x, y) = supp x ∪ supp y.
-/
section Products

namespace Prod

instance {X Y : Type} [PermSet 𝔸 X] [PermSet 𝔸 Y] : PermSet 𝔸 (X × Y) where

variable {X Y : Type} [NominalSet 𝔸 X] [NominalSet 𝔸 Y]

private lemma supp' {x : X} {y : Y} : PermSet.HasSupport 𝔸 (x, y) (NominalSet.supp 𝔸 x ∪ NominalSet.supp 𝔸 y) := by
  intro π hπ
  obtain ⟨hπ₁, hπ₂⟩ := PermSet.fix_union 𝔸 hπ
  exact Prod.mk_inj.mpr ⟨NominalSet.supp_is_support 𝔸 hπ₁, NominalSet.supp_is_support 𝔸 hπ₂⟩

instance : NominalSet 𝔸 (X × Y) where
  fin_supp := by
    intro ⟨x, y⟩; use NominalSet.supp 𝔸 x ∪ NominalSet.supp 𝔸 y
    exact ⟨supp' 𝔸, Set.Finite.union (NominalSet.supp_finite 𝔸) (NominalSet.supp_finite 𝔸)⟩

@[simp] lemma supp {x : X} {y : Y} : NominalSet.supp 𝔸 (x, y) = NominalSet.supp 𝔸 x ∪ NominalSet.supp 𝔸 y := by
  ext b; constructor
  · apply NominalSet.support_subset_supp
    · exact supp' 𝔸
    · exact Set.Finite.union (NominalSet.supp_finite 𝔸) (NominalSet.supp_finite 𝔸)
  · intro hb; cases hb with
    | inl hx =>
        intro S hS
        suffices PermSet.HasSupport 𝔸 x S by apply hx; exact ⟨this, hS.out.right⟩
        intro π hπ
        suffices π ∈ PermSet.fix' 𝔸 (x, y) by
          simp only [PermSet.fix', smul_mk, mk.injEq, Set.mem_setOf_eq] at this
          exact this.left
        exact (hS.out.left hπ).out
    | inr hy =>
        intro S hS
        suffices PermSet.HasSupport 𝔸 y S by apply hy; exact ⟨this, hS.out.right⟩
        intro π hπ
        suffices π ∈ PermSet.fix' 𝔸 (x, y) by
          simp only [PermSet.fix', smul_mk, mk.injEq, Set.mem_setOf_eq] at this
          exact this.right
        exact (hS.out.left hπ).out

@[simp] lemma fresh {a : 𝔸} {x : X} {y : Y} : a # (x, y) ↔ a # x ∧ a # y := by
  simp [NominalSet.IsFresh]

lemma supp_comm {x : X} {y : Y} : NominalSet.supp 𝔸 (x, y) = NominalSet.supp 𝔸 (y, x) := by
  simp; apply Set.union_comm

instance instFstEquivariant : Function.Equivariant 𝔸 (Prod.fst : X × Y → X) := by
  constructor; simp

instance instSndEquivariant : Function.Equivariant 𝔸 (Prod.snd : X × Y → Y) := by
  constructor; simp

end Prod

end Products

/-!
## Coproducts

Sums of two nominal sets also form a nominal set, where supp (inl x) = supp x and supp (inr y) = supp y.
-/
section Coproducts

namespace Sum

instance {X Y : Type} [PermSet 𝔸 X] [PermSet 𝔸 Y] : PermSet 𝔸 (X ⊕ Y) where

variable {X Y : Type} [NominalSet 𝔸 X] [NominalSet 𝔸 Y]

private lemma supp_inl' {x : X} : PermSet.HasSupport 𝔸 (.inl x : X ⊕ Y) (NominalSet.supp 𝔸 x) := by
  intro π hπ
  simp only [PermSet.fix', smul_inl, inl.injEq, Set.mem_setOf_eq]
  exact NominalSet.supp_is_support 𝔸 hπ

private lemma supp_inr' {y : Y} : PermSet.HasSupport 𝔸 (.inr y : X ⊕ Y) (NominalSet.supp 𝔸 y) := by
  intro π hπ
  simp only [PermSet.fix', smul_inr, inr.injEq, Set.mem_setOf_eq]
  exact NominalSet.supp_is_support 𝔸 hπ

instance : NominalSet 𝔸 (X ⊕ Y) where
  fin_supp := by
    intro xy
    cases xy with
    | inl x => use NominalSet.supp 𝔸 x; exact ⟨supp_inl' 𝔸, NominalSet.supp_finite 𝔸⟩
    | inr y => use NominalSet.supp 𝔸 y; exact ⟨supp_inr' 𝔸, NominalSet.supp_finite 𝔸⟩

@[simp] lemma supp_inl {x : X} : NominalSet.supp 𝔸 (.inl x : X ⊕ Y) = NominalSet.supp 𝔸 x := by
  ext b; constructor
  · apply NominalSet.support_subset_supp
    · exact supp_inl' 𝔸
    · exact NominalSet.supp_finite 𝔸
  · intro hb S hS
    suffices PermSet.HasSupport 𝔸 x S by apply hb; exact ⟨this, hS.out.right⟩
    intro π hπ
    suffices π ∈ PermSet.fix' 𝔸 (.inl x : X ⊕ Y) by
      simp only [PermSet.fix', smul_inl, inl.injEq, Set.mem_setOf_eq] at this
      exact this
    exact (hS.out.left hπ).out

@[simp] lemma supp_inr {y : Y} : NominalSet.supp 𝔸 (.inr y : X ⊕ Y) = NominalSet.supp 𝔸 y := by
  ext b; constructor
  · apply NominalSet.support_subset_supp
    · exact supp_inr' 𝔸
    · exact NominalSet.supp_finite 𝔸
  · intro hb S hS
    suffices PermSet.HasSupport 𝔸 y S by apply hb; exact ⟨this, hS.out.right⟩
    intro π hπ
    suffices π ∈ PermSet.fix' 𝔸 (.inr y : X ⊕ Y) by
      simp only [PermSet.fix', smul_inr, inr.injEq, Set.mem_setOf_eq] at this
      exact this
    exact (hS.out.left hπ).out

@[simp] lemma fresh_inl {a : 𝔸} {x : X} : a # (.inl x : X ⊕ Y) ↔ a # x := by
  simp [NominalSet.IsFresh]

@[simp] lemma fresh_inr {a : 𝔸} {y : Y} : a # (.inr y : X ⊕ Y) ↔ a # y := by
  simp [NominalSet.IsFresh]

instance instInlEquivariant : Function.Equivariant 𝔸 (Sum.inl : X → X ⊕ Y) := by
  constructor; simp

instance instInrEquivariant : Function.Equivariant 𝔸 (Sum.inr : Y → X ⊕ Y) := by
  constructor; simp

end Sum

end Coproducts

/-!
## Quotient Types

If ≈ is an equivariant relation on X × X, then X/≈ is a nominal set with `supp [x] = ⋂₀ {supp y | y ∈ [x]}`.
-/
section Quotient

namespace Quotient

variable {X : Type} [PermSet 𝔸 X] {setoid : Setoid X} [Rel.Equivariant 𝔸 setoid.r]

instance : SMul (FPerm 𝔸) (Quotient setoid) where
  smul π := Quotient.lift (fun x => ⟦π • x⟧) (by
    intro x y hxy
    simp only [Quotient.eq]
    rw [← Rel.equivariant 𝔸 setoid.r]
    exact hxy)

@[simp] lemma smul_quotient (π : FPerm 𝔸) (x : X) : (π • Quotient.mk setoid x) = ⟦π • x⟧ := rfl

instance : PermSet 𝔸 (Quotient setoid) where
  one_smul := by
    apply Quotient.ind
    intro x
    simp
  mul_smul := by
    intro π τ
    apply Quotient.ind
    simp [mul_smul]

instance instMkEquivariant : Function.Equivariant 𝔸 (Quotient.mk setoid) := by
  constructor; simp

def lift_equivariant
    {Y : Type} [PermSet 𝔸 Y] (f : X → Y) [Function.Equivariant 𝔸 f]
    (h : ∀ (a b : X), a ≈ b → f a = f b) :
    Function.Equivariant 𝔸 (Quotient.lift f h) := by
  constructor; intro x π
  obtain ⟨x', hx'⟩ := x.exists_rep
  simp [← hx', Function.equivariant 𝔸 f]

variable {X : Type} [NominalSet 𝔸 X] {setoid : Setoid X} [Rel.Equivariant 𝔸 setoid.r]

instance : NominalSet 𝔸 (Quotient setoid) :=
  NominalSet.ofEquivariantSurjective 𝔸 (Quotient.mk setoid) mk_surjective

lemma supp {x : X} : NominalSet.supp 𝔸 (Quotient.mk setoid x) = {a | ∀ (x' : X), x' ≈ x → a ∈ NominalSet.supp 𝔸 x'} := by
  ext a
  constructor
  · intro ha x' hx'
    rw [← Quotient.sound hx'] at ha
    have hsub : NominalSet.supp 𝔸 (Quotient.mk setoid x') ⊆ NominalSet.supp 𝔸 x' :=
      Function.Equivariant.supp_subset 𝔸
    exact hsub ha
  · contrapose; intro ha
    simp only [Set.mem_setOf_eq, Classical.not_imp, not_forall]
    have hfin₁ : (NominalSet.supp 𝔸 (Quotient.mk setoid x)).Finite := NominalSet.supp_finite 𝔸
    have hfin₂ : (NominalSet.supp 𝔸 x).Finite := NominalSet.supp_finite 𝔸
    -- Show that we can find a fresh c for which (a c) • ⟦x⟧ = ⟦x⟧
    obtain ⟨c, hc⟩ := Infinite.exists_notMem_finset (hfin₁.toFinset ∪ hfin₂.toFinset)
    simp only [Finset.mem_union, Set.Finite.mem_toFinset, not_or] at hc
    use ⸨a, c⸩ • x; constructor
    · apply Quotient.exact
      rw [← smul_quotient]
      exact NominalSet.fresh_swap_fix 𝔸 ha hc.left
    · rw [Function.equivariant 𝔸 (NominalSet.supp 𝔸), Set.mem_smul_set_iff_inv_smul_mem]
      simp only [FPerm.swap_inv, PermSet.atom_swap, Equiv.swap_apply_left]
      exact hc.right

/-- If an atom is not in the support of the equivalence class, one can always pick a representative where it is not in the least support. -/
lemma supp_pick_fresh {a : 𝔸} {x : X} : a ∉ NominalSet.supp 𝔸 (Quotient.mk setoid x) → ∃ (y : X), y ≈ x ∧ a ∉ NominalSet.supp 𝔸 y := by
  contrapose
  simp [not_exists, not_and, not_not, supp]

end Quotient

end Quotient

/-!
## Subtypes

If X is a nominal set, Y is a subtype of X, and π • x ∈ Y for all π ∈ FPerm 𝔸 and x ∈ X, then Y is
also a nominal set.
-/
section Subtype

namespace Subtype

variable {X : Type} [PermSet 𝔸 X] {S : Set X} [Set.Equivariant 𝔸 S]

instance : PermSet 𝔸 S where
  smul π x := ⟨π • x.val, by
    rw [← Set.equivariant 𝔸 S]
    exact x.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 𝔸) (x : S) : (π • x).val = π • x.val := rfl

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

variable {X : Type} [NominalSet 𝔸 X] {S : Set X} [Set.Equivariant 𝔸 S]

instance : NominalSet 𝔸 S where
  fin_supp x := by
    use NominalSet.supp 𝔸 x.val
    constructor
    · intro π hπ
      apply Subtype.eq
      exact NominalSet.supp_is_support 𝔸 hπ
    · exact NominalSet.supp_finite 𝔸

@[simp] lemma supp {x : S} : NominalSet.supp 𝔸 x = NominalSet.supp 𝔸 x.val := by
  symm
  apply Function.Equivariant.supp_eq_if_inj
  exact val_injective

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

end Subtype

end Subtype