"""
    Data

This module contains a collection of [`Possibility`](@ref) types, as well as some
useful utility functions for constructing them.

These currently include:

 * [`Integers{T}`](@ref), for producing integer types (except `BigInt`)
 * [`Floats{F}`](@ref), for producing floating point types (except `BigFloat`)
 * [`Booleans`](@ref), for producing booleans
 * [`Pairs`](@ref), for producing pairs of values `a => b`
 * [`Vectors`](@ref), for producing `Vector`s out of `Possibility`
 * [`Dicts`](@ref), for producing `Dict`s out of a `Possibility` for keys and one for values
 * [`AsciiCharacters`](@ref), for producing `Char`s that are `isascii`
 * [`Characters`](@ref), for producing all well-formed `Char`s, including invalid Unicode
 * [`UnicodeCharacters`](@ref), for producing ALL `Char`s, including invalid and malformed Unicode
 * [`Text`](@ref), for producing `String`s from `Possibility{Char}`
 * [`SampledFrom`](@ref), for producing a value from a given collection
 * [`Satisfying`](@ref), for filtering the values a `Possibility` produces through a predicate
 * [`Map`](@ref), for mapping a function over the values a `Possibility` produces
 * [`Just`](@ref), for producing a given fixed value
 * [`OneOf`](@ref), for choosing one of a number of given `Possibility` to produce from
 * [`Bind`](@ref), for binding a function that produces `Possibility` to the output of another `Possibility`
 * [`Recursive`](@ref), for creating recursive data structures using a basecase `Possibility` and a function that layers more `Possibility` around it
 * [`WeightedNumbers`](@ref), for sampling the numbers from `1:n` by giving `n` weights
 * [`WeightedSample`](@ref), for sampling a collection while biasing that sampling by a per-element weight

as well as these utility functions:

 * `map` for creating a `Map`
 * `filter` for creating a `Satisfying`
 * `|` for creating a `OneOf`
 * `bind` for creating a `Bind`
 * `recursive` for creating a `Recursive`
 * `Floats()` for producing all of `Float16`, `Float32` and `Float64` from the same `Possibility` at once
 * `BitIntegers()` for producing all of `Base.BitIntegers` from the same `Possibility` at once
"""
module Data

using Supposition
using Supposition: smootherstep, lerp, TestCase, choice!, weighted!, forced_choice!, reject
using RequiredInterfaces: @required
using StyledStrings: @styled_str
using Printf: format, @format_str

"""
    Possibility{T}

Abstract supertype for all generators.
The `T` type parameter describes the kinds of objects generated by this integrated shrinker.

Required methods:

  * `produce!(::TestCase, ::P) where P <: Possibility`

Fallback definitions:

  * `postype(::Possibility{T}) -> Type{T}`
  * `example(::Possibility{T}) -> T`
  * `==(::Possibility, ::Possibility) = false`
"""
abstract type Possibility{T} end
Base.:(==)(::Possibility, ::Possibility) = false

@required Possibility begin
    produce!(::TestCase, ::Possibility)
    Base.show(::IO, ::Possibility)
    Base.show(::IO, ::MIME"text/plain", ::Possibility)
end

"""
    |(::Possibility{T}, ::Possibility{S}) where {T,S} -> OneOf{Union{T,S}}

Combine two `Possibility` into one, sampling uniformly from either.

If either of the two arguments is a `OneOf`, the resulting object acts
as if all original non-`OneOf` had be given to `OneOf` instead.
That is, e.g. `OneOf(a, b) | c` will act like `OneOf(a,b,c)`.

See also [`OneOf`](@ref).
"""
Base.:(|)(a::Possibility, b::Possibility) = OneOf(a, b)

"""
    produce!(tc::TestCase, pos::Possibility{T}) -> T

Produces a value from the given `Possibility`, recording the required choices in the `TestCase` `tc`.

This needs to be implemented for custom `Possibility` objects, passing the given `tc` to any inner
requirements directly.

See also [`Supposition.produce!`](@ref)

!!! tip "Examples"
    You should not call this function when you have a `Possibility` and want to inspect what an object
    produced by that `Possibility` looks like - use [`example`](@ref) for that instead.
"""
function produce! end

"""
    postype(::Type{P<:Possibility})

Gives the type of objects this `Possibility` type will generate.
"""
postype(::Type{_P}) where {T, _P <: Possibility{T}} = T

"""
    postype(::P) where P <: Possibility

Gives the type of objects this `Possibility` object will generate.
"""
postype(::P) where P <: Possibility = postype(P)

"""
    Map(source::Possibility, f) <: Possibility

A `Possibility` representing mapping values from `source` through `f`.

Equivalent to calling `map(f, source)`.

The pre-calculated return type of `Map` is a *best effort* and may be wider than
necessary.

```julia-repl
julia> using Supposition

julia> makeeven(x) = (x √∑ 2) * 2

julia> pos = map(makeeven, Data.Integers{Int8}())

julia> all(iseven, example(pos, 10_000))
true
```
"""
struct Map{R, S <: Possibility, F} <: Possibility{R}
    source::S
    map::F
    Map(s::S, f::F) where {T, S <: Possibility{T}, F} = new{Base.promote_op(f, T), S, F}(s, f)
end
produce!(tc::TestCase, m::Map) = m.map(produce!(tc, m.source))

Base.:(==)(m1::Map, m2::Map) = m1.map == m2.map && m1.source == m2.source

function Base.show(io::IO, m::Map)
    print(io, Map, "(")
    show(io, m.source)
    print(io, ", ")
    show(io, m.map)
    print(io, ")")
    nothing
end

function Base.show(io::IO, ::MIME"text/plain", m::Map)
    func_str = string(m.map)
    padding = textwidth(func_str)+1
    print(io, styled"""
    {code,underline:$Map}:

        Map {code:$(m.map)}.

        $(" "^padding)‚ï≠‚îÄ‚îÄ {code:$(postype(m.source))} from {code:$(m.source)}
        {code:$func_str(x) -> $(postype(m))}
        $(" "^(padding+6))‚ï∞‚îÄ Result type""")

    obj = example(m.source)
    res = try
        Some(m.map(obj))
    catch
        nothing
    end

    if !isnothing(res)
        print(io, styled"\n\nE.g. {code:$(m.map)($obj) == $(something(res))}")
    end
end

"""
    map(f, pos::Possibility)
    map(f, pos::Vararg{Possibility})

Apply `f` to the result of calling `produce!` on `pos` (lazy mapping).

Equivalent to calling `Map(pos, f)`, or `MultiMap(pos, f)` for the `Vararg` case (where `pos` is the tuple of `Possibility`).
In the `Vararg` case, `f` is expected to take as many arguments as `Possibility` are passed to `map`.

See also [`Map`](@ref).
"""
function Base.map(f, p::Possibility, ps::Possibility...) 
    if isempty(ps)
        Map(p, f)
    else
        all_pos = (p, ps...)
        MultiMap(all_pos, f)
    end
end

struct MultiMap{R, N, S <: NTuple{N, Possibility}, F} <: Possibility{R}
    source::S
    map::F
    MultiMap(s::S, f::F) where {N, S <: NTuple{N, Possibility}, F} = new{multimap_eltype(f, S), N, S, F}(s, f)
end
produce!(tc::TestCase, mm::MultiMap) = mm.map(produce!.(Ref(tc), mm.source)...)

Base.:(==)(m1::MultiMap, m2::MultiMap) = m1.map == m2.map && m1.source == m2.source

function multimap_eltype(f, S::Type)
    # this is accessing internals again, but what can ya do :shrug:
    N = length(S.parameters)
    argtypes = foldr(Base.tuple_type_cons, ntuple(t -> Data.postype(S.parameters[t]), N); init=Tuple{})
    Base.promote_op(Base.splat(f), argtypes)
