"""
    SynchronousExecution

Functor object for using synchronous execution in a progressive-hedging algorithm (assuming multiple Julia cores are available). Create by supplying a [`Synchronous`](@ref) object through `execution` in the `ProgressiveHedgingSolver` factory function and then pass to a `StochasticPrograms.jl` model.

"""
struct SynchronousExecution{T <: AbstractFloat,
                            A <: AbstractVector,
                            PT <: AbstractPenaltyterm} <: AbstractProgressiveHedgingExecution
    subworkers::Vector{SubWorker{T,A,PT}}

    function SynchronousExecution(::Type{T}, ::Type{A}, ::Type{PT}) where {T <: AbstractFloat,
                                                                           A <: AbstractVector,
                                                                           PT <: AbstractPenaltyterm}
        return new{T,A,PT}(Vector{SubWorker{T,A,PT}}(undef, nworkers()))
    end
end

function initialize_subproblems!(ph::AbstractProgressiveHedging,
                                 execution::SynchronousExecution,
                                 scenarioproblems::DistributedScenarioProblems,
                                 penaltyterm::AbstractPenaltyterm)
    # Create subproblems on worker processes
    initialize_subproblems!(ph,
                            execution.subworkers,
                            scenarioproblems,
                            penaltyterm)
    # Initial reductions
    update_iterate!(ph)
    update_dual_gap!(ph)
    return nothing
end

function restore_subproblems!(::AbstractProgressiveHedging, execution::SynchronousExecution)
    restore_subproblems!(execution.subworkers)
    return nothing
end

function resolve_subproblems!(ph::AbstractProgressiveHedging, execution::SynchronousExecution{T}) where T <: AbstractFloat
    partial_solutions = Vector{SubproblemSolution{T}}(undef, nworkers())
    @sync begin
        for (i,w) in enumerate(workers())
            @async partial_solutions[i] = remotecall_fetch(resolve_subproblems!,
                                                           w,
                                                           execution.subworkers[w-1],
                                                           ph.ξ,
                                                           penalty(ph))
        end
    end
    return sum(partial_solutions)
end

function update_iterate!(ph::AbstractProgressiveHedging, execution::SynchronousExecution{T,A}) where {T <: AbstractFloat, A <: AbstractVector}
    partial_primals = Vector{A}(undef, nworkers())
    @sync begin
        for (i,w) in enumerate(workers())
            @async partial_primals[i] = remotecall_fetch(collect_primals,
                                                         w,
                                                         execution.subworkers[w-1],
                                                         length(ph.ξ))
        end
    end
    ξ_prev = copy(ph.ξ)
    ph.ξ .= sum(partial_primals)
    # Update δ₁
    ph.data.δ₁ = norm(ph.ξ - ξ_prev, 2) ^ 2
    return nothing
end

function update_subproblems!(ph::AbstractProgressiveHedging, execution::SynchronousExecution)
    # Update dual prices
    @sync begin
        for w in workers()
            @async remotecall_fetch(
                w,
                execution.subworkers[w-1],
                ph.ξ,
                penalty(ph)) do sw, ξ, r
                    subproblems = fetch(sw)
                    if length(subproblems) > 0
                        update_subproblems!(subproblems, ξ, r)
                    end
                end
        end
    end
    return nothing
end

function update_dual_gap!(ph::AbstractProgressiveHedging, execution::SynchronousExecution)
    return update_dual_gap!(ph, execution.subworkers)
end

function calculate_objective_value(ph::AbstractProgressiveHedging, execution::SynchronousExecution)
    return calculate_objective_value(execution.subworkers)
end

function scalar_subproblem_reduction(value::Function, execution::SynchronousExecution{T}) where T <: AbstractFloat
    partial_results = Vector{T}(undef, nworkers())
    @sync begin
        for (i,w) in enumerate(workers())
            @async partial_results[i] = remotecall_fetch(w, execution.subworkers[w-1], value) do sw, value
                subproblems = fetch(sw)
                return mapreduce(+, subproblems, init = zero(T)) do subproblem
                    π = subproblem.probability
                    return π * value(subproblem)
                end
            end
        end
    end
    return sum(partial_results)
end

function vector_subproblem_reduction(value::Function, execution::SynchronousExecution{T,A}, n::Integer) where {T <: AbstractFloat, A <: AbstractVector}
    partial_results = Vector{A}(undef, nworkers())
    @sync begin
        for (i,w) in enumerate(workers())
            @async partial_results[i] = remotecall_fetch(w, execution.subworkers[w-1], value, n) do sw, value, n
                subproblems = fetch(sw)
                return mapreduce(+, subproblems, init = zero(T, n)) do subproblem
                    π = subproblem.probability
                    return π * value(subproblem)
                end
            end
        end
    end
    return sum(partial_results)
end

# API
# ------------------------------------------------------------
function (execution::Synchronous)(::Type{T}, ::Type{A}, ::Type{PT}) where {T <: AbstractFloat,
                                                                           A <: AbstractVector,
                                                                           PT <: AbstractPenaltyterm}
    return SynchronousExecution(T,A,PT)
end

function str(::Synchronous)
    return "Synchronous "
end


"""
    SynchronousExecution

Functor object for using synchronous execution in an L-shaped algorithm (assuming multiple Julia cores are available). Create by supplying a [`Synchronous`](@ref) object through `execution` in the `LShapedSolver` factory function and then pass to a `StochasticPrograms.jl` model.

"""
struct SynchronousExecution{F <: AbstractFeasibility,
                            T <: AbstractFloat,
                            A <: AbstractVector,
                            S <: LQSolver} <: AbstractExecution
    subworkers::Vector{SubWorker{F,T,A,S}}
    subobjectives::A
    model_objectives::A
    metadata::Vector{MetaData}
    cutqueue::CutQueue{T}

    function SynchronousExecution(nscenarios::Integer, ::Type{F}, ::Type{T}, ::Type{A}, ::Type{S}) where {F <: AbstractFeasibility, T <: AbstractFloat, A <: AbstractVector, S <: LQSolver}
        return new{F,T,A,S}(Vector{SubWorker{F,T,A,S}}(undef, nworkers()),
                            A(),
                            A(),
                            Vector{MetaData}(undef,nworkers()),
                            RemoteChannel(() -> Channel{QCut{T}}(4*nworkers()*nscenarios)))
    end
end

function initialize_subproblems!(execution::SynchronousExecution,
                                 scenarioproblems::AbstractScenarioProblems,
                                 x::AbstractVector,
                                 subsolver::MPB.AbstractMathProgSolver)
    load_subproblems!(execution.subworkers, scenarioproblems, x, subsolver)
    return nothing
end

function finish_initilization!(lshaped::AbstractLShapedSolver, execution::SynchronousExecution)
    append!(execution.subobjectives, fill(1e10, nthetas(lshaped)))
    append!(execution.model_objectives, fill(-1e10, nthetas(lshaped)))
    for w in workers()
        execution.metadata[w-1] = RemoteChannel(() -> MetaChannel(), w)
    end
    return lshaped
end

function resolve_subproblems!(lshaped::AbstractLShapedSolver, execution::SynchronousExecution{F,T}) where {F <: AbstractFeasibility, T <: AbstractFloat}
    # Update metadata
    for w in workers()
        put!(execution.metadata[w-1], timestamp(lshaped), :gap, gap(lshaped))
    end
    @sync begin
        for (i,w) in enumerate(workers())
            worker_aggregator = remote_aggregator(lshaped.aggregation, scenarioproblems(lshaped.stochasticprogram), w)
            remotecall_fetch(resolve_subproblems!, w, execution.subworkers[w-1], lshaped.x, execution.cutqueue, worker_aggregator, timestamp(lshaped), execution.metadata[w-1])
        end
    end
    # Assume no cuts are added
    added = false
    # Collect incoming cuts
    while isready(execution.cutqueue)
        _, cut::SparseHyperPlane{T} = take!(execution.cutqueue)
        added |= add_cut!(lshaped, cut)
    end
    # Return current objective value and cut_added flag
    return current_objective_value(lshaped), added
end

function calculate_objective_value(lshaped::AbstractLShapedSolver, execution::SynchronousExecution)
    return get_obj(lshaped)⋅decision(lshaped) + eval_second_stage(execution.subworkers, decision(lshaped))
end

function fill_submodels!(lshaped::AbstractLShapedSolver, scenarioproblems, execution::SynchronousExecution)
    return fill_submodels!(execution.subworkers, decision(lshaped), scenarioproblems)
end

# API
# ------------------------------------------------------------
function (execution::Synchronous)(nscenarios::Integer, ::Type{F}, ::Type{T}, ::Type{A}, ::Type{S}) where {F <: AbstractFeasibility, T <: AbstractFloat, A <: AbstractVector, S <: LQSolver}
    return SynchronousExecution(nscenarios, F, T, A, S)
end

function str(::Synchronous)
    return "Synchronous "
end


__precompile__(true)

module ColorVectorSpace

using ColorTypes, FixedPointNumbers, Compat

import Base: ==, +, -, *, /, .+, .-, .*, ./, ^, .^, <, ~
import Base: abs, abs2, clamp, convert, copy, div, eps, isfinite, isinf,
    isnan, isless, length, mapreduce, norm, one, promote_array_type,
    promote_op, promote_rule, zero, trunc, floor, round, ceil, bswap,
    mod, rem, atan2, hypot, max, min, varm, real, histrange

if VERSION < v"0.5.0-dev"
    import Base.nan
else
    export nan
end

# The unaryOps
import Base:      conj, sin, cos, tan, sinh, cosh, tanh,
                  asin, acos, atan, asinh, acosh, atanh,
                  sec, csc, cot, asec, acsc, acot,
                  sech, csch, coth, asech, acsch, acoth,
                  sinc, cosc, cosd, cotd, cscd, secd,
                  sind, tand, acosd, acotd, acscd, asecd,
                  asind, atand, rad2deg, deg2rad,
                  log, log2, log10, log1p, exponent, exp,
                  exp2, expm1, cbrt, sqrt, erf,
                  erfc, erfcx, erfi, dawson,
                  significand, lgamma,
                  gamma, lfact, frexp, modf, airy, airyai,
                  airyprime, airyaiprime, airybi, airybiprime,
                  besselj0, besselj1, bessely0, bessely1,
                  eta, zeta, digamma

typealias AbstractGray{T} Color{T,1}
typealias TransparentRGB{C<:AbstractRGB,T}   TransparentColor{C,T,4}
typealias TransparentGray{C<:AbstractGray,T} TransparentColor{C,T,2}
typealias TransparentRGBFloat{C<:AbstractRGB,T<:AbstractFloat} TransparentColor{C,T,4}
typealias TransparentGrayFloat{C<:AbstractGray,T<:AbstractFloat} TransparentColor{C,T,2}
typealias TransparentRGBUFixed{C<:AbstractRGB,T<:UFixed} TransparentColor{C,T,4}
typealias TransparentGrayUFixed{C<:AbstractGray,T<:UFixed} TransparentColor{C,T,2}

typealias MathTypes{T,C} Union{AbstractRGB{T},TransparentRGB{C,T},AbstractGray{T},TransparentGray{C,T}}

# convert(RGB{Float32}, NaN) doesn't and shouldn't work, so we need to reintroduce nan
if VERSION >= v"0.5.0-dev"
    nan{T<:AbstractFloat}(::Type{T}) = convert(T, NaN)
end
nan{C<:MathTypes}(::Type{C}) = _nan(eltype(C), C)
_nan{T<:AbstractFloat,C<:AbstractGray}(::Type{T}, ::Type{C}) = (x = convert(T, NaN); C(x))
_nan{T<:AbstractFloat,C<:TransparentGray}(::Type{T}, ::Type{C}) = (x = convert(T, NaN); C(x,x))
_nan{T<:AbstractFloat,C<:AbstractRGB}(::Type{T}, ::Type{C}) = (x = convert(T, NaN); C(x,x,x))
_nan{T<:AbstractFloat,C<:TransparentRGB}(::Type{T}, ::Type{C}) = (x = convert(T, NaN); C(x,x,x,x))

## Generic algorithms
mapreduce(f, op::Base.ShortCircuiting, a::MathTypes) = f(a)  # ambiguity
mapreduce(f, op, a::MathTypes) = f(a)

for f in (:trunc, :floor, :round, :ceil, :eps, :bswap)
    @eval $f{T}(g::Gray{T}) = Gray{T}($f(gray(g)))
    @eval Compat.@dep_vectorize_1arg Gray $f
end
eps{T}(::Type{Gray{T}}) = Gray(eps(T))
Compat.@dep_vectorize_1arg AbstractGray isfinite
Compat.@dep_vectorize_1arg AbstractGray isinf
Compat.@dep_vectorize_1arg AbstractGray isnan
Compat.@dep_vectorize_1arg AbstractGray abs
Compat.@dep_vectorize_1arg AbstractGray abs2
for f in (:trunc, :floor, :round, :ceil)
    @eval $f{T<:Integer}(::Type{T}, g::Gray) = Gray{T}($f(T, gray(g)))
