text:
###
# Reflection tooling
##

using Base.Meta
using .Compiler: widenconst, argextype, Const, MethodMatchInfo,
    UnionSplitApplyCallInfo, UnionSplitInfo, ConstCallInfo,
    MethodResultPure, ApplyCallInfo

const is_return_type = Core.Compiler.is_return_type
const sptypes_from_meth_instance = Core.Compiler.sptypes_from_meth_instance
const may_invoke_generator = Base.may_invoke_generator
code_for_method(method, metharg, methsp, world, force=false) = Core.Compiler.specialize_method(method, metharg, methsp, force)

# https://github.com/JuliaLang/julia/pull/36318
_id(x) = x.id
_id(x::Core.Argument) = x.n
const SlotOrArgument = Union{Core.Argument,Core.Slot}

function unwrapconst(a)
    if isa(a, Const)
        a = Core.Typeof(a.val)
    elseif isa(a, Core.Compiler.PartialStruct)
        a = a.typ
    elseif isa(a, Core.Compiler.MaybeUndef)
        a = a.typ
    end
    return a
end

transform(::Val, callsite) = callsite
function transform(::Val{:CuFunction}, callsite, callexpr, CI, mi, slottypes; params=nothing, kwargs...)
    sptypes = sptypes_from_meth_instance(mi)
    tt = argextype(callexpr.args[4], CI, sptypes, slottypes)
    ft = argextype(callexpr.args[3], CI, sptypes, slottypes)
    isa(tt, Const) || return callsite
    return Callsite(callsite.id, CuCallInfo(callinfo(Tuple{widenconst(ft), tt.val.parameters...}, Nothing, params=params)), callsite.head)
end

function find_callsites(CI::Union{Core.CodeInfo, IRCode}, stmt_info::Union{Vector, Nothing}, mi::Core.MethodInstance, slottypes; params=current_params(), multichoose::Bool=false, kwargs...)
    sptypes = sptypes_from_meth_instance(mi)
    callsites = Callsite[]

    for (id, c) in enumerate(isa(CI, IRCode) ? CI.stmts.inst : CI.code)
        callsite = nothing
        isa(c, Expr) || continue
        if stmt_info !== nothing
            info = stmt_info[id]
            if info !== nothing
                rt = argextype(SSAValue(id), CI, sptypes, slottypes)
                was_return_type = false
                if isa(info, MethodResultPure)
                    # TODO: We could annotate this in the UI
                    continue
                end
                if isa(info, Core.Compiler.ReturnTypeCallInfo)
                    info = info.info
                    was_return_type = true
                end
                function process_info(info)
                    if isa(info, MethodMatchInfo)
                        if info.results === missing
                            return []
                        end

                        matches = info.results.matches
                        return map(match->MICallInfo(match, rt), matches)
                    elseif isa(info, UnionSplitInfo)
                        return mapreduce(process_info, vcat, info.matches)
                    elseif isa(info, UnionSplitApplyCallInfo)
                        return mapreduce(process_info, vcat, info.infos)
                    elseif isa(info, ApplyCallInfo)
                        # XXX: This could probably use its own info. For now,
                        # we ignore any implicit iterate calls.
                        return process_info(info.call)
                    elseif isa(info, ConstCallInfo)
                        if length(info.results) == 1
                            result = info.results[1]
                            return [ConstPropCallInfo(MICallInfo(result.linfo, rt), result)]
                        else
                            vcat(process_info(info.call),
                                map(result->ConstPropCallInfo(MICallInfo(result.linfo, rt), result),
                                    filter(!isnothing, info.results)))
                        end
                    elseif isdefined(Compiler, :OpaqueClosureCallInfo) && isa(info, Compiler.OpaqueClosureCallInfo)
                        return [OCCallInfo(Core.Compiler.specialize_method(info.match), rt)]
                    elseif isdefined(Compiler, :OpaqueClosureCreateInfo) && isa(info, Compiler.OpaqueClosureCreateInfo)
                        # TODO: Add ability to descend into OCs at creation site
                        return []
                    elseif info == false
                        return []
                    else
                        @show CI
                        @show c
                        @show info
                        error()
                    end
                end
                callinfos = process_info(info)
                if !was_return_type && isempty(callinfos)
                    continue
                end
                callsite = let
                    if length(callinfos) == 1
                        callinfo = callinfos[1]
                    else
                        types = mapany(arg -> widenconst(argextype(arg, CI, sptypes, slottypes)), c.args)
                        callinfo = MultiCallInfo(Core.Compiler.argtypes_to_type(types), rt, callinfos)
                    end
                    if was_return_type
                        callinfo = ReturnTypeCallInfo(callinfo)
                    end
                    Callsite(id, callinfo, c.head)
                end
            end
        end

        if callsite === nothing && c isa Expr
            if c.head === :(=)
                c = c.args[2]
                (c isa Expr) || continue
            end
            if c.head === :invoke
                rt = argextype(SSAValue(id), CI, sptypes, slottypes)
                at = argextype(c.args[2], CI, sptypes, slottypes)
                if isa(at, Const) && is_return_type(at.val)
                    callsite = process_return_type(id, c, rt)
                else
                    callsite = Callsite(id, MICallInfo(c.args[1], rt), c.head)
                end
            elseif c.head === :foreigncall
                # special handling of jl_new_task
                length(c.args) > 0 || continue
                if c.args[1] isa QuoteNode
                    cfunc = c.args[1].value
                    if cfunc === :jl_new_task
                        func = c.args[6]
                        ftype = widenconst(argextype(func, CI, sptypes, slottypes))
                        sig = Tuple{ftype}
                        callsite = Callsite(id, TaskCallInfo(callinfo(sig, nothing, params=params)), c.head)
                    end
                end
            end
        end

        if callsite !== nothing
            if callsite.info isa MICallInfo
                mi = get_mi(callsite)
                if nameof(mi.def.module) === :CUDAnative && mi.def.name === :cufunction
                    callsite = transform(Val(:CuFunction), callsite, c, CI, mi, slottypes; params=params, kwargs...)
                end
            end

            push!(callsites, callsite)
        end
    end
    return callsites
end

function dce!(ci, mi)
    argtypes = Core.Compiler.matching_cache_argtypes(mi, nothing)[1]
    ir = Compiler.inflate_ir(ci, sptypes_from_meth_instance(mi),
                             argtypes)
    dce!(ir, mi)
    Core.Compiler.replace_code_newstyle!(ci, ir, length(argtypes)-1)
end

function dce!(ir::IRCode, mi)
    compact = Core.Compiler.IncrementalCompact(ir, true)
    # Just run through the iterator without any processing
    Core.Compiler.foreach(x -> nothing, compact)
    ir = Core.Compiler.finish(compact)
end

function do_typeinf_slottypes(mi::Core.Compiler.MethodInstance, run_optimizer::Bool, interp::Core.Compiler.AbstractInterpreter)
    ccall(:jl_typeinf_begin, Cvoid, ())
    result = Core.Compiler.InferenceResult(mi)
    frame = Core.Compiler.InferenceState(result, false, interp)
    frame === nothing && return (nothing, Any, Any[])
    if Compiler.typeinf(interp, frame) && run_optimizer
        oparams = Core.Compiler.OptimizationParams(interp)
        opt = Compiler.OptimizationState(frame, oparams, interp)
        Compiler.optimize(interp, opt, oparams, result.result)
        opt.src.inferred = true
    end
    ccall(:jl_typeinf_end, Cvoid, ())
    frame.inferred || return (nothing, Any, Any[])
    return (frame.src, result.result, frame.slottypes)
end

function preprocess_ci!(ci::CodeInfo, mi, optimize, config::CthulhuConfig)
    if optimize && config.dead_code_elimination
        # if the optimizer hasn't run, the IR hasn't been converted
        # to SSA form yet and dce is not legal
        dce!(ci, mi)
        dce!(ci, mi)
    end
end

function preprocess_ci!(ir::IRCode, mi, optimize, config::CthulhuConfig)
    if optimize && config.dead_code_elimination
        # if the optimizer hasn't run, the IR hasn't been converted
        # to SSA form yet and dce is not legal
        ir = dce!(ir, mi)
        ir = dce!(ir, mi)
    end
end

const CompilerParams = Core.Compiler.NativeInterpreter
current_params() = CompilerParams()

function first_method_instance(@nospecialize(F), @nospecialize(TT); params=current_params())
    sig = Base.signature_type(F, TT)
    first_method_instance(sig; params=params)
end

function first_method_instance(@nospecialize(sig); params=current_params())
    ci = callinfo(sig, Any, 1, params=params)
    if ci isa Union{GeneratedCallInfo, FailedCallInfo, MultiCallInfo}
        if ci isa MultiCallInfo
            @warn "expected single CallInfo got multiple" ci
        end
        return nothing
    else
        @assert ci isa MICallInfo
        return get_mi(ci)
    end
end

function choose_method_instance(@nospecialize(F), @nospecialize(TT); params=current_params())
    sig = Base.signature_type(F, TT)
    choose_method_instance(sig; params=params)
end

function callinfo(sig, rt, max=-1; params=current_params())
    methds = Base._methods_by_ftype(sig, -1, params.world)
    (methds === false || length(methds) < 1) && return FailedCallInfo(sig, rt)
    callinfos = CallInfo[]
    for x in methds
        meth = x[3]
        atypes = x[1]
        sparams = x[2]
        if isdefined(meth, :generator) && !may_invoke_generator(meth, atypes, sparams)
            push!(callinfos, GeneratedCallInfo(sig, rt))
        else
            mi = code_for_method(meth, atypes, sparams, params.world)
            if mi !== nothing
                push!(callinfos, MICallInfo(mi, rt))
            else
                push!(callinfos, FailedCallInfo(sig, rt))
            end
        end
    end

    @assert length(callinfos) != 0
    length(callinfos) == 1 && return first(callinfos)
    return MultiCallInfo(sig, rt, callinfos)
end

function find_caller_of(callee::MethodInstance, caller::MethodInstance)
    interp = CthulhuInterpreter()
    do_typeinf!(interp, caller)
    params = current_params()
    locs = Tuple{Core.LineInfoNode,Int}[]
    for optimize in (true, false)
        (CI, rt, infos, slottypes) = lookup(interp, caller, optimize)
        preprocess_ci!(CI, caller, optimize, CONFIG)
        callsites = find_callsites(CI, infos, caller, slottypes; params=params)
        callsites = filter(cs->is_callsite(cs, callee), callsites)
        foreach(cs -> add_sourceline!(locs, CI, cs.id), callsites)
    end
    # Consolidate by method, but preserve the order
    prlookup = Dict{Tuple{Symbol,Symbol},Int}()
    ulocs = Pair{Tuple{Symbol,Symbol,Int},Vector{Int}}[]
    if !isempty(locs)
        for (loc, depth) in locs
            idx = get(prlookup, (loc.method, loc.file), nothing)
            if idx === nothing
                push!(ulocs, (loc.method, loc.file, depth) => Int[])
                prlookup[(loc.method, loc.file)] = idx = length(ulocs)
            end
            lines = ulocs[idx][2]
            line = loc.line
            if line ∉ lines
                push!(lines, line)
            end
        end
    end
    return ulocs
end

function add_sourceline!(locs, CI, stmtidx::Int)
    if isa(CI, IRCode)
        stack = Base.IRShow.compute_loc_stack(CI.linetable, CI.stmts.line[stmtidx])
        for (i, idx) in enumerate(stack)
            line = CI.linetable[idx]
            line.line == 0 && continue
            push!(locs, (CI.linetable[idx], i-1))
        end
    else
        push!(locs, (CI.linetable[CI.codelocs[stmtidx]], 0))
    end
    return locs
end


text:
using RollingFunctions