end

function Base.show(io::IO, mm::MultiMap)
    print(io, MultiMap, "(")
    show(io, mm.source)
    print(io, ", ")
    show(io, mm.map)
    print(io, ")")
    nothing
end

function Base.show(io::IO, ::MIME"text/plain", mm::MultiMap{R,N}) where {R,N}
    func_str = string(mm.map)
    padding = textwidth(func_str)
    print(io, styled"""
    {code,underline:$MultiMap}:

        MultiMap {code:$(mm.map)}.

        $(" "^padding)‚ï≠‚îÄ‚îÄ {code:$(postype.(mm.source))} from {code:$(mm.source)}
        {code:$func_str(}""")

    join(io, (styled"{code:x$i}" for i in 1:N), styled"{code:, }")

    #=
        N for the 'x'
        2*(N-1) for the ", "
        and clamps for the numbers 1..., 10..., up to 999
        if you pass more than that to `map`, the misalignment of printing is likely the least of your problems :)
    =#
    length_of_args = N + 2*(N-1)*(N>=2) + clamp(N, 0, 9) + 2*clamp(N-9, 0, 99) + 3*clamp(N-99, 0, 999)
    # length of function name + length of args + parentheses/fluff
    result_padding = padding + length_of_args + 10
    print(io, styled"""{code:) -> $(postype(mm))}
        $(" "^result_padding)‚ï∞‚îÄ Result type""")

    obj = example.(mm.source)
    res = try
        Some(mm.map(obj...))
    catch
        nothing
    end

    if !isnothing(res)
        # no need for parentheses, since we know `obj` to be a tuple of some kind which has parentheses in its printing
        print(io, styled"\n\nE.g. {code:$(mm.map)$obj == $(something(res))}")
    end
end

"""
    Satisfying(source::Possibility, pred) <: Possibility

A `Possibility` representing values from `source` fulfilling `pred`.

Equivalent to calling `filter(f, source)`.

```julia-repl
julia> using Supposition

julia> pos = filter(iseven, Data.Integers{Int8}())

julia> all(iseven, example(pos, 10_000))
true
```
"""
struct Satisfying{T, S <: Possibility{T}, P} <: Possibility{T}
    source::S
    predicate::P
    # flatten chained calls to `filter`
    Satisfying(s::Satisfying, pred) = Satisfying(s.source, pred ‚àò s.predicate)
    function Satisfying(s::S, pred::P) where {T, S <: Possibility{T}, P}
        new{T,S,P}(s, pred)
    end
end

Base.:(==)(s1::Satisfying, s2::Satisfying) = s1.predicate == s2.predicate && s1.source == s2.source

function Base.show(io::IO, s::Satisfying)
    print(io, Satisfying, "(")
    show(io, s.source)
    print(io, ", ")
    show(io, s.predicate)
    print(io, ")")
    nothing
end

function Base.show(io::IO, ::MIME"text/plain", s::Satisfying)
    obj = example(s.source)
    fulfill_pred = try
        Some(s.predicate(obj))
    catch
        nothing
    end
    print(io, styled"""
    {code,underline:$Satisfying}:

        Filter {code:$(s.source)} through the predicate {code:$(s.predicate)}.""")

    if !isnothing(fulfill_pred)
        fp = @something fulfill_pred
        isval = fp ? styled"{green:‚úî}" : styled"{red:‚ùå}"
        pred_descr = fp ? "fulfills the predicate" : "is rejected by the predicate"
        styled"E.g. {code:$obj} $pred_descr $isval"
    end
end

"""
    filter(f, pos::Possibility)

Filter the output of `produce!` on `pos` by applying the predicate
`f`.

!!! note "No stalling"
    In order not to stall generation of values, this will not
    try to produce a value from `pos` forever, but reject the
    testcase after some attempts.
"""
Base.filter(f, p::Possibility) = Satisfying(p, f)
satisfying(f, p::Possibility) = Satisfying(p, f)

function produce!(tc::TestCase, s::Satisfying)
    for _ in 1:3
        candidate = produce!(tc, s.source)
        if s.predicate(candidate)
            return candidate
        end
    end

    reject(tc)
end

"""
    Bind(source::Possibility, f)

Binds `f` to `source`, i.e., on `produce!(::Bind, ::TestCase)` this calls `produce!` on
`source`, the result of which is passed to `f`, the output of which will be used as input
to `produce!` again.

In other words, `f` takes a value `produce!`d by `source` and gives back a
`Possibility` that is then immediately `produce!`d from.

Equivalent to `bind(f, source)`.
"""
struct Bind{T, S <: Possibility{T}, M} <: Possibility{T}
    source::S
    map::M
end

Base.:(==)(b1::Bind, b2::Bind) = b1.map == b2.map && b1.source == b2.source

function Base.show(io::IO, b::Bind)
    print(io, Bind, "(")
    show(io, b.source)
    print(io, ", ")
    show(io, b.map)
    print(io, ")")
    nothing
end

function Base.show(io::IO, ::MIME"text/plain", b::Bind)
    print(io, styled"""
    {code,underline:$Bind}:

        Binds the output of {code:$(b.source)} through {code:$(b.map)} to a new {code:Possibility}.
    """)
end

"""
    bind(f, pos::Possibility)

Maps the output of `produce!` on `pos` through `f`, and calls `produce!` on
the result again. `f` is expected to take a value and return a `Possibility`.

Equivalent to calling `Bind(pos, f)`.

See also [`Bind`](@ref).
"""
bind(f, s::Possibility) = Bind(s, f)
function produce!(tc::TestCase, b::Bind)
    inner = produce!(tc, b.source)
    produce!(tc, b.map(inner))
end

## Possibilities of signed integers

"""
    Integers(minimum::T, maximum::T) <: Possibility{T <: Integer}
    Integers{T}() <: Possibility{T <: Integer}

A `Possibility` representing drawing integers from `[minimum, maximum]`.
The second constructors draws from the entirety of `T`.

Produced values are of type `T`.

```julia-repl
julia> using Supposition

julia> is = Data.Integers{Int}()

julia> example(is, 5)
5-element Vector{Int64}:
 -5854403925308846160
  4430062772779972974
    -9995351034504801
  2894734754389242339
 -6640496903289665416
```
"""
struct Integers{T<:Integer, U<:Unsigned} <: Possibility{T}
    minimum::T
    range::U
    function Integers(minimum::T, maximum::T) where T<:Integer
        minimum <= maximum || throw(ArgumentError("`minimum` must be `<= maximum`!"))
        new{T,unsigned(T)}(minimum, (maximum - minimum) % unsigned(T))
    end
    Integers{T}() where T <: Integer = new{T, unsigned(T)}(typemin(T), typemax(unsigned(T)))
end

Base.:(==)(i1::Integers{T}, i2::Integers{T}) where T = i1.minimum == i2.minimum && i1.range == i2.range

function Base.show(io::IO, i::Integers)
    print(io, Integers)
    if i.minimum == typemin(postype(i)) && i.range == typemax(i.range)
        # this was created/is equivalent to sampling everything
        print(io, "{", postype(i), "}()")
    else
        print(io, "(")
        show(io, i.minimum)
        print(io, ", ")
        show(io, (i.minimum + i.range) % postype(i))
        print(io, ")")
    end
    nothing
end

function Base.show(io::IO, ::MIME"text/plain", i::Integers)
    postype_str = styled"{code:$(postype(i))}"

    inner_str = if i.minimum == typemin(i.minimum) && i.range == typemax(i.range)
        # this was created/is equivalent to sampling everything
        styled"Produce an integer of type $postype_str"
    else
        lower = sprint(show, i.minimum % postype(i))
        upper = sprint(show, (i.minimum + i.range) % postype(i))
        styled"Produce an $postype_str from the closed interval {code:[$lower, $upper]}"
    end

    obj = styled"{code:$(sprint(show, example(i)))}"
    print(io, styled"""
    {code,underline:$Integers}:

        $inner_str

    E.g. $obj""")
    nothing
