#
# small sets
#

export SmallBitSet, bits, delete, pop, push, exchange, first_as_set, smallbitsettype

using Base: hasfastin

import Base: show, ==, hash, copy, convert,
    empty, isempty, in, first, last, iterate,
    length, issubset, ⊊, maximum, minimum, extrema,
    union, intersect, setdiff, symdiff, filter,
    count, any, all, replace

isinteger(x) = x isa Number && Base.isinteger(x)

"""
    SmallBitSet{U<:Unsigned} <: AbstractSet{Int}

    SmallBitSet{U}([iter])
    SmallBitSet([iter])

`SmallBitSet{U}` is an immutable set that can hold integers between `1` and the bit length of `U`.
The argument `iter` can be any iterator giving such values. If `U` is omitted, then `UInt` is taken.
Called without an argument, the constructor returns an empty set.

All non-mutating functions for sets are supported. The non-mutating analogs
[`push`](@ref push(::SmallBitSet, ::Vararg{Any})), [`pop`](@ref pop(::SmallBitSet)) and
[`delete`](@ref) of the corresponding `!`-functions are also provided.

# Examples
```jldoctest
julia> SmallBitSet((1, 3, 5))
SmallBitSet{$UInt} with 3 elements:
  1
  3
  5

julia> s = SmallBitSet{UInt8}(2*k for k in 1:3)
SmallBitSet{UInt8} with 3 elements:
  2
  4
  6

julia> first(s), last(s)
(2, 6)

julia> 4.0 in s
true

julia> SmallBitSet{UInt32}(s)
SmallBitSet{UInt32} with 3 elements:
  2
  4
  6
```
"""
struct SmallBitSet{U<:Unsigned} <: AbstractSet{Int}
    mask::U
    global _SmallBitSet(mask::U) where U = new{U}(mask)
end

function show(io::IO, s::SmallBitSet{U}) where U
    print(io, "SmallBitSet")
    if bitsize(U) > bitsize(UInt) || get(io, :typeinfo, Any) != SmallBitSet{U}
        print(io, '{', U, '}')
    end
    print(io, "([")
    join(io, s, ", ")
    print(io, "])")
end

MapStyle(::Type{SmallBitSet}, ::Type{OneTo{T}}) where T = iffasttypes(StrictStyle(), T)
MapStyle(::Type{SmallBitSet{U}}, ::Type{OneTo{T}}) where {U,T} = iffasttypes(StrictStyle(), U, T)

==(s::SmallBitSet, t::SmallBitSet) = s.mask == t.mask

copy(s::SmallBitSet) = s

"""
    bits(s::SmallBitSet{U}) where U -> U

Return the bit mask used internally to store the elements of the set `s`.

See also [`convert(::Type{SmallBitSet}, ::Integer)`](@ref).
"""
bits(s::SmallBitSet) = s.mask

"""
    capacity(::Type{<:SmallBitSet}) -> Int
    capacity(s::SmallBitSet) -> Int

Return the largest number that the given set or `SmallBitSet` type can store.
"""
capacity(::Type{<:SmallBitSet}),
capacity(::SmallBitSet)

capacity(::Type{SmallBitSet{U}}) where U = bitsize(U)
capacity(::Type{SmallBitSet}) = capacity(SmallBitSet{UInt})

"""
    smallbitsettype(::Val{N}) where N -> Type{<:SmallBitSet}

Returns a type `SmallBitSet{U}` where `U <: Unsigned` has at least `N` bits.

# Examples
```jldoctest
julia> smallbitsettype(Val(30))
SmallBitSet{UInt32}

julia> smallbitsettype(Val(200))
SmallBitSet{BitIntegers.UInt256}
```
"""
smallbitsettype(::Val{N}) where N = SmallBitSet{uinttype(Val(N))}

"""
    fasthash(s::SmallBitSet [, h0::UInt]) -> UInt

Return a hash for `s` that can be computed fast. This hash is consistent across
all `SmallBitSet`s, but it is not compatible with the `hash` used for sets.

See also `Base.hash`.

# Examples
```jldoctest
julia> s = SmallBitSet(1:3);

julia> fasthash(s)
0x828a4cc485149963

julia> fasthash(s) == hash(s)
false

julia> t = SmallBitSet{UInt16}(s);

julia> fasthash(s) == fasthash(t)
true
```
"""
fasthash(s::SmallBitSet, h0::UInt) = hash(bits(s), h0)