"""
    BasicMetrics{O<:AbstractTakeObjective, H<:ValueSelection} <: AbstractMetrics

`BasicMetrics` is a Type that stores useful information derived from consecutive search, during either a learning process or an evaluation process.
It is filled just after a search. The metrics are called in the `launch_experiment()` function and in the `evaluate()` function.

It satisfies the two AbstractMetrics requirements:
1) the constructor `metrics(model::CPModel, heuristic::ValueSelection)`.
2) the function `(::CustomMetrics)(model::CPmodel,dt::Float64)`.

# Fields

    heuristic::H                                        ->  The related heuristic for which the metrics stores the results.
    nodeVisited::Vector{Vector{Int64}}                  ->  contains the result of each search in term of node visited : number of nodes visited to
                                                            find every solution or infeasible case of an instance.
    meanNodeVisitedUntilfirstSolFound::Vector{Float32}  ->  contains the result of each search in term of node visited to find a first solution.
    meanNodeVisitedUntilEnd::Vector{Float32}            ->  contains the result of each search in term of node visited until the end of the search.
    timeneeded::Vector{Float32}                         ->  contains the computing time required to complete each search.
    scores::Union{Nothing,Vector{Vector{Float32}}}      ->  contains the result of each search in term of scores of every solution found
                                                            (only for problem that contains an objective)
    totalReward::Union{Nothing,Vector{Float32}}         ->  contains the total reward of each search (only if heuristic is a LearnedHeuristic).
    loss::Union{Nothing,Vector{Float32}}                ->  contains the total loss of each search (only if heuristic is a LearnedHeuristic).
    meanOver::Int64                                     ->  width of the windowspan for the moving average.
    nbEpisodes::Int64                                   ->  counts the number of search the metrics has been called on.
"""
mutable struct BasicMetrics{O<:AbstractTakeObjective, H<:ValueSelection} <: AbstractMetrics
    heuristic::H
    nodeVisited::Vector{Vector{Int64}}
    meanNodeVisitedUntilfirstSolFound::Vector{Float32}
    meanNodeVisitedUntilEnd::Vector{Float32}
    timeneeded::Vector{Float32}
    scores::Union{Nothing,Vector{Vector{Union{Nothing,Float32}}}}
    totalReward::Union{Nothing,Vector{Float32}}
    loss::Union{Nothing,Vector{Float32}}
    meanOver::Int64
    nbEpisodes::Int64

    BasicMetrics{O,H}(heuristic,meanOver) where {O,H}= new{O, H}(heuristic,Vector{Vector{Int64}}(),Float32[],Float32[], Float32[], O==TakeObjective ? Vector{Vector{Float32}}() : nothing, (H == BasicHeuristic) ? nothing : Float32[], (H == BasicHeuristic) ? nothing : Float32[], meanOver,0)
end

BasicMetrics(model::CPModel, heuristic::ValueSelection; meanOver=1) = BasicMetrics{(!isnothing(model.objective)) ? TakeObjective : DontTakeObjective ,typeof(heuristic)}(heuristic,meanOver)

"""
    function (metrics<:BasicMetrics)(model::CPModel,dt::Float64)

The function is called after a search on a Constraint Programming Model.
It updates all the metrics during the search.
"""
function (metrics::BasicMetrics{DontTakeObjective, <:BasicHeuristic})(model::CPModel,dt::Float64)
    metrics.nbEpisodes+=1
    push!(metrics.nodeVisited,copy(model.statistics.nodevisitedpersolution))
    push!(metrics.meanNodeVisitedUntilEnd,model.statistics.numberOfNodes)
    if ! isempty(model.statistics.nodevisitedpersolution)    #infeasible case
        push!(metrics.meanNodeVisitedUntilfirstSolFound,model.statistics.nodevisitedpersolution[1])
    end
    push!(metrics.timeneeded,dt)
    return
end

function (metrics::BasicMetrics{TakeObjective, <:BasicHeuristic})(model::CPModel,dt::Float64)
    metrics.nbEpisodes+=1
    push!(metrics.nodeVisited,copy(model.statistics.nodevisitedpersolution))
    push!(metrics.meanNodeVisitedUntilEnd,model.statistics.numberOfNodes)
    if ! isempty(model.statistics.nodevisitedpersolution)    #infeasible case
        push!(metrics.meanNodeVisitedUntilfirstSolFound,model.statistics.nodevisitedpersolution[1])
    end
    push!(metrics.timeneeded,dt)
    if ! isempty(model.statistics.objectives)
        push!(metrics.scores,copy(model.statistics.objectives))
    end

end

function (metrics::BasicMetrics{DontTakeObjective, <:LearnedHeuristic})(model::CPModel,dt::Float64)
    metrics.nbEpisodes+=1
    push!(metrics.nodeVisited,copy(model.statistics.nodevisitedpersolution))
    push!(metrics.meanNodeVisitedUntilEnd,model.statistics.numberOfNodes)
    if ! isempty(model.statistics.nodevisitedpersolution)    #infeasible case
        push!(metrics.meanNodeVisitedUntilfirstSolFound,model.statistics.nodevisitedpersolution[1])
    end
    push!(metrics.timeneeded,dt)
    push!(metrics.totalReward,last_episode_total_reward(metrics.heuristic.agent.trajectory))
    push!(metrics.loss,metrics.heuristic.agent.policy.learner.loss)
    return
end

function (metrics::BasicMetrics{TakeObjective, <:LearnedHeuristic})(model::CPModel,dt::Float64)
    metrics.nbEpisodes+=1
    push!(metrics.nodeVisited,copy(model.statistics.nodevisitedpersolution))
    push!(metrics.meanNodeVisitedUntilEnd,model.statistics.numberOfNodes)
    if ! isempty(model.statistics.nodevisitedpersolution)    #infeasible case
        push!(metrics.meanNodeVisitedUntilfirstSolFound,model.statistics.nodevisitedpersolution[1])
    end
    push!(metrics.timeneeded,dt)
    if ! isempty(model.statistics.objectives)
    push!(metrics.scores,copy(model.statistics.objectives))
    end
    push!(metrics.totalReward,last_episode_total_reward(metrics.heuristic.agent.trajectory))
    push!(metrics.loss,metrics.heuristic.agent.policy.learner.loss)
    return
end

"""
    function computemean!(metrics::BasicMetrics{O, H}) where{O, H<:ValueSelection}

At the end of an experiment, as the instances generated by the generator could be slightly different
(depending on the distribution of the generator), a moving average is applied on the vectors containing :
1) number of node visited by the heuristic to find a first solution along the training.
2) number of node visited by the heuristic to prove to optimality along the training.

this averaging smoothes the high frequency variations due to the differences between CPModel instances.
"""
function computemean!(metrics::BasicMetrics{O, H}) where {O, H<:ValueSelection}
    windowspan = min(metrics.meanOver,length(metrics.meanNodeVisitedUntilfirstSolFound))
    nodeVisitedUntilFirstSolFound = copy(metrics.meanNodeVisitedUntilfirstSolFound)
    metrics.meanNodeVisitedUntilfirstSolFound = rollmean(nodeVisitedUntilFirstSolFound, windowspan)
    nodeVisitedUntilOptimality = copy(metrics.meanNodeVisitedUntilEnd)
    metrics.meanNodeVisitedUntilEnd = rollmean(nodeVisitedUntilOptimality, windowspan)

    if isa(metrics,BasicMetrics{O,<:LearnedHeuristic})
        metrics.totalReward = rollmean(metrics.totalReward,windowspan)
    end
    return
end

ext: jl


text:
"""
    cudnnDropoutForward(x; dropout=0.5)
    cudnnDropoutForward(x, d::cudnnDropoutDescriptor)
    cudnnDropoutForward!(y, x; dropout=0.5)
    cudnnDropoutForward!(y, x, d::cudnnDropoutDescriptor)

Return a new array similar to `x` where approximately `dropout` fraction of the values are
replaced by a 0, and the rest are scaled by `1/(1-dropout)`.  Optionally `y` holds the
result and `d` specifies the operation. `y` should be similar to `x` if specified.

The user can set the global seed `cudnnDropoutSeed[]` to a positive number to always drop
the same values deterministically for debugging. Note that this slows down the operation by
about 40x.

The global constant `cudnnDropoutState::Dict` holds the random number generator state for
each CUDNN handle.
"""
cudnnDropoutForward, cudnnDropoutForward!, cudnnDropoutSeed, cudnnDropoutState

# Public methods
cudnnDropoutForward(x; o...)     = cudnnDropoutForwardWithDefaults(x; o...)
cudnnDropoutForward!(y, x; o...) = cudnnDropoutForwardWithDefaults(x; y, o...)
cudnnDropoutForward(x, d::cudnnDropoutDescriptor; o...)     = cudnnDropoutForwardWithDefaults(x; dropoutDesc=d, o...)
cudnnDropoutForward!(y, x, d::cudnnDropoutDescriptor; o...) = cudnnDropoutForwardWithDefaults(x; y, dropoutDesc=d, o...)

# Private method
function cudnnDropoutForwardWithDefaults(
    x;
    y = similar(x),
    dropout::Real = 0.5,
    dropoutDesc::cudnnDropoutDescriptor = cudnnDropoutDescriptor(Cfloat(dropout)),
    xDesc::cudnnTensorDescriptor = cudnnTensorDescriptor(x),
    yDesc::cudnnTensorDescriptor = xDesc,
    reserveSpace::CuArray = cudnnDropoutReserveSpace(xDesc)
)
    if cudnnDropoutSeed[] >= 0
        # This is a very expensive call (40x dropout), so only use for debugging
        @warn "CUDA.CUDNN.cudnnDropoutSeed[] >= 0: dropout operations will be deterministic but 40x more expensive" maxlog=1
        dropout, states, seed = cudnnGetDropoutDescriptor(dropoutDesc)
        hstate = cudnnDropoutState[handle()]
        @assert states == pointer(hstate)
        cudnnSetDropoutDescriptor(dropoutDesc, handle(), dropout, hstate, sizeof(hstate), cudnnDropoutSeed[])
    end
    cudnnDropoutForwardAD(x; xDesc, y, yDesc, dropoutDesc, reserveSpace)
end

function cudnnDropoutReserveSpace(td::cudnnTensorDescriptor)
    # reserveSpace is ~1/8 of tensor size and passes info between forw and back
    rss = Csize_t[0]; cudnnDropoutGetReserveSpaceSize(td, rss)
    return cudnnTempSpace(rss[1])
end

# AD method
function cudnnDropoutForwardAD(x; xDesc, y, yDesc, dropoutDesc, reserveSpace)
    cudnnDropoutForward(handle(), dropoutDesc, xDesc, x, yDesc, y, reserveSpace, sizeof(reserveSpace))
    return y
end

# Global RNG state: This should NOT be reallocated for each descriptor! However
# cudnnDropoutForward() doc says: "This function should not be running concurrently with
# another cudnnDropoutForward() function using the same states."  So I am going to assume
# using a single buffer per handle is ok.

const cudnnDropoutState = Dict{Ptr,CuArray}() # handle -> state

# Global dropout seed: To debug gradients set cudnnDropoutSeed[] >= 0 which makes all
# dropout operations deterministic but about 40x more expensive.

const cudnnDropoutSeed = Ref{Int}(-1)

# Helper for cudnnDropoutDescriptor constructor from float:
# Calls to cudnnDropoutDescriptor with identical Cfloats will return the same object thanks
# to caching. If the user wants to set the seed to replicate an experiment, that is taken
# care of during the forward call.

function cudnnSetDropoutDescriptorFromFloat(ptr::cudnnDropoutDescriptor_t, dropout::Real)
    hstate = get!(cudnnDropoutState, handle()) do
        cudnnTempSpace(cudnnDropoutGetStatesSize())
    end
    seed = floor(Culonglong,time())
    cudnnSetDropoutDescriptor(ptr, handle(), Cfloat(dropout), hstate, sizeof(hstate), seed)
end

function cudnnGetDropoutDescriptor(d::cudnnDropoutDescriptor)
    dropout, states, seed = Ref{Cfloat}(0), Ref{CuPtr{Nothing}}(0), Ref{Culonglong}(0)
    cudnnGetDropoutDescriptor(d::cudnnDropoutDescriptor, handle(), dropout, states, seed)
    return (dropout[], states[], seed[])
end

function cudnnDropoutGetStatesSize()
    ssize = Ref{Csize_t}(0)
    cudnnDropoutGetStatesSize(handle(), ssize)
    @assert ssize[] > 0
    return ssize[]
end


text:
# -------------------------------------------------------------------------
#   Saving and loading of (complete) observable
# -------------------------------------------------------------------------
"""
    saveobs(obs::Observable{T}[, filename::AbstractString, entryname::AbstractString])

Saves complete representation of the observable to JLD file.

Default filename is "Observables.jld" and default entryname is `name(obs)`.

See also [`loadobs`](@ref).
"""
function saveobs(obs::AbstractObservable, filename::AbstractString=obs.outfile,
                    entryname::AbstractString=(inmemory(obs) ? obs.group : obs.group*"/observable"))
    fileext(filename) == "jld" || error("\"$(filename)\" is not a valid JLD filename.")
    if !isfile(filename)
        save(filename, entryname, obs)
    else
        jldopen(filename, "r+") do f
            !HDF5.has(f.plain, entryname) || delete!(f, entryname)
            write(f, entryname, obs)
        end
    end
    nothing
end

"""
    loadobs(filename::AbstractString, entryname::AbstractString)

Load complete representation of an observable from JLD file.

See also [`saveobs`](@ref).
"""
function loadobs(filename::AbstractString, entryname::AbstractString)
    fileext(filename) == "jld" || error("\"$(filename)\" is not a valid JLD filename.")
    return load(filename, entryname)
end

struct ObservableResult{T,S}
    name::String
    count::Int64
    mean::T
    error::S
end

function _print_header(io::IO, r::ObservableResult{T, S}) where {T, S}
    print(io, "ObservableResult{$(T), $(S)}")
    nothing
end

function _println_body(io::IO, r::ObservableResult{T, S}) where {T, S}
    println(io)
    println(io, "| Name: ", r.name)
    print(io, "| Count: ", r.count)
    if r.count > 0 && ndims(r.error) == 0
        print(io, "\n| Mean: ", round.(r.mean, digits=5))
        print(io, "\n| StdError: ", round.(r.error, digits=5))
    end
    nothing