end

const BITINT_TYPES = (UInt8, Int8, UInt16, Int16, UInt32, Int32, UInt64, Int64, UInt128, Int128, )

"""
    BitIntegers() <: Possibility{$(Base.BitInteger)}

A `Possibility` for generating all possible bitintegers with fixed size.
"""
struct BitIntegers <: Possibility{Base.BitInteger} end

Base.:(==)(::BitIntegers, ::BitIntegers) = true

produce!(tc::TestCase, ::BitIntegers) =
    produce!(tc, OneOf((Integers{T}() for T in BITINT_TYPES)...))

Base.show(io::IO, ::BitIntegers) = print(io, "$BitIntegers()")

function Base.show(io::IO, ::MIME"text/plain", bi::BitIntegers)
    print(io,  styled"""
    {code,underline:$bi}:

        Produce any bitinteger of one of these types:

    """)

    join(io, (styled"        ‚àò {code:$T}" for T in BITINT_TYPES), '\n')

    o = example(bi)
    obj = styled"{code:$(sprint(show, o))}"
    print(io, styled"""


        I.e. all integer types from {code:Base} except for {code:BigInt}.

    E.g. $obj, a {code:$(typeof(o))}""")
end

function produce!(tc::TestCase, i::Integers{T}) where T
    offset = choice!(tc, i.range % UInt) % T
    return (i.minimum + offset) % T
end

function produce!(tc::TestCase, i::Integers{T}) where T <: Union{Int128, UInt128}
    # FIXME: this assumes a 64-bit architecture!
    upperbound = (i.range >> 64) % UInt
    lowerbound = i.range % UInt
    upper = choice!(tc, upperbound) % T
    lower = choice!(tc, lowerbound) % T
    offset = (upper << 64) | lower
    return (i.minimum + offset) % T
end

## Possibilities of vectors

"""
    Vectors(elements::Possibility{T}; min_size=0, max_size=10_000) <: Possibility{Vector{T}}

A `Possibility` representing drawing vectors with length `l` in `min_size <= l <= max_size`,
holding elements of type `T`.

`min_size` and `max_size` must be positive numbers, with `min_size <= max_size`.

```julia-repl
julia> using Supposition

julia> vs = Data.Vectors(Data.Floats{Float16}(); max_size=5)

julia> example(vs, 3)
3-element Vector{Vector{Float16}}:
 [9.64e-5, 9.03e3, 0.04172, -0.0003352]
 [9.793e-5, -2.893, 62.62, 0.0001961]
 [-0.007023, NaN, 3.805, 0.1943]
```
"""
struct Vectors{T, P <: Possibility{T}} <: Possibility{Vector{T}}
    elements::P
    min_size::UInt
    max_size::UInt

    function Vectors(elements::Possibility{T}; min_size=0, max_size=10_000) where T
        min_size <= max_size || throw(ArgumentError("`min_size` must be `<= max_size`!"))

        low = UInt(min_size)
        high = UInt(max_size)

        new{T,typeof(elements)}(elements, low, high)
    end
end

Base.:(==)(v1::Vectors{T,P}, v2::Vectors{T,P}) where {T, P} = v1.min_size == v2.min_size &&
                                                              v2.max_size == v2.max_size &&
                                                              v1.elements == v2.elements

function Base.show(io::IO, v::Vectors)
    print(io, Vectors, "(")
    show(io, v.elements)
    print(io, "; ")
    print(io, "min_size=", string(v.min_size; base=10), ", ")
    print(io, "max_size=", string(v.max_size; base=10), ")")
    nothing
end

function Base.show(io::IO, ::MIME"text/plain", v::Vectors)
    lengths = styled"{code:[$(Int128(v.min_size)), $(Int128(v.max_size))]}"
    print(io, styled"""
    {code,underline:$Vectors}:

        Produce a {code:Vector\{$(postype(v.elements))\}} with length in the interval $lengths.
        The elements are produced by {code:$(v.elements)}.""")
    obj = styled"""


    Elements of the vector may look like {code:$(sprint(show, example(v.elements)))}"""
    if length(obj) < 100
        print(io, obj)
    end
end

function produce!(tc::TestCase, v::Vectors{T}) where T
    result = T[]

    # it's VERY important to let the shrinker shrink this!
    # if we don't we get Invalids even when we shouldn't
    max_offset = choice!(tc, v.max_size - v.min_size)

    if tc.attempt.generation == tc.attempt.max_generation
        # if we're on the last try (should that exist)
        # guarantee that we're able to draw the maximum permissible size
        average_offset = (v.max_size√∑2 + v.min_size√∑2) - v.min_size
    else
        # otherwise, get an average according to a beta distribution
        raw_step = smootherstep(0.0, float(max(tc.attempt.max_generation√∑2, 5_000)), tc.attempt.generation)
        beta_param = lerp(0.5, 5.0, raw_step)
        average_offset = floor(UInt, max_offset*(beta_param/(beta_param+1.0)))
    end

    # give some hint to the amount of data we're going to need
    sizehint!(result, v.min_size+average_offset)

    # first, make sure we hit the minimum size
    for _ in 1:v.min_size
        push!(result, produce!(tc, v.elements))
    end

    # now for the fiddly bit to reaching `v.max_size`
    # the `min` here is important, otherwise we may _oversample_ if the
    # beta distribution drew too high after `max_offset` shrank!
    p_continue = _calc_p_continue(min(average_offset, max_offset), max_offset)

    # finally, draw with our targeted average until we're done
    for _ in 1:max_offset
        if weighted!(tc, p_continue)
            push!(result, produce!(tc, v.elements))
        else
            break
        end
    end

    return result
end

function _calc_p_continue(desired_avg, max_size)
    @assert desired_avg <= max_size "Require target <= max_size, not $desired_avg > $max_size"
    if desired_avg == max_size
        return 1.0
    end
    p_continue = 1.0 - 1.0/(1+desired_avg)
    if iszero(p_continue)
        @assert 0 <= p_continue < 1
        return p_continue
    end

    while _p_continue_to_avg(p_continue, max_size) > desired_avg
        p_continue -= 0.0001
        smallest_positive = nextfloat(0.0)
        if p_continue < smallest_positive
            p_continue = smallest_positive
            break
        end
    end

    hi = 1.0
    while desired_avg - _p_continue_to_avg(p_continue, max_size) > 0.01
        @assert 0 < p_continue < hi "Binary search failed: $p_continue, $hi"
        # this can't overflow, since the numbers are all in [0,1]
        mid = (p_continue + hi) / 2
        if _p_continue_to_avg(mid, max_size) <= desired_avg
            p_continue = mid
        else
            hi = mid
        end
    end
    @assert 0 < p_continue < 1 "Binary search faileD: $p_continue, $hi"
    @assert _p_continue_to_avg(p_continue, max_size) <= desired_avg "Found probability leads to higher-than-requested average"
    return p_continue
end

function _p_continue_to_avg(p_continue, max_size)
    p_continue >= 1 && return max_size
    return (1.0 / (1 - p_continue) - 1.0) * (1 - p_continue^max_size)
end

## Possibilities of pairs

"""
    Pairs(first::Possibility{T}, second::Possibility{S}) where {T,S} <: Possibility{Pair{T,S}}

A `Possibility` for producing `a => b` pairs. `a` is produced by `first`, `b` is produced by `second`.

```
julia> p = Data.Pairs(Data.Integers{UInt8}(), Data.Floats{Float64}());

julia> example(p, 4)
4-element Vector{Pair{UInt8, Float64}}:
 0x41 => 4.1183566661848205e-230
 0x48 => -2.2653631095108555e-119
 0x2a => -6.564396855333643e224
 0xec => 1.9330751262581671e-53
```
"""
struct Pairs{T,S} <: Possibility{Pair{T,S}}
    first::Possibility{T}
    second::Possibility{S}