"""
    convert(::Type{SmallBitSet{U}}, mask::Integer) where U -> SmallBitSet{U}
    convert(::Type{SmallBitSet}, mask::Integer) -> SmallBitSet{UInt}

Convert a bit mask to a `SmallBitSet` of the given type. This is the inverse operation to `bits`.

See also [`bits`](@ref).

# Examples
```jldoctest
julia> s = SmallBitSet{UInt16}([1, 5, 6]);

julia> u = bits(s)
0x0031

julia> convert(SmallBitSet, u)
SmallBitSet{$UInt} with 3 elements:
  1
  5
  6
```
"""
convert(::Type{SmallBitSet{U}}, ::Integer) where U <: Unsigned,
convert(::Type{SmallBitSet}, ::Integer)

convert(::Type{SmallBitSet{U}}, mask::Integer) where U = _SmallBitSet(U(mask))

convert(::Type{SmallBitSet}, mask::Integer) = convert(SmallBitSet{UInt}, mask)

@inline function _push(mask::U, iter) where U
    for n in iter
        @boundscheck if !isinteger(n) || n <= 0 || n > bitsize(U)
                error(LazyString(SmallBitSet{U}, " can only contain integers between 1 and ", bitsize(U)))
            end
        k = Integer(n)
        mask |= unsafe_shl(one(U), k-one(k))
    end
    _SmallBitSet(mask)
end

@propagate_inbounds SmallBitSet(args...) = SmallBitSet{UInt}(args...)

@inline function SmallBitSet{U}(s::SmallBitSet) where U
    mask = s.mask % U
    @boundscheck if mask != s.mask
        error(LazyString(SmallBitSet{U}, " can only contain integers between 1 and ", bitsize(U)))
    end
    _SmallBitSet(mask)
end

SmallBitSet{U}() where U = _SmallBitSet(zero(U))

@propagate_inbounds SmallBitSet{U}(iter) where U = _push(zero(U), iter)

@inline function SmallBitSet{U}(r::AbstractUnitRange{<:Integer}) where U
    r0, r1 = first(r), last(r)
    @boundscheck if r0 <= r1 && (r0 < 1 || r1 > bitsize(U))
        error(LazyString(SmallBitSet{U}, " can only contain integers between 1 and ", bitsize(U)))
    end
    m = one(U) << unsigned(r1-(r0-1)) - one(U)
    _SmallBitSet(unsafe_shl(m, r0-1))
end

isempty(s::SmallBitSet) = iszero(bits(s))

"""
    empty(s::S) where S <: SmallBitSet -> S

Return an empty `SmallBitSet` of the same type as `s`.
"""
empty(s::SmallBitSet)

empty(s::S) where S <: SmallBitSet = S()

default(::Type{S}) where S <: SmallBitSet = S()

length(s::SmallBitSet) = count_ones(bits(s))

# from https://discourse.julialang.org/t/faster-way-to-find-all-bit-arrays-of-weight-n/113658/12
iterate(s::SmallBitSet, m = bits(s)) =
    iszero(m) ? nothing : (trailing_zeros(m)+1, blsr(m))

@inline function first(s::SmallBitSet)
    @boundscheck isempty(s) && error("collection must be non-empty")
    trailing_zeros(bits(s))+1
end

@inline function last(s::SmallBitSet)
    @boundscheck isempty(s) && error("collection must be non-empty")
    top_set_bit(bits(s))
end

"""
    first_as_set(s::S) where S <: SmallBitSet -> S

Returns the smallest element of `s` as a singleton.
If `s` is empty, then an empty set is returned.

See also `Base.first`.

# Examples
```jldoctest
julia> s = SmallBitSet{UInt8}(3:5);

julia> first_as_set(s)
SmallBitSet{UInt8} with 1 element:
  3

julia> first_as_set(empty(s))
SmallBitSet{UInt8}([])
```
"""
first_as_set(s::SmallBitSet) = _SmallBitSet(blsi(s.mask))