end

# short version (shows up in arrays etc.)
Base.show(io::IO, r::ObservableResult{T, S}) where {T, S} = print(io, "ObservableResult{$(T), $(S)}()")
# verbose version (shows up in the REPL)
Base.show(io::IO, m::MIME"text/plain", r::ObservableResult) = (_print_header(io, r); _println_body(io, r))

function load_result(filename::AbstractString, obs::AbstractString)
    @assert isfile(filename) "File not found."

    p = joinpath("/", obs)

    local name, n, m, err
    jldopen(filename) do f
        name = read(f[joinpath(p, "name")])
        n = read(f[joinpath(p, "count")])
        m = read(f[joinpath(p, "mean")])
        err = read(f[joinpath(p, "error")])
    end

    return ObservableResult(name, n, m, err)
end

# function combined_mean_and_var(ors::ObservableResult{<:Number, <:Number}...)
#     ns = [r.count for r in ors]
#     μs = [r.mean for r in ors]
#     vs = [r.error^2 for r in ors]

#     return Helpers.combined_mean_and_var(ns, μs, vs)
# end

# function combined_mean_and_var(ors::ObservableResult{<:Number, <:AbstractArray}...)
#     ns = [r.count for r in ors]
#     μs = [r.mean for r in ors]
#     vs = [r.error^2 for r in ors]

#     return Helpers.combined_mean_and_var(ns, μs, vs)
# end

# personal wrapper
function getfrom(filename::AbstractString, obs::AbstractString, what::AbstractString)
    if !(what in ["ts", "ts_flat", "timeseries", "timeseries_flat"])
        d = joinpath("obs/", obs, what)
        return jldopen(filename) do f
            return read(f[d])
        end
    else
        grp = joinpath("obs/", obs)
        return occursin("flat", what) ? ts_flat(filename, grp) : ts(filename, grp)
    end
end

"""
List all observables in a given file and HDF5 group.
"""
function listobs(filename::AbstractString, group::AbstractString="obs/")
    s = Vector{String}()
    HDF5.h5open(filename, "r") do f
        if HDF5.has(f, group)
            for el in HDF5.names(f[group])
                # println(el)
                push!(s, el)
            end
        end
    end
    return s
end

"""
Remove an observable.
"""
function rmobs(filename::AbstractString, dset::AbstractString, group::AbstractString="obs/")
    HDF5.h5open(filename, "r+") do f
        HDF5.o_delete(f, joinpath(group,dset))
    end
end


text:
using Base.Test
using OPCSPs
using NearestNeighbors
using Distances
import OPCSPs: NearestStateMap, SimilarStateStruct, has_similar, get_nearest

rng = MersenneTwister(1)
op = gen_problem()
m = NearestStateMap(op, Int)

s1 = OPCSPState(1, IntSet([4, 5, 6]), 5.0, randn(rng, length(op)))
s2 = OPCSPState(1, IntSet([4, 5, 6]), 5.0, randn(rng, length(op)))

insert!(m, s1, 1)
insert!(m, s2, 2)

@test has_similar(m, s1)

near, dist = get_nearest(m, s1)
@test near == 1
@test dist == 0.0

s3 = deepcopy(s2)
s3.d[4] += 0.01
near, dist = get_nearest(m, s3)
@test near == 2
@test_approx_eq dist 0.01


text:
function run_hmmsearch(input_filenames; db_filename=joinpath(Pkg.dir("apfamet"),"db","Pfam-A.hmm"), cores="1", sampleIDs=[], project_table_file = "apfamet_project_table.txt", overwrite=false)
	#Check to see whether SampleIDs are present, if not then just number them starting from 1
	if sampleIDs == []
		sampleIDs = map(string,range(1,length(input_filenames)))
	end

	#If there are SampleIDs, check that there's enough of them
	if length(input_filenames) != length(sampleIDs)
		error("Number of Sample IDs and filenames do not match")
	end

	#Check to see whether hmmpress has been run, if not then run it
	if isfile(db_filename*".h3i") == false
		run(`hmmpress $db_filename`)
	end

	#Check to see whether project table file already exists, if it does and overwrite is set to the default false then stop the function
	if isfile(project_table_file)
		if overwrite==false
			error(project_table_file*" already exists, pick another name, delete the existing file, or specify overwrite=true")
		end
	end

	#Go through the input files and run hmmsearch, taking care not to overwrite any existing output files
	for filename in input_filenames
		output_filename = filename*".hmmsearch_output"
		if isfile(output_filename)
			if overwrite==false
				error(filename*".hmmsearch_output already exists, pick another name, delete the existing hmmsearch output, or specify overwrite=true")
			end
		end
		run(pipeline(`hmmsearch --tblout $output_filename --cpu $cores $db_filename $filename`,DevNull))
	end

	#Write project table ready for the next step
	output_table = open(project_table_file, "w")
	write(output_table, "SampleID\tFilename\n")
	for (i, filename) in enumerate(input_filenames)
		write(output_table, sampleIDs[i]*"\t"*filename*".hmmsearch_output"*"\n")
	end

	close(output_table)

end

max_stars_repo_name: ianpgm/apfamet

max_issues_repo_name: ianpgm/apfamet

max_forks_repo_name: ianpgm/apfamet


text:
"""
    @dowhile begin
        body
    end cond condition

Repeatedly evaluate a block, and end when condition is false.

# Examples
```jldoctest
julia> i = 1
1

julia> @dowhile begin
           println(i)
           i += 1
       end cond i <= 4
1
2
3
4
```
"""
macro dowhile(body, cond::Symbol, condition)
    cond ≠ :cond &&
        error("@dowhile expected `cond` got `$cond`")
    quote
        while true
            $body
            if ! $condition
                break
            end
        end
    end |> esc
end


text:
# Use baremodule to shave off a few KB from the serialized `.ji` file
baremodule TinyXML_jll
using Base
using Base: UUID
import JLLWrappers

JLLWrappers.@generate_main_file_header("TinyXML")
JLLWrappers.@generate_main_file("TinyXML", UUID("7dc7d583-5213-5900-9003-a99bbb4d6a76"))
end  # module TinyXML_jll


text:
if !@isdefined USE_GPU
    const USE_GPU = length(ARGS) > 0 && ARGS[1] == "gpu"
end
using CUDA
using LinearAlgebra
using DataFrames, CSV
using BenchmarkTools, Statistics
using ParallelStencil
using ParallelStencil.FiniteDifferences2D
if !ParallelStencil.is_initialized()
    @static if USE_GPU
        @init_parallel_stencil(CUDA, Float64, 2)
    else
        @init_parallel_stencil(Threads, Float64, 2)
    end
end
include("../scripts-part2/multigrid.jl")

const FILENAME = joinpath("benchmark-results", "bench_multigrid_$(USE_GPU ? "gpu" : "cpu").csv")
if isfile(FILENAME)
    bench_df = DataFrame(CSV.File(FILENAME))
else
    bench_df = DataFrame()
end
#for execution_policy in [parallel, parallel_shmem], l in 2:3, k in 7:10
for k in 4:13, l in 2:min(k-4, 8), solver in [jacobi, conjugate_gradient], execution_policy in [parallel, parallel_shmem]
    @info "Running with params" USE_GPU, execution_policy, l, k, solver
    n = (2^k)+1
    h = 1 / (n - 1)
    c = 0.0
    tol = 1e-6
    Nmax = 100

    opt = MGOpt()
    opt.execution_policy = execution_policy
    opt.coarse_solve_size = (2^l)+1
    opt.coarse_solver = solver

    for seed in 1:1
        try
            bench_res = @benchmark MGsolve_2DPoisson!(x, b, $h, $c, $tol, $Nmax, false; opt=$opt, verbose=false) setup=((x, b) = (@zeros($n, $n), @rand($n, $n))) evals=1
            bench_row = (
                 execution_policy = string(execution_policy),
                 coarse_solver = string(solver),
                 k = k,  # exponent for computing size of largest grid with nx = (2^k)+1
                 l = l,  # same as k but for the coarsest grid
                 median_time = median(bench_res).time / 1e9,
                 mean_time = mean(bench_res).time / 1e9,
                 std_time = std(bench_res).time / 1e9,
                 seed = seed,
                 use_gpu = USE_GPU,
                 nthreads = Threads.nthreads(),
                )
            push!(bench_df, bench_row)

            open(FILENAME, "w") do io
                CSV.write(io, bench_df)
            end
        catch e
            @warn "Failed with " e
        end
    end
end


text:
#!/usr/bin/env julia

using Pkg
Pkg.add("PackageCompiler")
using PackageCompiler
Pkg.activate(joinpath(@__DIR__, ".."))
using PackageCompiler
PackageCompiler.create_sysimage(
  :JuliaFormatter;
  precompile_execution_file = joinpath(@__DIR__, "precompile.jl"),
  sysimage_path = joinpath(
    @__DIR__,
    "..",
    "..",
    ".git",
    "hooks",
    "JuliaFormatterSysimage.so",
  ),
)


text:
using QuasiNewtonMethods, StrideArrays, Aqua
using Test

struct Rosenbrock end
function QuasiNewtonMethods.logdensity(::Rosenbrock, θ)
    s = zero(eltype(θ))
    N = length(θ) >> 1
    @inbounds @simd for i ∈ 1:N
        s -= 100(θ[i+N] - θ[i]^2)^2 + (1 - θ[i])^2
    end
    if isodd(length(θ))
        @inbounds δ = 1 - θ[end]
        -muladd(δ, δ, -s)
    else
        s
    end
end

function QuasiNewtonMethods.∂logdensity!(∇, ::Rosenbrock, θ)
    s = zero(eltype(θ))
    N = length(θ) >> 1
    @inbounds @simd for i ∈ 1:N
        s -= 100(θ[i+N] - θ[i]^2)^2 + (1 - θ[i])^2
        ∇[i] = 400(θ[i+N] - θ[i]^2)*θ[i] + 2(1 - θ[i])
        ∇[i+N] = 200(θ[i]^2 - θ[i+N])
    end
    if isodd(length(θ))
        @inbounds δ = 1 - θ[end]
        s = -muladd(δ, δ, -s)
        @inbounds ∇[end] = 2δ
    end
    s
end

@testset "QuasiNewtonMethods.jl" begin
    Aqua.test_all(QuasiNewtonMethods)
    for n ∈ 2:24
        @show n
        state = QuasiNewtonMethods.BFGSState{n}(undef);
        x = @StrideArray randn(StaticInt(n));
        # 2nd order
        @test abs(optimize!(state, Rosenbrock(), x)) < 2eps()
        @show QuasiNewtonMethods.optimum(state) .- 1;
        @test all(x -> x ≈ 1, QuasiNewtonMethods.optimum(state))
        @test maximum(abs, QuasiNewtonMethods.gradient(state)) < 1e-8
        # 3rd order
        @test abs(optimize!(state, Rosenbrock(), x, QuasiNewtonMethods.BackTracking{3}())) < eps()
        @show QuasiNewtonMethods.optimum(state) .- 1;
        @test all(x -> x ≈ 1, QuasiNewtonMethods.optimum(state))
        @test maximum(abs, QuasiNewtonMethods.gradient(state)) < 1e-8
    end
end


text:
include("adp/adp.jl")
include("irl/irl.jl")


text:
using Photon
using Photon.Layers
using Photon.Losses: CrossEntropyLoss

# Define a model with fully connected layers.
model = Sequential(
      Dense(256, :relu),
      Dense( 10),
      Softmax(),
)

# create the workout
workout = Workout(model, CrossEntropyLoss(); mover= SmartMover(Array{Float64}))

# create some dummy training data
X = [randn(25, 16) for i in 1:10]
Y = [rand(10, 16) for i in 1:10]

# perform the training
train!(workout, zip(X,Y))

# We can call the trained model directly to predict a minibatch of 16 samples
x = randn(25, 16)
y = model(x)

# Or use the predict function for a single sample.
predict(workout, randn(25))

println("### Done ###")


text:
"""
    list(config::Config; show_links=false)

Prints out all installed julia version and playgrounds.
"""
function list(config::Config; show_links=false)
	init(config)

    println("Julia Versions:")
    for j in readdir(config.bin)
        if j != p"playground"
            file_path = join(config.bin, j)
            if islink(file_path) && show_links
                println("\t$(j) -> $(readlink(file_path))")
            else
                println("\t$(j)")
            end
        end
    end

    println("\nPlaygrounds:")
    for p in readdir(config.share)
        file_path = join(config.share, p)
        if islink(file_path) && show_links
            println("\t$(string(p)) -> $(string(readlink(file_path)))")
        else
            println("\t$(string(p))")
        end
    end
end


text:
using HPAT

#HPAT.DomainPass.set_debug_level(3)