end

Base.:(==)(p1::Pairs{T,S}, p2::Pairs{T,S}) where {T, S} = p1.first == p2.first && p1.second == p2.second

function Base.show(io::IO, p::Pairs)
    print(io, Pairs, "(")
    show(io, p.first)
    print(io, ", ")
    show(io, p.second)
    print(io, ")")
    nothing
end

function Base.show(io::IO, ::MIME"text/plain", p::Pairs)
    print(io, styled"""
    {code,underline:$Pairs}:

        Produce a {code:$(postype(p))}:

        ‚ï≠‚îÄ‚îÄ‚îÄ From {code:$(sprint(show, p.first))}
        {code:a => b}
             ‚ï∞‚îÄ‚îÄ From {code:$(sprint(show, p.second))}

    E.g. {code:$(sprint(show, example(p)))}""")
end

pairs(a::Possibility, b::Possibility) = Pairs(a,b)
produce!(tc::TestCase, p::Pairs) = produce!(tc, p.first) => produce!(tc, p.second)

## Possibility of Just a value

"""
    Just(value::T) <: Possibility{T}

A `Possibility` that always produces `value`.

!!! warning "Mutable Data"
    The source object given to this `Just` is not copied when `produce!` is called.
    Be careful with mutable data!

```julia-repl
julia> using Supposition

julia> three = Data.Just(3)

julia> example(three, 3)
3-element Vector{Int64}:
 3
 3
 3
```
"""
struct Just{T} <: Possibility{T}
    value::T
end

Base.:(==)(j1::Just{T}, j2::Just{T}) where T = j1.value == j2.value

Base.show(io::IO, j::Just) = print(io, Just, "(", j.value, ")")

function Base.show(io::IO, ::MIME"text/plain", j::Just)
    print(io, styled"""
    {code,underline:$Just}:

        Produce the value {code:$(sprint(show, j.value))}
    """)
end

just(t) = Just(t)
produce!(_::TestCase, j::Just) = j.value

## Possibility of Nothing

produce!(tc::TestCase, ::Nothing) = reject(tc)

## Sampling one of N

"""
    OneOf(pos::Possibility...) <: Possibility

A `Possibility` able to generate any of the examples one of the given
`Possibility` can produce. The given `Possibility` are sampled from
uniformly.

At least one `Possibility` needs to be given to `OneOf`.

`postype(::OneOf)` is inferred as a _best effort_, and may be wider than necessary.

`OneOf` can also be constructed through use of `a | b` on `Possibility`. Constructed
in this way, if either `a` or `b` is already a `OneOf`, the resulting `OneOf`
acts as if it had been given the original `Possibility` objects in the first place.
That is, `OneOf(a, b) | c` acts like `OneOf(a, b, c)`.

See also [`WeightedNumbers`](@ref) and [`WeightedSample`](@ref).

```julia-repl
julia> of = Data.OneOf(Data.Integers{Int8}(), Data.Integers{UInt8}());

julia> Data.postype(of)
Union{Int8, UInt8}

julia> ex = map(of) do i
           (i, typeof(i))
       end;

julia> example(ex)
(-83, Int8)

julia> example(ex)
(0x9f, UInt8)
```
"""
struct OneOf{X, N} <: Possibility{X}
    strats::NTuple{N, Possibility}
    function OneOf(pos::Possibility...)
        isempty(pos) && throw(ArgumentError("Need at least one `Possibility` to draw from!"))
        new{Union{postype.(pos)...}, length(pos)}(pos)
    end
end

Base.:(==)(of1::OneOf{X,N}, of2::OneOf{X,N}) where {X,N} = all(Base.splat(==), zip(of1.strats, of2.strats))

function Base.show(io::IO, of::OneOf)
    print(io, OneOf, "(")
    for s in of.strats[1:end-1]
        show(io, s)
        print(io, ", ")
    end
    show(io, of.strats[end])
    print(io, ")")
    nothing
end

function Base.show(io::IO, ::MIME"text/plain", of::OneOf)
    print(io, styled"""
    {underline,code:$OneOf}:

        Produces an element from one of the following with equal probability:

    """)

    join(io, (styled"        ‚àò {code:$pos}" for pos in of.strats), '\n')
end

function produce!(tc::TestCase, @nospecialize(of::OneOf))
    strategy = produce!(tc, SampledFrom(of.strats))
    produce!(tc, strategy)::postype(of)
end

Base.:(|)(of::OneOf, b::Possibility) = OneOf(of.strats..., b)
Base.:(|)(a::Possibility, of::OneOf) = OneOf(a, of.strats...)
Base.:(|)(a::OneOf, b::OneOf) = OneOf(a.strats..., b.strats...)

## Recursion

"""
    Recursive(base::Possibility, extend; max_layers::Int=5) <: Possibility{T}

A `Possibility` for generating recursive data structures.
`base` is the basecase of the recursion. `extend` is a function returning a
new `Possibility` when given a `Possibility`, called to recursively
expand a tree starting from `base`. The returned `Possibility` is fed
back into `extend` again, expanding the recursion by one layer.

`max_layers` designates the maximum layers `Recursive` should
keep track of. This must be at least `1`, so that at least
the base case can always be generated. Note that this implies `extend`
will be used at most `max_layers-1` times, since the base case of
the recursion will not be wrapped.

Equivalent to calling `recursive(extend, base)`.

## Examples

```julia-repl
julia> base = Data.Integers{UInt8}()

julia> wrap(pos) = Data.Vectors(pos; min_size=2, max_size=3)

julia> rec = Data.recursive(wrap, base; max_layers=3);

julia> Data.postype(rec) # the result is formatted here for legibility
Union{UInt8,
      Vector{UInt8},
      Vector{Union{UInt8, Vector{UInt8}}}
}

julia> example(rec)
0x31

julia> example(rec)
2-element Vector{Union{UInt8, Vector{UInt8}}}:
     UInt8[0xa9, 0xb4]
 0x9b

julia> example(rec)
2-element Vector{UInt8}:
 0xbd
 0x25
```
"""
struct Recursive{T,F} <: Possibility{T}
    base::Possibility
    extend::F
    inner::Possibility{T}
    function Recursive(base::Possibility, extend; max_layers::Int=5)
        max_layers < 1 && throw(ArgumentError("Must be able to produce at least the base layer!"))
        strategies = Vector{Possibility}(undef, max_layers)
        strategies[1] = base
        for layer in 2:max_layers
            prev_layers = @view strategies[1:layer-1]
            strategies[layer] = extend(OneOf(prev_layers...))
        end
        inner = OneOf(strategies...)
        new{postype(inner), typeof(extend)}(base, extend, inner)
    end
end

Base.:(==)(r1::Recursive{T,F}, r2::Recursive{T,F}) where {T,F} = r1.inner == r2.inner

function Base.show(io::IO, r::Recursive)
    print(io, Recursive, "(")
    show(io, r.base)
    print(io, ", ", r.extend, "; ")
    print(io, "max_layers=", length(r.inner.strats), ")")
end

function Base.show(io::IO, ::MIME"text/plain", r::Recursive)
    print(io, styled"""
    {code,underline:$Recursive}:

        Wrap {code:$(r.base)} up to {code:$(length(r.inner.strats)-1)} times with {code:$(r.extend)}
    """)
end