end

for f in (:mod, :rem, :mod1)
    @eval $f(x::Gray, m::Gray) = Gray($f(gray(x), gray(m)))
end

# Real values are treated like grays
ColorTypes.gray(x::Real) = x

# Return types for arithmetic operations
multype{A,B}(::Type{A}, ::Type{B}) = coltype(typeof(zero(A)*zero(B)))
sumtype{A,B}(::Type{A}, ::Type{B}) = coltype(typeof(zero(A)+zero(B)))
divtype{A,B}(::Type{A}, ::Type{B}) = coltype(typeof(zero(A)/zero(B)))
powtype{A,B}(::Type{A}, ::Type{B}) = coltype(typeof(zero(A)^zero(B)))
multype(a::Colorant, b::Colorant) = multype(eltype(a),eltype(b))
sumtype(a::Colorant, b::Colorant) = sumtype(eltype(a),eltype(b))
divtype(a::Colorant, b::Colorant) = divtype(eltype(a),eltype(b))
powtype(a::Colorant, b::Colorant) = powtype(eltype(a),eltype(b))

coltype{T<:Fractional}(::Type{T}) = T
coltype{T}(::Type{T})             = Float64

# Scalar binary RGB operations require the same RGB type for each element,
# otherwise we don't know which to return
color_rettype{A<:AbstractRGB,B<:AbstractRGB}(::Type{A}, ::Type{B}) = _color_rettype(base_colorant_type(A), base_colorant_type(B))
color_rettype{A<:AbstractGray,B<:AbstractGray}(::Type{A}, ::Type{B}) = _color_rettype(base_colorant_type(A), base_colorant_type(B))
color_rettype{A<:TransparentRGB,B<:TransparentRGB}(::Type{A}, ::Type{B}) = _color_rettype(base_colorant_type(A), base_colorant_type(B))
color_rettype{A<:TransparentGray,B<:TransparentGray}(::Type{A}, ::Type{B}) = _color_rettype(base_colorant_type(A), base_colorant_type(B))
_color_rettype{A<:Colorant,B<:Colorant}(::Type{A}, ::Type{B}) = error("binary operation with $A and $B, return type is ambiguous")
_color_rettype{C<:Colorant}(::Type{C}, ::Type{C}) = C

color_rettype(c1::Colorant, c2::Colorant) = color_rettype(typeof(c1), typeof(c2))

## Math on Colors. These implementations encourage inlining and,
## for the case of UFixed types, nearly halve the number of multiplications (for RGB)

# Scalar RGB
copy(c::AbstractRGB) = c
(*)(f::Real, c::AbstractRGB) = base_colorant_type(c){multype(typeof(f),eltype(c))}(f*red(c), f*green(c), f*blue(c))
(*)(f::Real, c::TransparentRGB) = base_colorant_type(c){multype(typeof(f),eltype(c))}(f*red(c), f*green(c), f*blue(c), f*alpha(c))
function (*){T<:UFixed}(f::Real, c::AbstractRGB{T})
    fs = f*(1/reinterpret(one(T)))
    base_colorant_type(c){multype(typeof(f),T)}(fs*reinterpret(red(c)), fs*reinterpret(green(c)), fs*reinterpret(blue(c)))
end
function (*){T<:UFixed}(f::UFixed, c::AbstractRGB{T})
    fs = reinterpret(f)*(1/widen(reinterpret(one(T)))^2)
    base_colorant_type(c){multype(typeof(f),T)}(fs*reinterpret(red(c)), fs*reinterpret(green(c)), fs*reinterpret(blue(c)))
end
function (/){T<:UFixed}(c::AbstractRGB{T}, f::Real)
    fs = (one(f)/reinterpret(one(T)))/f
    base_colorant_type(c){divtype(typeof(f),T)}(fs*reinterpret(red(c)), fs*reinterpret(green(c)), fs*reinterpret(blue(c)))
end
function (/){T<:UFixed}(c::AbstractRGB{T}, f::Integer)
    fs = (1/reinterpret(one(T)))/f
    base_colorant_type(c){divtype(typeof(f),T)}(fs*reinterpret(red(c)), fs*reinterpret(green(c)), fs*reinterpret(blue(c)))
end
(+){S,T}(a::AbstractRGB{S}, b::AbstractRGB{T}) = color_rettype(a, b){sumtype(S,T)}(red(a)+red(b), green(a)+green(b), blue(a)+blue(b))
(-){S,T}(a::AbstractRGB{S}, b::AbstractRGB{T}) = color_rettype(a, b){sumtype(S,T)}(red(a)-red(b), green(a)-green(b), blue(a)-blue(b))
(+)(a::TransparentRGB, b::TransparentRGB) =
    color_rettype(a, b){sumtype(a,b)}(red(a)+red(b), green(a)+green(b), blue(a)+blue(b), alpha(a)+alpha(b))
(-)(a::TransparentRGB, b::TransparentRGB) =
    color_rettype(a, b){sumtype(a,b)}(red(a)-red(b), green(a)-green(b), blue(a)-blue(b), alpha(a)-alpha(b))
(*)(c::AbstractRGB, f::Real) = (*)(f, c)
(*)(c::TransparentRGB, f::Real) = (*)(f, c)
(.*)(f::Real, c::AbstractRGB) = (*)(f, c)
(.*)(f::Real, c::TransparentRGB) = (*)(f, c)
(.*)(c::AbstractRGB, f::Real) = (*)(f, c)
(.*)(c::TransparentRGB, f::Real) = (*)(f, c)
(/)(c::AbstractRGB, f::Real) = (one(f)/f)*c
(/)(c::TransparentRGB, f::Real) = (one(f)/f)*c
(/)(c::AbstractRGB, f::Integer) = (one(eltype(c))/f)*c
(/)(c::TransparentRGB, f::Integer) = (one(eltype(c))/f)*c
(./)(c::AbstractRGB, f::Real) = (/)(c, f)
(./)(c::TransparentRGB, f::Real) = (/)(c, f)

isfinite{T<:UFixed}(c::Colorant{T}) = true
isfinite{T<:AbstractFloat}(c::AbstractRGB{T}) = isfinite(red(c)) && isfinite(green(c)) && isfinite(blue(c))
isfinite(c::TransparentRGBFloat) = isfinite(red(c)) && isfinite(green(c)) && isfinite(blue(c)) && isfinite(alpha(c))
isnan{T<:UFixed}(c::Colorant{T}) = false
isnan{T<:AbstractFloat}(c::AbstractRGB{T}) = isnan(red(c)) || isnan(green(c)) || isnan(blue(c))
isnan(c::TransparentRGBFloat) = isnan(red(c)) || isnan(green(c)) || isnan(blue(c)) || isnan(alpha(c))
isinf{T<:UFixed}(c::Colorant{T}) = false
isinf{T<:AbstractFloat}(c::AbstractRGB{T}) = isinf(red(c)) || isinf(green(c)) || isinf(blue(c))
isinf(c::TransparentRGBFloat) = isinf(red(c)) || isinf(green(c)) || isinf(blue(c)) || isinf(alpha(c))
abs(c::AbstractRGB) = abs(red(c))+abs(green(c))+abs(blue(c)) # should this have a different name?
abs{T<:UFixed}(c::AbstractRGB{T}) = Float32(red(c))+Float32(green(c))+Float32(blue(c)) # should this have a different name?
abs(c::TransparentRGB) = abs(red(c))+abs(green(c))+abs(blue(c))+abs(alpha(c)) # should this have a different name?
abs{T<:UFixed}(c::TransparentRGB{T}) = Float32(red(c))+Float32(green(c))+Float32(blue(c))+Float32(alpha(c)) # should this have a different name?
abs2(c::AbstractRGB) = red(c)^2+green(c)^2+blue(c)^2
abs2{T<:UFixed}(c::AbstractRGB{T}) = Float32(red(c))^2+Float32(green(c))^2+Float32(blue(c))^2
abs2(c::TransparentRGB) = (ret = abs2(color(c)); ret + convert(typeof(ret), alpha(c))^2)
norm(c::AbstractRGB) = sqrt(abs2(c))
norm(c::TransparentRGB) = sqrt(abs2(c))

one{C<:AbstractRGB}(::Type{C})     = C(1,1,1)
one{C<:TransparentRGB}(::Type{C})  = C(1,1,1,1)
zero{C<:AbstractRGB}(::Type{C})    = C(0,0,0)
zero{C<:TransparentRGB}(::Type{C}) = C(0,0,0,0)
zero{C<:YCbCr}(::Type{C}) = C(0,0,0)
zero{C<:HSV}(::Type{C}) = C(0,0,0)
one(p::Colorant) = one(typeof(p))
zero(p::Colorant) = zero(typeof(p))

# Arrays
(+){CV<:AbstractRGB}(A::AbstractArray{CV}, b::AbstractRGB) = (.+)(A, b)
(+){CV<:AbstractRGB}(b::AbstractRGB, A::AbstractArray{CV}) = (.+)(b, A)
(-){CV<:AbstractRGB}(A::AbstractArray{CV}, b::AbstractRGB) = (.-)(A, b)
(-){CV<:AbstractRGB}(b::AbstractRGB, A::AbstractArray{CV}) = (.-)(b, A)
(*){T<:Number}(A::AbstractArray{T}, b::AbstractRGB) = A.*b
(*){T<:Number}(b::AbstractRGB, A::AbstractArray{T}) = A.*b
(.+){C<:AbstractRGB}(A::AbstractArray{C}, b::AbstractRGB) = plus(A, b)
(.+){C<:AbstractRGB}(b::AbstractRGB, A::AbstractArray{C}) = plus(b, A)
(.-){C<:AbstractRGB}(A::AbstractArray{C}, b::AbstractRGB) = minus(A, b)
(.-){C<:AbstractRGB}(b::AbstractRGB, A::AbstractArray{C}) = minus(b, A)
(.*){T<:Number}(A::AbstractArray{T}, b::AbstractRGB) = mul(A, b)
(.*){T<:Number}(b::AbstractRGB, A::AbstractArray{T}) = mul(b, A)

(+){CV<:TransparentRGB}(A::AbstractArray{CV}, b::TransparentRGB) = (.+)(A, b)
(+){CV<:TransparentRGB}(b::TransparentRGB, A::AbstractArray{CV}) = (.+)(b, A)
(-){CV<:TransparentRGB}(A::AbstractArray{CV}, b::TransparentRGB) = (.-)(A, b)
(-){CV<:TransparentRGB}(b::TransparentRGB, A::AbstractArray{CV}) = (.-)(b, A)
(*){T<:Number}(A::AbstractArray{T}, b::TransparentRGB) = A.*b
(*){T<:Number}(b::TransparentRGB, A::AbstractArray{T}) = A.*b
(.+){C<:TransparentRGB}(A::AbstractArray{C}, b::TransparentRGB) = plus(A, b)
(.+){C<:TransparentRGB}(b::TransparentRGB, A::AbstractArray{C}) = plus(b, A)
(.-){C<:TransparentRGB}(A::AbstractArray{C}, b::TransparentRGB) = minus(A, b)
(.-){C<:TransparentRGB}(b::TransparentRGB, A::AbstractArray{C}) = minus(b, A)
(.*){T<:Number}(A::AbstractArray{T}, b::TransparentRGB) = mul(A, b)
(.*){T<:Number}(b::TransparentRGB, A::AbstractArray{T}) = mul(b, A)

# Scalar Gray
copy(c::AbstractGray) = c
const unaryOps = (:~, :conj, :abs,
                  :sin, :cos, :tan, :sinh, :cosh, :tanh,
                  :asin, :acos, :atan, :asinh, :acosh, :atanh,
                  :sec, :csc, :cot, :asec, :acsc, :acot,
                  :sech, :csch, :coth, :asech, :acsch, :acoth,
                  :sinc, :cosc, :cosd, :cotd, :cscd, :secd,
                  :sind, :tand, :acosd, :acotd, :acscd, :asecd,
                  :asind, :atand, :rad2deg, :deg2rad,
                  :log, :log2, :log10, :log1p, :exponent, :exp,
                  :exp2, :expm1, :cbrt, :sqrt, :erf,
                  :erfc, :erfcx, :erfi, :dawson,
                  :significand, :lgamma,
                  :gamma, :lfact, :frexp, :modf, :airy, :airyai,
                  :airyprime, :airyaiprime, :airybi, :airybiprime,
                  :besselj0, :besselj1, :bessely0, :bessely1,
                  :eta, :zeta, :digamma)
for op in unaryOps
    @eval ($op)(c::AbstractGray) = $op(gray(c))
end