function minimum(s::SmallBitSet; init = missing)
    if !isempty(s)
        @inbounds first(s)
    elseif init !== missing
        init
    else
        error("collection must be non-empty unless `init` is given")
    end
end

function maximum(s::SmallBitSet; init = missing)
    if !isempty(s)
        @inbounds last(s)
    elseif init !== missing
        init
    else
        error("collection must be non-empty unless `init` is given")
    end
end

extrema(v::SmallBitSet; init::Tuple{Any,Any} = (missing, missing)) =
    (minimum(v; init = init[1]), maximum(v; init = init[2]))

# hasfastin(::Type{<:SmallBitSet}) = true
# this is the default for AbstractSet

function in(n, s::SmallBitSet{U}) where U <: Unsigned
    isinteger(n) && 1 <= n <= bitsize(U) && begin
        k = Integer(n)
        !iszero(s.mask & unsafe_shl(one(U), k-one(k)))
    end
end

issubset(s::SmallBitSet, t::SmallBitSet) = isempty(setdiff(s, t))

⊊(s::SmallBitSet, t::SmallBitSet) = issubset(s, t) && s != t

"""
    push(s::S, xs...) where S <: SmallBitSet -> S

Return the `SmallBitSet` obtained from `s` by adding the other arguments `xs`.

See also `Base.push!`, `BangBang.push!!`, [`$(@__MODULE__).unsafe_push`](@ref).
"""
@propagate_inbounds push(s::SmallBitSet, ns...) = _push(s.mask, ns)

"""
    $(@__MODULE__).unsafe_push(s::S, x::Integer) where S <: SmallBitSet -> S

A faster, but unsafe version of `push` where `x` is assumed to lie in the range `1:capacity(S)`.
The value `0` is also allowed for `x` to indicate that no element is added.

See also [`push`](@ref push(::SmallBitSet, ::Any)), [`capacity`](@ref capacity(::SmallBitSet)).

Examples
```jldoctest
julia> using $(@__MODULE__): unsafe_push

julia> unsafe_push(SmallBitSet(1:2), 5)
SmallBitSet{$UInt} with 3 elements:
  1
  2
  5

julia> unsafe_push(SmallBitSet(1:2), 0)
SmallBitSet{$UInt} with 2 elements:
  1
  2
```
"""
function unsafe_push(s::SmallBitSet{U}, n::Integer) where U
    m = ifelse(iszero(n), zero(U), unsafe_shl(one(U), n-1))
    _SmallBitSet(s.mask | m)
end

MapStyle(::typeof(unsafe_push), ::Type{SmallBitSet{U}}, ::Type{T}) where {U <: Unsigned, T <: Integer} = iffasttypes(StrictStyle(), U, T)

"""
    pop(s::S) where S <: SmallBitSet -> Tuple{S, Int}

Return the pair `(t, x)` where `x` is the smallest element from `s` and
`t` is the set `s` with `x` deleted. The set `s` must be non-empty.

See also `Base.pop!`, `BangBang.pop!!`.
"""
@inline function pop(s::SmallBitSet)
    @boundscheck isempty(s) && error("collection must be non-empty")
    n = last(s)
    unsafe_delete(s, n), n
end

"""
    pop(s::S, x) where S <: SmallBitSet -> Tuple{S, Int}

Return the pair `(t, Int(x))` where `t` is the set `s` with `x` deleted.
The set `s` must contain `x`.

See also `Base.pop!`, `BangBang.pop!!`.
"""
@inline function pop(s::SmallBitSet, n)
    @boundscheck n in s || keyerror(n)
    unsafe_delete(s, unsafe_int(n)), unsafe_int(n)
end

"""
    pop(s::S, x, default::T) where S <: SmallBitSet -> Tuple{S, Union{Int,T}}

If `s` contains `x`, return the pair `(t, Int(x))` where `t` is the set `s` with `x` deleted.
Otherwise return `(s, default)`

See also `Base.pop!`, `BangBang.pop!!`.
"""
function pop(s::SmallBitSet, n, default)
    n in s ? (unsafe_delete(s, unsafe_int(n)), unsafe_int(n)) : (s, default)
end