"""
    recursive(f, pos::Possibility; max_layers=5)

Recursively expand `pos` into deeper nested `Possibility` by repeatedly
passing `pos` itself to `f`. `f` returns a new `Possibility`, which is then
passed into `f` again until the maximum depth is achieved.

Equivalent to calling `Recursive(pos, f)`.

See also [`Recursive`](@ref).
"""
recursive(f, pos::Possibility; max_layers=5) = Recursive(pos, f; max_layers)

produce!(tc::TestCase, r::Recursive) = produce!(tc, r.inner)

## Possibility of Characters

"""
    Characters(;valid::Bool = false) <: Possibility{Char}

A `Possibility` of producing arbitrary well-formed `Char` instances.

!!! warning "Unicode"
    This will `produce!` ANY well-formed `Char` by default, not just valid unicode codepoints!
    Notably, this includes overlong `Char`. To only produce valid unicode codepoints, pass `valid=true` as a keyword argument.

Keyword arguments:

 * `valid`: Whether the produced `Char` must be valid, i.e. not malformed and not have Unicode category `Invalid`.

```julia-repl
julia> using Supposition

julia> chars = Data.Characters()

julia> example(chars, 5)
5-element Vector{Char}:
 '‚†∫': Unicode U+283A (category So: Symbol, other)
 '∞≥ç': Unicode U+30CCD (category Lo: Letter, other)
 '\\U6ec9c': Unicode U+6EC9C (category Cn: Other, not assigned)
 '\\U1a05c5': Unicode U+1A05C5 (category In: Invalid, too high)
 'ìÇ´': Unicode U+130AB (category Lo: Letter, other)
```
"""
struct Characters <: Possibility{Char}
    valid::Bool
    Characters(; valid=false) = new(valid)
end

Base.:(==)(c1::Characters, c2::Characters) = c1.valid == c2.valid

Base.show(io::IO, c::Characters) = print(io, Characters, "(; valid=", c.valid, ")")

function Base.show(io::IO, ::MIME"text/plain", c::Characters)
    ex = example(c)
    isval = isvalid(ex) ? styled"{green:‚úî}" : styled"{red:‚ùå}"
    obj = styled"{code:$(sprint(isvalid(ex) ? show : Base.show_invalid, ex))}"
    valid_str = c.valid ? styled"{green:always}" : styled"{yellow:maybe}"
    print(io, styled"""
    {code,underline:$Characters}:

        Produce well-formed {code:Char} which are {code:isvalid}: $valid_str

    E.g. $obj; {code:isvalid}: $isval""")
end

function produce!(tc::TestCase, c::Characters)
    # Ref. https://github.com/JuliaLang/julia/issues/44741#issuecomment-1079083216
    if c.valid
        sample = SampledFrom(typemin(Char):'\U0010ffff')
        s = filter(isvalid, sample)
    else
        s = SampledFrom(typemin(Char):"\xf7\xbf\xbf\xbf"[1])
    end
    produce!(tc, s)
end

"""
    UnicodeCharacters(;valid::Bool = false, malformed = true) <: Possibility{Char}

A `Possibility` of producing arbitrary `Char` instances.

!!! warning "Unicode"
    This will `produce!` ANY `Char` by default, not just valid or well-formed unicode codepoints!
    To only produce valid unicode codepoints, pass `valid=true` as a keyword argument.
    To produce well-formed unicode codepoints, pass `malformed=false` as a keyword argument.

Keyword arguments:

 * `malformed`: Whether produced `Char` are allowed to be malformed. This only has an effect when `valid=false`.
 * `valid`: Whether the produced `Char` must be valid, i.e. not malformed and not have Unicode category `Invalid`.

```julia-repl
julia> using Supposition

julia> chars = Data.UnicodeCharacters()

julia> example(chars, 5)
5-element Vector{Char}:
 '\\xd2\\x5a\\x96\\x07': Malformed UTF-8 (category Ma: Malformed, bad data)
 '\\x44\\x45\\xc5\\x64': Malformed UTF-8 (category Ma: Malformed, bad data)
 '\\x04': Unicode U+0004 (category Cc: Other, control)
 '\\x2b\\xe0\\x6a\\x89': Malformed UTF-8 (category Ma: Malformed, bad data)
 '\\xf5\\x9b\\x63\\x05': Malformed UTF-8 (category Ma: Malformed, bad data)
```
"""
struct UnicodeCharacters <: Possibility{Char}
    valid::Bool
    malformed::Bool
    UnicodeCharacters(; valid=false, malformed=true) = new(valid, malformed)
end

Base.:(==)(c1::UnicodeCharacters, c2::UnicodeCharacters) = c1.valid == c2.valid && c1.malformed == c2.malformed

function Base.show(io::IO, u::UnicodeCharacters)
    print(io, UnicodeCharacters, "(; ")
    print(io, "valid=", u.valid, ", ")
    print(io, "malformed=", u.malformed, ")")
    nothing
end

function Base.show(io::IO, ::MIME"text/plain", uc::UnicodeCharacters)
    ex = example(uc)
    affirm = styled"{green:‚úî}"
    deny = styled"{red:‚ùå}"
    isval = isvalid(ex) ? affirm : deny
    ismalf = Base.ismalformed(ex) ? affirm : deny
    obj = styled"{code:$(sprint(isvalid(ex) ? show : Base.show_invalid, ex))}"
    valid_str = uc.valid ? styled"{green:always}" : styled"{yellow:maybe}"
    malf_str = (!uc.valid && uc.malformed) ? styled"{yellow:maybe}" : styled"{red:no}"
    print(io, styled"""
    {code,underline:$UnicodeCharacters}:

        Produce {code:Char} which are
            ‚àò {code:$(Base.isvalid)}: $valid_str
            ‚àò {code:$(Base.ismalformed)}: $malf_str

    E.g. $obj; {code:isvalid}: $isval, {code:ismalformed}: $ismalf""")
end

function produce!(tc::TestCase, c::UnicodeCharacters)
    # Ref. https://github.com/JuliaLang/julia/issues/44741#issuecomment-1079083216
    if c.valid
        sample = SampledFrom(typemin(Char):'\U0010ffff')
        s = filter(isvalid, sample)
    elseif c.malformed
        s = map(Data.Integers{UInt32}()) do i
            reinterpret(Char, i)
        end
    else
        s = SampledFrom(typemin(Char):"\xf7\xbf\xbf\xbf"[1])
    end
    produce!(tc, s)
end

"""
    AsciiCharacters() <: Possibility{Char}

A `Possibility` of producing arbitrary `Char` instances that are `isascii`.
More efficient than filtering [`Characters`](@ref).

```julia-repl
julia> using Supposition

julia> ascii = Data.AsciiCharacters()

julia> example(ascii, 5)
5-element Vector{Char}:
 '8': ASCII/Unicode U+0038 (category Nd: Number, decimal digit)
 'i': ASCII/Unicode U+0069 (category Ll: Letter, lowercase)
 'R': ASCII/Unicode U+0052 (category Lu: Letter, uppercase)
 '\\f': ASCII/Unicode U+000C (category Cc: Other, control)
 '>': ASCII/Unicode U+003E (category Sm: Symbol, math)
```
"""
struct AsciiCharacters <: Possibility{Char} end

Base.:(==)(::AsciiCharacters, ::AsciiCharacters) = true

Base.show(io::IO, ::AsciiCharacters) = print(io, AsciiCharacters, "()")

function Base.show(io::IO, ::MIME"text/plain", ac::AsciiCharacters)
    obj = styled"{code:$(sprint(show, example(ac)))}"
    print(io, styled"""
    {code,underline:$AsciiCharacters}:

        Produce {code:Char} for which {code:isascii} returns {code:true}.

    E.g. $obj
    """)
end

function produce!(tc::TestCase, ::AsciiCharacters)
    s = SampledFrom(Char(0x0):Char(0x7f))
    produce!(tc, s)
end

## Possibility of Strings