(*)(f::Real, c::AbstractGray) = base_colorant_type(c){multype(typeof(f),eltype(c))}(f*gray(c))
(*)(f::Real, c::TransparentGray) = base_colorant_type(c){multype(typeof(f),eltype(c))}(f*gray(c), f*alpha(c))
(*)(c::AbstractGray, f::Real) = (*)(f, c)
(.*)(f::Real, c::AbstractGray) = (*)(f, c)
(.*)(c::AbstractGray, f::Real) = (*)(f, c)
(*)(c::TransparentGray, f::Real) = (*)(f, c)
(.*)(f::Real, c::TransparentGray) = (*)(f, c)
(.*)(c::TransparentGray, f::Real) = (*)(f, c)
(/)(c::AbstractGray, f::Real) = (one(f)/f)*c
(/)(n::Number, c::AbstractGray) = n/gray(c)
(/)(c::TransparentGray, f::Real) = (one(f)/f)*c
(/)(c::AbstractGray, f::Integer) = (one(eltype(c))/f)*c
(/)(c::TransparentGray, f::Integer) = (one(eltype(c))/f)*c
(./)(c::AbstractGray, f::Real) = c/f
(./)(n::Number, c::AbstractGray) = n/gray(c)
(./)(c::TransparentGray, f::Real) = c/f
(+){S,T}(a::AbstractGray{S}, b::AbstractGray{T}) = color_rettype(a,b){sumtype(S,T)}(gray(a)+gray(b))
(+)(a::TransparentGray, b::TransparentGray) = color_rettype(a,b){sumtype(eltype(a),eltype(b))}(gray(a)+gray(b),alpha(a)+alpha(b))
(-){S,T}(a::AbstractGray{S}, b::AbstractGray{T}) = color_rettype(a,b){sumtype(S,T)}(gray(a)-gray(b))
(-)(a::TransparentGray, b::TransparentGray) = color_rettype(a,b){sumtype(eltype(a),eltype(b))}(gray(a)-gray(b),alpha(a)-alpha(b))
(*){S,T}(a::AbstractGray{S}, b::AbstractGray{T}) = color_rettype(a,b){multype(S,T)}(gray(a)*gray(b))
(^){S}(a::AbstractGray{S}, b::Integer) = base_colorant_type(a){powtype(S,Int)}(gray(a)^convert(Int,b))
(^){S}(a::AbstractGray{S}, b::Real) = base_colorant_type(a){powtype(S,typeof(b))}(gray(a)^b)
(.^){S}(a::AbstractGray{S}, b) = a^b
(+)(c::AbstractGray) = c
(+)(c::TransparentGray) = c
(-)(c::AbstractGray) = typeof(c)(-gray(c))
(-)(c::TransparentGray) = typeof(c)(-gray(c),-alpha(c))
(/)(a::AbstractGray, b::AbstractGray) = gray(a)/gray(b)
div(a::AbstractGray, b::AbstractGray) = div(gray(a), gray(b))
(+)(a::AbstractGray, b::Number) = gray(a)+b
(-)(a::AbstractGray, b::Number) = gray(a)-b
(+)(a::Number, b::AbstractGray) = a+gray(b)
(-)(a::Number, b::AbstractGray) = a-gray(b)
(.+)(a::AbstractGray, b::Number) = gray(a)+b
(.-)(a::AbstractGray, b::Number) = gray(a)-b
(.+)(a::Number, b::AbstractGray) = a+gray(b)
(.-)(a::Number, b::AbstractGray) = a-gray(b)
max{T<:AbstractGray}(a::T, b::T) = T(max(gray(a),gray(b)))
max(a::AbstractGray, b::AbstractGray) = max(promote(a,b)...)
max(a::Number, b::AbstractGray) = max(promote(a,b)...)
max(a::AbstractGray, b::Number) = max(promote(a,b)...)
min{T<:AbstractGray}(a::T, b::T) = T(min(gray(a),gray(b)))
min(a::AbstractGray, b::AbstractGray) = min(promote(a,b)...)
min(a::Number, b::AbstractGray) = min(promote(a,b)...)
min(a::AbstractGray, b::Number) = min(promote(a,b)...)

isfinite{T<:AbstractFloat}(c::AbstractGray{T}) = isfinite(gray(c))
isfinite(c::TransparentGrayFloat) = isfinite(gray(c)) && isfinite(alpha(c))
isnan{T<:AbstractFloat}(c::AbstractGray{T}) = isnan(gray(c))
isnan(c::TransparentGrayFloat) = isnan(gray(c)) && isnan(alpha(c))
isinf{T<:AbstractFloat}(c::AbstractGray{T}) = isinf(gray(c))
isinf(c::TransparentGrayFloat) = isinf(gray(c)) && isnan(alpha(c))
norm(c::AbstractGray) = abs(gray(c))
abs(c::TransparentGray) = abs(gray(c))+abs(alpha(c)) # should this have a different name?
abs(c::TransparentGrayUFixed) = Float32(gray(c)) + Float32(alpha(c)) # should this have a different name?
abs2(c::AbstractGray) = gray(c)^2
abs2{T<:UFixed}(c::AbstractGray{T}) = Float32(gray(c))^2
abs2(c::TransparentGray) = gray(c)^2+alpha(c)^2
abs2(c::TransparentGrayUFixed) = Float32(gray(c))^2 + Float32(alpha(c))^2
atan2(x::Gray, y::Gray) = atan2(convert(Real, x), convert(Real, y))
hypot(x::Gray, y::Gray) = hypot(convert(Real, x), convert(Real, y))
norm(c::TransparentGray) = sqrt(abs2(c))

(<)(g1::AbstractGray, g2::AbstractGray) = gray(g1) < gray(g2)
(<)(c::AbstractGray, r::Real) = gray(c) < r
(<)(r::Real, c::AbstractGray) = r < gray(c)
isless(g1::AbstractGray, g2::AbstractGray) = isless(gray(g1), gray(g2))
isless(c::AbstractGray, r::Real) = isless(gray(c), r)
isless(r::Real, c::AbstractGray) = isless(r, gray(c))
Base.isapprox(x::AbstractGray, y::AbstractGray; kwargs...) = isapprox(gray(x), gray(y); kwargs...)
Base.isapprox(x::TransparentGray, y::TransparentGray; kwargs...) = isapprox(gray(x), gray(y); kwargs...) && isapprox(alpha(x), alpha(y); kwargs...)
Base.isapprox(x::AbstractRGB, y::AbstractRGB; kwargs...) = isapprox(red(x), red(y); kwargs...) && isapprox(green(x), green(y); kwargs...) && isapprox(blue(x), blue(y); kwargs...)
Base.isapprox(x::TransparentRGB, y::TransparentRGB; kwargs...) = isapprox(alpha(x), alpha(y); kwargs...) && isapprox(red(x), red(y); kwargs...) && isapprox(green(x), green(y); kwargs...) && isapprox(blue(x), blue(y); kwargs...)

function Base.isapprox{Cx<:MathTypes,Cy<:MathTypes}(x::AbstractArray{Cx},
                                                    y::AbstractArray{Cy};
                                                    rtol::Real=Base.rtoldefault(eltype(Cx),eltype(Cy)),
                                                    atol::Real=0,
                                                    norm::Function=vecnorm)
    d = norm(x - y)
    if isfinite(d)
        return d <= atol + rtol*max(norm(x), norm(y))
    else
        # Fall back to a component-wise approximate comparison
        return all(ab -> isapprox(ab[1], ab[2]; rtol=rtol, atol=atol), zip(x, y))
    end
end

zero{C<:TransparentGray}(::Type{C}) = C(0,0)
zero{C<:Gray}(::Type{C}) = C(0)
one{C<:TransparentGray}(::Type{C}) = C(1,1)
one{C<:Gray}(::Type{C}) = C(1)

 # Arrays
(+){CV<:AbstractGray}(A::AbstractArray{CV}, b::AbstractGray) = (.+)(A, b)
(+){CV<:AbstractGray}(b::AbstractGray, A::AbstractArray{CV}) = (.+)(b, A)
(-){CV<:AbstractGray}(A::AbstractArray{CV}, b::AbstractGray) = (.-)(A, b)
(-){CV<:AbstractGray}(b::AbstractGray, A::AbstractArray{CV}) = (.-)(b, A)
(*){T<:Number}(A::AbstractArray{T}, b::AbstractGray) = A.*b
(*){T<:Number}(b::AbstractGray, A::AbstractArray{T}) = A.*b
(/){C<:AbstractGray}(A::AbstractArray{C}, b::AbstractGray) = A./b
(.+){C<:AbstractGray}(A::AbstractArray{C}, b::AbstractGray) = plus(A, b)
(.+){C<:AbstractGray}(b::AbstractGray, A::AbstractArray{C}) = plus(b, A)
(.-){C<:AbstractGray}(A::AbstractArray{C}, b::AbstractGray) = minus(A, b)
(.-){C<:AbstractGray}(b::AbstractGray, A::AbstractArray{C}) = minus(b, A)
(.*){T<:Number}(A::AbstractArray{T}, b::AbstractGray) = mul(A, b)
(.*){T<:Number}(b::AbstractGray, A::AbstractArray{T}) = mul(b, A)
(./){C<:AbstractGray}(A::AbstractArray{C}, b::AbstractGray) = divd(A, b)

Compat.@dep_vectorize_2arg Gray max
Compat.@dep_vectorize_2arg Gray min
for f in (:min, :max)
    @eval begin
        @deprecate($f{T<:Gray}(x::Number, y::AbstractArray{T}),
                   @compat $f.(x, y))
        @deprecate($f{T<:Gray}(x::AbstractArray{T}, y::Number),
                   @compat $f.(x, y))
    end
end

(+){CV<:TransparentGray}(A::AbstractArray{CV}, b::TransparentGray) = (.+)(A, b)
(+){CV<:TransparentGray}(b::TransparentGray, A::AbstractArray{CV}) = (.+)(b, A)
(-){CV<:TransparentGray}(A::AbstractArray{CV}, b::TransparentGray) = (.-)(A, b)
(-){CV<:TransparentGray}(b::TransparentGray, A::AbstractArray{CV}) = (.-)(b, A)
(*){T<:Number}(A::AbstractArray{T}, b::TransparentGray) = A.*b
(*){T<:Number}(b::TransparentGray, A::AbstractArray{T}) = A.*b
(.+){C<:TransparentGray}(A::AbstractArray{C}, b::TransparentGray) = plus(A, b)
(.+){C<:TransparentGray}(b::TransparentGray, A::AbstractArray{C}) = plus(b, A)
(.-){C<:TransparentGray}(A::AbstractArray{C}, b::TransparentGray) = minus(A, b)
(.-){C<:TransparentGray}(b::TransparentGray, A::AbstractArray{C}) = minus(b, A)
(.*){T<:Number}(A::AbstractArray{T}, b::TransparentGray) = mul(A, b)
(.*){T<:Number}(b::TransparentGray, A::AbstractArray{T}) = mul(b, A)

varm{C<:AbstractGray}(v::AbstractArray{C}, s::AbstractGray; corrected::Bool=true) =
        varm(map(gray,v),gray(s); corrected=corrected)
real{C<:AbstractGray}(::Type{C}) = real(eltype(C))

# Called plus/minus instead of plus/sub because `sub` already has a meaning!
function plus(A::AbstractArray, b::Colorant)
    bT = convert(eltype(A), b)
    out = similar(A)
    plus!(out, A, bT)
end
plus(b::Colorant, A::AbstractArray) = plus(A, b)
function minus(A::AbstractArray, b::Colorant)
    bT = convert(eltype(A), b)
    out = similar(A)
    minus!(out, A, bT)
end
function minus(b::Colorant, A::AbstractArray)
    bT = convert(eltype(A), b)
    out = similar(A)
    minus!(out, bT, A)
end
function mul{T<:Number}(A::AbstractArray{T}, b::Colorant)
    bT = typeof(b*one(T))
    out = similar(A, bT)
    mul!(out, A, b)
end
mul{T<:Number}(b::Colorant, A::AbstractArray{T}) = mul(A, b)
function divd{C<:AbstractGray}(A::AbstractArray{C}, b::AbstractGray)
    bT = typeof(zero(C)/b)
    out = similar(A, bT)
    div!(out, A, b)
end

for (func, op) in ((:plus!, :+),
                   (:minus!, :-),
                   (:mul!, :*),
                   (:div!, :/))
    @eval begin
        function $func{T,N}(out, A::AbstractArray{T,N}, b)
            Rout, RA = eachindex(out), eachindex(A)
            if Rout == RA
                for I in RA
                    @inbounds out[I] = $op(A[I], b)
                end
            else
                for (Iout, IA) in zip(Rout, RA)
                    @inbounds out[Iout] = $op(A[IA], b)
                end
            end
            out
        end
    end
end

