text:
# Methods to help select or add a number of outputs given a new size as this problem apparently belongs to the class of FU-complete problems. And yes, I curse the day I conceived the idea for this project right now...

"""
    AbstractSelectionStrategy

Base type for how to select the exact inputs/outputs indices from a vertex given a size change.
"""
abstract type AbstractSelectionStrategy end

"""
    LogSelection <: AbstractSelectionStrategy
    LogSelection(msgfun::Function, andthen)
    LogSelection(level, msgfun::Function, andthen)

Logs output from function `msgfun` at LogLevel `level` and executes `AbstractSelectionStrategy andthen`.
"""
struct LogSelection{L,F,S <: AbstractSelectionStrategy} <: AbstractSelectionStrategy
    level::L
    msgfun::F
    andthen::S
end
LogSelection(msgfun::Function, andthen) = LogSelection(Logging.Info, msgfun, andthen)
LogSelectionFallback(nextstr, andthen; level=Logging.Warn) = LogSelection(level, v -> "Selection for vertex $(name(v)) failed! $nextstr", andthen)

"""
    SelectionFail <: AbstractSelectionStrategy

Throws an error.
"""
struct SelectionFail <: AbstractSelectionStrategy end

"""
    NoutRevert <: AbstractSelectionStrategy
    NoutRevert()

Reverts output size change for a vertex.
"""
struct NoutRevert <: AbstractSelectionStrategy end

"""
    SelectDirection <: AbstractSelectionStrategy
    SelectDirection()
    SelectDirection(s::AbstractSelectionStrategy)

Select indices for a vertex using `AbstractSelectionStrategy s` (default `OutSelect{Exact}`) in only the direction(s) in which the vertex has changed size.

Intended use it to reduce the number of constraints for a `AbstractJuMPSelectionStrategy` as only the parts of the graph which are changed will be considered.
"""
struct SelectDirection{S <: AbstractSelectionStrategy} <: AbstractSelectionStrategy
    strategy::S
end
SelectDirection() = SelectDirection(OutSelectExact())

"""
    ApplyAfter <: AbstractSelectionStrategy
    ApplyAfter()
    ApplyAfter(s::AbstractSelectionStrategy)
    ApplyAfter(apply::Function, s::AbstractSelectionStrategy)

Invokes `apply(v)` for each `AbstractVertex v` which was changed as a result of `AbstractSelectionStrategy s`.
"""
struct ApplyAfter{F, S} <: AbstractSelectionStrategy
    apply::F
    strategy::S
end
ApplyAfter() = ApplyAfter(OutSelectExact())
ApplyAfter(s::AbstractSelectionStrategy) = ApplyAfter(apply_mutation, s)

"""
    AbstractJuMPSelectionStrategy

Base type for how to select the exact inputs/outputs indices from a vertex given a size change using JuMP to handle the constraints.
"""
abstract type AbstractJuMPSelectionStrategy <: AbstractSelectionStrategy end

"""
    DefaultJuMPSelectionStrategy <: AbstractJuMPSelectionStrategy

Default strategy intended to be used when adding some extra constraints or objectives to a model on top of the default.
"""
struct DefaultJuMPSelectionStrategy <: AbstractJuMPSelectionStrategy end

fallback(::AbstractJuMPSelectionStrategy) = SelectionFail()

"""
    OutSelect{T} <: AbstractSelectionStrategy
    OutSelectExact()
    OutSelectRelaxed()
    OutSelect{T}(fallback)

If `T == Exact`, output indices for each vertex `v` are selected with the constraint that `nout(v)` shall not change.

If `T == Relaxed`, output indices for each vertex `v` are selected with the objective to deviate as little as possible from `nout(v)`.

Possible to set a `fallback AbstractSelectionStrategy` should it be impossible to select indices according to the strategy.
"""
struct OutSelect{T, S <: AbstractSelectionStrategy} <: AbstractJuMPSelectionStrategy
    fallback::S
end
OutSelectExact() = OutSelect{Exact}(LogSelectionFallback("Relaxing size constraint...", OutSelectRelaxed()))
OutSelectRelaxed() = OutSelect{Relaxed}(LogSelectionFallback("Reverting...", NoutRevert()))
OutSelect{T}(s::S) where {T,S} = OutSelect{T, S}(s)
fallback(s::OutSelect) = s.fallback

"""
    TruncateInIndsToValid{S} <: AbstractSelectionStrategy
    TruncateInIndsToValid()
    TruncateInIndsToValid(s::S)

Ensures that all selected input indices are within range of existing input indices after applying `s` (default `OutSelectExact`).

Not needed in normal cases, but certain structural mutations (e.g create_edge!) may cause this to happen due to how constraints are (not) created when original sizes do not align in conjunction with how result of selection is interpreted.

While this may be considered a flaw in the output selection procedure, it is rare enough so that in most cases when it happens it is the result of a user error or lower level bug. Therefore this strategy is left optional to be used only in cases when mismatches are expected.
"""
struct TruncateInIndsToValid{S <: AbstractSelectionStrategy} <: AbstractSelectionStrategy
    strategy::S
end
TruncateInIndsToValid() = TruncateInIndsToValid(OutSelectExact())

"""
    Δoutputs(g::CompGraph, valuefun::Function)
    Δoutputs(s::AbstractSelectionStrategy, g::CompGraph, valuefun::Function)
    Δoutputs(v::AbstractVertex, valuefun::Function)
    Δoutputs(s::AbstractSelectionStrategy, v::AbstractVertex, valuefun::Function)
    Δoutputs(d::Direction, v::AbstractVertex, valuefun::Function)
    Δoutputs(s::AbstractSelectionStrategy, d::Direction, v::AbstractVertex, valuefun::Function)

Change outputs of all vertices of graph `g` (or graph to which `v` is connected) according to the provided `AbstractSelectionStrategy s` (default `OutSelect{Exact}`).

Return true of operation was successful, false otherwise.

Argument `valuefun` provides a vector `value = valuefun(vx)` for any vertex `vx` in the same graph as `v` where `value[i] > value[j]` indicates that output index `i` shall be preferred over `j` for vertex `vx`.

If provided, `Direction d` will narrow down the set of vertices to evaluate so that only vertices which may change as a result of changing size of `v` are considered.
"""
Δoutputs(g::CompGraph, valuefun) = Δoutputs(OutSelectExact(), g, valuefun)
Δoutputs(s::AbstractSelectionStrategy, g::CompGraph, valuefun) = Δoutputs(s, vertices(g), valuefun)
Δoutputs(v::AbstractVertex, valuefun::Function) = Δoutputs(OutSelectExact(), v, valuefun)
Δoutputs(s::AbstractSelectionStrategy, v::AbstractVertex, valuefun::Function) = Δoutputs(s, all_in_graph(v), valuefun)
function Δoutputs(s::SelectDirection, v::AbstractVertex, valuefun::Function)
    nin_change = nin_org(v) != nin(v)
    nout_change = nout_org(v) != nout(v)
    if nout_change && nin_change
        return Δoutputs(s.strategy, Both(), v, valuefun)
    elseif nout_change
        return Δoutputs(s.strategy, Output(), v, valuefun)
    elseif nin_change
        return Δoutputs(s.strategy, Input(), v, valuefun)
    end
    return true
 end

Δoutputs(d::Direction, v::AbstractVertex, valuefun::Function) = Δoutputs(OutSelectExact(), d, v, valuefun)
Δoutputs(s::AbstractSelectionStrategy, d::Direction, v::AbstractVertex, valuefun::Function) = Δoutputs(s, all_in_Δsize_graph(v, d), valuefun)

function Δoutputs(s::ApplyAfter, vs::AbstractVector{<:AbstractVertex}, valuefun::Function)
    success = Δoutputs(s.strategy, vs, valuefun)
    # This is not 100% safe in all cases, as a strategy can chose to not do anything even in this case (although it probably should)
    foreach(s.apply, filter(v -> nout(v) != nout_org(v) || nin(v) != nin_org(v), vs))
    return success
end

function Δoutputs(s::AbstractSelectionStrategy, vs::AbstractVector{<:AbstractVertex}, valuefun::Function)
    success, ins, outs = solve_outputs_selection(s, vs, valuefun)
    if success
        Δoutputs(ins, outs, vs)
    end
    return success
end

function Δoutputs(s::TruncateInIndsToValid, vs::AbstractVector{<:AbstractVertex}, valuefun::Function)
    success, ins, outs = solve_outputs_selection(s.strategy, vs, valuefun)
    if success
        for (vv, ininds) in ins
            for innr in eachindex(ininds)
                ininds[innr] = aligntomax(nin_org(vv)[innr], ininds[innr])
            end
            newins, newouts = align_outs_to_ins(vv, ins[vv], outs[vv])
            ins[vv] = newins
            outs[vv] = newouts
        end
        Δoutputs(ins, outs, vs)
    end
    return success
end

aligntomax(maxval, ::Missing) = missing
function aligntomax(maxval, arr)
    arr = copy(arr)
    maxind = argmax(arr)
    while arr[maxind] > maxval
        newval = maxval
        while newval in arr && newval > -1
            newval -= 1
        end
        arr[maxind] = newval == 0 ? -1 : newval
        maxind = argmax(arr)
    end
    return arr
end

align_outs_to_ins(v, ins, outs) = align_outs_to_ins(trait(v), v, ins, outs)
align_outs_to_ins(t::DecoratingTrait, v, ins, outs) = align_outs_to_ins(base(t), v, ins, outs)
align_outs_to_ins(t, v, ins, outs) = ins,outs
function align_outs_to_ins(::SizeInvariant, v, ins::AbstractArray, outs)
    isempty(ins) && return ins, outs
    inds = ins[1]
    newins = repeat([inds], length(ins))
    newouts = ismissing(outs) ? outs : inds
    return newins, newouts