"""
    Text(alphabet::Possibility{Char}; min_len=0, max_len=10_000) <: Possibility{String}

A `Possibility` for producing `String`s containing `Char`s of a given alphabet.

```julia-repl
julia> using Supposition

julia> text = Data.Text(Data.AsciiCharacters(); max_len=15)

julia> example(text, 5)
5-element Vector{String}:
 "U\\x127lxf"
 "hm\\x172SJ-("
 "h`\\x03\\0\\x01[[il"
 "\\x0ep4"
 "9+Hk3 ii\\x1eT"
```
"""
struct Text <: Possibility{String}
    vectors::Vectors{Char}
    function Text(alphabet::A; min_len=0, max_len=10_000) where A <: Possibility{Char}
        vectors = Vectors(alphabet; min_size=min_len, max_size=max_len)
        new(vectors)
    end
end

Base.:(==)(t1::Text, t2::Text) = t1.vectors == t2.vectors

function Base.show(io::IO, t::Text)
    print(io, Text, "(")
    show(io, t.vectors.elements)
    print(io, "; min_len=", string(t.vectors.min_size; base=10))
    print(io, ", max_len=", string(t.vectors.max_size; base=10), ")")
    nothing
end

function Base.show(io::IO, ::MIME"text/plain", t::Text)
    str = sprint(show, example(t); context=:limit=>true)
    _, cols = displaysize(io)
    # these magic numbers are the rest of the text in that line
    available_space = (cols - 5 - 26 - 2)
    if textwidth(str) > available_space
        rem_len = textwidth(str) - available_space + 3
        str = str[begin:div(rem_len, 2)] * " _ " * str[(div(rem_len, 2)+3):end]
    end
    lower = string(t.vectors.min_size; base=10)
    upper = string(t.vectors.max_size; base=10)
    print(io, styled"""
    {code,underline:$Text}:

        Produce a {code:String} with length in the interval {code:[$lower, $upper]}
        The {code:Char} of the {code:String} are produced by {code:$(t.vectors.elements)}.

    E.g. {code:$str} (may be abbreviated here)
    """)
end

produce!(tc::TestCase, s::Text) = join(produce!(tc, s.vectors))

## Dictionaries

"""
    Dicts(keys::Possibility, values::Possibility; min_size=0, max_size=10_000)

A `Possibility` for generating `Dict` objects. The keys are drawn from `keys`,
    while the values are drawn from `values`. `min_size`/`max_size` control
    the number of objects placed into the resulting `Dict`, respectively.

```julia-repl
julia> dicts = Data.Dicts(Data.Integers{UInt8}(), Data.Integers{Int8}(); max_size=3);

julia> example(dicts)
Dict{UInt8, Int8} with 2 entries:
  0x54 => -29
  0x1f => -28
```
"""
struct Dicts{K,V} <: Possibility{Dict{K,V}}
    keys::Possibility{K}
    values::Possibility{V}
    min_size::Int
    max_size::Int
    function Dicts(keys::Possibility{K}, values::Possibility{V}; min_size=0, max_size=10_000) where {K,V}
        min_size <= max_size || throw(ArgumentError("`min_size` must be `<= max_size`!"))
        new{K,V}(keys, values, min_size, max_size)
    end
end

function Base.:(==)(d1::Dicts{K,V}, d2::Dicts{K,V}) where {K,V}
    d1.min_size == d2.min_size &&
    d1.max_size == d2.max_size &&
    d1.keys     == d2.keys     &&
    d1.values   == d2.values
end

function Base.show(io::IO, d::Dicts)
    print(io, Dicts, "(")
    show(io, d.keys)
    print(io, ", ")
    show(io, d.values)
    print(io, "; min_size=", d.min_size, ", max_size=", d.max_size, ")")
    nothing
end

function Base.show(io::IO, ::MIME"text/plain", d::Dicts)
    lengths = styled"{code:[$(Int128(d.min_size)), $(Int128(d.max_size))]}"
    print(io, styled"""
    {code,underline:$Dicts}:

        Produce a {code:Dict\{$(postype(d.keys)), $(postype(d.values))\}} with length in the interval $lengths.
        The keys are produced by {code:$(d.keys)}.
        The values are produced by {code:$(d.values)}.""")
    obj = styled"""


    Keys of the {code:Dict} may look like {code:$(sprint(show, example(d.keys)))}
    Values of the {code:Dict} may look like {code:$(sprint(show, example(d.values)))}"""
    if length(obj) < 100
        print(io, obj)
    end
end

function produce!(tc::TestCase, d::Dicts{K,V}) where {K,V}
    dict = Dict{K,V}()

    while true
        if length(dict) < d.min_size
            forced_choice!(tc, UInt(1))
        elseif (length(dict)+1) >= d.max_size
            forced_choice!(tc, UInt(0))
            break
        elseif !weighted!(tc, 0.9)
            break
        end
        k = produce!(tc, d.keys)
        v = produce!(tc, d.values)
        dict[k] = v
    end

    return dict
end

## Possibility of values from a collection

"""
    SampledFrom(collection) <: Possibility{eltype(collection)}

A `Possibility` for sampling uniformly from `collection`.

`collection`, as well as its `eachindex`, is assumed to be indexable.

!!! warning "Mutable Data"
    The source objects from the collection given to this `SampledFrom`
    is not copied when `produce!` is called. Be careful with mutable data!

!!! tip "Sampling from `String`"
    To sample from a `String`, you can `collect` the string first
    to get a `Vector{Char}`. This is necessary because `String`s
    use the variable-length UTF-8 encoding, which isn't arbitrarily indexable
    in constant time.

```julia-repl
julia> using Supposition

julia> sampler = Data.SampledFrom([1, 1, 1, 2])

julia> example(sampler, 4)
4-element Vector{Int64}:
 1
 1
 2
 1
```
"""
struct SampledFrom{T, C} <: Possibility{T}
    collection::C
    SampledFrom(col) = new{eltype(col), typeof(col)}(col)
end

Base.:(==)(sf1::SampledFrom{T,C}, sf2::SampledFrom{T,C}) where {T,C} = sf1.collection == sf2.collection

function Base.show(io::IO, sf::SampledFrom)
    print(io, SampledFrom, "(")
    show(IOContext(io, :compact=>true), sf.collection)
    print(io, ")")
    nothing
end

function Base.show(io::IO, ::MIME"text/plain", sf::SampledFrom)
    obj = example(sf)
    print(io, styled"""
    {code,underline:$SampledFrom}:

        Sample, with equal probability, an element from
        {code:$(sprint(show, sf.collection; context=:compact=>true))}

    E.g. {code:$obj}
    """)
end

function produce!(tc::TestCase, sf::SampledFrom)
    pos_indices = eachindex(sf.collection)
    idx = produce!(tc, Integers(firstindex(pos_indices), lastindex(pos_indices)))
    return sf.collection[pos_indices[idx]]
end

## Possibility of booleans

"""
    Booleans() <: Possibility{Bool}

A `Possibility` for sampling boolean values.

```julia-repl
julia> using Supposition

julia> bools = Data.Booleans()

julia> example(bools, 4)
4-element Vector{Bool}:
 0
 1
 0
 1
```
"""
struct Booleans <: Possibility{Bool} end

Base.:(==)(::Booleans, ::Booleans) = true

Base.show(io::IO, ::Booleans) = print(io, Booleans, "()")

function Base.show(io::IO, ::MIME"text/plain", b::Booleans)
    print(io, styled"""
    {code,underline:$Booleans}:

        Produce a {code:Bool} with a fair coin.
        Both {code:true} and {code:false} have a probability of 50%.

    E.g. {code:$(example(b))}
    """)
end

produce!(tc::TestCase, ::Booleans) = weighted!(tc, 0.5)

## Possibility of floating point values