# This needs separate implementation because we can take -b of unsigned types
function minus!{T,N}(out, b::Colorant, A::AbstractArray{T,N})
    Rout, RA = eachindex(out), eachindex(A)
    if Rout == RA
        for I in RA
            @inbounds out[I] = b - A[I]
        end
    else
        for (Iout, IA) in zip(Rout, RA)
            @inbounds out[Iout] = b - A[IA]
        end
    end
    out
end

#histrange for Gray type
Base.histrange{T}(v::AbstractArray{Gray{T}}, n::Integer) = histrange(convert(Array{Float32}, map(gray, v)), n)

# Promotions for reductions
if VERSION < v"0.5.0-dev+3701"
    Base.r_promote{T<:FixedPoint}(::Base.AddFun, c::MathTypes{T}) = convert(base_colorant_type(typeof(c)){Float64}, c)
    Base.r_promote{T<:FixedPoint}(::Base.MulFun, c::MathTypes{T}) = convert(base_colorant_type(typeof(c)){Float64}, c)
else
    Base.r_promote{T<:FixedPoint}(::typeof(+), c::MathTypes{T}) = convert(base_colorant_type(typeof(c)){Float64}, c)
    Base.r_promote{T<:FixedPoint}(::typeof(*), c::MathTypes{T}) = convert(base_colorant_type(typeof(c)){Float64}, c)
end

# To help type inference
promote_array_type{T<:Real,C<:MathTypes}(F, ::Type{T}, ::Type{C}) = base_colorant_type(C){Base.promote_array_type(F, T, eltype(C))}
promote_op{C<:MathTypes,T<:Real}(F, ::Type{C}, ::Type{T}) = typeof(F(zero(C), zero(T)))
promote_op{C<:MathTypes,T<:Real}(F, ::Type{T}, ::Type{C}) = typeof(F(zero(T), zero(C)))
promote_rule{C1<:Colorant,C2<:Colorant}(::Type{C1}, ::Type{C2}) = color_rettype(C1,C2){promote_type(eltype(C1), eltype(C2))}
promote_rule{T<:Real,C<:AbstractGray}(::Type{T}, ::Type{C}) = promote_type(T, eltype(C))

@deprecate sumsq abs2

end


function make_seq(panel_file::AbstractString, profile_file::AbstractString, output_folder::AbstractString, depth)
    profile = JSON.parsefile(profile_file)
    config = profile["config"]
    fill_default_config!(config)

    # load assembly
    assembly = load_assembly(config["assembly"])
    panel, panel_size = load_bed(panel_file, assembly)

    if !isdir(output_folder)
        mkpath(output_folder)
    end

    if depth<=0
        depth = config["depth"]
    end

    io = nothing
    filename = joinpath(output_folder, "SEQMAKER-$SEQMAKER_VERSION")
    filename *= "-" * basename(panel_file)
    filename *= "-" * basename(profile_file)

    if config["pair-end"]
        io = fastq_open_pair("$filename.R1.fq", "$filename.R2.fq", "w")
    else
        io = fastq_open("$filename.R1.fq", "w")
    end

    temp_max = config["template_len"]["max"]
    temp_min = config["template_len"]["min"]
    temp_len_mean = ( temp_min + temp_max )/2
    read_num = (depth * panel_size) / (temp_len_mean * config["duplication_rate"])
    read_num = Int(floor(read_num))
    rand_pos = rand(read_num)
    rand_temp_len = rand(read_num)
    for i in 1:length(rand_pos)
        panel_pos = Int(floor(rand_pos[i] * panel_size))
        # linear interpolation
        temp_len = Int(floor(rand_temp_len[i] * temp_min + (1.0-rand_temp_len[i]) * temp_max))
        chr, pos_in_chr = seek_in_panel(panel, panel_pos)
        cycles = 1
        if haskey(profile, "cnv")
            copy_num = simulate_copy_number(chr, pos_in_chr, profile["cnv"])
            cycles = Int(round((rand() * copy_num)/0.5))
        end
        for c in 1:cycles
            # vibration
            this_pos_in_chr = pos_in_chr + Int(round((rand() - 0.5) * 0.4 * temp_len))
            this_temp_len = temp_len + Int(round((rand() - 0.5) * 0.1 * temp_len))
            # sampling
            original_seq, start = sample(assembly, chr, this_pos_in_chr , this_temp_len)
            seq = simulate_snv(original_seq, chr, start, profile["snv"])
            seq = simulate_fusion(seq, chr, start, profile["fusion"], assembly)
            # simulate watson/crick strand
            if rand()> 0.5
                seq = ~seq
            end
            if config["pair-end"]
                reads = pair_end_seq(seq, config)
                for pair in reads
                    fastq_write_pair(io, pair)
                end
            end
        end
    end
    close(io)
end

function seek_in_panel(panel, panel_pos)
    # binary search to find the contig
    left = 1
    right = length(panel)
    contig = 0
    while left < right
        center = div(left+right, 2)
        rec = panel[center]
        if rec["start_in_panel"]<=panel_pos && rec["end_in_panel"]>=panel_pos
            contig = center
            break
        elseif rec["start_in_panel"]>panel_pos
            right = center - 1
        elseif rec["end_in_panel"]<panel_pos
            left = center + 1
        end
    end
    if contig == 0
        contig = left
    end
    rec = panel[contig]
    offset = panel_pos - rec["start_in_panel"]
    pos_in_chr = rec["from"] + offset
    return rec["chr"], pos_in_chr
end

function fill_default_config!(config)
    if !haskey(config, "duplication_rate")
        config["duplication_rate"] = 1.0
    end
    if !haskey(config, "random_index1")
        config["random_index1"] = false
    end
    if !haskey(config, "random_index2")
        config["random_index2"] = false
    end
end


using MATLAB

x = 1:10
y = rand(5,10)

mat"""
plot($x,$y)
"""


using Distributed
addprocs(2, exeflags="--project")

@everywhere begin
include("../../src/Constructors.jl")
include("../../src/RemoteBufferChannel.jl")
include("../../src/SelfPlay.jl")
include("../../src/ReplayBuffer.jl")
include("../../src/Learning.jl")
include("game.jl")
include("params.jl")
end

env=TicTacToe()
training_step = RemoteChannel(()->Channel{Int}(1))
num_played_games = RemoteChannel(()->Channel{Int}(1))
num_played_steps = RemoteChannel(()->Channel{Int}(1))
num_reanalysed_games = RemoteChannel(()->Channel{Int}(1))
total_samples = RemoteChannel(()->Channel{Int}(1))
remote_NNs = RemoteChannel(()->Channel{NamedTuple{(:representation, :prediction, :dynamics), Tuple{Any, Any, Any}}}(1))
remote_buffer = RemoteChannel(()->BufferChannel()) #TODO Start saving Buffer to disk at periodic intervals after a certain stage in the training process

put!(remote_NNs, (representation=init_representation(hyper), prediction=init_prediction(hyper), dynamics=init_dynamics(hyper)))
put!(training_step, 0)
put!(num_played_games, 0)
put!(num_played_steps, 0)
put!(num_reanalysed_games, 0)
put!(total_samples, 0)

sp = @spawnat :any self_play!(env,
training_step,
num_played_games,
num_played_steps,
total_samples,
remote_NNs,
remote_buffer)

learn = @spawnat :any learning!(num_played_games,
training_step,
remote_NNs,
remote_buffer)

# include("./src/MuZero.jl")
# using .MuZero
# include("./games/tictactoe/game.jl")
# ttt=TicTacToe()
# include("./games/tictactoe/params.jl")
# MuZero.self_play(hyper, ttt, MuZero.progress, MuZero.buffer)


############################################################################################
"""
$(SIGNATURES)
Return kernel that is used for jittering.

# Examples
```julia
```

"""
function jitterkernel(particles::SMCParticles, iter::Int64)
    return particles.kernel[iter]
end

############################################################################################
"""
$(SIGNATURES)
Resample particle ancestors, shuffle current particles and rejuvenate θ.

# Examples
```julia
```

"""
function resample!(
    _rng::Random.AbstractRNG,
    particles::SMCParticles,
    tune::SMCTune,
    data::D,
    temperature::F
) where {D, F<:AbstractFloat}
    ## Set resample back to false from previous iteration
    init!(tune.resample, false)
    ## Compute ESS
    ESS = BaytesCore.computeESS(particles.weights)
    resampled = BaytesCore.issmaller(ESS, tune.chains.Nchains * tune.chains.threshold)
    if resampled
        ## Record in tuning for diagnostics and propagation step
        init!(tune.resample, true)
        ## Resample ancestors ~ buffer already contain normalized weights from computeESS(weights) call
        resample!(
            _rng,
            tune.resample.method,
            particles.buffer.parameter.index,
            tune.iter.current,
            particles.weights.buffer,
            tune.chains.Nchains,
        )
        ## Equal weight normalized weights for next iteration memory
        Base.fill!(particles.weights.ℓweightsₙ, log(1.0 / tune.chains.Nchains))
        ## Reshuffle θ and rejuvenation particles - can be done inplace via buffer
        #!NOTE: Also set cumulative particle weights smc.buffer.weights back to correct position
        BaytesCore.shuffle!(particles.buffer.parameter, particles.kernel, particles.model, particles.buffer.cumweights)
        ## Rejuvenate Particles
        jitter!(_rng, particles, tune, data, temperature)
        ## Reweight ℓweights used for tempering so have correct index and temperature for next iteration
        reweight!(_rng, particles, tune, data, temperature)
    end
    return ESS, resampled
end

############################################################################################
"""
$(SIGNATURES)
Jitter θ particles with given kernels. This is performed in 2 stages:
1. kernel is updated with new data and shuffled particles, then one proposal step is performed.
2. If more than one jitterstep has to be performed, previous results might be captured depending on the kernel.

# Examples
```julia
```

"""
function jitter!(_rng::Random.AbstractRNG, particles::SMCParticles, tune::SMCTune, data::D, temperature::F) where {D, F<:AbstractFloat}
    ## Make first jitter step, update for new data and shuffled parameter
    Base.Threads.@threads for iter in eachindex(particles.kernel)
        ## Propose new parameter
        _, particles.buffer.jitterdiagnostics[iter] = propose!(
            _rng,
            jitterkernel(particles, iter),
            particles.model[iter],
            data,
            temperature,
            BaytesCore.UpdateTrue()
        )
        #!NOTE: has to be in loop for scoping rules
        particles.buffer.predictions[iter] = BaytesCore.get_prediction(
            particles.buffer.jitterdiagnostics[iter]
        )
    end
    # Calculate Correlation between old and new θ particles
    compute_ρ!(particles.buffer.correlation, particles.buffer.parameter.result, particles.kernel)
    # Check if jittering can be stopped
    jitter = BaytesCore.jitter!(tune.jitter, tune.jitterfun(particles.buffer.correlation.ρ))
    ## If more than one jitterstep has to be performed, previous results might be captured depending on the kernel.
    #!NOTE buffer.parameter.result.θᵤ at correct place because jitter! only happens after shuffle!
    while jitter
        ## Jitter Particles
        Base.Threads.@threads for iter in eachindex(particles.kernel)
            ## Propose new parameter
            #!TODO: This is only true in the first call though, if more than 1 step jittered, MCMC only kernels could be set to UpdateFalse()
            _, particles.buffer.jitterdiagnostics[iter] = propose!(
                _rng,
                jitterkernel(particles, iter),
                particles.model[iter],
                data,
                temperature,
                tune.capture
            )
            #!NOTE: has to be in loop for scoping rules
            particles.buffer.predictions[iter] = BaytesCore.get_prediction(
                particles.buffer.jitterdiagnostics[iter]
            )
        end
        ## Calculate Correlation between old and new θ particles
        compute_ρ!(particles.buffer.correlation, particles.buffer.parameter.result, particles.kernel)
        ## Check if jittering can be stopped
        jitter = jitter!(tune.jitter, tune.jitterfun(particles.buffer.correlation.ρ))
    end
    return nothing
end

############################################################################################
"""
$(SIGNATURES)
Propagate data forward over time. If (latent) data has to be extended, need to overload this function for specific smc kernel.

# Examples
```julia
```

"""
function propagate!(_rng::Random.AbstractRNG, particles::SMCParticles, tune::SMCTune, data::D, temperature::F) where {D, F<:AbstractFloat}
    return nothing
end

############################################################################################
"""
$(SIGNATURES)
Predict new data for each smc particle.

# Examples
```julia
```

"""
function predict!(_rng::Random.AbstractRNG, particles::SMCParticles, tune::SMCTune, data::D, temperature::F) where {D, F<:AbstractFloat}
    ## Propagate series forward with recent particle
    for iter in eachindex(particles.model)
        ## Predict new data point
        particles.buffer.predictions[iter] = predict(
            _rng,
            particles.kernel[iter],
            Objective(particles.model[iter], data, tune.tagged, temperature)
        )
    end
    #!NOTE: Here, particles.buffer.jitterdiagnostics will not be updated, and for SMCDiagnostics, last available jitterdiagnostics are provided.
    return nothing
