package gogo

// Set in which every element is identified
//
// Set is a pointer itself.
// if you want test someone is or not in the Set,
// you can just test set[elem].
// if you want know the length of Set,
// just len(set).
type Set[T comparable] map[T]bool

// NewSet make an empty set or for some elements.
//
// eg.
//     NewSet[int]() // new empty set
//     NewSet[int](1,2,3) // new for three elements
//     NewSet[int](intSlice...) // new from []int
func NewSet[T comparable](xs ...T) Set[T] {
	r := map[T]bool{}
	for _, e := range xs {
		r[e] = true
	}
	return r
}

// NewSetFromSet new from another set
func NewSetFromSet[T comparable](s Set[T]) Set[T] {
	r := NewSet[T]()
	for e := range s {
		r[e] = true
	}
	return r
}

// Add add some elements
func (s Set[T]) Add(xs ...T) {
	for _, x := range xs {
		s[x] = true
	}
}

// Slice get a slice version of set
func (s Set[T]) Slice() (r []T) {
	for x := range s {
		r = append(r, x)
	}
	return
}

// Del delete an elements from set
func (s Set[T]) Del(xs ...T) {
	for _, e := range xs {
		delete(s, e)
	}
}

// Belong whether s <- big
func Belong[T comparable](small, big Set[T]) bool {
	for e := range small {
		if !big[e] {
			return false
		}
	}
	return true
}

// Equal tests if a b are equal
func Equal[T comparable](a, b Set[T]) bool {
	return Belong(a, b) && Belong(b, a)
}

// Union add all elements from other to the set
func (s Set[T]) Union(other Set[T]) {
	for e := range other {
		s[e] = true
	}
}

// Union combine all elements
func Union[T comparable](ss ...Set[T]) Set[T] {
	r := NewSet[T]()
	for _, s := range ss {
		r.Union(s)
	}
	return r
}

// Inter return elements in both two sets
func Inter[T comparable](a, b Set[T]) Set[T] {
	if len(a) > len(b) {
		a, b = b, a
	}
	r := NewSet[T]()
	for e := range a {
		if b[e] {
			r[e] = true
		}
	}
	return r
}

// Sub remove all elements in other set
func (s Set[T]) Sub(other Set[T]) Set[T] {
	for b := range other {
		delete(s, b)
	}
	return s
}

// Sub remove all elements (which is in b) from a
func Sub[T comparable](a, b Set[T]) Set[T] {
	r := NewSet[T]()
	for e := range a {
		if !b[e] {
			r[e] = true
		}
	}
	return r
}