end

"""
    Δoutputs(ins::Dict outs::Dict, vertices::AbstractVector{<:AbstractVertex})

Set input and output indices of each `vi` in `vs` to `outs[vi]` and `ins[vi]` respectively.
"""
function Δoutputs(ins::Dict, outs::Dict, vs::AbstractVector{<:AbstractVertex})

    for vi in vs
        Δnin(OnlyFor(), vi, ins[vi]...)
        Δnout(OnlyFor(), vi, outs[vi])
    end

    for vi in vs
        after_Δnin(vi, ins[vi]...)
        after_Δnout(vi, outs[vi])
    end
end

function Δnin(::OnlyFor, v) end
function Δnin(::OnlyFor, v, inds::Missing) end
function Δnin(s::OnlyFor, v, inds::AbstractVector{<:Integer}...)
    any(inds .!= [1:insize for insize in nin(v)]) || return
    Δnin(s, trait(v), v, inds)
end

function Δnout(::OnlyFor, v, inds::Missing) end
function Δnout(s::OnlyFor, v, inds::AbstractVector{<:Integer})
    inds == 1:nout(v) && return
    Δnout(s, trait(v), v, inds)
end

function solve_outputs_selection(s::LogSelection, vertices::AbstractVector{<:AbstractVertex}, valuefun::Function)
    @logmsg s.level s.msgfun(vertices[1])
    return solve_outputs_selection(s.andthen, vertices, valuefun)
end

solve_outputs_selection(s::SelectionFail, vertices::AbstractVector{<:AbstractVertex}, valuefun::Function) = error("Selection failed for vertex $(name.(vertices))")

function NaiveNASlib.solve_outputs_selection(s::NoutRevert, vertices::AbstractVector{<:AbstractVertex}, valuefun::Function)
    for v in vertices
        Δnout(NaiveNASlib.OnlyFor(), v, nout_org(v) - nout(v))
        diff = nin_org(v) - nin(v)
        Δnin(NaiveNASlib.OnlyFor(), v, diff...)
    end

    return false, Dict(vertices .=> UnitRange.(1, nout.(vertices))), Dict(vertices .=> map(nins -> UnitRange.(1,nins), nin.(vertices)))
end

"""
    solve_outputs_selection(s::AbstractSelectionStrategy, vertices::AbstractVector{<:AbstractVertex}, valuefun::Function)

Returns a tuple `(success, nindict, noutdict)` where `nindict[vi]` are new input indices and `noutdict[vi]` are new output indices for each vertex `vi` in `vertices`.

The function generally tries to maximize `sum(valuefun(vi) .* selected[vi]) ∀ vi in vertices` where `selected[vi]` is all elements in `noutdict[vi]` larger than 0 (negative values in `noutdict` indicates a new output shall be inserted at that position). This however is up to the implementation of the `AbstractSelectionStrategy s`.

Since selection of outputs is not guaranteed to work in all cases, a flag `success` is also returned. If `success` is `false` then applying the new indices may (and probably will) fail.
"""
function solve_outputs_selection(s::AbstractJuMPSelectionStrategy, vertices::AbstractVector{<:AbstractVertex}, valuefun::Function)
    model = selectmodel(s, vertices, values)

    # The binary variables `outselectvars` tells us which existing output indices to select
    # The binary variables `outinsertvars` tells us where in the result we shall insert -1 where -1 means "create a new output (e.g. a neuron)
    # Thus, the result will consist of all selected indices with possibly interlaced -1s

    # TODO: Edge case when vertices are not all in graph (i.e there is at least one vertex for which its inputs or outputs are not part of vertices) and valuefun returns >= 0 might result in size inconsistency.
    # Either make sure values are > 0 for such vertices or add constraints to them!
    # Afaik this happens when only v's inputs are touched. Cleanest fix might be to separate "touch output" from "touch input". See https://github.com/DrChainsaw/NaiveNASlib.jl/issues/39
    outselectvars = Dict(vertices .=> map(v -> @variable(model, [1:length(valuefun(v))], binary=true), vertices))
    outinsertvars = Dict(vertices .=> map(v -> @variable(model, [1:nout(v)], binary=true), vertices))
    # Optimization: Init outinsertvars as empty and only add if needed: outinsertvars = Dict{eltype(outselectvars).types...}()

    objexpr = @expression(model, objective, 0)
    for v in vertices
        data = (model=model, outselectvars=outselectvars, outinsertvars=outinsertvars, objexpr = objexpr, valuefun = valuefun)
        vertexconstraints!(v, s, data)
        objexpr = selectobjective!(s, v, data)
    end

    @objective(model, Max, objexpr)

    JuMP.optimize!(model)

    !accept(s, model) && return solve_outputs_selection(fallback(s), vertices, valuefun)

    return true, extract_ininds_and_outinds(s, outselectvars, outinsertvars)...

end

accept(::AbstractJuMPSelectionStrategy, model::JuMP.Model) = JuMP.termination_status(model) != MOI.INFEASIBLE && JuMP.primal_status(model) == MOI.FEASIBLE_POINT # Beware: primal_status seems unreliable for Cbc. See MathOptInterface issue #822

selectmodel(::AbstractJuMPSelectionStrategy, v, values) = JuMP.Model(JuMP.optimizer_with_attributes(Cbc.Optimizer, "loglevel"=>0))

# First dispatch on traits to sort out things like immutable vertices
vertexconstraints!(v::AbstractVertex, s::AbstractJuMPSelectionStrategy, data) = vertexconstraints!(trait(v), v, s, data)
vertexconstraints!(t::DecoratingTrait, v, s::AbstractJuMPSelectionStrategy, data) = vertexconstraints!(base(t), v, s, data)
vertexconstraints!(t::MutationSizeTrait, v, s::AbstractJuMPSelectionStrategy, data) = vertexconstraints!(s, t, v, data) # Now dispatch on strategy (and trait)
function vertexconstraints!(::Immutable, v, s::AbstractJuMPSelectionStrategy, data)
     @constraint(data.model, data.outselectvars[v] .== 1)
     @constraint(data.model, data.outinsertvars[v] .== 0)
 end

function vertexconstraints!(s::AbstractJuMPSelectionStrategy, t::MutationSizeTrait, v, data)
    sizeconstraint!(s, t, v, data)
    compconstraint!(s, v, (data..., vertex=v))
    inoutconstraint!(s, t, v, data)
end

function sizeconstraint!(::OutSelect{Exact}, t, v, data)
    nselect = min(nout(v), nout_org(v))
    ninsert = max(0, nout(v) - nselect)
    @constraint(data.model, sum(data.outselectvars[v]) == nselect)
    @constraint(data.model, sum(data.outinsertvars[v]) == ninsert)
end

function sizeconstraint!(::OutSelect{Exact}, t::SizeStack, v, data)
    nselect = sum(min.(nin_org(v), nin(v)))
    ninsert = sum(max.(0, nin(v) - nin_org(v)))
    @constraint(data.model, sum(data.outselectvars[v]) == nselect)
    @constraint(data.model, sum(data.outinsertvars[v]) == ninsert)
end

function sizeconstraint!(::OutSelect{Relaxed}, t, v, data)
    @constraint(data.model, sum(data.outselectvars[v]) + sum(data.outinsertvars[v])  >= 1)

    # Handle insertions
    # The constraint that either there are no new outputs or the total number of outputs must be equal to the length of outinsertvars is a somewhat unfortunate result of the approach chosen to solve the problem.
    # If not enforced, we will end up in situations where some indices shall neither be selected nor have insertions. For example, the result might say "keep indices 1,2,3 and insert a new output at index 10".
    # If one can come up with a constraint to formulate "no gaps" (such as the gab above) instead of the current approach the chances of finding a feasible soluion would probably increase.
    # Maybe this https://cs.stackexchange.com/questions/12102/express-boolean-logic-operations-in-zero-one-integer-linear-programming-ilp in combination with this https://math.stackexchange.com/questions/2022967/how-to-model-a-constraint-of-consecutiveness-in-0-1-programming?rq=1
    outsel = data.outselectvars[v]
    outins = data.outinsertvars[v]
    model = data.model
    if length(outins) > length(outsel)
        @constraint(data.model, sum(outins) == length(outins) - sum(outsel))
    elseif length(outins) > 0
        @constraint(data.model, outins .== 0)
    end
end

function inoutconstraint!(s, ::SizeAbsorb, v, data) end
function inoutconstraint!(s, t::SizeTransparent, v, data)
    inoutconstraint!(s, t, v, data.model, data.outselectvars)
    inoutconstraint!(s, t, v, data.model, data.outinsertvars)
end

function inoutconstraint!(s, ::SizeStack, v, model, vardict::Dict)
    offs = 1
    var = vardict[v]
    for (i, vi) in enumerate(inputs(v))
        var_i = vardict[vi]
        # Sizes mismatch when vertex/edge was removed (or edge added)
        if nout_org(vi) == nin_org(v)[i] && offs+length(var_i)-1 <= length(var)
            @constraint(model, var_i .== var[offs:offs+length(var_i)-1])
        end
        offs += length(var_i)
    end
end