@acc hpat function q25(d_date, num_centroids, iterations, file_name)
    store_sales = DataSource(DataTable{:ss_customer_sk=Int64, :ss_ticket_number=Int64, :ss_sold_date_sk=Int64, :ss_net_paid=Float64}, HDF5, file_name)
    web_sales = DataSource(DataTable{:ws_bill_customer_sk=Int64, :ws_order_number=Int64, :ws_sold_date_sk=Int64, :ws_net_paid=Float64}, HDF5, file_name)

    store_sales = store_sales[:ss_sold_date_sk > d_date]
    web_sales   = web_sales[:ws_sold_date_sk > d_date]

    store_agg = aggregate(store_sales, :cid = :ss_customer_sk,
                                           :frequency = length(unique(:ss_ticket_number)),
                                           :most_recent_date = maximum(:ss_sold_date_sk),
                                           :amount = sum(:ss_net_paid))
    web_agg = aggregate(web_sales,  :cid = :ws_bill_customer_sk,
                                           :frequency = length(unique(:ws_order_number)),
                                           :most_recent_date = maximum(:ws_sold_date_sk),
                                           :amount = sum(:ws_net_paid))
    web_store_agg = [store_agg; web_agg]

    result = aggregate(web_store_agg, :cid,
                                      :recency = maximum(:most_recent_date),
                                      :frequency = sum(:frequency),
                                      :totalspend = sum(:amount))

    #return result[:cid], result[:recency], result[:frequency], result[:totalspend]
    recency = 37621 - result[:recency] .< 60;
    points = transpose(typed_hcat(Float64, result[:cid], recency, result[:frequency], result[:totalspend]))
    model = Kmeans(points, num_centroids, iterations)
    return model
end

println(q25(33000, 8, 20,"test_q25.hdf5"))


text:
# Autogenerated wrapper script for pandoc_crossref_jll for x86_64-linux-gnu
export pandoc_crossref