"""
    delete(s::S, x) where S <: SmallBitSet -> S

If `s` contains `x`, return the set obtained by deleting that element.
Otherwise return `s`.

See also `Base.delete!`, `BangBang.delete!!`, [`$(@__MODULE__).unsafe_delete`](@ref).
"""
function delete(s::SmallBitSet{U}, n) where U
    if isinteger(n)
        k = Integer(n)
        m = one(U) << (k-one(k))
        _SmallBitSet(s.mask & ~m)
    else
        s
    end
end

MapStyle(::typeof(delete), ::Type{SmallBitSet{U}}, ::Type{T}) where {U <: Unsigned, T <: Integer} = iffasttypes(StrictStyle(), U, T)

"""
    $(@__MODULE__).unsafe_delete(s::S, x::Integer) where S <: SmallBitSet -> S

A faster, but unsafe version of `delete` where `x` is assumed to lie in the range `1:capacity(S)`.

See also [`delete`](@ref delete(::SmallBitSet, ::Any)), [`capacity`](@ref capacity(::SmallBitSet)).
"""
function unsafe_delete(s::SmallBitSet{U}, n::Integer) where U
    m = unsafe_shl(one(U), n-1)
    _SmallBitSet(s.mask & ~m)
end

MapStyle(::typeof(unsafe_delete), ::Type{SmallBitSet{U}}, ::Type{T}) where {U <: Unsigned, T <: Integer} = iffasttypes(StrictStyle(), U, T)

@inline function replace(s::SmallBitSet{U}, ps::Vararg{Pair,N}) where {U,N}
    @boundscheck all(ps) do p
        p[1] isa Integer && p[2] isa Integer && 1 <= minimum(p) && maximum(p) <= bitsize(U)
    end || error(LazyString(SmallBitSet{U}, " can only contain integers between 1 and ", bitsize(U)))

    v, u = foldl(ps; init = (bits(s), zero(U))) do (v, u), p
        m = v & unsafe_shl(U(1), p[1]-1)
        v & ~m, u | bitrotate(m, p[2]-p[1])
    end
    convert(SmallBitSet{U}, v | u)
end

"""
    exchange(s::S, i::Integer, j::Integer) where S <: SmallBitSet -> S

Return the set `s` with the element `i`, if present, replaced by `j` and vice versa.
If `i` equals `j`, then the set is not modified.

This function is faster than the equivalent `replace(s, i => j, j => i)`.

See also [`invreplace`](@ref), `Base.replace`.

# Examples
```jldoctest
julia> s = SmallBitSet((1, 2)); exchange(s, 1, 2)
SmallBitSet{$UInt} with 2 elements:
  1
  2

julia> s = SmallBitSet((1, 2)); exchange(s, 2, 3)
SmallBitSet{$UInt} with 2 elements:
  1
  3

julia> s = SmallBitSet((1, 2)); exchange(s, 3, 4)
SmallBitSet{$UInt} with 2 elements:
  1
  2

julia> s = SmallBitSet((1, 2)); exchange(s, 1, 1)
SmallBitSet{$UInt} with 2 elements:
  1
  2
```
"""
@inline function exchange(s::SmallBitSet{U}, i::Integer, j::Integer) where U
    @boundscheck (1 <= i <= bitsize(U) && 1 <= j <= bitsize(U)) ||
        error(LazyString(SmallBitSet{U}, " can only contain integers between 1 and ", bitsize(U)))
    t = unsafe_shl(one(U), i-1) ⊻ unsafe_shl(one(U), j-1) # xor needed for the case i == j
    m = bits(s)
    convert(SmallBitSet{U}, ifelse(iszero(m & t) || m & t == t, m, m ⊻ t))
end