"""
    Floats{T <: Union{Float16,Float32,Float64}}(;infs=true, nans=true, minimum=-T(Inf), maximum=+T(Inf)) <: Possibility{T}

A `Possibility` for sampling floating point values.

The keyword `infs` controls whether infinities can be generated. `nans` controls whether
    any `NaN` (signaling & quiet) will be generated.

!!! warning "Inf, Nan"
    This possibility will generate *any* valid instance, including positive
    and negative infinities, signaling and quiet NaNs and every possible float.

```julia-repl
julia> using Supposition

julia> floats = Data.Floats{Float16}()

julia> example(floats, 5)
5-element Vector{Float16}:
  -8.3e-6
   1.459e4
   3.277
 NaN
  -0.0001688
```
"""
struct Floats{T <: Base.IEEEFloat} <: Possibility{T}
    nans::Bool
    infs::Bool
    minimum::T
    maximum::T
    function Floats{T}(; nans::Bool=true, infs::Bool=true, minimum=T(-Inf), maximum=T(Inf)) where T <: Base.IEEEFloat
        _min = convert(T, minimum)
        _max = convert(T, maximum)
        isnan(_min) && throw(ArgumentError("`minimum` must be non-NaN!"))
        isnan(_max) && throw(ArgumentError("`maximum` must be non-NaN!"))
        _min > _max && throw(ArgumentError("`minimum` must be `<= maximum`!"))
        infs = (isinf(minimum) | isinf(maximum)) & infs
        new{T}(nans, infs, _min, _max)
    end
end

Base.:(==)(f1::Floats{T}, f2::Floats{T}) where {T} = f1.nans == f2.nans && f1.infs == f2.infs && f1.minimum == f2.minimum && f1.maximum == f2.maximum

function Base.show(io::IO, f::Floats)
    print(io, typeof(f), "(; ")
    print(io, "nans=", f.nans, ", ")
    print(io, "infs=", f.infs)
    if !isinf(f.minimum)
        print(io, ", minimum=", f.minimum)
    end
    if !isinf(f.maximum)
        print(io, ", maximum=", f.maximum)
    end
    print(io, ")")
    nothing
end

function Base.show(io::IO, ::MIME"text/plain", f::Floats)
    ex = example(f)
    affirm = styled"{green:‚úî}"
    deny = styled"{red:‚ùå}"
    inf = isinf(ex) ? affirm : deny
    nan = isnan(ex) ? affirm : deny
    obj = sprint(show, ex)
    nan_str = f.nans ? styled"{yellow:maybe}" : styled"{red:never}"
    inf_str = f.infs ? styled"{yellow:maybe}" : styled"{red:never}"

    print(io, styled"""
    {code,underline:$Floats}:

        Produce a floating point value {code:x} of type {code:$(postype(f))}, which is
            * {code:isinf}: $inf_str
            * {code:isnan}: $nan_str
            * {code:$(f.minimum) <= x <= $(f.maximum)}

    E.g. {code:$obj}; {code:isinf}: $inf, {code:isnan}: $nan""")
end

function produce!(tc::TestCase, f::Floats{T}) where {T}
    iT = Supposition.uint(T)
    res = reinterpret(T, produce!(tc, Integers{iT}()))
    # early rejections
    !f.infs && isinf(res) && reject(tc)
    !f.nans && isnan(res) && reject(tc)

    # early acception
    f.nans && isnan(res) && return res
    f.minimum <= res <= f.maximum && return res

    return float_remap(res, f.minimum, f.maximum)
end

function float_remap(num::T, _min::T, _max::T) where T <: AbstractFloat
    # We're outside of the desired bounds, so use the mantissa
    # to resample the actual range
    range_size = min(_max - _min, floatmax(T))
    _, _, mantissa = Supposition.tear(num)
    max_mantissa = oftype(mantissa, (2^Supposition.fracsize(T)) - 1)
    num = _min + range_size * (mantissa / max_mantissa)

    # ensure the value is still in the desired range
    return clamp(convert(T, num), _min, _max)
end

struct AllFloats <: Data.Possibility{Union{Float16, Float32, Float64}}
    nans::Bool
    infs::Bool
    minimum::Float64
    maximum::Float64
    function AllFloats(nans, infs, minimum, maximum)
        _min = convert(Float64, minimum)
        _max = convert(Float64, maximum)
        isnan(_min) && throw(ArgumentError("`minimum` must be non-NaN!"))
        isnan(_max) && throw(ArgumentError("`maximum` must be non-NaN!"))
        _min > _max && throw(ArgumentError("`minimum` must be `<= maximum`!"))
        infs = (isinf(minimum) | isinf(maximum)) & infs
        new(convert(Bool, nans), convert(Bool, infs), _min, _max)
    end
end
AllFloats(nans, infs) = AllFloats(nans, infs, -Inf, Inf)
function AllFloats(; nans=true, infs=true,
                     minimum=-Inf, maximum=Inf)
    AllFloats(nans, infs, minimum, maximum)
end

Base.:(==)(af1::AllFloats, af2::AllFloats) = af1.nans == af2.nans && af1.infs == af2.infs && af1.minimum == af2.minimum && af1.maximum == af1.maximum

"""
    Floats(;nans=true, infs=true) <: Possibility{Union{Float64,Float32,Float16}}

A catch-all for generating instances of all three IEEE floating point types.
"""
Floats(;nans=true, infs=true, minimum=-Inf, maximum=Inf) = AllFloats(nans, infs, minimum, maximum)

function Base.show(io::IO, f::AllFloats)
    print(io, Floats, "(; ")
    print(io, "nans=", f.nans, ", ")
    print(io, "infs=", f.infs)
    if !isinf(f.minimum)
        print(io, ", minimum=", f.minimum)
    end
    if !isinf(f.maximum)
        print(io, ", maximum=", f.maximum)
    end
    print(io, ")")
    nothing
end

function Base.show(io::IO, ::MIME"text/plain", f::AllFloats)
    ex = example(f)
    affirm = styled"{green:‚úî}"
    deny = styled"{red:‚ùå}"
    inf = isinf(ex) ? affirm : deny
    nan = isnan(ex) ? affirm : deny
    obj = sprint(show, ex)
    nan_str = f.nans ? styled"{yellow:maybe}" : styled"{red:never}"
    inf_str = f.infs ? styled"{yellow:maybe}" : styled"{red:never}"

    print(io, styled"""
    {code,underline:$AllFloats}:

        Produce a floating point value of type {code:Float16}, {code:Float32} or {code:Float64}, which may be
            * {code:isinf}: $inf_str
            * {code:isnan}: $nan_str
            * {code:$(f.minimum) <= x <= $(f.maximum)}

    E.g. {code:$obj}, a {code:$(typeof(ex))}; {code:isinf}: $inf, {code:isnan}: $nan""")
end

function produce!(tc::TestCase, af::AllFloats)
    of = OneOf(
        Floats{Float16}(;af.nans,af.infs),
        Floats{Float32}(;af.nans,af.infs),
        Floats{Float64}(;af.nans,af.infs))
    num = produce!(tc, of)
    af.nans && isnan(num) && return num
    af.minimum <= num <= af.maximum && return num

    # Not all numbers in the given range are representable
    # by all types, so normalize the bounds to the closest
    # that are still in the interval for that type.
    # This ensures that we don't accidentally go outside of the
    # allowed range.
    _min = oftype(num, af.minimum)
    _max = oftype(num, af.maximum)
    while _min < af.minimum
        _min = nextfloat(_min)
    end
    while _max > af.maximum
        _max = prevfloat(_max)
    end

    # in this case, the interval doesn't contain values of the chosen type
    _min > _max && reject(tc)

    # and finally remap the value to a valid one
    return float_remap(num, _min, _max)
end

####
# Weighted dice
#
# Implementation based on https://www.keithschwarz.com/darts-dice-coins/
####