JLLWrappers.@generate_wrapper_header("pandoc_crossref")
JLLWrappers.@declare_executable_product(pandoc_crossref)
function __init__()
    JLLWrappers.@generate_init_header()
    JLLWrappers.@init_executable_product(
        pandoc_crossref,
        "bin/pandoc-crossref",
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
using Rubin
using Tests
using Elliptic
using HypergeometricFunctions
using Polylogarithms
using SpecialFunctions

@test integrate(x*(1+-1*x^2)^-1//2*arcsin(x), x) == :(x+-1*(1+-1*x^2)^(1/2)*arcsin(x))
@test integrate(arcsin((1+x)^(1/2)+-1*x^(1/2)), x) == :(-1*(3//8+x)*arcsin(x^(1/2)+-1*(1+x)^(1/2))+1//8*2^(1/2)*(-1x+x^(1/2)*(1+x)^(1/2))^(1/2)*(x^(1/2)+3*(1+x)^(1/2)))
@test integrate(log(1+x*(1+x^2)^(1/2)), x) == :(-2x+x*log(1+x*(1+x^2)^(1/2))+(2+2*5^(1/2))^(1/2)*arctan((-2+5^(1/2))^(1/2)*(x+(1+x^2)^(1/2)))+-1*(-2+2*5^(1/2))^(1/2)*arctanh((2+5^(1/2))^(1/2)*(x+(1+x^2)^(1/2))))
@test integrate((1+cos(x)^2+cos(x)^4)^-1//2*cos(x)^2, x) == :(1//3*x+1//3*arctan((1+(1+cos(x)^2+cos(x)^4)^(1/2)*cos(x)^2)^-1*(1+cos(x)^2)*cos(x)*sin(x)))
@test integrate((1+tan(x)^4)^(1/2)*tan(x), x) == :((1/2)*(1+tan(x)^4)^(1/2)+-1//2*arcsinh(tan(x)^2)+-1//2*2^(1/2)*arctanh((1/2)*2^(1/2)*(1+tan(x)^4)^-1//2*(1+-1*tan(x)^2)))
@test integrate((1+sec(x)^3)^-1//2*tan(x), x) == :(-2//3*arctanh((1+sec(x)^3)^(1/2)))
@test integrate((2+tan(x)^2+2*tan(x))^(1/2), x) == :(-1*(1/2+(1/2)*5^(1/2))^(1/2)*arctan((10+10*5^(1/2))^-1//2*(2+tan(x)^2+2*tan(x))^-1//2*(2*5^(1/2)+-1*(5+5^(1/2))*tan(x)))+-1*(-1//2+(1/2)*5^(1/2))^(1/2)*arctanh((-10+10*5^(1/2))^-1//2*(2+tan(x)^2+2*tan(x))^-1//2*(2*5^(1/2)+(5+-1*5^(1/2))*tan(x)))+arcsinh(1+tan(x)))
@test integrate(arctan((-1+sec(x))^(1/2))*sin(x), x) == :((1/2)*arctan((-1+sec(x))^(1/2))+(1/2)*(-1+sec(x))^(1/2)*cos(x)+-1*arctan((-1+sec(x))^(1/2))*cos(x))
@test integrate(x*(1+x^2)^-1//2*log(1+x^2)*log(x+(1+x^2)^(1/2)), x) == :(-2*arctan(x)+4x+-1*x*log(1+x^2)+-2*(1+x^2)^(1/2)*log(x+(1+x^2)^(1/2))+(1+x^2)^(1/2)*log(1+x^2)*log(x+(1+x^2)^(1/2)))
@test integrate(arctan(x+(1+-1*x^2)^(1/2)), x) == :(-1//2*arcsin(x)+-1//4*arctanh(x*(1+-1*x^2)^(1/2))+-1//8*log(1+x^4+-1*x^2)+x*arctan(x+(1+-1*x^2)^(1/2))+-1//4*3^(1/2)*arctan(1//3*3^(1/2)*(-1+2*x^2))+1//4*3^(1/2)*arctan((1+-1*x^2)^-1//2*(1+x*3^(1/2)))+1//4*3^(1/2)*arctan((1+-1*x^2)^-1//2*(-1+x*3^(1/2))))
@test integrate(x*(1+-1*x^2)^-1//2*arctan(x+(1+-1*x^2)^(1/2)), x) == :(-1//2*arcsin(x)+1//4*arctanh(x*(1+-1*x^2)^(1/2))+1//8*log(1+x^4+-1*x^2)+-1*(1+-1*x^2)^(1/2)*arctan(x+(1+-1*x^2)^(1/2))+-1//4*3^(1/2)*arctan(1//3*3^(1/2)*(-1+2*x^2))+1//4*3^(1/2)*arctan((1+-1*x^2)^-1//2*(1+x*3^(1/2)))+1//4*3^(1/2)*arctan((1+-1*x^2)^-1//2*(-1+x*3^(1/2))))
@test integrate((1+-1*x^2)^-3//2*log(x+(1+x^2)^(1/2)), x) == :(-1//2*arcsin(x^2)+x*(1+-1*x^2)^-1//2*log(x+(1+x^2)^(1/2)))
@test integrate((1+x^2)^-3//2*arcsin(x), x) == :(-1//2*arcsin(x^2)+x*(1+x^2)^-1//2*arcsin(x))
@test integrate((1+x^2)^-3//2*log(x+(-1+x^2)^(1/2)), x) == :(-1//2*arccosh(x^2)+x*(1+x^2)^-1//2*log(x+(-1+x^2)^(1/2)))
@test integrate(x^-2*(-1+x^2)^-1//2*log(x), x) == :(-1*arctanh(x*(-1+x^2)^-1//2)+x^-1*(-1+x^2)^(1/2)+x^-1*(-1+x^2)^(1/2)*log(x))
@test integrate(x^-1*(1+x^3)^(1/2), x) == :(-2//3*arctanh((1+x^3)^(1/2))+2//3*(1+x^3)^(1/2))
@test integrate(x*(-1+x^2)^-1//2*log(x+(-1+x^2)^(1/2)), x) == :(-1x+(-1+x^2)^(1/2)*log(x+(-1+x^2)^(1/2)))
@test integrate(x^3*(1+-1*x^4)^-1//2*arcsin(x), x) == :(1//4*arcsinh(x)+-1//2*(1+-1*x^4)^(1/2)*arcsin(x)+1//4*x*(1+x^2)^(1/2))
@test integrate(x*(1+x^2)^-1//2*arctan(x)*log(x+(1+x^2)^(1/2)), x) == :((1/2)*log(1+x^2)+-1//2*log(x+(1+x^2)^(1/2))^2+-1*x*arctan(x)+(1+x^2)^(1/2)*arctan(x)*log(x+(1+x^2)^(1/2)))
@test integrate(x*(1+-1*x^2)^-1//2*log(1+(1+-1*x^2)^(1/2)), x) == :((1+-1*x^2)^(1/2)+-1*log(1+(1+-1*x^2)^(1/2))+-1*(1+-1*x^2)^(1/2)*log(1+(1+-1*x^2)^(1/2)))
@test integrate(x*(1+x^2)^-1//2*log(x+(1+x^2)^(1/2)), x) == :(-1x+(1+x^2)^(1/2)*log(x+(1+x^2)^(1/2)))
@test integrate(x*(1+-1*x^2)^-1//2*log(x+(1+-1*x^2)^(1/2)), x) == :((1+-1*x^2)^(1/2)+(1/2)*2^(1/2)*arctanh(x*2^(1/2))+-1*(1+-1*x^2)^(1/2)*log(x+(1+-1*x^2)^(1/2))+-1//2*2^(1/2)*arctanh(2^(1/2)*(1+-1*x^2)^(1/2)))
@test integrate(x^-2*(1+-1*x^2)^-1//2*log(x), x) == :(-1*arcsin(x)+-1*x^-1*(1+-1*x^2)^(1/2)+-1*x^-1*(1+-1*x^2)^(1/2)*log(x))
@test integrate(x*(1+x^2)^-1//2*arctan(x), x) == :(-1*arcsinh(x)+(1+x^2)^(1/2)*arctan(x))
@test integrate(x^-2*(1+-1*x^2)^-1//2*arctan(x), x) == :(-1*arctanh((1+-1*x^2)^(1/2))+2^(1/2)*arctanh((1/2)*2^(1/2)*(1+-1*x^2)^(1/2))+-1*x^-1*(1+-1*x^2)^(1/2)*arctan(x))
@test integrate(x*(1+-1*x^2)^-1//2*arctan(x), x) == :(-1*arcsin(x)+2^(1/2)*arctan(x*2^(1/2)*(1+-1*x^2)^-1//2)+-1*(1+-1*x^2)^(1/2)*arctan(x))
@test integrate(x^-2*(1+x^2)^-1//2*arctan(x), x) == :(-1*arctanh((1+x^2)^(1/2))+-1*x^-1*(1+x^2)^(1/2)*arctan(x))
@test integrate(x^-2*(1+-1*x^2)^-1//2*arcsin(x), x) == :(-1*x^-1*(1+-1*x^2)^(1/2)*arcsin(x)+log(x))
@test integrate(x*(-1+x^2)^-1//2*log(x), x) == :(-1*(-1+x^2)^(1/2)+(-1+x^2)^(1/2)*log(x)+arctan((-1+x^2)^(1/2)))
@test integrate(x^-2*(1+x^2)^-1//2*log(x), x) == :(-1*x^-1*(1+x^2)^(1/2)+-1*x^-1*(1+x^2)^(1/2)*log(x)+arcsinh(x))
@test integrate(x*(-1+x^2)^-1//2*arcsec(x), x) == :((-1+x^2)^(1/2)*arcsec(x)+-1*x*(x^2)^-1//2*log(x))
@test integrate(x*(1+x^2)^-1//2*log(x), x) == :(-1*(1+x^2)^(1/2)+(1+x^2)^(1/2)*log(x)+arctanh((1+x^2)^(1/2)))
@test integrate((1+sin(x)^2)^-1*sin(x), x) == :(-1//2*2^(1/2)*arctanh((1/2)*2^(1/2)*cos(x)))
@test integrate((1+x^4)^-1//2*(1+-1*x^2)^-1*(1+x^2), x) == :((1/2)*2^(1/2)*arctanh(x*2^(1/2)*(1+x^4)^-1//2))
@test integrate((1+x^2)^-1*(1+x^4)^-1//2*(1+-1*x^2), x) == :((1/2)*2^(1/2)*arctan(x*2^(1/2)*(1+x^4)^-1//2))
@test integrate((1+sin(x))^-1*log(sin(x)), x) == :(-1x+-1*arctanh(cos(x))+-1*(1+sin(x))^-1*cos(x)*log(sin(x)))
@test integrate((1+sin(x))^(1/2)*log(sin(x)), x) == :(-4*arctanh((1+sin(x))^-1//2*cos(x))+4*(1+sin(x))^-1//2*cos(x)+-2*(1+sin(x))^-1//2*cos(x)*log(sin(x)))
@test integrate((-1+sec(x)^4)^-1//2*sec(x), x) == :(-1//2*2^(1/2)*arctanh((1/2)*2^(1/2)*(-1+sec(x)^4)^(1/2)*cos(x)*cot(x)))
@test integrate((1+tan(x)^4)^-1//2*tan(x), x) == :(-1//4*2^(1/2)*arctanh((1/2)*2^(1/2)*(1+tan(x)^4)^-1//2*(1+-1*tan(x)^2)))
@test integrate(((1+sec(x))^(1/2)+-1*(-1+sec(x))^(1/2))^(1/2), x) == :(2^(1/2)*(1+sec(x))^(1/2)*(-1+sec(x))^(1/2)*((-1+2^(1/2))^(1/2)*arctan((1/2)*(-2+2*2^(1/2))^(1/2)*((1+sec(x))^(1/2)+-1*(-1+sec(x))^(1/2))^-1//2*((1+sec(x))^(1/2)+-1*2^(1/2)+-1*(-1+sec(x))^(1/2)))+(-1+2^(1/2))^(1/2)*arctanh((2+2*2^(1/2))^(1/2)*((1+sec(x))^(1/2)+-1*(-1+sec(x))^(1/2))^(1/2)*(2^(1/2)+(1+sec(x))^(1/2)+-1*(-1+sec(x))^(1/2))^-1)+-1*(1+2^(1/2))^(1/2)*arctan((1/2)*(2+2*2^(1/2))^(1/2)*((1+sec(x))^(1/2)+-1*(-1+sec(x))^(1/2))^-1//2*((1+sec(x))^(1/2)+-1*2^(1/2)+-1*(-1+sec(x))^(1/2)))+-1*(1+2^(1/2))^(1/2)*arctanh((-2+2*2^(1/2))^(1/2)*((1+sec(x))^(1/2)+-1*(-1+sec(x))^(1/2))^(1/2)*(2^(1/2)+(1+sec(x))^(1/2)+-1*(-1+sec(x))^(1/2))^-1))*cot(x))
@test integrate(x*arctan(x)^2*log(1+x^2), x) == :(-3//2*arctan(x)^2+-3//2*log(1+x^2)+1//4*log(1+x^2)^2+3*x*arctan(x)+-1//2*x^2*arctan(x)^2+arctan(x)^2*(1/2+(1/2)*x^2)*log(1+x^2)+-1*x*arctan(x)*log(1+x^2))
@test integrate(arctan(x*(1+x^2)^(1/2)), x) == :((1/2)*arctan(3^(1/2)+-2*(1+x^2)^(1/2))+-1//2*arctan(3^(1/2)+2*(1+x^2)^(1/2))+x*arctan(x*(1+x^2)^(1/2))+-1//4*3^(1/2)*log(2+x^2+-1*3^(1/2)*(1+x^2)^(1/2))+1//4*3^(1/2)*log(2+x^2+3^(1/2)*(1+x^2)^(1/2)))
@test integrate(arcsin(x*(1+-1*x^2)^-1//2), x) == :(x*arcsin(x*(1+-1*x^2)^-1//2)+arctan((1+-2*x^2)^(1/2)))

max_stars_repo_path: test/0 Independent test suites/Charlwood Problems.jl


text:
# Fallback traits
# TODO: Should this instead be `NoChildren()`, `HasChild()`, etc. so we allow plural too, e.g. `HasChildren()`?

"""
    NodeTrait(context)
    NodeTrait(f, context)

Specifies the role of `context` in the context-tree.

The officially supported traits are:
- `IsLeaf`: `context` does not have any decendants.
- `IsParent`: `context` has a child context to which we often defer.
  Expects the following methods to be implemented:
  - [`childcontext`](@ref)
  - [`setchildcontext`](@ref)
"""
abstract type NodeTrait end
NodeTrait(_, context) = NodeTrait(context)

"""
    IsLeaf

Specifies that the context is a leaf in the context-tree.
"""
struct IsLeaf <: NodeTrait end
"""
    IsParent

Specifies that the context is a parent in the context-tree.
"""
struct IsParent <: NodeTrait end

"""
    childcontext(context)

Return the descendant context of `context`.
"""
childcontext

"""
    setchildcontext(parent::AbstractContext, child::AbstractContext)

Reconstruct `parent` but now using `child` is its [`childcontext`](@ref),
effectively updating the child context.

# Examples
```jldoctest
julia> ctx = SamplingContext();

julia> DynamicPPL.childcontext(ctx)
DefaultContext()

julia> ctx_prior = DynamicPPL.setchildcontext(ctx, PriorContext()); # only compute the logprior

julia> DynamicPPL.childcontext(ctx_prior)
PriorContext{Nothing}(nothing)
```
"""
setchildcontext

"""
    leafcontext(context)

Return the leaf of `context`, i.e. the first descendant context that `IsLeaf`.
"""
leafcontext(context) = leafcontext(NodeTrait(leafcontext, context), context)
leafcontext(::IsLeaf, context) = context
leafcontext(::IsParent, context) = leafcontext(childcontext(context))

"""
    setleafcontext(left, right)

Return `left` but now with its leaf context replaced by `right`.

Note that this also works even if `right` is not a leaf context,
in which case effectively append `right` to `left`, dropping the
original leaf context of `left`.

# Examples
```jldoctest
julia> using DynamicPPL: leafcontext, setleafcontext, childcontext, setchildcontext, AbstractContext

julia> struct ParentContext{C} <: AbstractContext
           context::C
       end

julia> DynamicPPL.NodeTrait(::ParentContext) = DynamicPPL.IsParent()

julia> DynamicPPL.childcontext(context::ParentContext) = context.context

julia> DynamicPPL.setchildcontext(::ParentContext, child) = ParentContext(child)

julia> Base.show(io::IO, c::ParentContext) = print(io, "ParentContext(", childcontext(c), ")")

julia> ctx = ParentContext(ParentContext(DefaultContext()))
ParentContext(ParentContext(DefaultContext()))

julia> # Replace the leaf context with another leaf.
       leafcontext(setleafcontext(ctx, PriorContext()))
PriorContext{Nothing}(nothing)

julia> # Append another parent context.
       setleafcontext(ctx, ParentContext(DefaultContext()))
ParentContext(ParentContext(ParentContext(DefaultContext())))
```
"""
function setleafcontext(left, right)
    return setleafcontext(
        NodeTrait(setleafcontext, left), NodeTrait(setleafcontext, right), left, right
    )
end
function setleafcontext(::IsParent, ::IsParent, left, right)
    return setchildcontext(left, setleafcontext(childcontext(left), right))
end
function setleafcontext(::IsParent, ::IsLeaf, left, right)
    return setchildcontext(left, setleafcontext(childcontext(left), right))
end
setleafcontext(::IsLeaf, ::IsParent, left, right) = right
setleafcontext(::IsLeaf, ::IsLeaf, left, right) = right

# Contexts
"""
    SamplingContext(rng, sampler, context)

Create a context that allows you to sample parameters with the `sampler` when running the model.
The `context` determines how the returned log density is computed when running the model.

See also: [`DefaultContext`](@ref), [`LikelihoodContext`](@ref), [`PriorContext`](@ref)
"""
struct SamplingContext{S<:AbstractSampler,C<:AbstractContext,R} <: AbstractContext
    rng::R
    sampler::S
    context::C
end
SamplingContext(sampler, context) = SamplingContext(Random.GLOBAL_RNG, sampler, context)
SamplingContext(context::AbstractContext) = SamplingContext(SampleFromPrior(), context)
SamplingContext(sampler::AbstractSampler) = SamplingContext(sampler, DefaultContext())
SamplingContext() = SamplingContext(SampleFromPrior())

NodeTrait(context::SamplingContext) = IsParent()
childcontext(context::SamplingContext) = context.context
function setchildcontext(parent::SamplingContext, child)
    return SamplingContext(parent.rng, parent.sampler, child)
end

"""
    struct DefaultContext <: AbstractContext end

The `DefaultContext` is used by default to compute log the joint probability of the data
and parameters when running the model.
"""
struct DefaultContext <: AbstractContext end
NodeTrait(context::DefaultContext) = IsLeaf()

"""
    struct PriorContext{Tvars} <: AbstractContext
        vars::Tvars
    end

The `PriorContext` enables the computation of the log prior of the parameters `vars` when
running the model.
"""
struct PriorContext{Tvars} <: AbstractContext
    vars::Tvars
end
PriorContext() = PriorContext(nothing)
NodeTrait(context::PriorContext) = IsLeaf()

"""
    struct LikelihoodContext{Tvars} <: AbstractContext
        vars::Tvars
    end

The `LikelihoodContext` enables the computation of the log likelihood of the parameters when
running the model. `vars` can be used to evaluate the log likelihood for specific values
of the model's parameters. If `vars` is `nothing`, the parameter values inside the `VarInfo` will be used by default.
"""
struct LikelihoodContext{Tvars} <: AbstractContext
    vars::Tvars
end
LikelihoodContext() = LikelihoodContext(nothing)
NodeTrait(context::LikelihoodContext) = IsLeaf()

"""
    struct MiniBatchContext{Tctx, T} <: AbstractContext
        context::Tctx
        loglike_scalar::T
    end

The `MiniBatchContext` enables the computation of
`log(prior) + s * log(likelihood of a batch)` when running the model, where `s` is the
`loglike_scalar` field, typically equal to `the number of data points / batch size`.
This is useful in batch-based stochastic gradient descent algorithms to be optimizing
`log(prior) + log(likelihood of all the data points)` in the expectation.
"""
struct MiniBatchContext{Tctx,T} <: AbstractContext
    context::Tctx
    loglike_scalar::T
end
function MiniBatchContext(context=DefaultContext(); batch_size, npoints)
    return MiniBatchContext(context, npoints / batch_size)
end
NodeTrait(context::MiniBatchContext) = IsParent()
childcontext(context::MiniBatchContext) = context.context
function setchildcontext(parent::MiniBatchContext, child)
    return MiniBatchContext(child, parent.loglike_scalar)
end

"""
    PrefixContext{Prefix}(context)

Create a context that allows you to use the wrapped `context` when running the model and
adds the `Prefix` to all parameters.

This context is useful in nested models to ensure that the names of the parameters are
unique.

See also: [`@submodel`](@ref)
"""
struct PrefixContext{Prefix,C} <: AbstractContext
    context::C
end
function PrefixContext{Prefix}(context::AbstractContext) where {Prefix}
    return PrefixContext{Prefix,typeof(context)}(context)
end

NodeTrait(context::PrefixContext) = IsParent()
childcontext(context::PrefixContext) = context.context
function setchildcontext(parent::PrefixContext{Prefix}, child) where {Prefix}
    return PrefixContext{Prefix}(child)
end

const PREFIX_SEPARATOR = Symbol(".")

function PrefixContext{PrefixInner}(
    context::PrefixContext{PrefixOuter}
) where {PrefixInner,PrefixOuter}
    if @generated
        :(PrefixContext{$(QuoteNode(Symbol(PrefixOuter, PREFIX_SEPARATOR, PrefixInner)))}(
            context.context
        ))
    else
        PrefixContext{Symbol(PrefixOuter, PREFIX_SEPARATOR, PrefixInner)}(context.context)
    end
end

function prefix(::PrefixContext{Prefix}, vn::VarName{Sym}) where {Prefix,Sym}
    if @generated
        return :(VarName{$(QuoteNode(Symbol(Prefix, PREFIX_SEPARATOR, Sym)))}(getlens(vn)))
    else
        VarName{Symbol(Prefix, PREFIX_SEPARATOR, Sym)}(getlens(vn))
    end
end

struct ConditionContext{Names,Values,Ctx<:AbstractContext} <: AbstractContext
    values::Values
    context::Ctx

    function ConditionContext{Values}(
        values::Values, context::AbstractContext
    ) where {names,Values<:NamedTuple{names}}
        return new{names,typeof(values),typeof(context)}(values, context)
    end
end

function ConditionContext(values::NamedTuple)
    return ConditionContext(values, DefaultContext())
end
function ConditionContext(values::NamedTuple, context::AbstractContext)
    return ConditionContext{typeof(values)}(values, context)
end

# Try to avoid nested `ConditionContext`.
function ConditionContext(
    values::NamedTuple{Names}, context::ConditionContext
) where {Names}
    # Note that this potentially overrides values from `context`, thus giving
    # precedence to the outmost `ConditionContext`.
    return ConditionContext(merge(context.values, values), childcontext(context))
end

function Base.show(io::IO, context::ConditionContext)
    return print(io, "ConditionContext($(context.values), $(childcontext(context)))")
end

NodeTrait(context::ConditionContext) = IsParent()
childcontext(context::ConditionContext) = context.context
setchildcontext(parent::ConditionContext, child) = ConditionContext(parent.values, child)

"""
    hasvalue(context, vn)

Return `true` if `vn` is found in `context`.
"""
hasvalue(context, vn) = false

function hasvalue(context::ConditionContext{vars}, vn::VarName{sym}) where {vars,sym}
    return sym in vars
end
function hasvalue(
    context::ConditionContext{vars}, vn::AbstractArray{<:VarName{sym}}
) where {vars,sym}
    return sym in vars
end

"""
    getvalue(context, vn)

Return value of `vn` in `context`.
"""
function getvalue(context::AbstractContext, vn)
    return error("context $(context) does not contain value for $vn")
end
getvalue(context::ConditionContext, vn) = get(context.values, vn)

"""
    hasvalue_nested(context, vn)

Return `true` if `vn` is found in `context` or any of its descendants.

This is contrast to [`hasvalue`](@ref) which only checks for `vn` in `context`,
not recursively checking if `vn` is in any of its descendants.
"""
function hasvalue_nested(context::AbstractContext, vn)
    return hasvalue_nested(NodeTrait(hasvalue_nested, context), context, vn)
end
hasvalue_nested(::IsLeaf, context, vn) = hasvalue(context, vn)
function hasvalue_nested(::IsParent, context, vn)
    return hasvalue(context, vn) || hasvalue_nested(childcontext(context), vn)
end
function hasvalue_nested(context::PrefixContext, vn)
    return hasvalue_nested(childcontext(context), prefix(context, vn))
end

"""
    getvalue_nested(context, vn)

Return the value of the parameter corresponding to `vn` from `context` or its descendants.

This is contrast to [`getvalue`](@ref) which only returns the value `vn` in `context`,
not recursively looking into its descendants.
"""
function getvalue_nested(context::AbstractContext, vn)
    return getvalue_nested(NodeTrait(getvalue_nested, context), context, vn)
end
function getvalue_nested(::IsLeaf, context, vn)
    return error("context $(context) does not contain value for $vn")
end
function getvalue_nested(context::PrefixContext, vn)
    return getvalue_nested(childcontext(context), prefix(context, vn))
end
function getvalue_nested(::IsParent, context, vn)
    return if hasvalue(context, vn)
        getvalue(context, vn)
    else
        getvalue_nested(childcontext(context), vn)
    end
end

"""
    condition([context::AbstractContext,] values::NamedTuple)
    condition([context::AbstractContext]; values...)

Return `ConditionContext` with `values` and `context` if `values` is non-empty,
otherwise return `context` which is [`DefaultContext`](@ref) by default.

See also: [`decondition`](@ref)
"""
condition(; values...) = condition(DefaultContext(), NamedTuple(values))
condition(values::NamedTuple) = condition(DefaultContext(), values)
condition(context::AbstractContext, values::NamedTuple{()}) = context
condition(context::AbstractContext, values::NamedTuple) = ConditionContext(values, context)
condition(context::AbstractContext; values...) = condition(context, NamedTuple(values))

"""
    decondition(context::AbstractContext, syms...)

Return `context` but with `syms` no longer conditioned on.

Note that this recursively traverses contexts, deconditioning all along the way.

See also: [`condition`](@ref)
"""
decondition(::IsLeaf, context, args...) = context
function decondition(::IsParent, context, args...)
    return setchildcontext(context, decondition(childcontext(context), args...))
end
decondition(context, args...) = decondition(NodeTrait(context), context, args...)
function decondition(context::ConditionContext)
    return decondition(childcontext(context))
end
function decondition(context::ConditionContext, sym)
    return condition(
        decondition(childcontext(context), sym), BangBang.delete!!(context.values, sym)
    )
end
function decondition(context::ConditionContext, sym, syms...)
    return decondition(
        condition(
            decondition(childcontext(context), syms...),
            BangBang.delete!!(context.values, sym),
        ),
        syms...,
    )
end

"""
    conditioned(context::AbstractContext)

Return `NamedTuple` of values that are conditioned on under context`.

Note that this will recursively traverse the context stack and return
a merged version of the condition values.
"""
function conditioned(context::AbstractContext)
    return conditioned(NodeTrait(conditioned, context), context)
end
conditioned(::IsLeaf, context) = ()
conditioned(::IsParent, context) = conditioned(childcontext(context))
function conditioned(context::ConditionContext)
    # Note the order of arguments to `merge`. The behavior of the rest of DPPL
    # is that the outermost `context` takes precendence, hence when resolving
    # the `conditioned` variables we need to ensure that `context.values` takes
    # precedence over decendants of `context`.
    return merge(context.values, conditioned(childcontext(context)))
end


text:
SUITE["Integrators"] = BenchmarkGroup()

SI = SUITE["Integrators"]

using DynamicalSystemsBase: CDS, DDS
using DynamicalSystemsBase.Systems: eom_towel, jacob_towel, eom_towel_iip, jacob_towel_iip
using DynamicalSystemsBase.Systems: loop, loop_jac, liip, liip_jac

u0 = [0, 10.0, 0]
p = [10, 28, 8/3]
u0t = ones(2)
pt = nothing

FUNCTIONS = (liip, liip_jac, loop, loop_jac, eom_towel_iip, jacob_towel_iip, eom_towel, jacob_towel)
INITCOD = (u0, u0t)
PARAMS = (p, pt)

combinations = ["C_I", "C_I_J", "C_O", "C_O_J", "D_I", "D_I_J", "D_O", "D_O_J"]

for i in 1:8
    SI[combinations[i]] = BenchmarkGroup()
    sysindx = i < 5 ? 1 : 2
    if i < 5
        if isodd(i)
            ds = CDS(FUNCTIONS[i], INITCOD[sysindx], PARAMS[sysindx])
        else
            ds = CDS(FUNCTIONS[i-1], INITCOD[sysindx], PARAMS[sysindx], FUNCTIONS[i])
        end
    else
        if isodd(i)
            ds = DDS(FUNCTIONS[i], INITCOD[sysindx], PARAMS[sysindx])
        else
            ds = DDS(FUNCTIONS[i-1], INITCOD[sysindx], PARAMS[sysindx], FUNCTIONS[i])
        end
    end

    tinteg = tangent_integrator(ds); step!(tinteg)

    if isodd(i)
        integ = integrator(ds)
        pinteg = parallel_integrator(ds, [INITCOD[sysindx], INITCOD[sysindx]])
        step!(pinteg); step!(integ)
        SI[combinations[i]]["integ"] = @benchmarkable step!($integ)
        SI[combinations[i]]["pinteg"] = @benchmarkable step!($pinteg)
    end
    SI[combinations[i]]["tinteg"] = @benchmarkable step!($tinteg)
end


text:
import UnicodePlots
using StaticArrays

const CountType = UInt32
struct UnicodePlotter{T<:Number} <: Plotter
    fps::T
    counter::Base.RefValue{CountType}
    skip::CountType
end

function UnicodePlotter(; fps = 20, counter = 0, skip = 1 )
    UnicodePlotter{typeof(fps)}(
        fps,
        Ref(convert(CountType, counter)),
        convert(CountType, skip)
        )
end # function

function plot( ::UnicodePlotter, world::World, time )
    active_dogs = filter( d->(d.state==ACTIVE), world.dogs.member )
    sleepy_dogs = typeof(world.dogs.member)[]
    occupied_dogs = typeof(world.dogs.member)[]
    if length( active_dogs ) < length( world.dogs )
        sleepy_dogs = filter( d->(d.state==SLEEPY), world.dogs.member )
        occupied_dogs = filter( d->(d.state==OCCUPIED), world.dogs.member )
    end
    p = UnicodePlots.scatterplot( Float64[],
            xlim = @SVector([0, world.boxsize[]]),
            ylim = @SVector([0, world.boxsize[]])
            )
    for data in ( (active_dogs,:green), (sleepy_dogs,:red), (occupied_dogs,:white) )
        positions = getproperty.( data[1], :position )
        xs = map( p->p[1], positions )
        ys = map( p->p[2], positions )
        p = UnicodePlots.scatterplot!( p, xs, ys, color = data[2] )
    end # for
    sheeps = collect( world.sheeps )
    sheep_xs = map( s->s[2][1], sheeps )
    sheep_ys = map( s->s[2][2], sheeps )
    UnicodePlots.scatterplot!( p, sheep_xs, sheep_ys )
    return p
end # function

function plot!( io, p, plotter::UnicodePlotter, world::World, time )
    if plotter.counter[] % plotter.skip == zero(CountType)
        for _ in 1:(UnicodePlots.nrows( p.graphics )+p.margin)
            print(io, "\e[2K\e[1F")
        end # for
        p = plot( plotter, world, time )
        show(IOContext(io, :color => true), p)
        print(io,"\n\t\tPress Enter twice to end simulation.")
        print(stdout, String(take!(io)))
        sleep( 1 / plotter.fps )
    end # if
    plotter.counter[] += 1
    return nothing
end # function


text:
using SpecialFunctions
using LinearAlgebra
using ProgressMeter

"""
    Helper functions
"""
# power functions
@inline p2(x) = x*x
@inline p3(x) = x*x*x

"""
    kr_fitting_function(x::Real,
                             a0::Real, a1::Real, a2::Real, a3::Real, a4::Real)

Helper function to use the fitting function from KR07.
"""
@inline function kr_fitting_function(x::Real,
                             a0::Real, a1::Real, a2::Real, a3::Real, a4::Real)
    mm = x - 1.0

	return ( a0 + a1*mm + a2*p2(mm) + a3*p3(mm) + a4*p2(p2(mm)) ) / p2(p2(x))
end

"""
    Efficiency functions
"""

"""
    KR07_acc(M::Float64)

Efficiency model from Kang, Ryu, Cen, Ostriker 2007, http://arxiv.org/abs/0704.1521v1
"""
function KR07_acc(M::Real)
    if M < 2.0
        return 0.0
    elseif M <= 2.0
        return 1.96e-3*(M^2 - 1.)             # eq. A3
    else
        return kr_fitting_function(M, 5.46, -9.78, 4.17, -0.334, 0.57)
    end
end

"""
    KR13_acc(M::Real)

Efficiency model from Kang&Ryu 2013, doi:10.1088/0004-637X/764/1/95
"""
function KR13_acc(M::Real)

    if M < 2.0
        return 0.0
    elseif 2.0 <= M <= 5.0
        return -0.0005950569221922047 + 1.880258286365841e-5 * M^5.334076006529829
    elseif 5.0 < M <= 15.0
        return kr_fitting_function(M, -2.8696966498579606, 9.667563166507879,
                                  -8.877138312318019, 1.938386688261113, 0.1806112438315771)
    else
        return 0.21152
    end
end

"""
    Ryu19_acc(M::Real)

Efficiency model from Ryu et al. 2019, https://arxiv.org/abs/1905.04476
Values for 2.25 < M <= 5.0 extrapolated to entire range
"""
function Ryu19_acc(M::Real)

    if M < 2.25
        return 0.0
    elseif M <= 34.0
        return kr_fitting_function(M, -1.5255114554627316, 2.4026049650156693,
                 -1.2534251472776456, 0.22152323784680614, 0.0335800899612107)
    else
        return 0.0348
    end
end

"""
    CS14_acc(M::Real)

Efficiency from Caprioli&Spitkovsky 2015, doi: 10.1088/0004-637x/783/2/91
Same simplified approach as Vazza+12 -> is roughly half the efficiency of Kang&Ryu 2013.
"""
function CS14_acc(M::Real)
    vazza_factor = 0.5
    return vazza_factor * KR13_acc(M)
end

"""
    P16_acc(M::Real)

Constant efficiency as in Pfrommer+ 2016, doi: 10.1093/mnras/stw2941
"""
function P16_acc(M::Real)
    return 0.5
end

"""
    cre_spec_norm_particle(M::Real, eff_function::Function)

Computes the CR electron norm of the particles.
This depends on the Mach number `M` and the acceleration efficiency given by `eff_function`.
"""
function cre_spec_norm_particle(M::Real, eff_function::Function)
    norm = eff_function(M)
    return norm / 7.6e14 # Donnert et al 2016, eq. 40, p_0 = 0.1 me c
end

"""
    dsa_spectral_index(M::Real)

Spectral index given by standard Diffuse-Shock-Acceleration (DSA).
"""
function dsa_spectral_index(M::Real)

    # diverges for Mach smaller 1.01
    if M <= 1.01
        return maxintfloat(Float64)
    else
        return 2.0 * ( M^2 + 1.0 ) / ( M^2 - 1 )
    end
end

"""
    EpsNtherm(rho_cgs::Real, T_K::Real)

Thermal energy density in cgs.
"""
function EpsNtherm(rho_cgs::Real, T_K::Real; xH::Real=0.76)
    u_mol = (4.0/(5.0*xH+3.0))
    return (rho_cgs / (m_p * u_mol) * k_B * T_K)
end

"""
    analytic_synchrotron_emission( rho_cgs::Array{<:Real}, B_cgs::Array{<:Real},
                                   T_K::Array{<:Real}, Mach::Array{<:Real};
                                   xH::Real=0.76, dsa_model::Integer=1, ν0::Real=1.44e9,
                                   integrate_pitch_angle::Bool=true )

Computes the analytic synchrotron emission with the simplified approach described in Longair Eq. 8.128.
Returns J_ν in units [erg/cm^3/Hz/s].

# Arguments
- `rho_cgs::Array{<:Real}`: Density in ``g/cm^3``.
- `B_cgs::Array{<:Real}`:   Magnetic field in Gauss.
- `T_K::Array{<:Real}`:     Temperature in Kelvin.
- `Mach::Array{<:Real}`:    Mach number.

# Keyword Arguments
- `xH::Float64 = 0.76`:               Hydrogen fraction of the simulation, if run without chemical model.
- `dsa_model::Integer=1`:             Diffuse-Shock-Acceleration model. Takes values `0...4`, see next section.
- `ν0::Real=1.44e9`:                  Observation frequency in ``Hz``.
- `K_ep::Real=0.01`:                  Ratio of CR proton to electron energy density.
- `integrate_pitch_angle::Bool=true`: Integrates over the pitch angle as in Longair Eq. 8.87.
- `convert_to_mJy::Bool=false`:       Convert the result from ``[erg/cm^3/Hz/s]`` to ``mJy/cm``.

# DSA models
- `0`: [`KR07_acc`](@ref). Efficiency model from Kang, Ryu, Cen, Ostriker 2007, http://arxiv.org/abs/0704.1521v1.
- `1`: [`KR13_acc`](@ref). Efficiency model from Kang&Ryu 2013, doi:10.1088/0004-637X/764/1/95 .
- `2`: [`Ryu19_acc`](@ref). Efficiency model from Ryu et al. 2019, https://arxiv.org/abs/1905.04476 .
- `3`: [`CS14_acc`](@ref). Efficiency model from Caprioli&Spitkovsky 2015, doi: 10.1088/0004-637x/783/2/91 .
- `4`: [`P16_acc`](@ref). Constant efficiency as in Pfrommer+ 2016, doi: 10.1093/mnras/stw2941 .

"""
function analytic_synchrotron_emission(rho_cgs::Array{<:Real}, B_cgs::Array{<:Real},
                                       T_K::Array{<:Real}, Mach::Array{<:Real};
                                       xH::Real=0.76, dsa_model::Integer=1, ν0::Real=1.44e9,
                                       K_ep::Real=0.01,
                                       integrate_pitch_angle::Bool=true,
                                       convert_to_mJy::Bool=false)

    Npart = length(T_K)

    # default ratio for electron to protons from Donnert+16.
    K_ep_default = 0.01

    # allocate storage array
    J_ν = Vector{Float64}(undef, Npart)

    if length(B_cgs[1,:]) == 1
        B_1dim = true
    else
        B_1dim = false
    end

    # select DSA model
    if dsa_model == 0
        acc_function = KR07_acc
    elseif dsa_model == 1
        acc_function = KR13_acc
    elseif dsa_model == 2
        acc_function = Ryu19_acc
    elseif dsa_model == 3
        acc_function = CS14_acc
    elseif dsa_model == 4
        acc_function = P16_acc
    else
        error("Invalid DSA model selection!")
    end

    # bracket of Longair eq. 8128
    nufac = (3q_e)/(p3(m_e) * c_light * c_light * c_light * c_light * c_light * 2π * ν0)

    @inbounds for i = 1:Npart
        if B_1dim
            B = B_cgs[i]
        else
            B  = sqrt( B_cgs[1,i]^2 + B_cgs[2,i]^2 + B_cgs[3,i]^2)
        end
        n0 = K_ep / K_ep_default * cre_spec_norm_particle(Mach[i], acc_function) * EpsNtherm(rho_cgs[i], T_K[i], xH=xH)
        s  = dsa_spectral_index(Mach[i])

        if n0 > 0.0

            # Longair eq. 8.129
            a_p = gamma(s / 4.0 + 19.0 / 12.0) * gamma(s / 4.0 - 1.0 / 12.0)

            if integrate_pitch_angle
                # Longair eq 8.87
                a_p *= 0.5 * sqrt(π) * gamma(s / 4.0 + 5.0 / 4.0 ) /
                                        gamma(s / 4.0 + 7.0 / 4.0 )
            end

            # Longair eq 8.128 prefactor
            prefac = √(3)*q_e^3/(m_e * p2(c_light) * (s + 1.0))

            # Longair eq 8.128
            J_ν[i] = prefac * n0 *
                     nufac^( 0.5 * ( s - 1.0 ) ) * B^( 0.5 * ( s + 1.0 ) ) *
                     a_p
        else
            J_ν[i] = 0.0
        end
    end

    if convert_to_mJy
        J_ν .*= mJy_factor
    end

    return J_ν
end


text:
using Test
import ClimaCore: Domains, Meshes, Topologies
import ClimaCore.Geometry: Cartesian2DPoint, Cartesian123Point
using ClimaCore.Meshes:
    equispaced_rectangular_mesh,
    cube_panel_mesh,
    AbstractWarp,
    EquidistantSphereWarp,
    EquiangularSphereWarp,
    NoWarp,
    Mesh2D
using StaticArrays
using IntervalSets

function cube_panel_topology(ne, ::Type{FT}) where {FT <: AbstractFloat}
    domain = Domains.CubePanelDomain{FT}()
    mesh = Mesh2D(domain, ne)
    grid_topology = Topologies.Grid2DTopology(mesh)
    return (domain, mesh, grid_topology)
end

FT = Float64
_, _, grid_topology = cube_panel_topology(1, FT)

@testset "simple cube surface grid opposing face" begin

    @testset "assert correct element numbering" begin
        _, _, grid_topology = cube_panel_topology(1, FT)
        @test_throws AssertionError Topologies.opposing_face(
            grid_topology,
            0,
            1,
        )
        @test_throws AssertionError Topologies.opposing_face(
            grid_topology,
            7,
            1,
        )
    end

    @testset "6 element mesh with 1 element per panel" begin
        _, _, grid_topology = cube_panel_topology(1, FT)
        @test Topologies.opposing_face(grid_topology, 1, 1) == (4, 3, false)
        @test Topologies.opposing_face(grid_topology, 1, 2) == (2, 3, false)
        @test Topologies.opposing_face(grid_topology, 1, 3) == (5, 3, false)
        @test Topologies.opposing_face(grid_topology, 1, 4) == (3, 3, false)

        @test Topologies.opposing_face(grid_topology, 2, 1) == (5, 2, false)
        @test Topologies.opposing_face(grid_topology, 2, 2) == (3, 2, false)
        @test Topologies.opposing_face(grid_topology, 2, 3) == (1, 2, false)
        @test Topologies.opposing_face(grid_topology, 2, 4) == (6, 2, false)

        @test Topologies.opposing_face(grid_topology, 3, 1) == (4, 2, false)
        @test Topologies.opposing_face(grid_topology, 3, 2) == (2, 2, false)
        @test Topologies.opposing_face(grid_topology, 3, 3) == (1, 4, false)
        @test Topologies.opposing_face(grid_topology, 3, 4) == (6, 4, false)

        @test Topologies.opposing_face(grid_topology, 4, 1) == (5, 1, false)
        @test Topologies.opposing_face(grid_topology, 4, 2) == (3, 1, false)
        @test Topologies.opposing_face(grid_topology, 4, 3) == (1, 1, false)
        @test Topologies.opposing_face(grid_topology, 4, 4) == (6, 1, false)

        @test Topologies.opposing_face(grid_topology, 5, 1) == (4, 1, false)
        @test Topologies.opposing_face(grid_topology, 5, 2) == (2, 1, false)
        @test Topologies.opposing_face(grid_topology, 5, 3) == (1, 3, false)
        @test Topologies.opposing_face(grid_topology, 5, 4) == (6, 3, false)

        @test Topologies.opposing_face(grid_topology, 6, 1) == (4, 4, false)
        @test Topologies.opposing_face(grid_topology, 6, 2) == (2, 4, false)
        @test Topologies.opposing_face(grid_topology, 6, 3) == (5, 4, false)
        @test Topologies.opposing_face(grid_topology, 6, 4) == (3, 4, false)

        # 6 faces, 4 vertices per face, 8 global vertices, so each vertex should be part of 3 elements
        @test Topologies.vertex_coordinates(grid_topology, 1)[1] isa
              Cartesian123Point
        # check that all vertices appear as part of 3 elements
        for vert in Topologies.vertices(grid_topology)
            @test length(vert) == 3
        end
    end

    @testset "24 element mesh with 4 elements per panel" begin
        _, _, grid_topology = cube_panel_topology(2, FT)
        @test Topologies.opposing_face(grid_topology, 1, 1) == (13, 3, false)
        @test Topologies.opposing_face(grid_topology, 1, 2) == (2, 1, false)
        @test Topologies.opposing_face(grid_topology, 1, 3) == (17, 3, false)
        @test Topologies.opposing_face(grid_topology, 1, 4) == (3, 3, false)
    end
end

@testset "cube surface grid interior faces iterator" begin
    @testset "all faces should be interior faces" begin
        _, _, grid_topology = cube_panel_topology(1, FT)
        @test length(Topologies.interior_faces(grid_topology)) ==
              grid_topology.mesh.nfaces
        _, _, grid_topology = cube_panel_topology(2, FT)
        @test length(Topologies.interior_faces(grid_topology)) ==
              grid_topology.mesh.nfaces
    end

    @testset "6 element mesh with 1 element per panel" begin
        _, _, grid_topology = cube_panel_topology(1, FT)
        faces = collect(Topologies.interior_faces(grid_topology))
        @test faces[1] == (1, 3, 5, 3, false)
        @test faces[2] == (1, 4, 3, 3, false)
        @test faces[3] == (5, 4, 6, 3, false)
        @test faces[4] == (6, 4, 3, 4, false)
        @test faces[5] == (1, 1, 4, 3, false)
        @test faces[6] == (1, 2, 2, 3, false)
        @test faces[7] == (4, 4, 6, 1, false)
        @test faces[8] == (6, 2, 2, 4, false)
        @test faces[9] == (4, 1, 5, 1, false)
        @test faces[10] == (5, 2, 2, 1, false)
        @test faces[11] == (4, 2, 3, 1, false)
        @test faces[12] == (3, 2, 2, 2, false)
    end
end

@testset "simple cube surface grid boundary faces iterator" begin
    @testset "24 element mesh with 4 elements per panel (no boundary faces for this topology)" begin
        _, _, grid_topology = cube_panel_topology(2, FT)
        @test length(Topologies.boundary_faces(grid_topology, 1)) == 0
        @test length(Topologies.boundary_faces(grid_topology, 2)) == 0
        @test length(Topologies.boundary_faces(grid_topology, 3)) == 0
        @test length(Topologies.boundary_faces(grid_topology, 4)) == 0
    end
end

@testset "sphere mesh" begin
    FT = Float64
    @testset "4 elements per edge, equidistant spherical mesh of radius 10; Float type = Float64" begin
        radius = FT(10)
        domain = Domains.SphereDomain(radius)
        mesh = Mesh2D(domain, EquidistantSphereWarp(), 4)
        crad = abs.(sqrt.(sum(mesh.coordinates .^ 2, dims = 2)) .- radius)
        @test maximum(crad) ≤ 100 * eps(FT)
    end
    @testset "4 elements per edge, equiangular spherical mesh of radius 10; Float type = Float64" begin
        radius = FT(10)
        domain = Domains.SphereDomain(radius)
        mesh = Mesh2D(domain, EquiangularSphereWarp(), 4)
        crad = abs.(sqrt.(sum(mesh.coordinates .^ 2, dims = 2)) .- radius)
        @test maximum(crad) ≤ 100 * eps(FT)
    end

    FT = BigFloat
    @testset "4 elements per edge, equidistant spherical mesh of radius 10; Float type = BigFloat" begin
        radius = FT(10)
        domain = Domains.SphereDomain(radius)
        mesh = Mesh2D(domain, EquidistantSphereWarp(), 4)
        crad = abs.(sqrt.(sum(mesh.coordinates .^ 2, dims = 2)) .- radius)
        @test maximum(crad) ≤ 100 * eps(FT)
    end
    @testset "4 elements per edge, equiangular spherical mesh of radius 10; Float type = BigFloat" begin
        radius = FT(10)
        domain = Domains.SphereDomain(radius)
        mesh = Mesh2D(domain, EquiangularSphereWarp(), 4)
        crad = abs.(sqrt.(sum(mesh.coordinates .^ 2, dims = 2)) .- radius)
        @test maximum(crad) ≤ 100 * eps(FT)
    end
end


text:
@testset "utils.jl" begin
    @testset "size_function" begin
        y = ones(15)
        X = randn(15, 2)
        s = 2
        model = structural(y, s; X = X)

        n, p, m, r = size(model)
        @test n == 15
        @test p == 1
        @test m == s + size(X, 2) + 1
        @test r == 3

        y = ones(15)
        X = randn(15, 15)
        s = 2
        model = structural(y, s; X = X)
        n, p, m, r = size(model)
        @test n == 15
        @test p == 1
        @test m == s + size(X, 2) + 1
        @test r == 3
    end

    @testset "ztr" begin
        y = ones(15)
        model = local_level(y)
        Z, T, R = ztr(model)
        @test all(Z .== [1.0])
        @test T == [1.0][:, :]
        @test R == [1.0][:, :]
    end

    @testset "gram" begin
        A = [1.0][:, :]
        @test A == SSM.gram(A)
        B = [1 1.0;0 1]
        @test B*B' == SSM.gram(B)
    end

    @testset "gram_in_time" begin
        C = Array{Float64, 3}(undef, 2, 2, 5)
        B = [1 1.0;0 1]
        C[:, :, 1:5] .= B
        @test B*B' == SSM.gram_in_time(C)[:, :, 1]
        @test SSM.gram_in_time(C)[:, :, 2] == SSM.gram_in_time(C)[:, :, 1]
    end

    @testset "ensure_pos_sym!" begin
        A = randn(3, 3, 2)
        SSM.ensure_pos_sym!(A, 1)
        @test issymmetric(A[:, :, 1])
        SSM.ensure_pos_sym!(A, 2)
        @test issymmetric(A[:, :, 2])
    end

    @testset "show" begin
        y = ones(15)
        model = local_level(y)
        ss = statespace(model)
        @test show(model) == nothing
        @test show(ss) == nothing
    end

    @testset "statespace_recursion" begin
        # Unknowns error
        y = ones(15)
        model = local_level(y)
        initial_a = [1.0][:, :]
        @test_throws ErrorException("StateSpaceModel has unknown parameters.") statespace_recursion(model, initial_a)

        # initial_a dimension mismatch
        y = ones(15)
        model = local_level(y)
        model.H .= [1e-8]
        model.Q .= [1e-8]
        initial_a = [1.0 1.0]
        @test_throws ErrorException("intial_a must be a 1 by 1 matrix.") statespace_recursion(model, initial_a)

        # univariate local_level
        y = ones(15)
        model = local_level(y)
        model.H .= [1e-8]
        model.Q .= [1e-8]
        initial_a = [1.0][:, :]
        y, a = statespace_recursion(model, initial_a)
        @test y ≈ ones(15) atol = 1e-3 rtol = 1e-3
        @test a ≈ ones(15) atol = 1e-3 rtol = 1e-3

        # univariate linear_trend
        y = collect(1.0:15.0)
        model = linear_trend(y)
        model.H .= [1e-8]
        model.Q .= [1e-8 0
                    0 1e-8]
        initial_a = [1.0 0.5]
        y, a = statespace_recursion(model, initial_a)
        @test y ≈ collect(1.0:0.5:8.0) atol = 1e-3 rtol = 1e-3
        @test a ≈ [collect(1.0:0.5:8.0) 0.5*ones(15)] atol = 1e-3 rtol = 1e-3

        # multivariate local_level
        y = ones(15, 4)
        model = local_level(y)
        model.H .= Diagonal(1e-8*ones(4))
        model.Q .= Diagonal(1e-8*ones(4))
        initial_a = ones(1, 4)
        y, a = statespace_recursion(model, initial_a)
        @test y ≈ ones(15, 4) atol = 1e-3 rtol = 1e-3
        @test a ≈ ones(15, 4) atol = 1e-3 rtol = 1e-3

        # univariate linear_trend
        y_c = collect(1.0:15.0)
        y = [y_c y_c y_c y_c]
        model = linear_trend(y)
        model.H .= Diagonal(1e-8*ones(4))
        model.Q .= Diagonal(1e-8*ones(8))
        initial_a = [1.0 0.5 1.0 0.5 1.0 0.5 1.0 0.5]
        y, a = statespace_recursion(model, initial_a)
        @test y[:, 1] ≈ collect(1.0:0.5:8.0) atol = 1e-3 rtol = 1e-3
        @test y[:, 2] ≈ collect(1.0:0.5:8.0) atol = 1e-3 rtol = 1e-3
        @test y[:, 3] ≈ collect(1.0:0.5:8.0) atol = 1e-3 rtol = 1e-3
        @test y[:, 4] ≈ collect(1.0:0.5:8.0) atol = 1e-3 rtol = 1e-3
        @test a[:, [1, 2]] ≈ [collect(1.0:0.5:8.0) 0.5*ones(15)] atol = 1e-3 rtol = 1e-3
        @test a[:, [3, 4]] ≈ [collect(1.0:0.5:8.0) 0.5*ones(15)] atol = 1e-3 rtol = 1e-3
        @test a[:, [5, 6]] ≈ [collect(1.0:0.5:8.0) 0.5*ones(15)] atol = 1e-3 rtol = 1e-3
        @test a[:, [7, 8]] ≈ [collect(1.0:0.5:8.0) 0.5*ones(15)] atol = 1e-3 rtol = 1e-3
    end
end

function compare_forecast_simulation(ss::StateSpace, N::Int, S::Int, rtol::Float64)
    sim = simulate(ss, N, S)
    forec, dist = forecast(ss, N)

    @test forec ≈ mean(sim, dims = 3) rtol = rtol
end


text:
export
    # methods that must be overloaded
    jld_fn, group_keys,

    # convenience methods
    groupname, groupname!, with_group, @with_group, with_file, @with_file,
    groups, init_jld,

    # core api methods
    write_jld, write_jld!, read_jld, delete_jld!, names_jld, delete_group!,
    groups, haskey_jld

# just define _function_ here so others can define methods
function jld_fn end
function group_keys end

"""
```julia
init_jld(x, [force=false])
```

Initialize the jld_file for objects of type `typeof(x)`. This creates the file
(if needed) and ensures that a `groups` dataset exists and constructs a group
for the instance `x`.

If `force` is `true` any existing file is overwritten and the above operations
are performed.

This should only be called by users when they want to completely reconstruct
the jld file.
"""
function init_jld(x, force::Bool=false)
    if !isfile(jld_fn(x)) || force
        # create the file and the map and first group
        jldopen(jld_fn(x), "w") do f
            grps = Dict(group_keys(x) => "group1")
            write(f, "groups", grps)
            g_create(f, "group1")

            # write a value.... somehow this prevents bugs?
            write(f, "group1/visted", true)
        end

        return
    end

    # we know the file exists, now we just need to make sure the map and
    # group for x exist. That's what groupname!(x) is for!
    groupname!(x)
    return
end

"""
```julia
groupname(x::T)
```

Returns a tuple `(name, group_map_exists, group_exists)` specifying the `name`
of the group for x, whether a map from `group_keys(x)` to a group name exists
in the jld file, and whether or not a group with `name` exists in the jld file.
Mostly an internal function
"""
function groupname(x)
    # NOTE: can't use `groups` method here because we might get stuck in an
    #       infinite recursion when `init_jld` calls this routine (via a call
    #       to `groupname!`)
    grps = jldopen(jld_fn(x), "r+") do f
        if exists(f, "groups")
            read(f, "groups")
        else
            # make sure group1 doesn't already exist
            if exists(f, "group1")
                delete!(f, "group1")
            end
            return "group1", false, false
        end
    end

    # HACK to return if grps isn't a Dict
    if isa(grps, Tuple)
        return grps
    end

    if haskey(grps, group_keys(x))
        name = grps[group_keys(x)]
        group_map_exists = true

        group_exists = jldopen(jld_fn(x), "r") do f
            exists(f, name)
        end

        return name, group_map_exists, group_exists
    else
        # get new unique name
        group_nums = Int[parse(Int, x[6:end]) for x in values(grps)]
        name = string("group", maximum(group_nums) + 1)

        # the mapping entry didn't exist, and neither does the group
        return name, false, false
    end
end

"""
```julia
groupname!(x::T)
```

Returns the same as `groupname(x)`, but also creates the group map (if
`group_map_exists == false`) and the group (if `group_exists == false`).
"""
function groupname!(x)
    name, group_map_exists, group_exists = groupname(x)

    # if either the map or group doesn't exist, we need to do _something_ with
    # the file
    if !group_map_exists || !group_exists
        # NOTE: can't use with_file here because we might get stuck in an
        #       infinite recursion due to the call of this routine from within
        #       init_jld
        jldopen(jld_fn(x), "r+") do f

            if !group_map_exists
                if exists(f, "groups")
                    # get current groups
                    groups = read(f, "groups")

                    # add this name
                    groups[group_keys(x)] = name

                    # delete the groups object, then update it with the new one
                    delete!(f, "groups")
                else
                    groups = Dict(group_keys(x) => name)
                end
                write(f, "groups", groups)
            end

            if !group_exists
                # create the group
                g_create(f, name)
                write(f, "$(name)/visted", true)
            end
        end
    end

    # group_map_exists, group_exists are both true now
    return name, true, true
end

"""
```julia
with_group(func!::Function, x)
```

With the jld group corresponding to the instance `x` as `g`, call `func!(g)`.
This is a core method that is used to interact with the group for an object.
"""
function with_group(func!::Function, x)
    init_jld(x)
    jldopen(jld_fn(x), "r+") do f
        g = g_open(f, groupname!(x)[1])
        local out
        try
            out = func!(g)
        finally
            close(g)
        end
        out
    end
end

"""
```julia
with_file(func!::Function, x)
```

With the jld file object for objects of type `typeof(x)` as `f`, call
`func!(f)`. This is a core method that is used to interact with the jld file
as a whole.
"""
function with_file(func!::Function, x)
    init_jld(x)
    jldopen(jld_fn(x), "r+") do f
        func!(f)
    end
end

function _get_f_args_body(ex::Expr)
    # try one line function definitino version
    @capture(ex, f_(args__) = body_)
    if any(map(x->x==nothing, (f, args, body)))

    # now try full `function name(args...) ` version
    @capture(ex, function f_(args__) body_ end)
    if any(map(x->x==nothing, (f, args, body)))
        error("I couldn't understand your function definition")
        end
    end

    return f, args, body
end

"""
```julia
macro with_group(func)
```

Wrap the body of a function `func` in a call to `with_group`. This is best
understood by example.

```julia
@with_group write_jld(x, nm, obj) = write(g, nm, obj)
```

is equivalent to

```julia
write_jld(x, nm, obj) = with_group(g -> write(g, nm, obj), x)
```
"""
macro with_group(ex)
    f, args, body = _get_f_args_body(ex)
    :($(esc(f))($(args...)) = with_group(g -> $body, $(args[1])))
end

"""
```julia
macro with_group(func)
```

Wrap the body of a function `func` in a call to `with_file`. This is best
understood by example.

```julia
@with_file foobar(x, y) = foobar(f, y)
```

is equivalent to

```julia
foobar(x, nm, obj) = with_file(f -> foobar(y), x)
```
"""
macro with_file(ex)
    f, args, body = _get_f_args_body(ex)
    :($(esc(f))($(args...)) = with_file(f -> $body, $(args[1])))
end

function _write_jld(overrwrite::Bool, x; kwargs...)
    grp_nm = groupname!(x)[1]
    with_file(x) do f
        for (_nm, obj) in kwargs
            nm = "$grp_nm/$(_nm)"
            if overrwrite
                exists(f, nm) && delete!(f, nm)
            end
            write(f, nm, obj)
        end
    end
end

function _write_jld(overrwrite::Bool, x, _nm, obj)
    grp_nm = groupname!(x)[1]
    with_file(x) do f
        nm = "$grp_nm/$(_nm)"
        if overrwrite
            exists(f, nm) && delete!(f, nm)
        end
        write(f, nm, obj)
    end
end

"""
```julia
write_jld(x, nm, obj)
```

Write the object `obj` `x`'s jld group under name `nm`
"""
write_jld(x, nm, obj) = _write_jld(false, x, nm, obj)
write_jld(x; kwargs...) = _write_jld(false, x; kwargs...)

"""
```julia
write_jld!(x, nm, obj)
```

Write the object `obj` `x`'s jld group under name `nm`. This routine will
overwrite any existing object with name `nm` in the group.
"""
write_jld!(x, nm, obj) = _write_jld(true, x, nm, obj)
write_jld!(x; kwargs...) = _write_jld(true, x; kwargs...)

"""
```julia
read_jld(x, nm)
```

Read the object with `nm` from  `x`'s jld group
"""
@with_group read_jld(x, dataset) = read(g, dataset)

"""
```julia
delete_jld!(x, dataset)
```

Delete the object with name `dataset` from `x`'s jld group
"""
@with_group delete_jld!(x, dataset) = delete!(g, dataset)

"""
```julia
haskey_jld(x, dataset)
```

Check if the jld group for `x` has an object named `dataset`
"""
@with_group haskey_jld(x, dataset) = exists(g, dataset)

"""
```julia
names_jld(x)
```

List all the names in `x`'s jld group
"""
@with_group names_jld(x) = names(g)

"""
```julia
groups(x)
```

Return the dict mapping keys to jld group names for all objects of type
`typeof(x)`. If the dict doesn't yet exist, it is created, along with a group
for the instance `x`.
"""
@with_file groups(x) = read(f, "groups")

"""
```julia
delete_group!(x)
```

Delete the jld group associated with the instance of `x` and remove its entry
from the group map.
"""
function delete_group!(x)
    # NOTE: we don't use the macro here so that we can avoid opening the file if
    #       it isn't necessary.
    name, group_map_exists, group_exists = groupname(x)
    if !group_map_exists && !group_exists
        # nothing to do
        return
    end

    # we need to do something
    with_file(x) do f
        # need to delete the entry in the group map
        if group_map_exists
            # read existing map
            groups = read(f, "groups")

            # remove this entry
            pop!(groups, group_keys(x))

            # delete existing map
            delete!(f, "groups")

            # write updated map
            write(f, "groups", groups)
        end

        # if the group exists, delete it
        if group_exists
            delete!(f, name)
        end
    end

end