"""
    any(f::Function, s::SmallBitSet{U}; [style::MapStyle]) where U
    all(f::Function, s::SmallBitSet{U}; [style::MapStyle]) where U
    count(f, s::SmallBitSet{U}; init = 0, [style::MapStyle]) where U
    filter(f, s::SmallBitSet{U}; [style::MapStyle]) where U

With a `SmallBitSet` `s` as second argument, these functions accept
the additional keyword argument `style`. If it equals `LazyStyle()`, then the
function `f` is only evaluated for elements of `s`. For any
other value of `style`, `f` is evaluated on all numbers between `1` and the
bit size of `U`. This is often faster for simple functions.

As discussed under `MapStyle`, the default value for `style` is based on a list
of known functions.

See also [`$(@__MODULE__).MapStyle`](@ref).
"""
any(::Function, ::SmallBitSet),
all(::Function, ::SmallBitSet),
count(::Any, ::SmallBitSet),
filter(::Any, ::SmallBitSet)

filter(f, s::SmallBitSet; style::MapStyle = MapStyle(f, Int)) = smallbitset_filter(style, f, s)

function smallbitset_filter(::MapStyle, f::F, s::SmallBitSet{U}) where {F,U}
    N = bitsize(U)
    v = map(f, FixedVector{N}(1:N))
    eltype(v) == Bool || error("given function must return Bool values")
    _SmallBitSet((bits(v) % U) & bits(s))
end

function smallbitset_filter(::LazyStyle, f::F, s::SmallBitSet) where F
    m = bits(s)
    q = zero(m)
    while !iszero(m)
        p = blsr(m)
        n = trailing_zeros(m)+1
        if f(n)
            q |= m ⊻ p
        end
        m = p
    end
    _SmallBitSet(q)
end

function filter(f::Fix2{typeof(<=),T}, s::SmallBitSet{U}) where {T <: Real, U <: Unsigned}
    k = T <: Integer ? f.x : round(Int, f.x, RoundDown)
    s ∩ @inbounds SmallBitSet{U}(1:min(k, bitsize(U)))
end

function filter(f::Fix2{typeof(<),T}, s::SmallBitSet{U}) where {T <: Real, U <: Unsigned}
    k = T <: Integer ? f.x : round(Int, f.x, RoundUp)
    setdiff(s, @inbounds SmallBitSet{U}(max(k, 1):bitsize(U)))
end

function filter(f::Fix2{typeof(>=),T}, s::SmallBitSet{U}) where {T <: Real, U <: Unsigned}
    k = T <: Integer ? f.x : round(Int, f.x, RoundUp)
    s ∩ @inbounds SmallBitSet{U}(max(k, 1):bitsize(U))
end

function filter(f::Fix2{typeof(>),T}, s::SmallBitSet{U}) where {T <: Real, U <: Unsigned}
    k = T <: Integer ? f.x : round(Int, f.x, RoundDown)
    setdiff(s, @inbounds SmallBitSet{U}(1:min(k, bitsize(U))))
end

function any(f::F, s::SmallBitSet; style::MapStyle = MapStyle(f, Int)) where F <: Function
    if style isa LazyStyle
        invoke(any, Tuple{F, AbstractSet}, f, s)
    else
        !isempty(filter(f, s; style))
    end
end

function all(f::F, s::SmallBitSet; style::MapStyle = MapStyle(f, Int)) where F <: Function
    if style isa LazyStyle
        invoke(all, Tuple{F, AbstractSet}, f, s)
    else
        filter(f, s; style) == s
    end
end

count(f, s::SmallBitSet; init = 0, style::MapStyle = MapStyle(f, Int)) = smallbitset_count(style, f, s, init)

smallbitset_count(::LazyStyle, f::F, s, init) where F = invoke(count, Tuple{Any, AbstractSet}, f, s; init)

smallbitset_count(::MapStyle, f::F, s, init::T) where {F,T} = init + (length(filter(f, s)) % T)

union(s::SmallBitSet, t::SmallBitSet) = _SmallBitSet(s.mask | t.mask)

union(s::SmallBitSet, ts::SmallBitSet...) = foldl(union, ts; init = s)

intersect(s::SmallBitSet{U}, t::SmallBitSet) where U <: Unsigned = _SmallBitSet(s.mask & (t.mask % U))

function intersect(s::SmallBitSet{U}, t) where U <: Unsigned
    u = _SmallBitSet(zero(U))
    for n in (hasfastin(t) ? s : t)
        if n in (hasfastin(t) ? t : s)
            @inbounds u = push(u, n)
        end
    end
    u
end

intersect(s::SmallBitSet, ts...) = foldl(intersect, ts; init = s)