end

############################################################################################
"""
$(SIGNATURES)
Compute new cumulative and incremental particle weights.

# Examples
```julia
```

"""
function weight!(
    _rng::Random.AbstractRNG,
    particles::SMCParticles,
    tune::SMCTune,
    data::D,
    temperature::F
) where {D, F<:AbstractFloat}
    ## Compute incremental weight for resampling step, and cumulative weight for i) next iteration and ii) temperature adjustment
    @inbounds for iter in eachindex(particles.weights.ℓweights)
        objective = Objective(particles.model[iter], data, tune.tagged, temperature)
        #!NOTE: particles.buffer.cumweights is cumulative weight at previous iteration, accounting for jittering step.
        particles.buffer.cumweights[iter], particles.weights.ℓweights[iter] = SMCweight(
            _rng,
            objective, particles.kernel[iter],
            particles.buffer.cumweights[iter]
        )
    end
    BaytesFilters.normalize!(particles.weights)
    return nothing
end

############################################################################################
"""
$(SIGNATURES)
Compute cumulative and incremental weight of objective at time/iteration t, given weight at t-1.
Incremental weight will be used as particle weight for resampling.
Cumulative weight will be used to adapt temperature.

If temperature is constant, this function can be overloaded with your ModelName if incremental weight can be computed independent of previous weight,
which speeds up computation. `cumweightsₜ` is not needed in this case.

# Examples
```julia
```

"""
function SMCweight(_rng::Random.AbstractRNG, objective::Objective, algorithm, cumweightsₜ₋₁)
    cumweightsₜ = objective.temperature * objective(objective.model.val)
    #!NOTE: cumweightsₜ₋₁ already has correct temperature at t-1
    return cumweightsₜ, cumweightsₜ - cumweightsₜ₋₁
end

############################################################################################
"""
$(SIGNATURES)
Compute new cumulative particle weights, accounting for jittering steps. This is only needed for
next iteration's weight calculation, and will not adjust current incremental and normalized weight.

# Examples
```julia
```

"""
function reweight!(
    _rng::Random.AbstractRNG,
    particles::SMCParticles,
    tune::SMCTune,
    data::D,
    temperature::F
) where {D, F<:AbstractFloat}
    ## Compute incremental weight for resampling step, and cumulative weight for i) next iteration and ii) temperature adjustment
    @inbounds for iter in eachindex(particles.weights.ℓweights)
        objective = Objective(particles.model[iter], data, tune.tagged, temperature)
        #!NOTE: particles.buffer.cumweights is cumulative weight at previous iteration, accounting for jittering step.
        particles.buffer.cumweights[iter], _ = SMCreweight(
            _rng,
            objective,
            particles.kernel[iter],
            particles.buffer.cumweights[iter]
        )
    end
    return nothing
end

############################################################################################
"""
$(SIGNATURES)
Computes particle weights after jiterring step. Defaults to `SMCweight` function.

# Examples
```julia
```

"""
function SMCreweight(_rng::Random.AbstractRNG, objective::Objective, algorithm, cumweightsₜ₋₁)
    return SMCweight(_rng, objective, algorithm, cumweightsₜ₋₁)
end

############################################################################################
# Export
export
    resample!,
    jitter!,
    propagate!,
    weight!,
    reweight!,
    SMCweight,
    SMCreweight


# types ------------------------------------------------------------------------

# type corresponding to a solution
mutable struct tSolution{T}
    x :: Vector{T}                # vector variables x (1..n)
    y :: Vector{T}                # vector outcomes  y (1..p)
end

# type corresponding to a point generator
mutable struct tGenerateur
    sRel :: tSolution{Float64}    # initial relaxed solution
    sInt :: tSolution{Int64}      # integer solution
    sPrj :: tSolution{Float64}    # projected solution
    sFea :: Bool                  # indicate if sInt is feasible or not
end

# type of a point (x,y) in the objective space
mutable struct tPoint
    x :: Float64
    y :: Float64
end

# type grouping the lists of points for dysplaying purposes
mutable struct tListDisplay
    xLf1  :: Vector{Float64};  yLf1  :: Vector{Float64} # liste des points (x,y) relaches
    xLf2  :: Vector{Float64};  yLf2  :: Vector{Float64} # liste des points (x,y) relaches
    xL    :: Vector{Float64};  yL    :: Vector{Float64} # liste des points (x,y) relaches
    XInt  :: Vector{Int64};    YInt  :: Vector{Int64}   # liste des points (x,y) entiers
    XProj :: Vector{Float64};  YProj :: Vector{Float64} # liste des points (x,y) projetes
    XFeas :: Vector{Int64};    YFeas :: Vector{Int64}   # liste des points (x,y) admissibles
    XPert :: Vector{Int64};    YPert :: Vector{Int64}   # liste des points (x,y) perturbes
end

# ==============================================================================
# Initialisation structure donnees contenant tous les generateurs

function allocateDatastructure(nbgen::Int64, nbvar::Int64, nbobj::Int64)

    verbose ? println("\n  → Allocation memoire pour ",nbgen," generateurs\n") : nothing

    vg = Vector{tGenerateur}(undef, nbgen)
    for k = 1:nbgen
        vg[k] = tGenerateur( tSolution{Float64}(zeros(Float64,nbvar),zeros(Float64,nbobj)),
                              tSolution{Int64}(zeros(Int64,nbvar),zeros(Int64,nbobj)),
                              tSolution{Float64}(zeros(Float64,nbvar),zeros(Float64,nbobj)),
                              false
                            )
    end
    return vg
end


include("deepmap.jl")
include("fileio.jl")
include("preprocess.jl")
include("resnet.jl"); export ResNet
include("mobilenet.jl"); export MobileNet
include("efficientnet.jl"); export EfficientNet


using FastHistograms
using FastHistograms: BinType, FixedWidth, VariableWidth, UnboundedWidth
using FastHistograms: BinSearchAlgorithm, Arithmetic, BinarySearch, HashFunction
using FastHistograms: HistogramParallelization, NoParallelization, SIMD, PrivateThreads
using Test
using Random
import StatsBase
using BenchmarkTools
using Downloads

const run_1D_tests = true
const run_2D_tests = true

function get_test_resources()
    path1 = joinpath(@__DIR__, "pride_and_prejudice.txt")
    if !isfile(path1)
        Downloads.download("https://www.gutenberg.org/files/1342/1342-0.txt", path1)
    end
end

get_test_resources()

include("test_fixed_width.jl")
include("test_variable_width.jl")
include("test_text_hist.jl")

invalid_combination(::BinType, ::BinSearchAlgorithm, ::HistogramParallelization) = false
# SIMD binary search needs a special implementation that I don't have yet
invalid_combination(::BinType, ::BinarySearch, ::SIMD) = true
# Arithmetic only works for fixed width
invalid_combination(::Union{VariableWidth,UnboundedWidth}, ::Arithmetic, ::HistogramParallelization) = true
# Text histograms can only use UnboundedWidth, HashFunction, and NoParallelization or PrivateThreads. SIMD not implemented yet.
invalid_combination(::UnboundedWidth, ::HashFunction, ::SIMD) = true

@testset "FastHistograms.jl" begin
    bin_types = [FixedWidth(), VariableWidth()]
    search_algorithms = [Arithmetic(), BinarySearch()]
    parallelizations = [NoParallelization(), SIMD(), PrivateThreads()]

    @testset "histogram computations BinType=$(bin_type)" for bin_type in bin_types
        @testset "BinSearchAlgorithm=$(search_algorithm)" for search_algorithm in search_algorithms
            @testset "HistogramParallelization=$(parallelization)" for parallelization in parallelizations
                if invalid_combination(bin_type, search_algorithm, parallelization)
                    continue
                end

                test_parameterized_hist(bin_type, search_algorithm, parallelization)
            end
        end
    end

    @testset "zero bins" begin
        h = create_fast_histogram(FixedWidth(), Arithmetic(), NoParallelization(), [(0x00, 0xff, 16), (0x00, 0xff, 16)])

        increment_bins!(h, rand(UInt8, 10, 10), rand(UInt8, 10, 10))

        @test any(counts(h) .!= 0)

        zero!(h)
        @test all(counts(h) .== 0)
    end

    @testset "default hist bin type" begin
        h = create_fast_histogram(FixedWidth(), Arithmetic(), NoParallelization(), [(0x00, 0xff, 16), (0x00, 0xff, 16)])

        @test eltype(h) == UInt8
    end

    @testset "custom hist bin type" begin
        h = create_fast_histogram(
            FixedWidth(),
            Arithmetic(),
            NoParallelization(),
            [(Int32(0), Int32(16), 16), (Int32(0), Int32(16), 16)],
        )

        @test eltype(h) == Int32
    end

    @testset "AHTL fixed data" begin
        # Generate the data in the same way the AHTL test does
        data = rand(Float32, 100000000) .* (512 - Float32(0.01))

        img_vec = vec(data)
        sb_counts = StatsBase.fit(
            StatsBase.Histogram,
            img_vec,
            range(Float32(0.0); stop = Float32(512.0), length = 513),
        ).weights

        h = create_fast_histogram(
            FastHistograms.FixedWidth(),
            FastHistograms.Arithmetic(),
            FastHistograms.PrivateThreads(),
            [(Float32(0.0), Float32(512.0), 512)],
        )
        increment_bins!(h, data)
        fh_counts = counts(h)
        @test fh_counts == sb_counts
    end

    let bin_type = UnboundedWidth(), search_algorithm = HashFunction()
        @testset "HistogramParallelization=$(parallelization)" for parallelization in parallelizations
            if invalid_combination(bin_type, search_algorithm, parallelization)
                continue
            end

            test_text_hist(parallelization)
        end
    end
end


""
    ispangram(input)

Return `true` if `input` contains every alphabetic character (case insensitive).