function inoutconstraint!(s, ::SizeInvariant, v, model, vardict::Dict)
    var = vardict[v]
    for (i, vi) in enumerate(inputs(v))
        # Sizes mismatch when vertex/edge was removed (or edge added)
        nout_org(vi) == nout_org(v) || continue
        var_i = vardict[vi]
        @constraint(model, var_i .== var)
    end
end

function selectobjective!(s::AbstractJuMPSelectionStrategy, v, data)
    value = valueobjective!(s, v, data)
    insertlast = insertlastobjective!(s, v, data)
    return @expression(data.model, data.objexpr + value + insertlast)
end

function selectobjective!(::OutSelect{Relaxed}, v, data)

    # No thought behind scaling other than wanting to have roughly same order of magnitude
    scale = max(0, maximum(data.valuefun(v)))
    sizediff = @expression(data.model, sum(data.outselectvars[v]) + sum(data.outinsertvars[v]) - nout(v) + count(<(0), data.valuefun(v)))
    sizediffnorm = norm!(ScaleNorm(scale, MaxNormLinear()), data.model, sizediff)

    default = selectobjective!(DefaultJuMPSelectionStrategy(), v, data)
    return @expression(data.model, default - sizediffnorm)
end

valueobjective!(::AbstractJuMPSelectionStrategy, v, data) = @expression(data.model, sum(data.valuefun(v) .* data.outselectvars[v]))

function insertlastobjective!(s, v, data)
    insvars = data.outinsertvars[v]
    preferend = collect(1:length(insvars))
    return @expression(data.model, 0.05*sum(insvars .* preferend))
end

function extract_ininds_and_outinds(s, outselectvars::Dict, outinsertvars::Dict)
    outinds = Dict([v => extract_inds(s, outselectvars[v], outinsertvars[v]) for v in keys(outselectvars)])
    ininds = Dict([v => getall(outinds, inputs(v)) for v in keys(outselectvars)])
    return ininds, outinds
end

function extract_inds(::AbstractJuMPSelectionStrategy, selectvars::T, insertvars::T) where T <: AbstractVector{JuMP.VariableRef}
    # insertvar is 1.0 at indices where a new output shall be added and 0.0 where an existing one shall be selected
    result = -round.(Int, JuMP.value.(insertvars))
    selected = findall(xi -> xi > 0, JuMP.value.(selectvars))

    # TODO: Needs investigation
    sum(result) == 0 && return selected

    j = 1
    for i in eachindex(result)
        if result[i] == 0
            result[i] = selected[j]
            j += 1
        end
    end

    return result
end


text:
using PyPlot
using Printf

pi5(x, y) = x > y ? (x-y)^5/120 : zero(x)
w1 = 1/6
w2 = 5/6
a = 1/sqrt(5)

K(y) = (1-y)^6/720 - ( w2*pi5(-a,y) + w2*pi5(a,y) + w1*pi5(1,y) )

figure(1)
y = range(-1, stop=1, length=201)
plot(y, K.(y))
grid(true)
xlabel(L"$y$")
ylabel(L"$K(y)$")
savefig("quadrature_kernel.pdf")

E2(n) = 2//(2n+1) - (1+5//5^n)/3
for n = 0:3
    err = E2(n)
    @printf("%6d  %0d / %0d\n", n, err.num, err.den)
end

ext: jl

max_stars_repo_name: billmclean/ComputationalMathsNotes

max_issues_repo_name: billmclean/ComputationalMathsNotes

max_forks_repo_name: billmclean/ComputationalMathsNotes


text:
"""
    fill_nans(x, intp::Linear)

Fill `NaN` values in `x` using linear interpolation. `NaN`s at
the left or right edges of `x` are not interpolated, but left
as they are. See also [`fill_nans!`](@ref).

## Example

```julia
x = [NaN, 1.4, 5.6, NaN, 4.3, 3.1, NaN, NaN, 5.6, NaN]
fill_nans(x, Linear())
```
"""
function fill_nans(x, intp::Linear)
    if isnan(x[1]) && isnan(x[end])
        nan_chunks = findall_nan_chunks(x)[2:end-1]
    elseif isnan(x[1])
        nan_chunks = findall_nan_chunks(x)[2:end]
    elseif isnan(x[end])
        nan_chunks = findall_nan_chunks(x)[1:end-1]
    else
        nan_chunks = findall_nan_chunks(x)
    end

    x_filled = copy(x)

    for (start_idx, stop_idx) in nan_chunks
        x1 = start_idx - 1
        x2 = stop_idx + 1

        y1 = x[start_idx - 1]
        y2 = x[stop_idx + 1]

        slope = (y2 - y1) / (x2 - x1)

        intcpt = y1 - slope*x1

        for (i, idx) in enumerate(start_idx:stop_idx)
            x_filled[idx] = y1 + i*slope
        end
    end

    return x_filled
end

"""
    fill_nans!(x, intp::Linear)

Fill `NaN` values in `x` in-place using linear interpolation. `NaN`s at
the left or right edges of `x` are not interpolated, but left
as they are. See also [`fill_nans`](@ref).

## Example

```julia
x = [NaN, 1.4, 5.6, NaN, 4.3, 3.1, NaN, NaN, 5.6, NaN]
fill_nans!(x, Linear())
```
"""
function fill_nans!(x, intp::Linear)
    if isnan(x[1]) && isnan(x[end])
        nan_chunks = findall_nan_chunks(x)[2:end-1]
    elseif isnan(x[1])
        nan_chunks = findall_nan_chunks(x)[2:end]
    elseif isnan(x[end])
        nan_chunks = findall_nan_chunks(x)[1:end-1]
    else
        nan_chunks = findall_nan_chunks(x)
    end

    for (start_idx, stop_idx) in nan_chunks
        x1 = start_idx - 1
        x2 = stop_idx + 1

        y1 = x[start_idx - 1]
        y2 = x[stop_idx + 1]

        slope = (y2 - y1) / (x2 - x1)

        intcpt = y1 - slope*x1

        for (i, idx) in enumerate(start_idx:stop_idx)
            x[idx] = y1 + i*slope
        end
    end

    return x
end

export fill_nans, fill_nans!


text:
using PMDPs
using RandomNumbers.Xorshifts
using DrWatson
using BSON

using POMDPSimulators # load histories
using StaticArrays, Distributions # load

using LightGraphs
using GraphPlot
using Cairo, Compose

N_individ=100
pps = [
    # Dict(pairs((nᵣ=3, c=3, T=10, expected_res=3., res_budget_μ=5.))),
    # Dict(pairs((nᵣ=6, c=5, T=100, expected_res=60., res_budget_μ=5.))),
    # Dict(pairs((nᵣ=10, c=5, T=100, expected_res=100., res_budget_μ=5.))),
    # Dict(pairs((nᵣ=10, c=40, T=1000, expected_res=Float64(800), res_budget_μ=5.))),
    # Dict(pairs((nᵣ=50, c=40, T=1000, expected_res=Float64(4000), res_budget_μ=5.)))
    # Dict(pairs((nᵣ=6, c=5, T=100, expected_res=60., res_budget_μ=5., objective=:utilization))),
    Dict(pairs((nᵣ=3, c=3, T=10, expected_res=3., res_budget_μ=5., objective=:utilization))),
    Dict(pairs((nᵣ=10, c=5, T=100, expected_res=100., res_budget_μ=5., objective=:utilization))),
    Dict(pairs((nᵣ=10, c=40, T=1000, expected_res=Float64(800), res_budget_μ=5., objective=:utilization))),
    Dict(pairs((nᵣ=50, c=40, T=1000, expected_res=Float64(4000), res_budget_μ=5., objective=:utilization)))
]
name = "linear_problem"
for pp_params in pps
    display("Generating $name with $pp_params")
    pp = PMDPs.linear_pp(;pp_params...)
    mg = PMDPs.PMDPg(pp)

    rnd = Xorshift128Plus(1)
    traces = [PMDPs.simulate_trace(mg, rnd) for i in 1:N_individ]

    sname = savename("traces_lp", pp_params,  "bson")
    @tagsave(datadir("traces", sname), @dict(name, pp, pp_params, traces))
end

N_individ=100
gpps = [
    Dict(pairs((NV=5, NE=8, seed=1, NP=20, c=5, T=100, expected_res=Float64(80), res_budget_μ=5., objective=:utilization))),
    Dict(pairs((NV=8, NE=20, seed=1, NP=50, c=10, T=1000, expected_res=Float64(400), res_budget_μ=5., objective=:utilization))),
    Dict(pairs((NV=15, NE=30, seed=1, NP=100, c=10, T=1000, expected_res=Float64(600), res_budget_μ=5., objective=:utilization))),
    Dict(pairs((NV=30, NE=45, seed=1, NP=100, c=10, T=1000, expected_res=Float64(900), res_budget_μ=5., objective=:utilization)))
]
name = "graph_problem"
for pp_params in gpps
    display("Generating $name with $pp_params")
    pp, g = PMDPs.graph_pp(;pp_params...)
    display(gplot(g, nodelabel=1:nv(g)))

    #  https://juliapackages.com/p/graphplot
    draw(PNG(plotsdir(savename("graph_", pp_params, "png")), 16cm, 16cm), gplot(g))
    # mg = PMDPs.PMDPg(pp)

    # rnd = Xorshift128Plus(1)
    # traces = [PMDPs.simulate_trace(mg, rnd) for i in 1:N_individ]

    # sname = savename("traces_gp", pp_params,  "bson")
    # @tagsave(datadir("traces", sname), @dict(name, pp, pp_params, traces, g))
end