"""
    WeightedNumbers(weights::Vector{Float64}) <: Possibility{Int}

Sample the numbers from `1:length(weights)`, each with a weight of `weights[i]`.

The weights may be any number > 0.0.

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

```julia-repl
julia> using Supposition

julia> bn = Data.WeightedNumbers([1.0, 1.0, 3.0]);

julia> example(Data.Vectors(bn; min_size=3, max_size=15), 5)
5-element Vector{Vector{Int64}}:
 [3, 2, 3, 3, 2, 3, 3]
 [1, 1, 1, 2, 1, 3, 1, 3]
 [2, 3, 3, 3, 3, 3, 3, 1, 1, 3, 3, 3]
 [3, 3, 2, 3, 3]
 [1, 3, 3, 3, 2, 2]
```
"""
struct WeightedNumbers <: Possibility{Int}
    table::Vector{Tuple{Int,Int,Float64}}
    function WeightedNumbers(weights::Vector{Float64})
        n = length(weights)
        iszero(n) && throw(ArgumentError("Weights must not be empty!"))
        any(<=(0.0), weights) && throw(ArgumentError("All weights must be `>= 0.0`!"))
        total = sum(weights)
        probabilities = weights ./ total

        table = UInt64[ (-1 % UInt64) for _ in 1:n, _ in 1:3 ]
        table[:, 1] .= UInt64.(1:n)

        small = Vector{UInt64}()
        large = Vector{UInt64}()
        sizehint!(small, n)
        sizehint!(large, n)

        scaled_probabilities = Vector{Float64}(undef, n)
        for (i,pi) in enumerate(probabilities)
            scaled = pi*n
            scaled_probabilities[i] = scaled
            if scaled < 1.0
                push!(small, i % UInt64)
            else
                push!(large, i % UInt64)
            end
        end

        while !isempty(small) && !isempty(large)
            lo = pop!(small)
            hi = pop!(large)
            @assert lo != hi
            @assert scaled_probabilities[hi] >= 1.0
            @assert table[lo, 2] == -1 % UInt64
            table[lo, 2] = hi
            table[lo, 3] = reinterpret(UInt64, scaled_probabilities[lo])
            scaled_probabilities[hi] = (scaled_probabilities[hi] + scaled_probabilities[lo]) - 1.0

            if scaled_probabilities[hi] < 1.0
                push!(small, hi)
            else
                push!(large, hi)
            end
        end

        while !isempty(large)
            g = pop!(large)
            table[g, 3] = reinterpret(UInt64, 0.0)
        end

        while !isempty(small)
            l = pop!(small)
            table[l, 3] = reinterpret(UInt64, 0.0)
        end

        res = Vector{Tuple{Int,Int,Float64}}(undef, n)
        for (ubase, ualternate, uchance) in eachrow(table)
            base = ubase % Int64
            alternate = ualternate % Int64
            alternate_chance = reinterpret(Float64, uchance)
            if alternate == -1
                res[base] = (base, base, reinterpret(Float64, alternate_chance))
            elseif alternate < base
                res[base] = (alternate, base, 1.0 - alternate_chance)
            else
                res[base] = (base, alternate, alternate_chance)
            end
        end

        new(res)
    end
end

Base.:(==)(wn1::WeightedNumbers, wn2::WeightedNumbers) = wn1.table == wn2.table

function recover_weights(wn::WeightedNumbers)
    probs = zeros(Float64, size(wn.table))
    n = length(wn.table)
    for l in wn.table
        probs[l[1]] += l[3] / n
        probs[l[2]] += (1.0 - l[3])/n
    end
    probs
end

function Base.show(io::IO, wn::WeightedNumbers)
    print(io, WeightedNumbers, "(")
    show(IOContext(io, :compact=>true), recover_weights(wn))
    print(io, ")")
    nothing
end

function Base.show(io::IO, ::MIME"text/plain", wn::WeightedNumbers)
    weights = recover_weights(wn)
    print(io, styled"""
    {code,underline:$WeightedNumbers}:

        Select the numbers from {code:1:$(length(weights))} with some probability.
        The probabilities are chosen based on some given weights.

    """)

    # this format ends up printing like `xx.yy%`, regardless of the actual number.
    # leading characters are padded with spaces
    percent_format = format"%5.2f%%"
    for (i,w) in enumerate(weights)
        print(io, "        ‚àò ")
        format(io, percent_format, 100.0*w)
        print(io, " : ", i, '\n')
    end

    obj = example(wn)

    print(io, styled"""

    E.g. {code:$obj}, selected with probability """)
    format(io, percent_format, 100.0*weights[obj])

    nothing
end

function Data.produce!(tc::TestCase, bn::WeightedNumbers)
    base, alternate, alternate_chance = choice!(tc, bn.table)
    use_base = weighted!(tc, alternate_chance)
    use_base ? base : alternate
end

"""
    WeightedSample{T}(colllection, weights::Vector{Float64}) <: Possibility{T}

Draw an element from the indexable `collection`, biasing the 
drawing process by assigning each index `i` of `col` the weight at
`weights[i]`.

`length` of `col` and `weights` must be equal and `eachindex(col)` must
be indexable.

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

!!! tip "Sampling from `String`"
    To sample from a `String`, you can `collect` the string first
    to get a `Vector{Char}`. This is necessary because `String`s
    use the variable-length UTF-8 encoding, which isn't arbitrarily indexable
    in constant time.

```
julia> bs = Data.WeightedSample(["foo", "bar", "baz"], [3.0, 1.0, 1.0]);

julia> example(bs, 10)
10-element Vector{String}:
 "foo"
 "foo"
 "bar"
 "baz"
 "baz"
 "foo"
 "bar"
 "foo"
 "foo"
 "bar"
```
"""
struct WeightedSample{T,C} <: Possibility{T}
    idx::WeightedNumbers
    col::C
    function WeightedSample(col, weights::Vector{Float64})
        len_col = length(eachindex(col))
        len_weights = length(weights)
        len_weights == len_col || throw(ArgumentError("Length mismatch - need $len_col weights, but $len_weights given!"))
        bn = WeightedNumbers(weights)
        new{eltype(col),typeof(col)}(bn, col)
    end
end

Base.:(==)(wn1::WeightedSample{T,C}, wn2::WeightedSample{T,C}) where {T,C} = wn1.col == wn2.col && wn1.idx == wn2.idx

function Base.show(io::IO, ws::WeightedSample)
    print(io, WeightedSample, "(")
    show(IOContext(io, :compact=>true), ws.col)
    print(io, ", ")
    show(IOContext(io, :compact=>true), recover_weights(ws.idx))
    print(io, ")")
    nothing
end

function Base.show(io::IO, ::MIME"text/plain", ws::WeightedSample)
    weights = recover_weights(ws.idx)
    print(io, styled"""
    {code,underline:$WeightedSample}:

        Select an object from the given collection with some probability.
        The probabilities per element are chosen based on some given weights.

    """)

    # this format ends up printing like `xx.yy%`, regardless of the actual number.
    # leading characters are padded with spaces
    percent_format = format"%5.2f%%"
    for (i,w) in enumerate(weights)
        print(io, "        ‚àò ")
        format(io, percent_format, 100.0*w)
        print(io, " : ", ws.col[eachindex(ws.col)[i]], '\n')
    end

    idx = example(ws.idx)
    prob = weights[idx]
    obj = sprint(show, ws.col[idx])

    print(io, styled"""

    E.g. {code:$obj}, selected with probability """)
    format(io, percent_format, 100.0*prob)

    nothing
end

function produce!(tc::TestCase, bs::WeightedSample)
    idx = produce!(bs.idx)
    indices = eachindex(bs.col)
    bs.col[indices[idx]]
end

end # data module