"""
ispangram(input) = prod([ch ∈ lowercase(input) for ch in ['a':'z';]])


# Tools for dealing with PCI
const DRV_IS_PCI_VENDOR_ID_INTEL = 0x8086
const VENDOR_ID_MASK = 0x0000_FFFF
const DEVICE_ID_MASK = 0xFFFF_0000
const DEVICE_ID_BITSHIFT = 16
const PCI_ENABLE = 0x8000_0000

# These are found in uncore performance monitoring guide
# Under: "Uncore Performance Monitoring State in PCICFG space"
# Device and Function IDS for Skylake based IMCs
const SKYLAKE_IMC_REGISTERS = (
    # IMC 0 - Channels 0, 1, 2
    ((device = 10, fn = 2), (device = 10, fn = 6), (device = 11, fn = 2)),
    # IMC 1 - Channels 0, 1, 2
    ((device = 12, fn = 2), (device = 12, fn = 6), (device = 13, fn = 2)),
)

const IMC_DEVICE_IDS = (0x2042, 0x2046, 0x204A)

# This is taken from the PCM code
#
# TODO: Dynamically find number of sockets
#
# The general idea is that we know the device and function numbers of IMC devices
# Furthermore, we know the device IDS
#
# So, we enumerate all of the PCI busses until
#   1. We find a valid path
#   2. We read the vendor ID and device ID from the device
#       - Vendor ID must match Intel (0x8086)
#       - Device ID must match one of the iMC device IDS
#
# When we've found this, we've found the bus for the socket.
#
# NOTE: I'm assuming that lower bus numbers correspond to lower sockets.
# This appears consistent with what's happening in PCM.
#
# NOTE: There's a lot of business with PCI group numbers.
# I'm ignoring that for now because we only have group 0 on our system.
#
# TODO: We can abstract this for different architectures using dispatch, but I'm not
# too worried about that at the moment.
function findbusses()
    socket_to_bus = UInt[]

    bus_numbers = 0:255
    device = first(first(SKYLAKE_IMC_REGISTERS)).device
    fn = first(first(SKYLAKE_IMC_REGISTERS)).fn

    for bus in bus_numbers
        # Check to see if the path exists
        path = pcipath(bus, device, fn)
        ispath(path) || continue

        pci = Handle(path)

        # Read the first value from the bus - compare against vendor
        seek(pci, IndexZero(0))
        value = read(pci, UInt32)
        vendor_id = value & VENDOR_ID_MASK
        device_id = (value & DEVICE_ID_MASK) >> DEVICE_ID_BITSHIFT

        # Check if this is run by Intel
        vendor_id == DRV_IS_PCI_VENDOR_ID_INTEL || continue
        in(device_id, IMC_DEVICE_IDS) || continue
        push!(socket_to_bus, bus)

        close(pci)
    end

    return socket_to_bus
end


function SolutionCallback()
    return SolutionCallback(())
end

function get_solution(obj::SolutionCallback)
    return jcall(obj, "getSolution", PointValuePair, ())
end

function is_solution_optimal(obj::SolutionCallback)
    return jcall(obj, "isSolutionOptimal", jboolean, ())
end


using Test
using SingleColumnModels.DomainDecompositions
DD = DomainDecompositions

@testset "DomainSubSet, unit tests" begin

    @test DD.n_subdomains(DD.GridMean{1}) == 1
    @test DD.n_subdomains(DD.Environment{1}) == 1
    @test DD.n_subdomains(DD.Updraft{5}) == 5

    @test DD.gridmean(DomainSubSet(gm = true)) == true
    @test DD.environment(DomainSubSet(gm = true)) == false
    @test DD.updraft(DomainSubSet(gm = true)) == false

    @test DD.gridmean(DomainSubSet(en = true)) == false
    @test DD.environment(DomainSubSet(en = true)) == true
    @test DD.updraft(DomainSubSet(en = true)) == false

    @test DD.gridmean(DomainSubSet(ud = true)) == false
    @test DD.environment(DomainSubSet(ud = true)) == false
    @test DD.updraft(DomainSubSet(ud = true)) == true

    domain_set = DomainSet(gm = 1, en = 1, ud = 4)
    @test DD.n_subdomains(domain_set, DomainSubSet(gm = true)) == (1, 0, 0)
    @test DD.n_subdomains(domain_set, DomainSubSet(en = true)) == (0, 1, 0)
    @test DD.n_subdomains(domain_set, DomainSubSet(ud = true)) == (0, 0, 4)

    @test DD.gridmean(domain_set, DomainSubSet(gm = true)) == 1
    @test DD.environment(domain_set, DomainSubSet(en = true)) == 1
    @test DD.updraft(domain_set, DomainSubSet(ud = true)) == 4

    @test DD.gridmean(domain_set, DomainSubSet(en = true)) == 0
    @test DD.environment(domain_set, DomainSubSet(gm = true)) == 0
    @test DD.updraft(domain_set, DomainSubSet(gm = true)) == 0

    @test DD.sum(domain_set, DomainSubSet(gm = true)) == 1
    @test DD.sum(domain_set, DomainSubSet(en = true)) == 1
    @test DD.sum(domain_set, DomainSubSet(ud = true)) == 4
end

@testset "DomainSubSet, single domain" begin
    domain_set, domain_subset = DomainSet(gm = 1), DomainSubSet(gm = true)
    idx = DomainIdx(domain_set)
    idx_ss = DomainIdx(domain_set, domain_subset)
    a_map = DD.get_sv_a_map(idx, idx_ss)
    gm, en, ud, sd, al = allcombinations(idx)
    @test gm == 1
    @test en == 0
    @test ud == (0,)
    @test idx_ss == idx

    domain_set, domain_subset = DomainSet(en = 1), DomainSubSet(en = true)
    idx = DomainIdx(domain_set)
    idx_ss = DomainIdx(domain_set, domain_subset)
    a_map = DD.get_sv_a_map(idx, idx_ss)
    gm, en, ud, sd, al = allcombinations(idx)
    @test en == 1
    @test gm == 0
    @test ud == (0,)
    @test idx_ss == idx

    domain_set, domain_subset = DomainSet(ud = 1), DomainSubSet(ud = true)
    idx = DomainIdx(domain_set)
    idx_ss = DomainIdx(domain_set, domain_subset)
    a_map = DD.get_sv_a_map(idx, idx_ss)
    gm, en, ud, sd, al = allcombinations(idx)
    @test en == 0
    @test gm == 0
    @test ud == (1,)
    @test idx_ss == idx

    domain_set, domain_subset = DomainSet(ud = 2), DomainSubSet(ud = true)
    idx = DomainIdx(domain_set)
    idx_ss = DomainIdx(domain_set, domain_subset)
    a_map = DD.get_sv_a_map(idx, idx_ss)
    gm, en, ud, sd, al = allcombinations(idx)
    @test en == 0
    @test gm == 0
    @test ud == (1, 2)
    @test idx_ss == idx
end

@testset "DomainSubSet, two domains" begin
    domain_set, domain_subset =
        DomainSet(gm = 1, en = 1), DomainSubSet(gm = true, en = true)
    idx = DomainIdx(domain_set)
    idx_ss = DomainIdx(domain_set, domain_subset)
    a_map = DD.get_sv_a_map(idx, idx_ss)
    gm, en, ud, sd, al = allcombinations(idx)
    @test gm == 2
    @test en == 1
    @test ud == (0,)
    @test idx_ss == idx

    domain_set, domain_subset =
        DomainSet(gm = 1, ud = 3), DomainSubSet(gm = true, ud = true)
    idx = DomainIdx(domain_set)
    idx_ss = DomainIdx(domain_set, domain_subset)
    a_map = DD.get_sv_a_map(idx, idx_ss)
    gm, en, ud, sd, al = allcombinations(idx)
    @test gm == 4
    @test en == 0
    @test ud == (1, 2, 3)
    @test idx_ss == idx

    domain_set, domain_subset =
        DomainSet(en = 1, ud = 3), DomainSubSet(en = true, ud = true)
    idx = DomainIdx(domain_set)
    idx_ss = DomainIdx(domain_set, domain_subset)
    a_map = DD.get_sv_a_map(idx, idx_ss)
    gm, en, ud, sd, al = allcombinations(idx)
    @test gm == 0
    @test en == 4
    @test ud == (1, 2, 3)
    @test idx_ss == idx

    domain_set, domain_subset =
        DomainSet(gm = 1, ud = 1), DomainSubSet(gm = true, ud = true)
    idx = DomainIdx(domain_set)
    idx_ss = DomainIdx(domain_set, domain_subset)
    a_map = DD.get_sv_a_map(idx, idx_ss)
    gm, en, ud, sd, al = allcombinations(idx)
    @test gm == 2
    @test en == 0
    @test ud == (1,)
    @test idx_ss == idx

    domain_set, domain_subset =
        DomainSet(en = 1, ud = 1), DomainSubSet(en = true, ud = true)
    idx = DomainIdx(domain_set)
    idx_ss = DomainIdx(domain_set, domain_subset)
    a_map = DD.get_sv_a_map(idx, idx_ss)
    gm, en, ud, sd, al = allcombinations(idx)
    @test gm == 0
    @test en == 2
    @test ud == (1,)
    @test idx_ss == idx
end

@testset "DomainIdx, all domains" begin
    domain_set, domain_subset = DomainSet(gm = 1, en = 1, ud = 4),
    DomainSubSet(gm = true, en = true, ud = true)
    idx = DomainIdx(domain_set)
    idx_ss = DomainIdx(domain_set, domain_subset)
    a_map = DD.get_sv_a_map(idx, idx_ss)
    gm, en, ud, sd, al = allcombinations(idx)
    @test gm == 6
    @test en == 5
    @test ud == (1, 2, 3, 4)
    @test idx_ss == idx
end

@testset "DomainIdx, utilizing DomainSubSet" begin
    domain_set = DomainSet(gm = 1, en = 1, ud = 4)
    idx = DomainIdx(domain_set)

    gm, en, ud, sd, al = allcombinations(idx)
    @test gm == 6
    @test en == 5
    @test ud == (1, 2, 3, 4)

    domain_subset = DomainSubSet(gm = true)
    idx_ss = DomainIdx(domain_set, domain_subset)
    a_map = DD.get_sv_a_map(idx, idx_ss)

    gm, en, ud, sd, al = allcombinations(idx_ss)
    @test gm == 1
    @test en == 0
    @test ud == (0,)

    domain_subset = DomainSubSet(en = true)
    idx_ss = DomainIdx(domain_set, domain_subset)
    a_map = DD.get_sv_a_map(idx, idx_ss)

    gm, en, ud, sd, al = allcombinations(idx_ss)
    @test gm == 0
    @test en == 1
    @test ud == (0,)

    domain_subset = DomainSubSet(ud = true)
    idx_ss = DomainIdx(domain_set, domain_subset)
    a_map = DD.get_sv_a_map(idx, idx_ss)

    gm, en, ud, sd, al = allcombinations(idx_ss)
    @test gm == 0
    @test en == 0
    @test ud == (1, 2, 3, 4)
end

@testset "Test DomainSubSet indexing, multiple domains" begin
    domain_set, domain_subset =
        DomainSet(gm = 1, en = 1, ud = 4), DomainSubSet(gm = true)

    idx = DomainIdx(domain_set)
    idx_ss = DomainIdx(domain_set, domain_subset)
    a_map = DD.get_sv_a_map(idx, idx_ss)
    gm, en, ud, sd, al = allcombinations(idx)

    idx_ss = DomainIdx(domain_set, domain_subset)
    a_map = DD.get_sv_a_map(idx, idx_ss)
    @test DD.updraft(idx_ss) == (0,)
    @test DD.environment(idx_ss) == 0
    @test DD.gridmean(idx_ss) == 1

    domain_subset = DomainSubSet(en = true)
    idx_ss = DomainIdx(domain_set, domain_subset)
    a_map = DD.get_sv_a_map(idx, idx_ss)
    @test DD.updraft(idx_ss) == (0,)
    @test DD.environment(idx_ss) == 1
    @test DD.gridmean(idx_ss) == 0

    domain_subset = DomainSubSet(ud = true)
    idx_ss = DomainIdx(domain_set, domain_subset)
    a_map = DD.get_sv_a_map(idx, idx_ss)
    @test DD.updraft(idx_ss) == (1, 2, 3, 4)
    @test DD.environment(idx_ss) == 0
    @test DD.gridmean(idx_ss) == 0

    domain_subset = DomainSubSet(gm = true, en = true)
    idx_ss = DomainIdx(domain_set, domain_subset)
    a_map = DD.get_sv_a_map(idx, idx_ss)
    @test DD.updraft(idx_ss) == (0,)
    @test DD.environment(idx_ss) == 1
    @test DD.gridmean(idx_ss) == 2

    domain_subset = DomainSubSet(gm = true, ud = true)
    idx_ss = DomainIdx(domain_set, domain_subset)
    a_map = DD.get_sv_a_map(idx, idx_ss)
    @test DD.updraft(idx_ss) == (1, 2, 3, 4)
    @test DD.environment(idx_ss) == 0
    @test DD.gridmean(idx_ss) == 5

    domain_subset = DomainSubSet(en = true, ud = true)
    idx_ss = DomainIdx(domain_set, domain_subset)
    a_map = DD.get_sv_a_map(idx, idx_ss)
    @test DD.updraft(idx_ss) == (1, 2, 3, 4)
    @test DD.environment(idx_ss) == 5
    @test DD.gridmean(idx_ss) == 0

    domain_subset = DomainSubSet(gm = true, en = true, ud = true)
    idx_ss = DomainIdx(domain_set, domain_subset)
    a_map = DD.get_sv_a_map(idx, idx_ss)
    @test DD.updraft(idx_ss) == (1, 2, 3, 4)
    @test DD.environment(idx_ss) == 5
    @test DD.gridmean(idx_ss) == 6
end

@testset "Test global indexing, many updrafts" begin

    vars = (
        (:ρ_0, DomainSubSet(gm = true)),
        (:a, DomainSubSet(gm = true, en = true, ud = true)),
        (:tke, DomainSubSet(en = true, ud = true)),
        (:K_h, DomainSubSet(ud = true)),
    )

    D = Dict([x => y for (x, y) in vars])

    domain_set = DomainSet(gm = 1, en = 1, ud = 4)
    idx = DomainIdx(domain_set)
    @test DD.n_subdomains(domain_set) == (1, 1, 4)
    @test DD.has_gridmean(idx) == true
    @test DD.has_environment(idx) == true
    @test DD.has_updraft(idx) == true

    sd = subdomains(idx)
    al = alldomains(idx)
    gm, en, ud = eachdomain(idx)
    gm, en, ud, sd, al = allcombinations(idx)

    @test sum(domain_set) == 6

    sd_unmapped = DD.get_sd_unmapped(vars, idx, domain_set)
    sd_mapped = DD.get_sd_mapped(vars, idx, domain_set)

    @test sd_unmapped[:a] == Int[6, 5, 1, 2, 3, 4]
    @test sd_unmapped[:tke] == Int[5, 1, 2, 3, 4]
    @test sd_unmapped[:K_h] == Int[1, 2, 3, 4]
    @test sd_unmapped[:ρ_0] == Int[6]

    @test sd_mapped[:a] == Int[6, 5, 1, 2, 3, 4]
    @test sd_mapped[:tke] == Int[0, 5, 1, 2, 3, 4]
    @test sd_mapped[:K_h] == Int[0, 0, 1, 2, 3, 4]
    @test sd_mapped[:ρ_0] == Int[1, 0, 0, 0, 0]

    vm, dss_per_var, var_names = DD.get_var_mapper(vars, domain_set)

    idx_ss_per_var = Dict([name => DomainIdx(domain_set, dss_per_var[name]) for
    name in var_names]...)
    a_map = Dict([name => DD.get_sv_a_map(idx, idx_ss_per_var[name]) for
    name in var_names]...)
    gm, en, ud, sd, al = allcombinations(idx)

    @test DD.var_string(vm, idx, idx_ss_per_var[:a], :a, gm) == "a_gm"
    @test DD.var_string(vm, idx, idx_ss_per_var[:a], :a, en) == "a_en"
    @test DD.var_string(vm, idx, idx_ss_per_var[:a], :a, ud[1]) == "a_ud_1"
    @test DD.var_string(vm, idx, idx_ss_per_var[:a], :a, ud[2]) == "a_ud_2"
    @test DD.var_string(vm, idx, idx_ss_per_var[:a], :a, ud[3]) == "a_ud_3"
    @test DD.var_string(vm, idx, idx_ss_per_var[:a], :a, ud[4]) == "a_ud_4"
    @test DD.var_suffix(vm, idx, idx_ss_per_var[:a], :a, gm) == "_gm"
    @test DD.var_suffix(vm, idx, idx_ss_per_var[:a], :a, en) == "_en"
    @test DD.var_suffix(vm, idx, idx_ss_per_var[:a], :a, ud[1]) == "_ud_1"
    @test DD.var_suffix(vm, idx, idx_ss_per_var[:a], :a, ud[2]) == "_ud_2"
    @test DD.var_suffix(vm, idx, idx_ss_per_var[:a], :a, ud[3]) == "_ud_3"
    @test DD.var_suffix(vm, idx, idx_ss_per_var[:a], :a, ud[4]) == "_ud_4"
    @test_throws BoundsError DD.var_suffix(
        vm,
        idx,
        idx_ss_per_var[:a],
        :a,
        1000,
    )

    @test DD.get_i_var(a_map[:ρ_0], gm) == 1
    @test DD.get_i_var(a_map[:a], ud[1]) == 1
    @test DD.get_i_var(a_map[:a], ud[2]) == 2
    @test DD.get_i_var(a_map[:a], ud[3]) == 3
    @test DD.get_i_var(a_map[:a], ud[4]) == 4
    @test DD.get_i_var(a_map[:a], en) == 5
    @test DD.get_i_var(a_map[:a], gm) == 6

    @test DD.get_i_state_vec(vm, a_map[:ρ_0], :ρ_0, gm) == 1
    @test DD.get_i_state_vec(vm, a_map[:a], :a, ud[1]) == 2
    @test DD.get_i_state_vec(vm, a_map[:a], :a, ud[2]) == 3
    @test DD.get_i_state_vec(vm, a_map[:a], :a, ud[3]) == 4
    @test DD.get_i_state_vec(vm, a_map[:a], :a, ud[4]) == 5
    @test DD.get_i_state_vec(vm, a_map[:a], :a, en) == 6
    @test DD.get_i_state_vec(vm, a_map[:a], :a, gm) == 7

    @test DD.get_i_state_vec(vm, a_map[:tke], :tke, ud[1]) == 8
    @test DD.get_i_state_vec(vm, a_map[:tke], :tke, ud[2]) == 9
    @test DD.get_i_state_vec(vm, a_map[:tke], :tke, ud[3]) == 10
    @test DD.get_i_state_vec(vm, a_map[:tke], :tke, ud[4]) == 11
    @test DD.get_i_state_vec(vm, a_map[:tke], :tke, en) == 12

    @test DD.get_i_state_vec(vm, a_map[:K_h], :K_h, ud[1]) == 13
    @test DD.get_i_state_vec(vm, a_map[:K_h], :K_h, ud[2]) == 14
    @test DD.get_i_state_vec(vm, a_map[:K_h], :K_h, ud[3]) == 15
    @test DD.get_i_state_vec(vm, a_map[:K_h], :K_h, ud[4]) == 16

end


#JSON Parser
#Adapted from http://www.mathworks.com/matlabcentral/fileexchange/23393
#Original BSD Licence, (c) 2011, François Glineur

function parse_json(strng::String)
    pos = 1
    len = length(strng)
    # String delimiters and escape characters are identified beforehand to improve speed
    # esc = regexp(str, "[\"\\\\]"); index_esc = 1; len_esc = length(esc);  #TODO Enable for speed

    function parse_object()
        parse_char('{')
        object = Dict{String, Any}()
        if next_char() != '}'
            while true
                str = parse_string()
                if isempty(str)
                    error("Name of value at position $pos cannot be empty")
                end
                parse_char(':')
                val = parse_value()
                object[str] = val
                if next_char() == '}'
                    break
                end
                parse_char(',')
            end
        end
        parse_char('}')
        return object
    end

    function parse_array()
        parse_char('[')
        object = Set()
        if next_char != ']'
            while true
                val = parse_value()
                push!(object, val)
                if next_char() == ']'
                    break
                end
                parse_char(',')
            end
        end
        parse_char(']')
        return object
    end

    function parse_char(c::Char)
        skip_whitespace()
        if pos > len || strng[pos] != c
            error("Expected $c at position $pos")
        else
            pos = pos + 1
            skip_whitespace()
        end
    end

    function next_char()
        skip_whitespace()
        if pos > len
            c = '\0'
        else
            c = strng[pos]
        end
    end

    function skip_whitespace()
        while pos <= len && isspace(strng[pos])
            pos = pos + 1
        end
    end

    function parse_string()
        if strng[pos] != '"'
            error("String starting with quotation expected at position $pos")
        else
            pos = pos + 1
        end
        str = ""
        while pos <= len
            # while index_esc <= len_esc && esc(index_esc) < pos
            #     index_esc = index_esc + 1
            # end
            # if index_esc > len_esc
            #     str = string(str, strng[pos:end])
            #     pos = len + 1
            #     break
            # else
            #     str = string(str, strng[pos:esc(index_esc)-1])
            #     pos = esc(index_esc)
            # end
            nc = strng[pos]
            if nc == '"'
                pos = pos + 1
                return string(str)
            elseif nc ==  '\\'
                if pos+1 > len
                    error_pos("End of file reached right after escape character")
                end
                pos = pos + 1
                anc = strng[pos]
                if anc == '"' || anc == '\\' || anc == '/'
                    str = string(str, strng[pos])
                    pos = pos + 1
                elseif anc ==  'b' || anc == 'f'|| anc == 'n' || anc == 'r' || anc == 't'
                    str = string(str, '\\', string[pos])
                    pos = pos + 1
                elseif anc == 'u'
                    if pos+4 > len
                        error_pos("End of file reached in escaped unicode character")
                    end
                    str = string(str, strng[pos-1:pos+4])
                    pos = pos + 5
                end
            else # should never happen
                str = string(str,strng[pos])
                pos = pos + 1
            end
        end
        error("End of file while expecting end of string")
    end

    function parse_number()
        num_regex = r"^[\w]?[-+]?[0-9]*\.?[0-9]+([eE][-+]?[0-9]+)?[\w]?"
        m = match(num_regex, strng[pos:min(len,pos+20)])
        if m === nothing
            error("Error reading number at position $pos")
        end
        delta = m.offset + length(m.match)
        pos = pos + delta -1
        return float(m.match)
    end

    function  parse_value()
        nc = strng[pos]
        if nc == '"'
            val = parse_string()
            return val
        elseif nc == '['
            val = parse_array()
            return val
        elseif nc == '{'
            val = parse_object()
            return val
        elseif nc == '-' || nc == '0' || nc == '1' || nc == '2' || nc == '3' || nc == '4' || nc == '5' || nc == '6' || nc == '7' || nc == '8' || nc == '9'
            val = parse_number()
            return val
        elseif nc == 't'
            if pos+3 <= len && strng[pos:pos+3] == "true"
                val = true
                pos = pos + 4
                return val
            end
        elseif nc == 'f'
            if pos+4 <= len && strng[pos:pos+4] == "false"
                val = false
                pos = pos + 5
                return val
            end
        elseif nc == 'n'
            if pos+3 <= len && strng[pos:pos+3] == "null"
                val = []
                pos = pos + 4
                return val
            end
        end
        error("Value expected at position $pos")
    end

    if pos <= len
        nc = next_char()
        if nc == '{'
            return parse_object()
        elseif nc ==  '['
            return parse_array()
        else
            error("Outer level structure must be an object or an array")
        end
    end
end


# # Battle Royale
#
# ```@raw html
# <video width="auto" controls autoplay loop>
# <source src="../battle.mp4" type="video/mp4">
# </video>
# ```
#
# This example illustrates how to leverage higher dimensions of a `GridSpace` to identify
# the distance from neighbors not just spatially, but also categorically. We'll also use
# the [`walk!`](@ref) function extensively.

# The `Models` module includes this example as [`Models.battle`](@ref).

# ## Rules of Engagement
# Agents wander around the map looking for opponents. When a grid space is occupied by
# two or more agents there will be a battle. With experience gained from the fight, the
# victor searches for more opponents to crush and losers scurry away defeated or possibly
# even die. This process repeats until there is a single, definitive winner.
#
# For this battle ground to exist, the following rules must be followed:
# - Agents have an experience level, starting at level 1 up to a maximum of 10.
# - Agents will search for the nearest worthy opponent (one with equal or ±1 experience level) and move towards them to attack, so long as something more important doesn't happen, which could be
#   - A tougher opponent (with experience level +2 or higher) is nearby: run!
#   - There are no worthy opponents available, but there are weak ones (with experience level -2 or lower): chase them down.
#   - Capture and taunt a weaker opponent, then kill them.
#   - Notice a tough opponent is occupied, sneak up and kill them.
#   - There is no-one worthy to fight, but also no-one left to taunt. All bets are off: THERE CAN BE ONLY ONE.

# Battles are won by weighted chance - a higher level gives an agent a larger chance of
# winning, but does not guarantee it. When a victor is chosen
#   - The difference in experience between opponents is swapped.
#   - If an agents experience reaches 0, they die.
#
# Captured opponents will be killed once taunted. The captor will gain half of their
# experience. If an opportunist manages to take the captor by surprise, they can gain
# up to half of the captor's experience. This means a level 1 agent may eliminate a
# level 10 captor and jump straight to level 6.
#
# Once all rules of engagement have been exhausted, the final showdown begins. Opponents
# fight their closest adversary regardless of experience level. Winner takes all.

# ## Model Setup

using Random # hide
using Agents
using InteractiveDynamics
using CairoMakie
CairoMakie.activate!() # hide

mutable struct Fighter <: AbstractAgent
    id::Int
    pos::Dims{3}
    has_prisoner::Bool
    capture_time::Int
    shape::Symbol # For plotting
end

# As you can see, the properties of out agent are very simple and contain only two
# parameters that are needed to store context from one time step to the next. All
# other properties needed are stored in the space. `pos` is three-dimensional, two
# for the actual space agents move within, and a third categorical dimension representing
# their level. `shape` is used solely for plotting (well, used once just for convenience).

# Now let's set up the battle field:
function battle(; fighters = 50)
    model = ABM(
        Fighter,
        GridSpace((100, 100, 10); periodic = false);
        scheduler = Schedulers.randomly,
    )

    n = 0
    while n != fighters
        pos = (rand(model.rng, 1:100, 2)..., 1) # Start at level 1
        if isempty(pos, model)
            add_agent!(pos, model, false, 0, :diamond)
            n += 1
        end
    end

    return model
end

Random.seed!(6547) # hide
model = battle()

# 50 opponents positioned randomly on a 100x100 grid, with no escape
# (`periodic = false`). To leverage categorical dimensions fully, non-periodic chebyshev
# space is necessary.

# ## Game Dynamics
#
# To implement the rules of engagement, only an `agent_step!` function is required,
# along with a few helper functions.

space(agent) = agent.pos[1:2]
level(agent) = agent.pos[3]

# `space` allows us to invoke a number of helpful utilities provided by Agents.jl
# but only operate on our spatial dimensions, `level` is a wrapper to access the agent's
# experience easily.

# Nearest agents that satisfy our search criteria can be identified via Euclidean
# distance solely on the spatial dimensions of our `GridSpace`.

function closest_target(agent::Fighter, ids::Vector{Int}, model::ABM)
    if length(ids) == 1
        closest = ids[1]
    else
        close_id = argmin(map(id -> edistance(space(agent), space(model[id]), model), ids))
        closest = ids[close_id]
    end
    return model[closest]
end

# Since our battles are only between opponents with equal, or as much as one level apart,
# the odds can be set explicitly. Stronger opponents have twice the capacity of winning a
# match.

function battle!(one::Fighter, two::Fighter, model)
    if level(one) == level(two)
        ## Odds are equivalent
        one_winner = rand(model.rng) < 0.5
    elseif level(one) > level(two)
        ## Odds are in favor of one
        one_winner = 2 * rand(model.rng) > rand(model.rng)
    else
        ## Odds are in favor of two
        one_winner = rand(model.rng) > 2 * rand(model.rng)
    end

    one_winner ? (up = one; down = two) : (up = two; down = one)

    new_lvl_up = min(level(up) + 1, 10)
    new_pos_up =
        clamp.(rand(model.rng, -1:1, 2) .+ space(up), [1, 1], size(model.space)[1:2])
    move_agent!(up, (new_pos_up..., new_lvl_up), model)
    new_lvl_down = level(down) - 1
    if new_lvl_down == 0
        kill_agent!(down, model)
    else
        move_agent!(down, (space(down)..., new_lvl_down), model)
    end
end

# If an agent has a prisoner, it will taunt it for a time, then kill it, so long as
# an opportunist doesn't sneak up on them first!
# Here we use the tuple constructor with [`nearby_ids`](@ref) to look for agents at
# the same position as the captor `(0, 0)`, and *any* level `(..., 10)`. We could also
# use the range constructor in this instance
# `nearby_ids(agent, model, [(1, 0:0), (2, 0:0)])`, meaning which is more performant but not as
# readable.

function captor_behavior!(agent, model)
    close_ids = collect(nearby_ids(agent, model, (0, 0, 10)))
    if length(close_ids) == 1
        ## Taunt prisoner or kill it
        prisoner = model[close_ids[1]]
        if prisoner.capture_time > 10
            agent.shape = :rect
            gain = ceil(Int, level(prisoner) / 2)
            new_lvl = min(level(agent) + gain, 10)
            kill_agent!(prisoner, model)
            agent.has_prisoner = false
            move_agent!(agent, (space(agent)..., new_lvl), model)
        end
    else
        ## Someone is here to kill the captor. Could be more than one opponent
        prisoner = [model[id] for id in close_ids if model[id].capture_time > 0][1]
        exploiter = rand(
            model.rng,
            [
                model[id]
                for
                id in close_ids if
                model[id].capture_time == 0 && model[id].has_prisoner == false
            ],
        )
        exploiter.shape = :rect
        gain = ceil(Int, level(agent) / 2)
        new_lvl = min(level(agent) + rand(model.rng, 1:gain), 10)
        kill_agent!(agent, model)
        move_agent!(exploiter, (space(exploiter)..., new_lvl), model)
        ## Prisoner runs away in the commotion
        prisoner.shape = :utriangle
        prisoner.capture_time = 0
        walk!(prisoner, (rand(model.rng, -1:1, 2)..., 0), model)
    end
end

# When there are only few fighters standing, the stakes are higher. Prior experience is
# paramount since there is no gain, and fights are to the death.

function endgame!(agent, model)
    origin = space(agent)
    end_ids = collect(Iterators.filter(
        id -> model[id].shape == :circle && id != agent.id,
        allids(model),
    ))
    agent.shape = :circle
    if !isempty(end_ids)
        opponent = closest_target(agent, end_ids, model)
        target = space(opponent)
        if origin == target
            ## Battle
            agent.shape = :rect
            opponent.shape = :rect
            showdown!(agent, opponent, model)
        else
            walk!(agent, (sign.(target .- origin)..., 0), model)
        end
    end
end

function showdown!(one::Fighter, two::Fighter, model)
    if level(one) == level(two)
        ## Odds are equivalent
        one_winner = rand(model.rng) < 0.5
    elseif level(one) > level(two)
        ## Odds are in favor of one
        one_winner = level(one) - level(two) * rand(model.rng) > rand(model.rng)
    else
        ## Odds are in favor of two
        one_winner = rand(model.rng) > level(two) - level(one) * rand(model.rng)
    end

    one_winner ? kill_agent!(two, model) : kill_agent!(one, model)
end

# The rest of our interactions flow down a hierarchy, so we'll place them directly in the
# `agent_step!` function. We use the tuple search for `occupied_ids` here, as we did with
# `close_ids` above. The rest of the searches however use the range search to provide a
# more precise criteria.
#
# The easiest context to explore is `worthy_ids`: all we want to do is find an agent with
# a similar level. If we used the tuple search here, we would have to search
# `(100, 100, 1)` - even though we are not at all interested in the spatial location of
# the neighbors at this time. `(3, -1:1)` is therefore more accurate representation.
#
# A more complex example is that of `strong_ids`. We are looking for agents with a level
# 2-4 points higher withing a distance of `(5, 5)`. The range search becomes a little
# verbose, but precise. An equivalent tuple search is not completely possible however.
# The closest solution is `(5, 5, 4)`, which also looks for *weaker* opponents and must
# be filtered to the correct neighbor set after the fact. In this instance the range
# search has significant performance gains.

function agent_step!(agent, model)
    if agent.capture_time > 0
        ## Captured agents are powerless, but we need to keep track of how long
        ## they have been in this state
        agent.capture_time += 1
    elseif agent.has_prisoner
        captor_behavior!(agent, model)
    else
        origin = space(agent)
        ## Find agents that have captives, they are not focused
        occupied_ids = collect(Iterators.filter(
            id -> model[id].has_prisoner,
            nearby_ids(agent, model, (7, 7, 10)),
        ))
        if !isempty(occupied_ids)
            ## Sneak up behind them
            target = space(closest_target(agent, occupied_ids, model))
            agent.shape = :pentagon
            walk!(agent, (sign.(target .- origin)..., 0), model)
        else
            ## Opponents that are greatly higher in rank that the current agent
            strong_ids = collect(nearby_ids(agent, model, [(1, -5:5), (2, -5:5), (3, 2:4)]))
            if !isempty(strong_ids)
                ## Run away from nearest
                target = space(closest_target(agent, strong_ids, model))
                agent.shape = :utriangle
                walk!(agent, (sign.(origin .- target)..., 0), model)
            else
                ## There are no distractions. Search for the closest worthy opponent
                worthy_ids = collect(nearby_ids(agent, model, [(3, -1:1)]))
                if !isempty(worthy_ids)
                    opponent = closest_target(agent, worthy_ids, model)
                    target = space(opponent)
                    if origin == target
                        ## Battle
                        agent.shape = :rect
                        opponent.shape = :rect
                        battle!(agent, opponent, model)
                    else
                        ## Move towards worthy opponent
                        agent.shape = :diamond
                        walk!(agent, (sign.(target .- origin)..., 0), model)
                    end
                else
                    ## Find any weak targets in the vicinity
                    weak_ids = collect(nearby_ids(
                        agent,
                        model,
                        [(1, -10:10), (2, -10:10), (3, -4:-2)],
                    ))
                    if !isempty(weak_ids)
                        prisoner = closest_target(agent, weak_ids, model)
                        target = space(prisoner)
                        if origin == target
                            ## Capture and taunt target
                            agent.has_prisoner = true
                            agent.shape = :vline
                            prisoner.capture_time += 1
                            prisoner.shape = :hline
                        else
                            ## Chase down nearest (can move 2 steps at a time!)
                            agent.shape = :star4
                            walk!(agent, (2 .* sign.(target .- origin)..., 0), model)
                        end
                    else
                        ## Abandon honour. This is the end
                        endgame!(agent, model)
                    end
                end
            end
        end
    end

    return nothing
end

# ## Let the Battle Begin
# Plotting is relatively straightforward. [`plotabm`](@ref) cannot be used explicitly (yet)
# since it expects our categorical dimension is actually a third spatial one.
# We start with some custom legends to easier understand the dynamics.

label_action = ["Battle", "Run", "Showdown", "Sneak", "Duel", "Captor", "Prisoner", "Chase"]
actions = [:rect, :utriangle, :circle, :pentagon, :diamond, :vline, :hline, :star4]
group_action = [
    MarkerElement(
        marker = marker,
        color = :black,
        strokecolor = :transparent,
        markersize = 15,
    ) for marker in actions
]
group_level = [
    PolyElement(color = color, strokecolor = :transparent) for color in cgrad(:tab10)[1:10]
]
nothing #hide

# And some complex internals that will be hidden away in the near future

e = size(model.space.s)[1:2] .+ 2
o = zero.(e) .- 2
clr(agent) = cgrad(:tab10)[level(agent)]
mkr(a) = a.shape
colors = Observable(to_color.([clr(model[id]) for id in Schedulers.by_id(model)]))
markers = Observable([mkr(model[id]) for id in Schedulers.by_id(model)])
pos = Observable([model[id].pos for id in Schedulers.by_id(model)])
stepper = InteractiveDynamics.ABMStepper(
    clr,
    mkr,
    15,
    nothing,
    Schedulers.by_id,
    pos,
    colors,
    Observable(15),
    markers,
    nothing,
    nothing
)
nothing #hide

# Finally, the plot:

f = Figure(resolution = (600, 700))
ax = f[1, 1] = Axis(f, title = "Battle Royale")
hidedecorations!(ax)
ax.xgridvisible = true
ax.ygridvisible = true
f[2, 1] = Legend(
    f,
    [group_action, group_level],
    [label_action, string.(1:10)],
    ["Action", "Level"],
    orientation = :horizontal,
    tellheight = true,
    tellwidth = false,
    nbanks = 5,
)

scatter!(ax, pos; color = colors, markersize = 15, marker = markers, strokewidth = 0.0)
xlims!(ax, o[1], e[1])
ylims!(ax, o[2], e[2])
record(f, "battle.mp4", 0:225; framerate = 10) do i
    Agents.step!(stepper, model, agent_step!, dummystep, 1)
end
nothing # hide
# ```@raw html
# <video width="auto" controls autoplay loop>
# <source src="../battle.mp4" type="video/mp4">
# </video>
# ```

# Some interesting behaviour emerges: sometimes you see a group of diamonds chasing one triangle.
# What ends up happening here is usually a close pair that wishes to fight gets caught
# out by the weaker one of the two running away from an even stronger opponent. Problem
# is that this stronger opponent is chasing the stronger of the pair, but since the
# weakest of the pair is still closer to the newcomer, there is a stalemate.
# This is usually resolved by hitting a boundary or other opponents.


mutable struct FunctionCallingAffect{funcFunc, funcatType, funcatCacheType}
    func::funcFunc
    funcat::funcatType
    funcat_cache::funcatCacheType
    func_everystep::Bool
    func_start::Bool
    funciter::Int
end

function (affect!::FunctionCallingAffect)(integrator,force_func = false)
    # see OrdinaryDiffEq.jl -> integrator_utils.jl, function funcvalues!
    while !isempty(affect!.funcat) && integrator.tdir*top(affect!.funcat) <= integrator.tdir*integrator.t # Perform funcat
        affect!.funciter += 1
        curt = pop!(affect!.funcat) # current time
        if curt != integrator.t # If <t, interpolate
            if typeof(integrator) <: ODEIntegrator
                # Expand lazy dense for interpolation
                DiffEqBase.addsteps!(integrator)
            end
            if typeof(integrator.u) <: Union{Number,SArray}
                curu = integrator(curt)
            else
                curu = first(get_tmp_cache(integrator))
                integrator(curu,curt) # inplace since func_func allocates
            end
            affect!.func(curu, curt, integrator)
        else # ==t, just func
            affect!.func(integrator.u, integrator.t, integrator)
        end
    end
    if affect!.func_everystep || force_func
        affect!.funciter += 1
        affect!.func(integrator.u, integrator.t, integrator)
    end
    u_modified!(integrator, false)
end

function functioncalling_initialize(cb, u, t, integrator)
    if cb.affect!.funciter != 0
        if integrator.tdir > 0
            cb.affect!.funcat = BinaryMinHeap(cb.affect!.funcat_cache)
        else
            cb.affect!.funcat = BinaryMaxHeap(cb.affect!.funcat_cache)
        end
        cb.affect!.funciter = 0
    end
    cb.affect!.func_start && cb.affect!(integrator)
end

"""
    FunctionCallingCallback(func;
                    funcat=Vector{Float64}(),
                    func_everystep=isempty(funcat),
                    func_start = true,
                    tdir=1)

A `DiscreteCallback` applied after every step to call `func(u,t,integrator)`
If `func_everystep`, every step of the integrator gives a `func` call.
If `funcat` is specified, the function is called at the given times, using
interpolation if necessary.
If the time `tdir` direction is not positive, i.e. `tspan[1] > tspan[2]`,
`tdir = -1` has to be specified.
"""
function FunctionCallingCallback(func;
                        funcat=Vector{Float64}(),
                        func_everystep=isempty(funcat),
                        func_start = true,
                        tdir=1)
    # funcat conversions, see OrdinaryDiffEq.jl -> integrators/type.jl
    funcat_vec = collect(funcat)
    if tdir > 0
        funcat_internal = BinaryMinHeap(funcat_vec)
    else
        funcat_internal = BinaryMaxHeap(funcat_vec)
    end
    affect! = FunctionCallingAffect(func, funcat_internal,
                                    funcat_vec, func_everystep, func_start, 0)
    condtion = (u, t, integrator) -> true
    DiscreteCallback(condtion, affect!;
                     initialize = functioncalling_initialize,
                     save_positions=(false,false))
end

export FunctionCallingCallback