const N = 100 # number of traces
# Linear instances
for expected_res in 50:50:1200
    pp_params = Dict(pairs((nᵣ=10, c=40, T=1000, expected_res=Float64(expected_res), res_budget_μ=5.)))
    name = "linear_problem"
    display("Generating $name with $pp_params")
    pp = PMDPs.linear_pp(;pp_params...)
    mg = PMDPs.PMDPg(pp)

    rnd = Xorshift128Plus(1)
    traces = [PMDPs.simulate_trace(mg, rnd) for i in 1:N]

    sname = savename("traces_lp", pp_params,  "bson")
    @tagsave(datadir("traces", sname), @dict(name, pp, pp_params, traces))
end

# Graph instances
for seed in 1:10
    for expected_res in 25:25:600
        pp_params = Dict(pairs((NV=8, NE=20,seed=seed, NP=50, c=10, T=1000, expected_res=Float64(expected_res), res_budget_μ=5.)))
        name = "graph_problem"
        display("Generating $name with $pp_params")
        pp, g = PMDPs.graph_pp(;pp_params...)
        mg = PMDPs.PMDPg(pp)

        rnd = Xorshift128Plus(1)
        traces = [PMDPs.simulate_trace(mg, rnd) for i in 1:N]

        sname = savename("traces_gp", pp_params,  "bson")
        @tagsave(datadir("traces", sname), @dict(name, pp, pp_params, traces, g))
    end
end


text:
"""
Utilities to evaluate players against one another.

Typically, between each training iteration, different players
that possibly depend on the current neural network
compete against a set of baselines.
"""
module Benchmark

import ..Network, ..MinMax, ..GI
import ..Env, ..MCTS, ..MctsParams, ..TwoPlayers
import ..simulate, ..Simulator, ..rewards_and_redundancy, ..record_trace
import ..ColorPolicy, ..ALTERNATE_COLORS
import ..AbstractPlayer, ..EpsilonGreedyPlayer, ..NetworkPlayer, ..MctsPlayer
import ..PlayerWithTemperature, ..ConstSchedule

using ProgressMeter
using Statistics: mean

"""
    Benchmark.DuelOutcome

The outcome of a duel between two players.

# Fields
- `player` and `baseline` are `String` fields containing the names of
    both players involved in the duel
- `avgr` is the averagereward collected by `player`
- `rewards` is the sequence of rewards collected by `player` (one per game)
- `redundancy` is the ratio of duplicate positions encountered during the
   evaluation, not counting the initial position. If this number is too high,
   you may want to increase the move selection temperature.
- `time` is the computing time spent running the duel, in seconds
"""
struct DuelOutcome
  player :: String
  baseline :: String
  avgr :: Float64
  rewards :: Vector{Float64}
  redundancy :: Float64
  time :: Float64
end

"""
    Benchmark.Report = Vector{Benchmark.DuelOutcome}

A benchmark report is a vector of [`Benchmark.DuelOutcome`](@ref) objects.
"""
const Report = Vector{DuelOutcome}

"""
    Benchmark.Player

Abstract type to specify a player that can be featured in a benchmark duel.

Subtypes must implement the following functions:
  - `Benchmark.instantiate(player, nn)`: instantiate the player specification
      into an [`AbstractPlayer`](@ref) given a neural network
  - `Benchmark.name(player)`: return a `String` describing the player
"""
abstract type Player end

# instantiate(::Player, nn)
function instantiate end

# name(::Player) :: String
function name end

"""
    Benchmark.Duel(player, baseline; num_games)

Specify a duel that consists in `num_games` games between
`player` and `baseline`, each of them of type [`Benchmark.Player`](@ref).

# Optional keyword arguments
- `reset_every`: if set, the MCTS tree is reset every `reset_mcts_every` games
    to avoid running out of memory
- `color_policy` has type [`ColorPolicy`](@ref) and is `ALTERNATE_COLORS`
    by default
"""
struct Duel
  num_games :: Int
  num_workers :: Int
  use_gpu :: Bool
  reset_every :: Union{Nothing, Int}
  flip_probability :: Float64
  color_policy :: ColorPolicy
  player :: Player
  baseline :: Player
  function Duel(player, baseline;
      num_games, num_workers, use_gpu=false, reset_every=nothing,
      color_policy=ALTERNATE_COLORS, flip_probability=0.)
    return new(num_games, num_workers, use_gpu, reset_every,
      flip_probability, color_policy, player, baseline)
  end
end

"""
    Benchmark.run(env::Env, duel::Benchmark.Duel, progress=nothing)

Run a benchmark duel and return a [`Benchmark.DuelOutcome`](@ref).

If a `progress` is provided, `next!(progress)` is called
after each simulated game.
"""
function run(env::Env{G}, duel::Duel, progress=nothing) where G
  net() = Network.copy(env.bestnn, on_gpu=duel.use_gpu, test_mode=true)
  simulator = Simulator(net, record_trace) do net
    player = instantiate(duel.player, net)
    baseline = instantiate(duel.baseline, net)
    return TwoPlayers(player, baseline)
  end
  samples, elapsed = @timed simulate(
    simulator,
    num_games=duel.num_games,
    num_workers=duel.num_workers,
    game_simulated=(() -> next!(progress)),
    reset_every=duel.reset_every,
    flip_probability=duel.flip_probability,
    color_policy=duel.color_policy)
  gamma = env.params.self_play.mcts.gamma
  rewards, redundancy = rewards_and_redundancy(samples, gamma=gamma)
  avgr = mean(rewards)
  return DuelOutcome(
    name(duel.player), name(duel.baseline), avgr, rewards, redundancy, elapsed)
end

#####
##### Statistics for games with ternary rewards
#####

struct TernaryOutcomeStatistics
  num_won  :: Int
  num_draw :: Int
  num_lost :: Int
end

function TernaryOutcomeStatistics(rewards::AbstractVector{<:Number})
  num_won  = count(==(1), rewards)
  num_draw = count(==(0), rewards)
  num_lost = count(==(-1), rewards)
  @assert num_won + num_draw + num_lost == length(rewards)
  return TernaryOutcomeStatistics(num_won, num_draw, num_lost)
end

function TernaryOutcomeStatistics(outcome::DuelOutcome)
  return TernaryOutcomeStatistics(outcome.rewards)
end

#####
##### Standard players
#####

"""
    Benchmark.MctsRollouts(params) <: Benchmark.Player

Pure MCTS baseline that uses rollouts to evaluate new positions.

Argument `params` has type [`MctsParams`](@ref).
"""
struct MctsRollouts <: Player
  params :: MctsParams
end

name(p::MctsRollouts) = "MCTS ($(p.params.num_iters_per_turn) rollouts)"

function instantiate(p::MctsRollouts, nn::MCTS.Oracle{G}) where G
  return MctsPlayer(MCTS.RolloutOracle{G}(), p.params)
end

"""
    Benchmark.Full(params) <: Benchmark.Player

Full AlphaZero player that combines MCTS with the learnt network.

Argument `params` has type [`MctsParams`](@ref).
"""
struct Full <: Player
  params :: MctsParams
  Full(params) = new(params)
end

name(::Full) = "AlphaZero"

function instantiate(p::Full, nn)
  return MctsPlayer(nn, p.params)
end

"""
    Benchmark.NetworkOnly(;τ=1.0) <: Benchmark.Player

Player that uses the policy output by the learnt network directly,
instead of relying on MCTS.
"""
struct NetworkOnly <: Player
  τ :: Float64
  NetworkOnly(;τ=1.0) = new(τ)
end

name(::NetworkOnly) = "Network Only"

function instantiate(p::NetworkOnly, nn)
  player = NetworkPlayer(nn)
  return PlayerWithTemperature(player, ConstSchedule(p.τ))
end

"""
    Benchmark.MinMaxTS(;depth, τ=0.) <: Benchmark.Player

Minmax baseline, which relies on [`MinMax.Player`](@ref).
"""
struct MinMaxTS <: Player
  depth :: Int
  amplify_rewards :: Bool
  τ :: Float64
  MinMaxTS(;depth, amplify_rewards, τ=0.) = new(depth, amplify_rewards, τ)
end

name(p::MinMaxTS) = "MinMax (depth $(p.depth))"

function instantiate(p::MinMaxTS, ::MCTS.Oracle{G}) where G
  return MinMax.Player{G}(
    depth=p.depth, amplify_rewards=p.amplify_rewards, τ=p.τ)
end

"""
    Benchmark.Solver(;ϵ) <: Benchmark.Player

Perfect solver that plays randomly with probability `ϵ`.
"""
struct Solver <: Player
  ϵ :: Float64
  Solver(;ϵ) = new(ϵ)
end

# Return the type of the perfect player for a given type of game
# PerfectPlayer(::Type{<:GI.AbstractGame})
function PerfectPlayer end

name(p::Solver) = "Perfect Player ($(round(Int, 100 * p.ϵ))% random)"

function instantiate(p::Solver, nn::MCTS.Oracle{G}) where G
  return EpsilonGreedyPlayer(PerfectPlayer(G)(), p.ϵ)
end

end


