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_neq (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.discrete (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

@[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.discrete 𝔸 Unit

/-- We consider Bool a discrete nominal set. -/
instance : NominalSet.Discrete 𝔸 Bool := NominalSet.discrete 𝔸 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

lemma fst_equivariant : Equivariant 𝔸 (Prod.fst : X × Y → X) := by
  simp [Equivariant]

lemma snd_equivariant : Equivariant 𝔸 (Prod.snd : X × Y → Y) := by
  simp [Equivariant]

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]

lemma inl_equivariant : Equivariant 𝔸 (Sum.inl : X → X ⊕ Y) := by
  simp [Equivariant]

lemma inr_equivariant : Equivariant 𝔸 (Sum.inr : Y → X ⊕ Y) := by
  simp [Equivariant]

end Sum

end Coproducts