setdiff(s::SmallBitSet{U}, t::SmallBitSet) where U <: Unsigned = _SmallBitSet(s.mask & ~(t.mask % U))

function setdiff(s::SmallBitSet, t)
    if hasfastin(t)
        u = s
        for n in s
            if n in t
                u = delete(u, n)
            end
        end
        return u
    else
        foldl(delete, t; init = s)
    end
end

setdiff(s::SmallBitSet, ts...) = foldl(setdiff, ts; init = s)

symdiff(s::SmallBitSet, t::SmallBitSet) = _SmallBitSet(s.mask ⊻ t.mask)

symdiff(s::SmallBitSet, ts::SmallBitSet...) = foldl(symdiff, ts; init = s)

Random.rand(rng::AbstractRNG, ::SamplerType{SmallBitSet{U}}) where U <: Unsigned = _SmallBitSet(rand(rng, U))
Random.rand(rng::AbstractRNG, ::SamplerType{SmallBitSet}) = rand(rng, SmallBitSet{UInt})

# ordering and sorting

using Base: Order, Sort

"""
    isless(s::SmallBitSet, t::SmallBitSet) -> Bool

Compares `s` and `t` in the colexicographic ordering.

Currently, sorting with the default ordering given by `isless` is significantly
faster than specifiying one of the other orderings (including `Colex`).

See also [`$(@__MODULE__).Colex`](@ref).
"""
Base.isless(s::SmallBitSet, t::SmallBitSet) = Order.lt(Colex, s, t)

"""
    $(@__MODULE__).Colex <: Base.Order.Ordering

The colexicographic ordering for `SmallBitSet`.

See also [`$(@__MODULE__).Lex`](@ref), [`$(@__MODULE__).GradedColex`](@ref).
"""
const Colex = Order.ord(isless, bits, false)

"""
    $(@__MODULE__).GradedColex <: Base.Order.Ordering

The graded colexicographic ordering for `SmallBitSet`.
A smaller set either has smaller size,
or for equal size it is smaller in the colexicographic ordering.

See also [`$(@__MODULE__).GradedLex`](@ref), [`$(@__MODULE__).Colex`](@ref).
"""
const GradedColex = Order.ord(identity, false) do s::SmallBitSet, t::SmallBitSet
    (length(s) < length(t)) | ((length(s) == length(t)) & Order.lt(Colex, s, t))
end

"""
    $(@__MODULE__).Lex <: Base.Order.Ordering

The lexicographic ordering for `SmallBitSet`. This is the slowest among the available orderings.

See also [`$(@__MODULE__).Colex`](@ref), [`$(@__MODULE__).GradedLex`](@ref).
"""
const Lex = Order.ord(bits, false) do x::Unsigned, y::Unsigned
    a = blsi(x ⊻ y)
    b = ~blsmsk(x ⊻ y)
    (!iszero(x & a) & !iszero(y & b)) | (iszero(x & b) & !iszero(y & a))
end

"""
    $(@__MODULE__).GradedLex <: Base.Order.Ordering

The graded lexicographic ordering for `SmallBitSet`.
A smaller set either has smaller size,
or for equal size it is smaller in the lexicographic ordering.

See also [`$(@__MODULE__).Lex`](@ref), [`$(@__MODULE__).GradedColex`](@ref).
"""
const GradedLex = Order.ord(identity, false) do s::SmallBitSet, t::SmallBitSet
    u = first_as_set(symdiff(s, t))
    (length(s) < length(t)) | ((length(s) == length(t)) & !isempty(s ∩ u))
end

# The following definitions don't work for custom orderings, see julia#60184
Sort.uint_map(s::SmallBitSet, ::Order.ForwardOrdering) = bits(s)
Sort.uint_unmap(::Type{SmallBitSet{U}}, u::U, ::Order.ForwardOrdering) where U <: Unsigned = _SmallBitSet(u)
Sort.UIntMappable(::Type{SmallBitSet{U}}, ::Order.ForwardOrdering) where U <: Unsigned = U

# HWType conversion

from_hwvalue(::Type{SmallBitSet{U}}, x::U) where U <: Unsigned = _SmallBitSet(x)