text:
function get_number_in_each_fold(num_samples::Integer, num_folds::Integer)::Vector{Int}
    if num_samples < num_folds
        throw(ArgumentError("`num_samples` must be greater than or equal to `num_folds`"))
    end
    result::Vector{Int} = zeros(Int, num_folds)
    if num_samples % num_folds == 0
        result .= Int(num_samples//num_folds)
    else
        min_number_per_fold = round(Int, floor(num_samples/num_folds))
        remaining_number = num_samples - min_number_per_fold*num_folds
        result[1:remaining_number] .= min_number_per_fold + 1
        result[(remaining_number+1):end] .= min_number_per_fold
    end
    @assert sum(result) == num_samples
    return result
end

function get_indices_in_each_fold(all_indices::AbstractVector{<:Integer}, num_folds::Integer)::Vector{Vector{Int}}
    all_indices_deepcopy = deepcopy(all_indices)
    unique!(all_indices_deepcopy)
    sort!(all_indices_deepcopy)
    num_samples = length(all_indices_deepcopy)
    number_in_each_fold = get_number_in_each_fold(num_samples, num_folds)
    result::Vector{Vector{Int}} = Vector{Vector{Int}}(undef, num_folds)
    lower::Int = 1
    upper::Int = lower + number_in_each_fold[1] - 1
    result[1] = all_indices_deepcopy[lower:upper]
    @assert length(result[1]) == length(unique(result[1]))
    @assert length(result[1]) == number_in_each_fold[1]
    for fold = 2:num_folds
        lower = upper + 1
        upper = lower + number_in_each_fold[fold] - 1
        result[fold] = all_indices_deepcopy[lower:upper]
        # @assert length(result[fold]) == length(unique(result[fold]))
        # @assert length(result[fold]) == number_in_each_fold[fold]
    end
    @assert all(sort(unique(all_indices_deepcopy)) .== sort(unique(vcat(result...))))
    return result
end

function get_leavein_indices(all_indices::AbstractVector{<:Integer}, num_folds::Integer, fold_choice::Integer)::Vector{Int}
    indices_in_each_fold::Vector{Vector{Int}} = get_indices_in_each_fold(all_indices, num_folds)
    result::Vector{Int} = reduce(vcat, indices_in_each_fold[(1:num_folds).!=fold_choice])
    return result
end

function get_leaveout_indices(all_indices::AbstractVector{<:Integer}, num_folds::Integer, fold_choice::Integer)::Vector{Int}
    indices_in_each_fold::Vector{Vector{Int}} = get_indices_in_each_fold(all_indices, num_folds)
    result::Vector{Int} = indices_in_each_fold[fold_choice]
    return result
end

function vector_to_ranges(vector::T1)::Vector{UnitRange{Int}} where
        T1 <: AbstractVector{T2} where
        T2 <: Integer
    result::Vector{UnitRange{Int}} = Vector{UnitRange{Int}}(undef, 0)
    vector_deepcopy = deepcopy(vector)
    unique!(vector_deepcopy)
    sort!(vector_deepcopy)
    if !isempty(vector_deepcopy)
        lower::Int = vector_deepcopy[1]
        upper::Int = vector_deepcopy[1]
        for x in vector_deepcopy
            if (x) > (upper + 1)
                push!(result, lower:upper)
                lower = x
                upper = x
            else
                upper = x
            end
        end
        push!(result, lower:upper)
    end
    unique!(result)
    sort!(result)
    return result
end

function ranges_to_vector(ranges::T1)::Vector{Int} where
        T1 <: AbstractVector{T2} where
        T2 <: AbstractRange{T3} where
        T3 <: Integer
    ranges_deepcopy = deepcopy(ranges)
    unique!(ranges_deepcopy)
    sort!(ranges_deepcopy)
    result = Vector{Int}(undef, 0)
    for r in ranges_deepcopy
        append!(result, collect(r))
    end
    unique!(result)
    sort!(result)
    return result
end

function Base.:(==)(x::CrossValidation{T}, y::CrossValidation{T})::Bool where T
   result::Bool = (x.leaveout == y.leaveout) && (x.leavein == y.leavein)
   return result
end

function CrossValidation{T}(
        ;
        num_folds_per_level,
        all_indices,
        )::CrossValidation{T} where T
    all_indices_deepcopy = deepcopy(all_indices)
    unique!(all_indices_deepcopy)
    sort!(all_indices_deepcopy)
    if length(num_folds_per_level) < 1
        leaveout = Vector{Vector{T}}(undef, 1)
        leaveout[1] = all_indices_deepcopy
        leavein = Vector{CrossValidation{T}}(undef, 0)
    else
        num_folds_this_level = num_folds_per_level[1]
        remaining_num_folds_per_level = num_folds_per_level[2:end]
        leaveout = Vector{Vector{T}}(undef, num_folds_this_level)
        leavein = Vector{CrossValidation{T}}(undef, num_folds_this_level)
        for fold_choice = 1:num_folds_this_level
            leaveout[fold_choice] = get_leaveout_indices(all_indices, num_folds_this_level, fold_choice)
            leavein[fold_choice] = CrossValidation{T}(
                ;
                num_folds_per_level = remaining_num_folds_per_level,
                all_indices = get_leavein_indices(all_indices, num_folds_this_level, fold_choice)
                )
        end
    end
    result::CrossValidation{T} = CrossValidation{T}(leavein, leaveout)
    return result
end

function CrossValidation{T}(
        cv::CrossValidation{R},
        )::CrossValidation{T} where
        T where
        R <: AbstractUnitRange{S} where
        S <: Integer
    if isleaf(cv)
        leaveout = Vector{Vector{T}}(undef, 1)
        leaveout[1] = ranges_to_vector(get_leavein_indices(cv))
        leavein = Vector{CrossValidation{T}}(undef, 0)
    else
        num_folds = get_top_level_num_folds(cv)
        leaveout = Vector{Vector{T}}(undef, num_folds)
        leavein = Vector{CrossValidation{T}}(undef, num_folds)
        for fold_choice = 1:num_folds
            leaveout[fold_choice] = ranges_to_vector(get_leaveout_indices(cv, fold_choice))
            leavein[fold_choice] = CrossValidation{T}(get_leavein_cv(cv, fold_choice))
        end
    end
    result::CrossValidation{T} = CrossValidation{T}(leavein, leaveout)
    return result
end

function CrossValidation{T}(
        cv::CrossValidation{<:Integer},
        )::CrossValidation{T} where T
    if isleaf(cv)
        leaveout = Vector{Vector{T}}(undef, 1)
        leaveout[1] = vector_to_ranges(get_leavein_indices(cv))
        leavein = Vector{CrossValidation{T}}(undef, 0)
    else
        num_folds = get_top_level_num_folds(cv)
        leaveout = Vector{Vector{T}}(undef, num_folds)
        leavein = Vector{CrossValidation{T}}(undef, num_folds)
        for fold_choice = 1:num_folds
            leaveout[fold_choice] = vector_to_ranges(get_leaveout_indices(cv, fold_choice))
            leavein[fold_choice] = CrossValidation{T}(get_leavein_cv(cv, fold_choice))
        end
    end
    result::CrossValidation{T} = CrossValidation{T}(leavein, leaveout)
    return result
end

function isleaf(cv::CrossValidation{T})::Bool where T
    result::Bool = length(cv.leavein) < 1
    return result
end

function get_leavein_cv(cv::CrossValidation{T}, fold_choice)::CrossValidation{T} where T
    if isleaf(cv)
        throw(ArgumentError("`cv` is a leaf, so `get_leavein_cv(cv, fold_choice)` is not defined"))
    end
    result::CrossValidation{T} = cv.leavein[fold_choice]
    return result
end

function get_leavein_indices(cv::CrossValidation{T})::Vector{T} where T
    if !isleaf(cv)
        throw(ArgumentError("`cv` is not a leaf, so `get_leavein_indices(cv)` is not defined"))
    end
    result::Vector{T} = cv.leaveout[1]
    return result
end

function get_leaveout_indices(cv::CrossValidation{T}, fold_choice)::Vector{T} where T
    if isleaf(cv)
        throw(ArgumentError("`cv` is a leaf, so `get_leaveout_indices(cv, fold_choice)` is not defined"))
    end
    result::Vector{T} = cv.leaveout[fold_choice]
    return result
end

function get_all_indices(cv::CrossValidation{T})::Vector{T} where T
    result::Vector{T} = sort(
        unique(
            reduce(vcat, cv.leaveout)
            )
        )
    return result
end

function get_top_level_num_folds(cv::CrossValidation{T})::Int where T
    if isleaf(cv)
        result = 0
    else
        if length(cv.leaveout) != length(cv.leavein)
            throw(DimensionMismatch("length(cv.leaveout) != length(cv.leavein)"))
        end
        result = length(cv.leaveout)
    end
    return result
end


text:
@noinline function find_first_tests(x, collection)
    @testset "find first test - $x, $(typeof(collection))" begin
        @test @inferred(catch_nothing(find_first(>(x), collection))) == catch_nothing(findfirst(i -> i > x,  collection))
        @test @inferred(catch_nothing(find_first(>=(x), collection))) == catch_nothing(findfirst(i -> i >= x, collection))
        @test @inferred(catch_nothing(find_first(<(x), collection))) == catch_nothing(findfirst(i -> i < x, collection))
        @test @inferred(catch_nothing(find_first(<=(x), collection))) == catch_nothing(findfirst(i -> i <= x, collection))
        @test @inferred(catch_nothing(find_first(==(x), collection))) == catch_nothing(findfirst(i -> i == x, collection))
    end
end

@noinline function find_last_tests(x, collection)
    @testset "find first test - $x, $(typeof(collection))" begin
        @test @inferred(catch_nothing(find_last(>(x), collection))) == catch_nothing(findlast(i -> i > x,  collection))
        @test @inferred(catch_nothing(find_last(>=(x), collection))) == catch_nothing(findlast(i -> i >= x, collection))
        @test @inferred(catch_nothing(find_last(<(x), collection))) == catch_nothing(findlast(i -> i < x, collection))
        @test @inferred(catch_nothing(find_last(<=(x), collection))) == catch_nothing(findlast(i -> i <= x, collection))
        @test @inferred(catch_nothing(find_last(==(x), collection))) == catch_nothing(findlast(i -> i == x, collection))
    end
end

@noinline function find_all_tests(x, collection)
    @testset "find first test - $x, $(typeof(collection))" begin
        @test @inferred(catch_nothing(find_all(>(x), collection))) == catch_nothing(findall(i -> i > x,  collection))
        @test @inferred(catch_nothing(find_all(>=(x), collection))) == catch_nothing(findall(i -> i >= x, collection))
        @test @inferred(catch_nothing(find_all(<(x), collection))) == catch_nothing(findall(i -> i < x, collection))
        @test @inferred(catch_nothing(find_all(<=(x), collection))) == catch_nothing(findall(i -> i <= x, collection))
        @test @inferred(catch_nothing(find_all(==(x), collection))) == catch_nothing(findall(i -> i == x, collection))
    end
end

@noinline function filter_tests(x, collection)
    @testset "filter test - $x, $(typeof(collection))" begin
        @test @inferred(filter(>(x), collection)) ==
              filter(i -> i > x, collection) ==
              filter(>(x), collection)

        @test @inferred(filter(>=(x), collection)) ==
              filter(i -> i >= x, collection) ==
              filter(>=(x), collection)

        @test @inferred(filter(<(x), collection)) ==
              filter(i -> i < x, collection) ==
              filter(<(x), collection)

        @test @inferred(filter(<=(x), collection)) ==
              filter(i -> i <= x, collection) ==
              filter(<=(x), collection)

        @test @inferred(filter(==(x), collection)) ==
              filter(i -> i == x, collection) ==
              filter(==(x), collection)
    end
end

@noinline function count_tests(x, collection)
    @testset "count test - $x, $(typeof(collection))" begin
        @test @inferred(count(>(x), collection)) == count(i -> i > x, collection)
        @test @inferred(count(>=(x), collection)) == count(i -> i >= x, collection)
        @test @inferred(count(<(x), collection)) == count(i -> i < x, collection)
        @test @inferred(count(<=(x), collection)) == count(i -> i <= x, collection)
        @test @inferred(count(==(x), collection)) == count(i -> i == x, collection)
    end
end

const range_list = (DynamicAxis(10),
                    #IdOffsetRange(OneTo(10), 2),
                    1:10,
                    #IdOffsetRange(1:10, 2),
                    mutable(UnitRange(1.0, 10.0)),
                    static(StepRange(1, 2, 10)),
                    StepRange(10, -2, 1),
                    LinRange(1, 5, 10),
                    LinRange(5, 1, 10),
                    mrange(1.0, step=0.25, stop=10),
                    mrange(10.0, step=-0.25, stop=1.0))

@testset "find_first" begin
    for x in (-1, -0.5, 0, 0.5, 1)
        for r in range_list
            find_first_tests(x, r)
        end
    end
end

@testset "find_last" begin
    for x in (-1, -0.5, 0, 0.5, 1)
        for r in range_list
            find_last_tests(x, r)
        end
    end
end

@testset "find_all" begin
    for x in (-1, -0.5, 0, 0.5, 1)
        for r in range_list
            find_all_tests(x, r)
        end
    end
end

@testset "find_all(in)" begin
    for (x, y, z) in ((1:10, 1:2:10, 1:5),
                      (1:2:20, 1:8:20, 1:3),
                      (1:2:20, 1:10, 1:2:9),
                      (1:2:10, 1.1:1:10.1, []),
                      (UnitRange(1.1, 10.1), UnitRange(2.1, 8.1), 1:7),
                      (UnitRange(2.1, 8.1), UnitRange(1.1, 10.1), 2:8))
        @testset "find_all(in($x), $y)" begin
            @test find_all(in(x), y) == z
        end
    end
end

@testset "find_all(in(x), r)" begin
    r = @inferred(find_all(in(OneTo(10)), static(OneTo(8))))
    @test r == 1:8

    r = @inferred(find_all(in(OneTo(10)), DynamicAxis(8)))
    @test r == 1:8

    r = @inferred(find_all(in(static(OneTo(8))), static(OneTo(10))))
    @test r == 1:8

    r = @inferred(find_all(in(UnitRange(1,10)), static(UnitRange(1,8))))
    @test r == 1:8
    @test isa(r, UnitRange)

    r = @inferred(find_all(in(OneTo(10)), mutable(UnitRange(1, 8))))
    @test r == OneTo(8)
    @test isa(r, UnitRange) == true

    r = @inferred(find_all(in(static(UnitRange(1, 8))), static(UnitRange(1, 10))))
    @test r == static(UnitRange(1, 8))

    @test find_all(in(collect(1:10)), 1:20) == find_all(in(1:10), 1:20)
    @test find_all(in(1:10), collect(1:20)) == 1:10

    @testset "steps match but no overlap" begin
        r = @inferred(find_all(in(1:3), 4:5))
        @test r == 1:0
        @test isa(r, UnitRange)
    end

    @test find_all(in([1, 2, 3]), collect(1:10)) == [1, 2, 3]
end

@testset "find_all(in(::IntervalSets), r)" begin
    for (i,t) in ((Interval{:closed,:closed}(1, 10), 1:10),
                  (Interval{:open,:closed}(1, 10), 2:10),
                  (Interval{:closed,:open}(1, 10), 1:9),
                  (Interval{:open,:open}(1, 10), 2:9))
        @test find_all(in(i), 1:10) == t
    end
end

#=
f1 = <(1)
f2 = <(3)
bitop = or
=#

function test_find_chained_fix(m, s, b)
    for i1 in (m[1] - step(m), m[1], m[4], m[5] + 2step(m))
        for i2 in (m[2], m[3], m[5], m[5] + step(m))
            for f in (<, >, <=, >=, ==)
               for f2 in (<, >, <=, >=, ==)
                    for bitop in (and, or)
                        @testset "find_all($bitop($f($i1), $f2($i2)), $b)" begin
                            @test @inferred(find_all(bitop(f(i1),  f2(i2)), m)) == find_all(x -> bitop(f(i1)(x), f2(i2)(x)), b)
                        end

                        @testset "find_all($bitop($f($i2), $f2($i1)), $b)" begin
                            @test @inferred(find_all(bitop(f(i2),  f2(i1)), m)) == findall(x -> bitop(f(i2)(x), f2(i1)(x)), b)
                        end
                    end
                end
            end
        end
    end
end

@testset "findall(::ChainedFix,...)" begin
    test_find_chained_fix(DynamicAxis(5), static(OneTo(5)), OneTo(5))
    test_find_chained_fix(mutable(UnitRange(2, 6)), static(UnitRange(2, 6)), UnitRange(2, 6))
    test_find_chained_fix(mutable(StepRangeLen(1, 3, 5)), static(StepRangeLen(1, 3, 5)), StepRangeLen(1, 3, 5))
end

@testset "find not equal" begin
    for (m,s,b) in ((DynamicAxis(5), static(OneTo(5)), OneTo(5)),
                    (mutable(UnitRange(2, 6)), static(UnitRange(2, 6)), UnitRange(2, 6)),
                    (mutable(StepRangeLen(1, 3, 5)), static(StepRangeLen(1, 3, 5)), StepRangeLen(1, 3, 5))
                   )
        for i in (m[1] - step(m), m[1], m[4], m[5] + 2step(m))
            @testset "find_all(!=($i), $b)" begin
                @test @inferred(findall(!=(i), m)) == @inferred(findall(!=(i), b))
                @test @inferred(findall(!=(i), s)) == @inferred(findall(!=(i), b))
            end
        end
    end
end

@testset "chained finds" begin
    for (m,s,b) in ((DynamicAxis(5), static(OneTo(5)), OneTo(5)),
                    (mutable(UnitRange(2, 6)), static(UnitRange(2, 6)), UnitRange(2, 6)),
                    (mutable(StepRange(1, 2, 11)), static(StepRange(1, 2, 11)), StepRange(1, 2, 11)),
                   )
        @testset "Type: $(typeof(b))" begin
            for i1 in (m[1] - step(m), m[1], m[4], m[5] + 2step(m))
                for i2 in (m[2], m[3], m[5], m[5] + step(m))
                    for f in (<, >, <=, >=, ==)
                        @testset "Comparison: $f" begin
                            if !isa(b, LinRange)
                                for f2 in (<, >, <=, >=, ==)
                                    for bitop in (and, or)
                                        @testset "filter($f($i1) $bitop $f2($i2), $b)" begin
                                            @test filter(bitop(f(i1),  f2(i2)), m) == filter(x -> bitop(f(i1)(x), f2(i2)(x)), b)
                                            @test filter(bitop(f(i2),  f2(i1)), m) == filter(x -> bitop(f(i2)(x), f2(i1)(x)), b)
                                        end
                                    end
                                end
                            end
                            for f2 in (<, >, <=, >=, ==)
                                for bitop in (and, or)
                                    @testset "find_all($f($i1) $bitop $f2($i2), $b)" begin
                                        @test find_all(bitop(f(i1),  f2(i2)), m) == find_all(x -> bitop(f(i1)(x), f2(i2)(x)), b)
                                        @test find_all(bitop(f(i2),  f2(i1)), m) == find_all(x -> bitop(f(i2)(x), f2(i1)(x)), b)
                                    end
                                end
                            end
                        end
                    end
                end
            end
        end
    end
end

@testset "filter non numerics" begin
    x = Second(1):Second(1):Second(10)
    @test @inferred(find_all(and(>=(Second(1)), <=(Second(3))), x)) == 1:3
    @test @inferred(find_all(in(Second(1):Second(1):Second(3)), x)) == 1:3

    x = Second(10):Second(-1):Second(1)
    @test @inferred(find_all(in(Second(1):Second(1):Second(3)), x)) == 10:1:9
end

@testset "filter" begin
    for (m,s,b) in ((DynamicAxis(5), static(OneTo(5)), OneTo(5)),
                    (mutable(UnitRange(2, 6)), static(UnitRange(2, 6)), UnitRange(2, 6)),
                    (mutable(StepRange(1, 2, 11)), static(StepRange(1, 2, 11)), StepRange(1, 2, 11)),
                    (mutable(StepRange(11, -2, 1)), static(StepRange(11, -2, 1)), StepRange(11, -2, 1)),
                    (mutable(LinRange(1, 10, 5)), static(LinRange(1, 10, 5)), LinRange(1, 10, 5)),
                    (mutable(StepRangeLen(1, 3, 5)), static(StepRangeLen(1, 3, 5)), StepRangeLen(1, 3, 5))
                   )
        @testset "Type: $(typeof(b))" begin
            for i1 in (m[1] - step(m), m[1], m[4], m[5] + 2step(m))
                for i2 in (m[2], m[3], m[5], m[5] + step(m))
                    for f in (<, >, <=, >=, ==)
                        @testset "Comparison: $f" begin
                            # FIXME - filter AbstractLinRange doesn't come out as
                            # expected because getindex results in inexact values
                            # in both base and AbstractLinRange.
                            if !isa(b, LinRange)
                                for f2 in (<, >, <=, >=, ==)
                                    for bitop in (and, or)
                                        @testset "filter($f($i1) $bitop $f2($i2), $b)" begin
                                            @test @inferred(filter(bitop(f(i1),  f2(i2)), m)) == filter(x -> bitop(f(i1)(x), f2(i2)(x)), b)
                                            @test @inferred(filter(bitop(f(i2),  f2(i1)), m)) == filter(x -> bitop(f(i2)(x), f2(i1)(x)), b)
                                        end
                                    end
                                end
                            end
                       end
                    end
                end
            end
        end
    end

    #=
    @testset "Filter with empty range" begin
        for i in (0, 1)
            @testset "Number: $i" begin
                for f in (<, >, <=, >=, ==)
                    # FIXME - filter AbstractLinRange doesn't come out as
                    # expected because getindex results in inexact values
                    # in both base and AbstractLinRange.
                    @testset "filter" begin
                        @test filter(f(i), m) == filter(f(i), b)
                        @test filter(f(i), s) == filter(f(i), b)
                    end
                end
            end
        end
    end
    =#

    for (m,s,b) in ((DynamicAxis(5), static(OneTo(5)), OneTo(5)),
                    (mutable(UnitRange(2, 6)), static(UnitRange(2, 6)), UnitRange(2, 6)),
                    (mutable(StepRange(1, 2, 11)), static(StepRange(1, 2, 11)), StepRange(1, 2, 11)),
                    (mutable(StepRange(11, -2, 1)), static(StepRange(11, -2, 1)), StepRange(11, -2, 1)),
                    (mutable(LinRange(1, 10, 5)), static(LinRange(1, 10, 5)), LinRange(1, 10, 5)),
                    (mutable(StepRangeLen(1, 3, 5)), static(StepRangeLen(1, 3, 5)), StepRangeLen(1, 3, 5))
                   )
        for i in (m[1] - step(m), m[1], m[4], m[5] + 2step(m))
            @testset "filter(!=($i), $b)" begin
                @test filter(!=(i), m) == filter(!=(i), b)
                @test filter(!=(i), s) == filter(!=(i), b)
            end
        end
    end
end

for collection in range_list
    @testset "$collection" begin
        for x in (first(collection) - step(collection),
            first(collection) - step(collection) / 2,
            first(collection),
            first(collection) + step(collection) / 2,
            first(collection) + step(collection),
            last(collection) - step(collection),
            last(collection) - step(collection) / 2,
            last(collection),
            last(collection) + step(collection) / 2,
            last(collection) + step(collection))
            find_first_tests(x, collection)
            find_last_tests(x, collection)
            find_all_tests(x, collection)
            # FIXME - filter AbstractLinRange doesn't come out as
            # expected because getindex results in inexact values
            # in both base and AbstractLinRange.
            if !isa(collection, LinRange)
                filter_tests(x, collection)
            end
            count_tests(x, collection)
        end
    end
end

#=
for x in (-1, -0.5, 0, 0.5, 1, 9, 9.5, 10, 10.5, 11)
    find_first_tests(x, collect(1:10))
    find_last_tests(x, collect(1:10))
    find_all_tests(x, collect(1:10))
end
=#


text:
int_rules_1_1_2_x = @theory begin

    #= ::Subsection::Closed:: =#
    #= 1.1.2.0.x*P(x)*(a+b*x^2)^p =#
    @apply_utils Antiderivative(~Pq * (~a + ~(b') * (~x) ^ 2) ^ ~(p'), ~x) => Antiderivative(ExpandIntegrand(~Pq * (~a + ~b * (~x) ^ 2) ^ ~p, ~x), ~x) <-- FreeQ([~a, ~b], ~x) && (PolyQ(~Pq, ~x) && IGtQ(~p, -2))
    @apply_utils Antiderivative(~Pq * (~a + ~(b') * (~x) ^ 2) ^ ~p, ~x) => Antiderivative(~x * PolynomialQuotient(~Pq, ~x, ~x) * (~a + ~b * (~x) ^ 2) ^ ~p, ~x) <-- FreeQ([~a, ~b, ~p], ~x) && (PolyQ(~Pq, ~x) && (EqQ(Coeff(~Pq, ~x, 0), 0) && Not(MatchQ(~Pq, (~x) ^ ~(m') * ~(u') <-- IntegerQ(m)))))
    @apply_utils Antiderivative(~Px * (~a + ~(b') * (~x) ^ 2) ^ ~(p'), ~x) => Antiderivative(PolynomialQuotient(~Px, ~a + ~b * (~x) ^ 2, ~x) * (~a + ~b * (~x) ^ 2) ^ (~p + 1), ~x) <-- FreeQ([~a, ~b, ~p], ~x) && (PolyQ(~Px, ~x) && EqQ(PolynomialRemainder(~Px, ~a + ~b * (~x) ^ 2, ~x), 0))
    @apply_utils Antiderivative(~Pq * (~a + ~(b') * (~x) ^ 2) ^ ~p, ~x) => With([A = Coeff(~Pq, ~x, 0), Q = PolynomialQuotient(~Pq - Coeff(~Pq, ~x, 0), (~x) ^ 2, ~x)], (A * ~x * (~a + ~b * (~x) ^ 2) ^ (~p + 1)) / ~a + (1 / ~a) * Antiderivative((~x) ^ 2 * (~a + ~b * (~x) ^ 2) ^ ~p * (~a * Q - A * ~b * (2 * ~p + 3)), ~x)) <-- FreeQ([~a, ~b], ~x) && (PolyQ(~Pq, (~x) ^ 2) && (ILtQ(~p + 1 / 2, 0) && LtQ(Expon(~Pq, ~x) + 2 * ~p + 1, 0)))
    @apply_utils Antiderivative(~Pq * (~a + ~(b') * (~x) ^ 2) ^ ~p, ~x) => With([Q = PolynomialQuotient(~Pq, ~a + ~b * (~x) ^ 2, ~x), f = Coeff(PolynomialRemainder(~Pq, ~a + ~b * (~x) ^ 2, ~x), ~x, 0), g = Coeff(PolynomialRemainder(~Pq, ~a + ~b * (~x) ^ 2, ~x), ~x, 1)], ((~a * g - ~b * f * ~x) * (~a + ~b * (~x) ^ 2) ^ (~p + 1)) / (2 * ~a * ~b * (~p + 1)) + (1 / (2 * ~a * (~p + 1))) * Antiderivative((~a + ~b * (~x) ^ 2) ^ (~p + 1) * ExpandToSum(2 * ~a * (~p + 1) * Q + f * (2 * ~p + 3), ~x), ~x)) <-- FreeQ([~a, ~b], ~x) && (PolyQ(~Pq, ~x) && LtQ(~p, -1))
    @apply_utils Antiderivative(~Pq * (~a + ~(b') * (~x) ^ 2) ^ ~p, ~x) => With([q = Expon(~Pq, ~x), e = Coeff(~Pq, ~x, Expon(~Pq, ~x))], (e * (~x) ^ (q - 1) * (~a + ~b * (~x) ^ 2) ^ (~p + 1)) / (~b * (q + 2 * ~p + 1)) + (1 / (~b * (q + 2 * ~p + 1))) * Antiderivative((~a + ~b * (~x) ^ 2) ^ ~p * ExpandToSum((~b * (q + 2 * ~p + 1) * ~Pq - ~a * e * (q - 1) * (~x) ^ (q - 2)) - ~b * e * (q + 2 * ~p + 1) * (~x) ^ q, ~x), ~x)) <-- FreeQ([~a, ~b, ~p], ~x) && (PolyQ(~Pq, ~x) && Not(LeQ(~p, -1)))
end

max_stars_repo_path: src/rules/1 Algebraic functions/1.1 Binomial products/1.1.2 Quadratic/1.1.2.x P(x) (a+b x^2)^p.jl


text:
s = BenchmarkGroup()
suite["persistence"] = s

fname = joinpath(bench_dir,"data","benchgraph_")
Random.seed!(17)
for G in GLIST
    for n in [100]
        g = G(n, 5n, seed=1)
        s["write","gt","$g"] = @benchmarkable writegraph($(fname*"$n.gt"), $g)
        s["write","net","$g"] = @benchmarkable writegraph($(fname*"$n.net"), $g)
        s["write","gml","$g"] = @benchmarkable writegraph($(fname*"$n.gml"), $g)
        s["write","dot","$g"] = @benchmarkable writegraph($(fname*"$n.dot"), $g)
        s["write","graphml","$g"] = @benchmarkable writegraph($(fname*"$n.graphml"), $g)
        s["write","gexf","$g"] = @benchmarkable writegraph($(fname*"$n.gexf"), $g)

        s["read","gt","$g"] = @benchmarkable readgraph($(fname*"$n.gt"), $G)
        s["read","net","$g"] = @benchmarkable readgraph($(fname*"$n.net"), $G)
        # s["read","gml","$g"] = @benchmarkable readgraph($(fname*"$n.gml"), $G) #SLOW
        # s["read","dot","$g"] = @benchmarkable readgraph($(fname*"$n.dot"), $G) #SLOW
        s["read","graphml","$g"] = @benchmarkable readgraph($(fname*"$n.graphml"), $G)
        s["read","gexf","$g"] = @benchmarkable readgraph($(fname*"$n.gexf"), $G)
    end
end


text:
using Statistics
using Dates
using Insolation

using CLIMAParameters
using CLIMAParameters.Planet
struct EarthParameterSet <: AbstractEarthParameterSet end
const param_set = EarthParameterSet()

atol = 1e-6
rtol = 1e-2

## Test zero insolation at night
# sunrise at equator
date = Dates.DateTime(2020, 1, 1, 6, 0, 0)
lon, lat = [0.0, 0.0]
sza, azi, d = instantaneous_zenith_angle(date, lon, lat, param_set)
F = insolation(sza, d, param_set)
@test F ≈ 0.0 atol=atol

S, mu = solar_flux_and_cos_sza(date, lon, lat, param_set)
@test S ≈ tot_solar_irrad(param_set) rtol=4e-2
@test mu ≈ 0.0 atol=atol

# polar night NH 1
date = Dates.DateTime(2020, 12, 20, 11, 0, 0)
lon, lat = [0.0, 80.0]
sza, azi, d = instantaneous_zenith_angle(date, lon, lat, param_set)
F = insolation(sza, d, param_set)
@test F ≈ 0.0 atol=atol

S, mu = solar_flux_and_cos_sza(date, lon, lat, param_set)
@test S ≈ tot_solar_irrad(param_set) rtol=4e-2
@test mu ≈ 0.0 atol=atol

# polar night NH 2
date = Dates.DateTime(2020, 12, 20, 23, 0, 0)
lon, lat = [0.0, 80.0]
sza, azi, d = instantaneous_zenith_angle(date, lon, lat, param_set)
F = insolation(sza, d, param_set)
@test F ≈ 0.0 atol=atol

S, mu = solar_flux_and_cos_sza(date, lon, lat, param_set)
@test S ≈ tot_solar_irrad(param_set) rtol=4e-2
@test mu ≈ 0.0 atol=atol

## Test globally averaged insolation ≈ TSI
ndays, nlats = [365, 181]
d_arr = Array{Int}(round.(collect(range(0, stop = 365, length = ndays))))
l_arr = collect(range(-90, stop = 90, length = nlats))
F_arr = zeros(ndays, nlats)

for (i, d) in enumerate(d_arr)
    for (j, lat) in enumerate(l_arr)
        datei = Dates.DateTime(2020,1,1) + Dates.Day(d)
        θ, dist = daily_zenith_angle(datei, lat, param_set)
        F_arr[i, j] = insolation(θ, dist, param_set)
    end
end

zonal_mean_insol = mean(F_arr, dims=1)
area_fac = abs.(cosd.(l_arr))
global_mean_insol = sum(zonal_mean_insol * area_fac) / sum(area_fac)
@test global_mean_insol ≈ tot_solar_irrad(param_set) / 4 rtol=rtol

## Test symmetry of insolation at equinox
date = Dates.DateTime(2021,3,20,9,34)
F_arr = zeros(nlats)
for (i, lat) in enumerate(l_arr)
    θ, dist = daily_zenith_angle(date, lat, param_set)
    F_arr[i] = insolation(θ, dist, param_set)
end
F_NH = sort(F_arr[l_arr .>= 0])
F_SH = sort(F_arr[l_arr .<= 0])
@test F_NH ≈ F_SH rtol=rtol

## Test invariance of zonal-mean insolation under rotation of ϖ
CLIMAParameters.Planet.lon_perihelion(::EarthParameterSet) = lon_perihelion_epoch(param_set) + π

ndays, nlats = [365, 181]
d_arr = Array{Int}(round.(collect(range(0, stop = 365, length = ndays))))
l_arr = collect(range(-90, stop = 90, length = nlats))
F_arr = zeros(ndays, nlats)

for (i, d) in enumerate(d_arr)
    for (j, lat) in enumerate(l_arr)
        datei = Dates.DateTime(2020,1,1) + Dates.Day(d)
        θ, dist = daily_zenith_angle(datei, lat, param_set)
        F_arr[i, j] = insolation(θ, dist, param_set)
    end
end

zonal_mean_insol_rotate = mean(F_arr, dims=1)
@test zonal_mean_insol_rotate ≈ zonal_mean_insol rtol=rtol


text:
# This file is a part of TypeDBClient.  License is MIT: https://github.com/Humans-of-Julia/TypeDBClient.jl/blob/main/LICENSE

mutable struct TypeDBOptions
    infer::Optional{Bool}
    trace_inference::Optional{Bool}
    explain::Optional{Bool}
    parallel::Optional{Bool}
    batch_size::Optional{Int}
    prefetch::Optional{Bool}
    session_idle_timeout_millis::Optional{Int}
    schema_lock_acquire_timeout_millis::Optional{Int}
end

function TypeDBOptions()
    return TypeDBOptions(nothing,nothing,nothing,nothing,nothing,nothing,nothing,nothing)
end

function typedb_options_core()
    return TypeDBOptions()
end


text:
using Poirot, Poirot.LAX, Test

double(x) = xla(() -> x+x)

@test double(21) == 42
@test double(4.5) == 9.0

let x = 2, y = 2.0
  @test_broken xla(() -> x+y) == 4.0
end

add(a, b) = a+b
@test @code_xla(add(2, 3)) isa Poirot.Abstract.IR

@test xla(() -> 2+2) == 4

let x = 5
  @test xla(() -> 3x^(1+1) + (2x + 1)) == 86
end

relu(x) = xla(() -> x > 0 ? x : 0)

@test relu(5) == 5
@test relu(-5) == 0
@test_broken relu(5.0) == 5.0

let x = rand(3), y = rand(3)
  @test collect(xla(() -> x+y)) == x+y
end

function updatezero!(env)
  if env[:x] < 0
    env[:x] = 0
  end
end

function relu(x)
  env = Dict()
  env[:x] = x
  updatezero!(env)
  return env[:x]
end

@test xla(() -> relu(5)) == 5
@test xla(() -> relu(-5)) == 0


text:
## install package for Gaia.Clustering

using Pkg
ENV["PYTHON"] = ""

Pkg.add("PyCall")
Pkg.build("PyCall")

Pkg.add("DataFrames")
Pkg.add("Clustering")
Pkg.add("Statistics")
Pkg.add("Distributions")
Pkg.add("Random")
Pkg.add("CSV")
Pkg.add("PyPlot")
Pkg.add("Formatting")
Pkg.add("Query")
Pkg.add("TSne")
Pkg.add("Printf")
Pkg.add("Interpolations")
Pkg.add("Images")
Pkg.add("LightGraphs")
Pkg.add("SimpleWeightedGraphs")
Pkg.add("Distances")
Pkg.add("Glob")
Pkg.add("StatsBase")
Pkg.add("FFTW")
Pkg.add("LaTeXStrings")
Pkg.add("MultivariateStats")
Pkg.add("ArgParse")
Pkg.add("VoronoiCells")
Pkg.add("GeometryBasics")

println("## Package installation for GaiaClustering done...")

max_stars_repo_name: bosscha/gaia-shock

max_issues_repo_name: bosscha/gaia-shock

max_forks_repo_name: bosscha/gaia-shock


text:
using Clines, SimpleDrawing, Plots

function apollo_work(C1::Circle, C2::Circle, C3::Circle, min_radius)
    S = soddy(C1, C2, C3)
    if radius(S) < min_radius
        return
    end
    draw(S, linewidth = 0.25)
    apollo_work(C1, C2, S, min_radius)
    apollo_work(C1, C3, S, min_radius)
    apollo_work(C2, C3, S, min_radius)
end

function apollo(C1::Circle, C2::Circle, C3::Circle, min_radius::Real = 0.01)
    for C in [C1, C2, C3]
        draw(C, linewidth = 0.25)
    end
    apollo_work(C1, C2, C3, min_radius)
    finish()
end

function apollo(z1::Number, z2::Number, z3::Number, min_radius::Real = 0.01)
    newdraw()
    CC = kiss(z1, z2, z3)
    apollo(CC..., min_radius)
end

function apollo(min_radius::Real = 0.01)
    angles = [(2 * pi * k / 3 + pi / 6) * im for k = 1:3]
    zz = exp.(angles)
    apollo(zz..., min_radius)
end