text:
function run!(model, grid, output_freq, ncfilename)

    etime = 0.0
    sim_time = grid.nt * grid.dt

    if output_freq > 0.0
        nsnaps = floor(Int, sim_time/output_freq) + 1
        snapshots = Array{Float64}(undef, 4, grid.nx, grid.nz, nsnaps)
        etimes = Vector{Float64}(undef, nsnaps)
        # output initial state
        snapshots[:,:,:,1] = createSnapshot(model, grid)
    end

    direction_switch = true
    output_counter = 0.0
    counter = 1
    @showprogress for i in 1:grid.nt
        if output_freq > 0.0 && output_counter >= output_freq
            counter += 1
            output_counter = output_counter - output_freq
            snapshots[:,:,:,counter] = createSnapshot(model, grid)
            etimes[counter] = etime
        end

        perform_timestep!(model, grid, direction_switch)
        direction_switch = !direction_switch
        etime += grid.dt
        output_counter += grid.dt
        #output_gif!(model, etime, grid, anim);
    end

    if output_freq > 0
        output(snapshots, grid, etimes, ncfilename)
    end

end

#Performs a single dimensionally split time step using a simple low-storate three-stage Runge-Kutta time integrator
#The dimensional splitting is a second-order-accurate alternating Strang splitting in which the
#order of directions is alternated each time step.
#The Runge-Kutta method used here is defined as follows:
# q*     = q[n] + dt/3 * rhs(q[n])
# q^    = q[n] + dt/2 * rhs(q*  )
# q[n+1] = q[n] + dt/1 * rhs(q^ )
function perform_timestep!(model, grid, direction_switch)
    if direction_switch
        #x-direction first
        semi_discrete_step!(model, grid, DIR_X, 1)
        semi_discrete_step!(model, grid, DIR_X, 2)
        semi_discrete_step!(model, grid, DIR_X, 3)
        #z-direction second
        semi_discrete_step!(model, grid, DIR_Z, 1)
        semi_discrete_step!(model, grid, DIR_Z, 2)
        semi_discrete_step!(model, grid, DIR_Z, 3)
    else
        #z-direction second
        semi_discrete_step!(model, grid, DIR_Z, 1)
        semi_discrete_step!(model, grid, DIR_Z, 2)
        semi_discrete_step!(model, grid, DIR_Z, 3)
        #x-direction first
        semi_discrete_step!(model, grid, DIR_X, 1)
        semi_discrete_step!(model, grid, DIR_X, 2)
        semi_discrete_step!(model, grid, DIR_X, 3)
    end
end

#Perform a single semi-discretized step in time with the form:
#state_out = state_init + dt * rhs(state_forcing)
#Meaning the step starts from state_init, computes the rhs using state_forcing, and stores the result in state_out
function semi_discrete_step!(model, grid, dir, mode)
    # mode=1 sets model.state=state_init & state_forcing and model.state_tmp=state_out
    # mode=2 sets model.state=initial-state and model.state_tmp=state_forcing & state_out
    # mode=3 sets model.state=initial state & state_out and model.state_tmp=state_forcing
    # time step depends on mode
    dt = grid.dt / (4.0 - mode)
    if dir == DIR_X
        #Set the halo values for this MPI task's fluid state in the x-direction
        if mode == 1
            set_halo_values_x!(
                model.state,
                model.hy_dens_cell,
                model.hy_dens_theta_cell,
                grid,
                model.weather_type,
            )
            compute_tendencies_x!(
                model.state,
                model.flux,
                model.tend,
                model.hy_dens_cell,
                model.hy_dens_theta_cell,
                grid,
            )
        elseif mode == 2 || mode == 3
            set_halo_values_x!(
                model.state_tmp,
                model.hy_dens_cell,
                model.hy_dens_theta_cell,
                grid,
                model.weather_type,
            )
            compute_tendencies_x!(
                model.state_tmp,
                model.flux,
                model.tend,
                model.hy_dens_cell,
                model.hy_dens_theta_cell,
                grid,
            )
        else
            throw(ArgumentError("mode must be either 1, 2 or 3"))
        end
        #Compute the time tendencies for the fluid state in the x-direction
    elseif dir == DIR_Z
        #Set the halo values for this MPI task's fluid state in the z-direction
        if mode == 1
            set_halo_values_z!(model.state, grid, model.weather_type)
            compute_tendencies_z!(
                model.state,
                model.flux,
                model.tend,
                model.hy_dens_int,
                model.hy_dens_theta_int,
                model.hy_pressure_int,
                grid,
            )
        elseif mode == 2 || mode == 3
            set_halo_values_z!(model.state_tmp, grid, model.weather_type)
            compute_tendencies_z!(
                model.state_tmp,
                model.flux,
                model.tend,
                model.hy_dens_int,
                model.hy_dens_theta_int,
                model.hy_pressure_int,
                grid,
            )
        else
            throw(ArgumentError("mode must be either 1, 2 or 3"))
        end
        #Compute the time tendencies for the fluid state in the z-direction
    else
        throw(ArgumentError("dir must be either $DIR_X or $DIR_Z, got $dir"))
    end

    #################################################
    ## TODO: THREAD ME
    #################################################
    #Apply the tendencies to the fluid state
    #state_out = zeros(size(model.state))
    state_out = similar(model.state)
    #    state_out = zeros(grid.nx + 2 * hs, grid.nz + 2 * hs, NUM_VARS)
    for ll = 1:NUM_VARS
        for k = 1:grid.nz, i = 1:grid.nx
            state_out[i+hs, k+hs, ll] =
                model.state[i+hs, k+hs, ll] + dt * model.tend[i, k, ll]
        end
    end
    if mode == 1 || mode == 2
        model.state_tmp .= state_out
    elseif mode == 3
        model.state .= state_out
    else
        throw(ArgumentError("mode must be either 1, 2 or 3"))
    end
end

#Compute the time tendencies of the fluid state using forcing in the x-direction
#Since the halos are set in a separate routine, this will not require MPI
#First, compute the flux vector at each cell interface in the x-direction (including hyperviscosity)
#Then, compute the tendencies using those fluxes
function compute_tendencies_x!(state, flux, tend, hy_dens_cell, hy_dens_theta_cell, grid)
    #    flux = zeros(grid.nx + 1, grid.nz + 1, NUM_VARS)
    #    tend = zeros(grid.nx, grid.nz, NUM_VARS)
    d3_vals = zeros(NUM_VARS)
    vals = zeros(NUM_VARS)
    stencil = zeros(4)
    #Compute the hyperviscosity coeficient
    hv_coef = -hv_beta * grid.dx / (16 * grid.dt)

    #################################################
    ## TODO: THREAD ME
    #################################################
    #Compute fluxes in the x-direction for each cell
    for k = 1:grid.nz, i = 1:grid.nx+1
        #Use fourth-order interpolation from four cell averages to compute the value at the interface in question
        for ll = 1:NUM_VARS
            for s = 1:sten_size
                stencil[s] = state[i+s-1, k+hs, ll]
            end
            #Fourth-order-accurate interpolation of the state
            vals[ll] =
                -stencil[1] / 12 + 7 * stencil[2] / 12 + 7 * stencil[3] / 12 -
                stencil[4] / 12
            #First-order-accurate interpolation of the third spatial derivative of the state (for artificial viscosity)
            d3_vals[ll] = -stencil[1] + 3 * stencil[2] - 3 * stencil[3] + stencil[4]
        end

        #Compute density, u-wind, w-wind, potential temperature, and pressure (r,u,w,t,p respectively)
        r = vals[ID_DENS] + hy_dens_cell[k+hs]
        u = vals[ID_UMOM] / r
        w = vals[ID_WMOM] / r
        t = (vals[ID_RHOT] + hy_dens_theta_cell[k+hs]) / r
        p = C0 * (r * t)^gamma

        #Compute the flux vector
        flux[i, k, ID_DENS] = r * u - hv_coef * d3_vals[ID_DENS]
        flux[i, k, ID_UMOM] = r * u * u + p - hv_coef * d3_vals[ID_UMOM]
        flux[i, k, ID_WMOM] = r * u * w - hv_coef * d3_vals[ID_WMOM]
        flux[i, k, ID_RHOT] = r * u * t - hv_coef * d3_vals[ID_RHOT]
    end

    ######
    # TODO: THREAD ME
    #####
    #Use the fluxes to compute tendencies for each cell
    for ll = 1:NUM_VARS
        for k = 1:grid.nz, i = 1:grid.nx
            tend[i, k, ll] = -(flux[i+1, k, ll] - flux[i, k, ll]) / grid.dx
        end
    end
end

#Compute the time tendencies of the fluid state using forcing in the z-direction
#Since the halos are set in a separate routine, this will not require MPI
#First, compute the flux vector at each cell interface in the z-direction (including hyperviscosity)
#Then, compute the tendencies using those fluxes
function compute_tendencies_z!(
    state,
    flux,
    tend,
    hy_dens_int,
    hy_dens_theta_int,
    hy_pressure_int,
    grid,
)
    #    flux = zeros(grid.nx + 1, grid.nz + 1, NUM_VARS)
    #    tend = zeros(grid.nx, grid.nz, NUM_VARS)
    d3_vals = zeros(NUM_VARS)
    vals = zeros(NUM_VARS)
    stencil = zeros(4)
    #Compute the hyperviscosity coeficient
    hv_coef = -hv_beta * grid.dz / (16 * grid.dt)
    ####
    ## TODO: THREAD ME
    ####
    #Compute fluxes in the x-direction for each cell
    for k = 1:grid.nz+1, i = 1:grid.nx
        #Use fourth-order interpolation from four cell averages to compute the value at the interface in question
        for ll = 1:NUM_VARS
            for s = 1:sten_size
                stencil[s] = state[i+hs, k-1+s, ll]
            end
            #Fourth-order-accurate interpolation of the state
            vals[ll] =
                -stencil[1] / 12 + 7 * stencil[2] / 12 + 7 * stencil[3] / 12 -
                stencil[4] / 12
            #First-order-accurate interpolation of the third spatial derivative of the state
            d3_vals[ll] = -stencil[1] + 3 * stencil[2] - 3 * stencil[3] + stencil[4]
        end

        #Compute density, u-wind, w-wind, potential temperature, and pressure (r,u,w,t,p respectively)
        r = vals[ID_DENS] + hy_dens_int[k]
        u = vals[ID_UMOM] / r
        w = vals[ID_WMOM] / r
        t = (vals[ID_RHOT] + hy_dens_theta_int[k]) / r
        p = C0 * (r * t)^gamma - hy_pressure_int[k]
        #Enforce vertical boundary condition and exact mass conservation
        if k == 1 || k == grid.nz + 1
            w = 0
            d3_vals[ID_DENS] = 0
        end

        #Compute the flux vector with hyperviscosity
        flux[i, k, ID_DENS] = r * w - hv_coef * d3_vals[ID_DENS]
        flux[i, k, ID_UMOM] = r * w * u - hv_coef * d3_vals[ID_UMOM]
        flux[i, k, ID_WMOM] = r * w * w + p - hv_coef * d3_vals[ID_WMOM]
        flux[i, k, ID_RHOT] = r * w * t - hv_coef * d3_vals[ID_RHOT]
    end

    ####
    ## TODO: THREAD ME
    ####
    #Use the fluxes to compute tendencies for each cell
    for ll = 1:NUM_VARS
        for k = 1:grid.nz, i = 1:grid.nx
            tend[i, k, ll] = -(flux[i, k+1, ll] - flux[i, k, ll]) / grid.dz
            if ll == ID_WMOM
                tend[i, k, ll] -= state[i+hs, k+hs, ID_DENS] * grav
            end
        end
    end
end

#Set this MPI task's halo values in the x-direction. This routine will require MPI
function set_halo_values_x!(state, hy_dens_cell, hy_dens_theta_cell, grid, weather_type)
    #real(rp), intent(inout) :: state(1-hs:nx+hs,1-hs:nz+hs,NUM_VARS)
    ###
    # TODO: EXCHANGE HALO VALUES WITH NEIGHBORING MPI TASKS
    # (1) give    state(1:hs,1:nz,1:NUM_VARS)       to   my left  neighbor
    # (2) receive state(1-hs:0,1:nz,1:NUM_VARS)     from my left  neighbor
    # (3) give    state(nx-hs+1:nx,1:nz,1:NUM_VARS) to   my right neighbor
    # (4) receive state(nx+1:nx+hs,1:nz,1:NUM_VARS) from my right neighbor
    ###

    ###
    # DELETE THE SERIAL CODE BELOW AND REPLACE WITH MPI
    ###
    for ll = 1:NUM_VARS
        for k = 1:grid.nz
            state[1, hs+k, ll] = state[grid.nx+hs-1, hs+k, ll]
            state[2, hs+k, ll] = state[grid.nx+hs, hs+k, ll]
            state[grid.nx+hs+1, hs+k, ll] = state[hs+1, hs+k, ll]
            state[grid.nx+hs+2, hs+k, ll] = state[hs+2, hs+k, ll]
        end
    end
    ###

    if (weather_type == "injection")
        #if (myrank == 0)
        for k = 1:grid.nz
            z = (k_beg - 1 + k - 0.5) * grid.dz
            if (abs(z - 3 * zlen / 4) <= zlen / 16)
                state[1:2, hs+k, ID_UMOM] =
                    (state[1:2, hs+k, ID_DENS] + hy_dens_cell[k+hs]) * 50.0
                state[1:2, hs+k, ID_RHOT] =
                    (state[1:2, hs+k, ID_DENS] + hy_dens_cell[k+hs]) * 298.0 -
                    hy_dens_theta_cell[k+hs]
            end
        end
        #end
    end
end

#Set this MPI task's halo values in the z-direction. This does not require MPI because there is no MPI
#decomposition in the vertical direction
function set_halo_values_z!(state, grid, weather_type)
    #real(rp), intent(inout) :: state(1-hs:nx+hs,1-hs:nz+hs,NUM_VARS)
    mnt_width = xlen / 8
    ###
    # TODO: THREAD ME
    ###
    for ll = 1:NUM_VARS
        for i = 1:grid.nx+2*hs
            if (ll == ID_WMOM)
                state[i, 1, ll] = 0
                state[i, 2, ll] = 0
                state[i, grid.nz+hs+1, ll] = 0
                state[i, grid.nz+hs+2, ll] = 0
                #Impose the vertical momentum effects of an artificial cos^2 mountain at the lower boundary
                if (weather_type == "mountain_waves")
                    x = (i_beg - 1 + i - 0.5) * grid.dx
                    if (abs(x - xlen / 4) < mnt_width)
                        xloc = (x - (xlen / 4)) / mnt_width
                        #Compute the derivative of the fake mountain
                        mnt_deriv =
                            -pi * cos(pi * xloc / 2) * sin(pi * xloc / 2) * 10 / grid.dx
                        #w = (dz/dx)*u
                        state[i, 1, ID_WMOM] = mnt_deriv * state[i, 3, ID_UMOM]
                        state[i, 2, ID_WMOM] = mnt_deriv * state[i, 3, ID_UMOM]
                    end
                end
            else
                state[i, 1, ll] = state[i, 3, ll]
                state[i, 2, ll] = state[i, 3, ll]
                state[i, grid.nz+3, ll] = state[i, grid.nz+2, ll]
                state[i, grid.nz+4, ll] = state[i, grid.nz+2, ll]
            end
        end
    end
end


text:
using DistStat, Pkg, Test

type=[Float32 ,Float64]

if get(ENV,"JULIA_MPI_TEST_ARRAYTYPE","") == "CuArray"
    using CUDA
    A = CuArray
else
    A = Array
end

for T in type
    x=A(zeros(T,5,13)); x_dist=distribute(x)
    cols=x_dist.partitioning[DistStat.Rank()+1][2]

    y=A{T}(reshape(collect(-6:6),1,13))
    y_dist=distribute(y)

    x_dist .+= 8.0; x .+= 8.0

    @test isapprox(x_dist.localarray,x[:,cols])

    x_dist .*= 1.5; x .*= 1.5

    @test isapprox(x_dist.localarray,x[:,cols])

    x_dist .+= y_dist; x .+= y

    @test isapprox(x_dist.localarray, x[:, cols])

    z=A{T}([1,2,3,4,5])

    x_dist .+= z; x .+= z

    @test isapprox(x_dist.localarray,x[:,cols])

    x_dist .= log.(x_dist); x .= log.(x)

    @test isapprox(x_dist.localarray,x[:,cols])

end

ext: jl


text:
module Sch

using Distributed
import MemPool: DRef, poolset
import Statistics: mean
import Random: randperm

import ..Dagger
import ..Dagger: Context, Processor, Thunk, WeakThunk, ThunkFuture, ThunkFailedException, Chunk, OSProc, AnyScope
import ..Dagger: order, free!, dependents, noffspring, istask, inputs, unwrap_weak_checked, affinity, tochunk, timespan_start, timespan_finish, unrelease, procs, move, capacity, chunktype, processor, default_enabled, get_processors, get_parent, execute!, rmprocs!, addprocs!, thunk_processor, constrain, cputhreadtime

const OneToMany = Dict{Thunk, Set{Thunk}}

include("util.jl")
include("fault-handler.jl")
include("dynamic.jl")

mutable struct ProcessorCacheEntry
    gproc::OSProc
    proc::Processor
    next::ProcessorCacheEntry

    ProcessorCacheEntry(gproc::OSProc, proc::Processor) = new(gproc, proc)
end
Base.isequal(p1::ProcessorCacheEntry, p2::ProcessorCacheEntry) =
    p1.proc === p2.proc
function Base.show(io::IO, entry::ProcessorCacheEntry)
    entries = 1
    next = entry.next
    while next !== entry
        entries += 1
        next = next.next
    end
    print(io, "ProcessorCacheEntry(pid $(entry.gproc.pid), $(entry.proc), $entries entries)")
end

const Signature = Vector{Any}

"""
    ComputeState

The internal state-holding struct of the scheduler.

Fields:
- `uid::UInt64` - Unique identifier for this scheduler instance
- `waiting::OneToMany` - Map from downstream `Thunk` to upstream `Thunk`s that still need to execute
- `waiting_data::Dict{Union{Thunk,Chunk},Set{Thunk}}` - Map from input `Chunk`/upstream `Thunk` to all unfinished downstream `Thunk`s, to retain caches
- `ready::Vector{Thunk}` - The list of `Thunk`s that are ready to execute
- `cache::WeakKeyDict{Thunk, Any}` - Maps from a finished `Thunk` to it's cached result, often a DRef
- `running::Set{Thunk}` - The set of currently-running `Thunk`s
- `running_on::Dict{Thunk,OSProc}` - Map from `Thunk` to the OS process executing it
- `thunk_dict::Dict{Int, WeakThunk}` - Maps from thunk IDs to a `Thunk`
- `node_order::Any` - Function that returns the order of a thunk
- `worker_pressure::Dict{Int,Dict{Processor,UInt64}}` - Cache of worker pressure
- `worker_capacity::Dict{Int,Dict{Processor,UInt64}}` - Maps from worker ID to capacity
- `worker_loadavg::Dict{Int,NTuple{3,Float64}}` - Worker load average
- `worker_chans::Dict{Int, Tuple{RemoteChannel,RemoteChannel}}` - Communication channels between the scheduler and each worker
- `procs_cache_list::Base.RefValue{Union{ProcessorCacheEntry,Nothing}}` - Cached linked list of processors ready to be used
- `function_cost_cache::Dict{Signature,UInt64}` - Cache of estimated CPU time required to compute the given signature
- `transfer_rate::Ref{UInt64}` - Estimate of the network transfer rate in bytes per second
- `halt::Base.Event` - Event indicating that the scheduler is halting
- `lock::ReentrantLock` - Lock around operations which modify the state
- `futures::Dict{Thunk, Vector{ThunkFuture}}` - Futures registered for waiting on the result of a thunk.
- `errored::WeakKeyDict{Thunk,Bool}` - Indicates if a thunk's result is an error.
- `chan::RemoteChannel{Channel{Any}}` - Channel for receiving completed thunks.
"""
struct ComputeState
    uid::UInt64
    waiting::OneToMany
    waiting_data::Dict{Union{Thunk,Chunk},Set{Thunk}}
    ready::Vector{Thunk}
    cache::WeakKeyDict{Thunk, Any}
    running::Set{Thunk}
    running_on::Dict{Thunk,OSProc}
    thunk_dict::Dict{Int, WeakThunk}
    node_order::Any
    worker_pressure::Dict{Int,Dict{Processor,UInt64}}
    worker_capacity::Dict{Int,Dict{Processor,UInt64}}
    worker_loadavg::Dict{Int,NTuple{3,Float64}}
    worker_chans::Dict{Int, Tuple{RemoteChannel,RemoteChannel}}
    procs_cache_list::Base.RefValue{Union{ProcessorCacheEntry,Nothing}}
    function_cost_cache::Dict{Signature,UInt64}
    transfer_rate::Ref{UInt64}
    halt::Base.Event
    lock::ReentrantLock
    futures::Dict{Thunk, Vector{ThunkFuture}}
    errored::WeakKeyDict{Thunk,Bool}
    chan::RemoteChannel{Channel{Any}}
end

function start_state(deps::Dict, node_order, chan)
    state = ComputeState(rand(UInt64),
                         OneToMany(),
                         deps,
                         Vector{Thunk}(undef, 0),
                         Dict{Thunk, Any}(),
                         Set{Thunk}(),
                         Dict{Thunk,OSProc}(),
                         Dict{Int, WeakThunk}(),
                         node_order,
                         Dict{Int,Dict{Type,UInt64}}(),
                         Dict{Int,Dict{Type,UInt64}}(),
                         Dict{Int,NTuple{3,Float64}}(),
                         Dict{Int, Tuple{RemoteChannel,RemoteChannel}}(),
                         Ref{Union{ProcessorCacheEntry,Nothing}}(nothing),
                         Dict{Signature,UInt64}(),
                         Ref{UInt64}(1_000_000),
                         Base.Event(),
                         ReentrantLock(),
                         Dict{Thunk, Vector{ThunkFuture}}(),
                         WeakKeyDict{Thunk,Bool}(),
                         chan)

    for k in sort(collect(keys(deps)), by=node_order)
        if istask(k)
            waiting = Set{Thunk}(Iterators.filter(istask, inputs(k)))
            if isempty(waiting)
                push!(state.ready, k)
            else
                state.waiting[k] = waiting
            end
        end
    end
    state
end

"""
    SchedulerOptions

Stores DAG-global options to be passed to the Dagger.Sch scheduler.

# Arguments
- `single::Int=0`: Force all work onto worker with specified id. `0` disables
this option.
- `proclist=nothing`: Force scheduler to use one or more processors that are
instances/subtypes of a contained type. Alternatively, a function can be
supplied, and the function will be called with a processor as the sole
argument and should return a `Bool` result to indicate whether or not to use
the given processor. `nothing` enables all default processors.
- `allow_errors::Bool=true`: Allow thunks to error without affecting
non-dependent thunks.
- `checkpoint=nothing`: If not `nothing`, uses the provided function to save
the final result of the current scheduler invocation to persistent storage, for
later retrieval by `restore`.
- `restore=nothing`: If not `nothing`, uses the provided function to return the
(cached) final result of the current scheduler invocation, were it to execute.
If this returns a `Chunk`, all thunks will be skipped, and the `Chunk` will be
returned.  If `nothing` is returned, restoring is skipped, and the scheduler
will execute as usual. If this function throws an error, restoring will be
skipped, and the error will be displayed.
- `round_robin::Bool=false`: Whether to schedule in round-robin mode, which
spreads load instead of the default behavior of filling processors to capacity.
"""
Base.@kwdef struct SchedulerOptions
    single::Int = 0
    proclist = nothing
    allow_errors::Bool = false
    checkpoint = nothing
    restore = nothing
    round_robin::Bool = false
end

"""
    ThunkOptions

Stores Thunk-local options to be passed to the Dagger.Sch scheduler.

# Arguments
- `single::Int=0`: Force thunk onto worker with specified id. `0` disables this
option.
- `proclist=nothing`: Force thunk to use one or more processors that are
instances/subtypes of a contained type. Alternatively, a function can be
supplied, and the function will be called with a processor as the sole
argument and should return a `Bool` result to indicate whether or not to use
the given processor. `nothing` enables all default processors.
- `procutil::Dict{Type,Any}=Dict{Type,Any}()`: Indicates the maximum expected
processor utilization for this thunk. Each keypair maps a processor type to
the utilization, where the value can be a real (approx. the number of processors
of this type utilized), or `MaxUtilization()` (utilizes all processors of this
type). By default, the scheduler assumes that this thunk only uses one
processor.
- `allow_errors::Bool=true`: Allow this thunk to error without affecting
non-dependent thunks.
- `checkpoint=nothing`: If not `nothing`, uses the provided function to save
the result of the thunk to persistent storage, for later retrieval by
`restore`.
- `restore=nothing`: If not `nothing`, uses the provided function to return the
(cached) result of this thunk, were it to execute.  If this returns a `Chunk`,
this thunk will be skipped, and its result will be set to the `Chunk`.  If
`nothing` is returned, restoring is skipped, and the thunk will execute as
usual. If this function throws an error, restoring will be skipped, and the
error will be displayed.
"""
Base.@kwdef struct ThunkOptions
    single::Int = 0
    proclist = nothing
    procutil::Dict{Type,Any} = Dict{Type,Any}()
    allow_errors::Bool = false
    checkpoint = nothing
    restore = nothing
end

# Eager scheduling
include("eager.jl")

"""
    Base.merge(sopts::SchedulerOptions, topts::ThunkOptions) -> ThunkOptions

Combine `SchedulerOptions` and `ThunkOptions` into a new `ThunkOptions`.
"""
function Base.merge(sopts::SchedulerOptions, topts::ThunkOptions)
    single = topts.single != 0 ? topts.single : sopts.single
    allow_errors = sopts.allow_errors || topts.allow_errors
    proclist = topts.proclist !== nothing ? topts.proclist : sopts.proclist
    ThunkOptions(single, proclist, topts.procutil, allow_errors, topts.checkpoint, topts.restore)
end
Base.merge(sopts::SchedulerOptions, ::Nothing) =
    ThunkOptions(sopts.single, sopts.proclist, Dict{Type,Any}())

function isrestricted(task::Thunk, proc::OSProc)
    if (task.options !== nothing) && (task.options.single != 0) &&
       (task.options.single != proc.pid)
        return true
    end
    return false
end

function cleanup(ctx)
end

const WORKER_MONITOR_LOCK = Threads.ReentrantLock()
const WORKER_MONITOR_TASKS = Dict{Int,Task}()
const WORKER_MONITOR_CHANS = Dict{Int,Dict{UInt64,RemoteChannel}}()
function init_proc(state, p, log_sink)
    ctx = Context(Int[]; log_sink)
    timespan_start(ctx, :init_proc, p.pid, 0)
    # Initialize pressure and capacity
    gproc = OSProc(p.pid)
    lock(state.lock) do
        state.worker_pressure[p.pid] = Dict{Processor,UInt64}()
        #state.worker_capacity[p.pid] = Dict{Processor,UInt64}()
        state.worker_loadavg[p.pid] = (0.0, 0.0, 0.0)
        for proc in get_processors(gproc)
            state.worker_pressure[p.pid][proc] = 0
            #state.worker_capacity[p.pid][proc] = capacity(gproc, proc) * UInt64(1e9)
        end
        state.worker_pressure[p.pid][gproc] = 0
        #state.worker_capacity[p.pid][OSProc] = 0
    end
    #=
    cap = remotecall(capacity, p.pid)
    @async begin
        cap = fetch(cap) * UInt64(1e9)
        lock(state.lock) do
            state.worker_capacity[p.pid] = cap
        end
    end
    =#
    lock(WORKER_MONITOR_LOCK) do
        wid = p.pid
        if !haskey(WORKER_MONITOR_TASKS, wid)
            t = @async begin
                try
                    # Wait until this connection is terminated
                    remotecall_fetch(sleep, wid, typemax(UInt64))
                catch err
                    if err isa ProcessExitedException
                        lock(WORKER_MONITOR_LOCK) do
                            d = WORKER_MONITOR_CHANS[wid]
                            for uid in keys(d)
                                put!(d[uid], (wid, OSProc(wid), nothing, (ProcessExitedException(wid), nothing)))
                            end
                            empty!(d)
                            delete!(WORKER_MONITOR_CHANS, wid)
                        end
                    end
                end
            end
            WORKER_MONITOR_TASKS[wid] = t
            WORKER_MONITOR_CHANS[wid] = Dict{UInt64,RemoteChannel}()
        end
        WORKER_MONITOR_CHANS[wid][state.uid] = state.chan
    end

    # Setup worker-to-scheduler channels
    inp_chan = RemoteChannel(p.pid)
    out_chan = RemoteChannel(p.pid)
    lock(state.lock) do
        state.worker_chans[p.pid] = (inp_chan, out_chan)
    end

    # Setup dynamic listener
    dynamic_listener!(ctx, state, p.pid)

    timespan_finish(ctx, :init_proc, p.pid, 0)
end
function _cleanup_proc(uid, log_sink)
    empty!(CHUNK_CACHE) # FIXME: Should be keyed on uid!
end
function cleanup_proc(state, p, log_sink)
    ctx = Context(Int[]; log_sink)
    timespan_start(ctx, :cleanup_proc, p.pid, 0)
    lock(WORKER_MONITOR_LOCK) do
        wid = p.pid
        if haskey(WORKER_MONITOR_CHANS, wid)
            delete!(WORKER_MONITOR_CHANS[wid], state.uid)
            remote_do(_cleanup_proc, wid, state.uid, log_sink)
        end
    end
    timespan_finish(ctx, :cleanup_proc, p.pid, 0)
end

"Process-local condition variable (and lock) indicating task completion."
const TASK_SYNC = Threads.Condition()

"Process-local set of running task IDs."
const TASKS_RUNNING = Set{Int}()

"Process-local dictionary tracking per-processor total utilization."
const PROC_UTILIZATION = Dict{UInt64,Dict{Processor,Ref{UInt64}}}()

# TODO: "Process-local count of actively-executing Dagger tasks per processor type."

"""
    MaxUtilization

Indicates a thunk that uses all processors of a given type.
"""
struct MaxUtilization end

function compute_dag(ctx, d::Thunk; options=SchedulerOptions())
    if options === nothing
        options = SchedulerOptions()
    end
    ctx.options = options
    if options.restore !== nothing
        try
            result = options.restore()
            if result isa Chunk
                return result
            elseif result !== nothing
                throw(ArgumentError("Invalid restore return type: $(typeof(result))"))
            end
        catch err
            report_catch_error(err, "Scheduler restore failed")
        end
    end

    chan = RemoteChannel(()->Channel(1024))
    deps = dependents(d)
    ord = order(d, noffspring(deps))

    node_order = x -> -get(ord, x, 0)
    state = start_state(deps, node_order, chan)

    master = OSProc(myid())

    timespan_start(ctx, :scheduler_init, 0, master)
    try
        scheduler_init(ctx, state, d, options, deps)
    finally
        timespan_finish(ctx, :scheduler_init, 0, master)
    end

    value, errored = try
        scheduler_run(ctx, state, d, options)
    finally
        # Always try to tear down the scheduler
        timespan_start(ctx, :scheduler_exit, 0, master)
        try
            scheduler_exit(ctx, state, options)
        catch err
            @error "Error when tearing down scheduler" exception=(err,catch_backtrace())
        finally
            timespan_finish(ctx, :scheduler_exit, 0, master)
        end
    end

    if errored
        throw(value)
    end
    return value
end

function scheduler_init(ctx, state::ComputeState, d::Thunk, options, deps)
    # setup thunk_dict mappings
    for node in filter(istask, keys(deps))
        state.thunk_dict[node.id] = WeakThunk(node)
        for dep in deps[node]
            state.thunk_dict[dep.id] = WeakThunk(dep)
        end
    end

    # Initialize workers
    @sync for p in procs_to_use(ctx)
        @async begin
            try
                init_proc(state, p, ctx.log_sink)
            catch err
                @error "Error initializing worker $p" exception=(err,catch_backtrace())
                remove_dead_proc!(ctx, state, p)
            end
        end
    end

    # Listen for new workers
    @async begin
        try
            monitor_procs_changed!(ctx, state)
        catch err
            @error "Error assigning workers" exception=(err,catch_backtrace())
        end
    end
end

function scheduler_run(ctx, state::ComputeState, d::Thunk, options)
    safepoint(state)

    # Loop while we still have thunks to execute
    while !isempty(state.ready) || !isempty(state.running)
        if !isempty(state.ready)
            # Nothing running, so schedule up to N thunks, 1 per N workers
            schedule!(ctx, state)
        end

        check_integrity(ctx)

        isempty(state.running) && continue
        timespan_start(ctx, :take, 0, 0)
        chan_value = take!(state.chan) # get result of completed thunk
        timespan_finish(ctx, :take, 0, 0)
        if chan_value isa RescheduleSignal
            continue
        end
        pid, proc, thunk_id, (res, metadata) = chan_value
        gproc = OSProc(pid)
        safepoint(state)
        lock(state.lock) do
            thunk_failed = false
            if res isa Exception
                if unwrap_nested_exception(res) isa ProcessExitedException
                    @warn "Worker $(pid) died, rescheduling work"

                    # Remove dead worker from procs list
                    timespan_start(ctx, :remove_procs, 0, 0)
                    remove_dead_proc!(ctx, state, gproc)
                    timespan_finish(ctx, :remove_procs, 0, 0)

                    timespan_start(ctx, :handle_fault, 0, 0)
                    handle_fault(ctx, state, gproc)
                    timespan_finish(ctx, :handle_fault, 0, 0)
                    return # effectively `continue`
                else
                    if ctx.options.allow_errors || unwrap_weak_checked(state.thunk_dict[thunk_id]).options.allow_errors
                        thunk_failed = true
                    else
                        throw(res)
                    end
                end
            end
            node = unwrap_weak_checked(state.thunk_dict[thunk_id])
            if metadata !== nothing
                state.worker_pressure[pid][proc] = metadata.pressure
                state.worker_loadavg[pid] = metadata.loadavg
                sig = signature(node, state)
                state.function_cost_cache[sig] = (metadata.threadtime + get(state.function_cost_cache, sig, 0)) ÷ 2
                if metadata.transfer_rate !== nothing
                    state.transfer_rate[] = (state.transfer_rate[] + metadata.transfer_rate) ÷ 2
                end
            end
            state.cache[node] = res
            state.errored[node] = thunk_failed
            if node.options !== nothing && node.options.checkpoint !== nothing
                try
                    node.options.checkpoint(node, res)
                catch err
                    report_catch_error(err, "Thunk checkpoint failed")
                end
            end

            timespan_start(ctx, :finish, thunk_id, (;thunk_id))
            finish_task!(ctx, state, node, thunk_failed)
            timespan_finish(ctx, :finish, thunk_id, (;thunk_id))
        end

        safepoint(state)
    end

    # Final value is ready
    value = state.cache[d]
    errored = get(state.errored, d, false)
    if !errored
        if options.checkpoint !== nothing
            try
                options.checkpoint(value)
            catch err
                report_catch_error(err, "Scheduler checkpoint failed")
            end
        end
    end
    return value, errored
end
function scheduler_exit(ctx, state::ComputeState, options)
    close(state.chan)
    notify(state.halt)
    @sync for p in procs_to_use(ctx)
        @async cleanup_proc(state, p, ctx.log_sink)
    end

    # Let the context procs handler clean itself up
    lock(ctx.proc_notify) do
        notify(ctx.proc_notify)
    end
end

function procs_to_use(ctx, options=ctx.options)
    return if options.single !== 0
        @assert options.single in vcat(1, workers()) "Sch option `single` must specify an active worker ID."
        OSProc[OSProc(options.single)]
    else
        procs(ctx)
    end
end

check_integrity(ctx) = @assert !isempty(procs_to_use(ctx)) "No suitable workers available in context."

struct SchedulingException <: Exception
    reason::String
end
function Base.show(io::IO, se::SchedulingException)
    print(io, "SchedulingException ($(se.reason))")
end

const CHUNK_CACHE = Dict{Chunk,Dict{Processor,Any}}()

function schedule!(ctx, state, procs=procs_to_use(ctx))
    lock(state.lock) do
        safepoint(state)
        @assert length(procs) > 0

        populate_processor_cache_list!(state, procs)

        # Schedule tasks
        to_fire = Dict{Tuple{OSProc,<:Processor},Vector{Tuple{Thunk,<:Any}}}()
        failed_scheduling = Thunk[]

        # Select a new task and get its options
        task = nothing
        @label pop_task
        if task !== nothing
            timespan_finish(ctx, :schedule, task.id, (;thunk_id=task.id))
        end
        if isempty(state.ready)
            @goto fire_tasks
        end
        task = pop!(state.ready)
        timespan_start(ctx, :schedule, task.id, (;thunk_id=task.id))
        @assert !haskey(state.cache, task)
        opts = merge(ctx.options, task.options)
        sig = signature(task, state)

        # Calculate scope
        scope = if task.f isa Chunk
            task.f.scope
        else
            AnyScope()
        end
        for input in task.inputs
            input = unwrap_weak_checked(input)
            chunk = if istask(input)
                state.cache[input]
            elseif input isa Chunk
                input
            else
                nothing
            end
            chunk isa Chunk || continue
            scope = constrain(scope, chunk.scope)
            if scope isa Dagger.InvalidScope
                ex = SchedulingException("Scopes are not compatible: $(scope.x), $(scope.y)")
                state.cache[task] = ex
                state.errored[task] = true
                set_failed!(state, task)
                @goto pop_task
            end
        end

        fallback_threshold = 1024 # TODO: Parameterize this threshold
        if length(procs) > fallback_threshold
            @goto fallback
        end
        local_procs = unique(vcat([Dagger.get_processors(gp) for gp in procs]...))
        if length(local_procs) > fallback_threshold
            @goto fallback
        end

        local_procs, costs = estimate_task_costs(state, local_procs, task)
        scheduled = false

        # Move our corresponding ThreadProc to be the last considered
        if length(local_procs) > 1
            sch_threadproc = Dagger.ThreadProc(myid(), Threads.threadid())
            sch_thread_idx = findfirst(proc->proc==sch_threadproc, local_procs)
            if sch_thread_idx !== nothing
                deleteat!(local_procs, sch_thread_idx)
                push!(local_procs, sch_threadproc)
            end
        end

        for proc in local_procs
            gproc = get_parent(proc)
            if can_use_proc(task, gproc, proc, opts, scope)
                has_cap, cap, extra_util = has_capacity(state, proc, gproc.pid, opts.procutil, sig)
                if has_cap
                    # Schedule task onto proc
                    extra_util = extra_util isa MaxUtilization ? cap : extra_util
                    push!(get!(()->Vector{Tuple{Thunk,<:Any}}(), to_fire, (gproc, proc)), (task, extra_util))
                    state.worker_pressure[gproc.pid][proc] += extra_util
                    @goto pop_task
                end
            end
        end
        state.cache[task] = SchedulingException("No processors available, try making proclist more liberal")
        state.errored[task] = true
        set_failed!(state, task)
        @goto pop_task

        # Fast fallback algorithm, used when the smarter cost model algorithm
        # would be too expensive
        @label fallback
        selected_entry = nothing
        entry = state.procs_cache_list[]
        cap, extra_util = nothing, nothing
        procs_found = false
        # N.B. if we only have one processor, we need to select it now
        if can_use_proc(task, entry.gproc, entry.proc, opts, scope)
            has_cap, cap, extra_util = has_capacity(state, entry.proc, entry.gproc.pid, opts.procutil, sig)
            if has_cap
                selected_entry = entry
            else
                procs_found = true
                entry = entry.next
            end
        else
            entry = entry.next
        end
        while selected_entry === nothing
            if entry === state.procs_cache_list[]
                # Exhausted all procs
                if procs_found
                    push!(failed_scheduling, task)
                else
                    state.cache[task] = SchedulingException("No processors available, try making proclist more liberal")
                    state.errored[task] = true
                    set_failed!(state, task)
                end
                @goto pop_task
            end

            if can_use_proc(task, entry.gproc, entry.proc, opts, scope)
                has_cap, cap, extra_util = has_capacity(state, entry.proc, entry.gproc.pid, opts.procutil, sig)
                if has_cap
                    # Select this processor
                    selected_entry = entry
                else
                    # We could have selected it otherwise
                    procs_found = true
                    entry = entry.next
                end
            else
                # Try next processor
                entry = entry.next
            end
        end
        @assert selected_entry !== nothing

        # Schedule task onto proc
        gproc, proc = entry.gproc, entry.proc
        extra_util = extra_util isa MaxUtilization ? cap : extra_util
        push!(get!(()->Vector{Tuple{Thunk,<:Any}}(), to_fire, (gproc, proc)), (task, extra_util))

        # Proceed to next entry to spread work
        if !ctx.options.round_robin
            @warn "Round-robin mode is always on"
        end
        state.procs_cache_list[] = state.procs_cache_list[].next
        @goto pop_task

        # Fire all newly-scheduled tasks
        @label fire_tasks
        for gpp in keys(to_fire)
            fire_tasks!(ctx, to_fire[gpp], gpp, state)
        end

        append!(state.ready, failed_scheduling)
    end
end

"""
Monitors for workers being added/removed to/from `ctx`, sets up or tears down
per-worker state, and notifies the scheduler so that work can be reassigned.
"""
function monitor_procs_changed!(ctx, state)
    # Load current set of procs
    old_ps = procs_to_use(ctx)

    while !state.halt.set
        # Wait for the notification that procs have changed
        lock(ctx.proc_notify) do
            wait(ctx.proc_notify)
        end

        timespan_start(ctx, :assign_procs, 0, 0)

        # Load new set of procs
        new_ps = procs_to_use(ctx)

        # Initialize new procs
        diffps = setdiff(new_ps, old_ps)
        for p in diffps
            init_proc(state, p, ctx.log_sink)

            # Empty the processor cache list and force reschedule
            lock(state.lock) do
                state.procs_cache_list[] = nothing
            end
            put!(state.chan, RescheduleSignal())
        end

        # Cleanup removed procs
        diffps = setdiff(old_ps, new_ps)
        for p in diffps
            cleanup_proc(state, p, ctx.log_sink)

            # Empty the processor cache list
            lock(state.lock) do
                state.procs_cache_list[] = nothing
            end
        end

        timespan_finish(ctx, :assign_procs, 0, 0)
        old_ps = new_ps
    end
end

function remove_dead_proc!(ctx, state, proc, options=ctx.options)
    @assert options.single !== proc.pid "Single worker failed, cannot continue."
    rmprocs!(ctx, [proc])
    delete!(state.worker_pressure, proc.pid)
    #delete!(state.worker_capacity, proc.pid)
    delete!(state.worker_loadavg, proc.pid)
    delete!(state.worker_chans, proc.pid)
    state.procs_cache_list[] = nothing
end

function pop_with_affinity!(ctx, tasks, proc)
    # TODO: use the size
    parent_affinity_procs = Vector(undef, length(tasks))
    # parent_affinity_sizes = Vector(undef, length(tasks))
    for i=length(tasks):-1:1
        t = tasks[i]
        aff = affinity(t)
        aff_procs = first.(aff)
        if proc in aff_procs
            if !isrestricted(t,proc)
                deleteat!(tasks, i)
                return t
            end
        end
        parent_affinity_procs[i] = aff_procs
    end
    for i=length(tasks):-1:1
        # use up tasks without affinities
        # let the procs with the respective affinities pick up
        # other tasks
        aff_procs = parent_affinity_procs[i]
        if isempty(aff_procs)
            t = tasks[i]
            if !isrestricted(t,proc)
                deleteat!(tasks, i)
                return t
            end
        end
        if all(!(p in aff_procs) for p in procs(ctx))
            # no proc is ever going to ask for it
            t = tasks[i]
            if !isrestricted(t,proc)
                deleteat!(tasks, i)
                return t
            end
        end
    end
    return nothing
end

function finish_task!(ctx, state, node, thunk_failed)
    pop!(state.running, node)
    delete!(state.running_on, node)
    if thunk_failed
        set_failed!(state, node)
    end
    if node.cache
        node.cache_ref = state.cache[node]
    end
    schedule_dependents!(state, node, thunk_failed)
    fill_registered_futures!(state, node, thunk_failed)

    to_evict = cleanup_inputs!(state, node)
    if node.f isa Chunk
        # FIXME: Check the graph for matching chunks
        push!(to_evict, node.f)
    end
    if haskey(state.waiting_data, node) && isempty(state.waiting_data[node])
        delete!(state.waiting_data, node)
    end
    evict_all_chunks!(ctx, to_evict)
end

function evict_all_chunks!(ctx, to_evict)
    if !isempty(to_evict)
        @sync for w in map(p->p.pid, procs_to_use(ctx))
            @async remote_do(evict_chunks!, w, ctx.log_sink, to_evict)
        end
    end
end
function evict_chunks!(log_sink, chunks::Set{Chunk})
    # Need worker id or else Context might use Processors which user does not want us to use.
    # In particular workers which have not yet run using Dagger will cause the call below to throw an exception
    ctx = Context([myid()];log_sink)
    for chunk in chunks
        timespan_start(ctx, :evict, myid(), (;data=chunk))
        haskey(CHUNK_CACHE, chunk) && delete!(CHUNK_CACHE, chunk)
        timespan_finish(ctx, :evict, myid(), (;data=chunk))
    end
    nothing
end

fire_task!(ctx, thunk::Thunk, p, state; util=10^9) =
    fire_task!(ctx, (thunk, util), p, state)
fire_task!(ctx, (thunk, util)::Tuple{Thunk,<:Any}, p, state) =
    fire_tasks!(ctx, [(thunk, util)], p, state)
function fire_tasks!(ctx, thunks::Vector{<:Tuple}, (gproc, proc), state)
    to_send = []
    for (thunk, util) in thunks
        push!(state.running, thunk)
        state.running_on[thunk] = gproc
        if thunk.cache && thunk.cache_ref !== nothing
            # the result might be already cached
            data = unrelease(thunk.cache_ref) # ask worker to keep the data around
                                              # till this compute cycle frees it
            if data !== nothing
                # cache hit
                state.cache[thunk] = data
                thunk_failed = get(state.errored, thunk, false)
                finish_task!(ctx, state, thunk, thunk_failed)
                continue
            else
                # cache miss
                thunk.cache_ref = nothing
            end
        end
        if thunk.options !== nothing && thunk.options.restore !== nothing
            try
                result = thunk.options.restore(thunk)
                if result isa Chunk
                    state.cache[thunk] = result
                    state.errored[thunk] = false
                    finish_task!(ctx, state, thunk, false)
                    continue
                elseif result !== nothing
                    throw(ArgumentError("Invalid restore return type: $(typeof(result))"))
                end
            catch err
                report_catch_error(err, "Thunk restore failed")
            end
        end

        ids = Int[0]
        data = Any[thunk.f]
        for (idx, x) in enumerate(thunk.inputs)
            x = unwrap_weak_checked(x)
            push!(ids, istask(x) ? x.id : -idx)
            push!(data, istask(x) ? state.cache[x] : x)
        end
        toptions = thunk.options !== nothing ? thunk.options : ThunkOptions()
        options = merge(ctx.options, toptions)
        propagated = get_propagated_options(thunk)
        @assert (options.single == 0) || (gproc.pid == options.single)
        # TODO: Set `sch_handle.tid.ref` to the right `DRef`
        sch_handle = SchedulerHandle(ThunkID(thunk.id, nothing), state.worker_chans[gproc.pid]...)

        # TODO: De-dup common fields (log_sink, uid, etc.)
        push!(to_send, Any[util, thunk.id, fn_type(thunk.f), data, thunk.get_result,
                           thunk.persist, thunk.cache, thunk.meta, options,
                           propagated, ids,
                           (log_sink=ctx.log_sink, profile=ctx.profile),
                           sch_handle, state.uid])
    end
    # N.B. We don't batch these because we might get a deserialization
    # error due to something not being defined on the worker, and then we don't
    # know which task failed.
    tasks = Task[]
    for ts in to_send
        # TODO: errormonitor
        @async begin
            timespan_start(ctx, :fire, gproc.pid, 0)
            try
                remotecall_wait(do_tasks, gproc.pid, proc, state.chan, [ts])
            catch err
                bt = catch_backtrace()
                put!(state.chan, (gproc.pid, proc, ts[2], (CapturedException(err, bt), nothing)))
            finally
                timespan_finish(ctx, :fire, gproc.pid, 0)
            end
        end
    end
end

"""
    do_tasks(to_proc, chan, tasks)

Executes a batch of tasks on `to_proc`.
"""
function do_tasks(to_proc, chan, tasks)
    for task in tasks
        should_launch = lock(TASK_SYNC) do
            # Already running; don't try to re-launch
            if !(task[2] in TASKS_RUNNING)
                push!(TASKS_RUNNING, task[2])
                true
            else
                false
            end
        end
        should_launch || continue
        @async begin
            try
                result = do_task(to_proc, task)
                put!(chan, (myid(), to_proc, task[2], result))
            catch ex
                bt = catch_backtrace()
                put!(chan, (myid(), to_proc, task[2], (CapturedException(ex, bt), nothing)))
            end
        end
    end
end
"Executes a single task on `to_proc`."
function do_task(to_proc, comm)
    extra_util, thunk_id, Tf, data, send_result, persist, cache, meta, options, propagated, ids, ctx_vars, sch_handle, uid = comm
    ctx = Context(Processor[]; log_sink=ctx_vars.log_sink, profile=ctx_vars.profile)

    from_proc = OSProc()
    Tdata = Any[]
    for x in data
        push!(Tdata, x isa Chunk ? chunktype(x) : x)
    end
    f = isdefined(Tf, :instance) ? Tf.instance : nothing
    f_chunk = first(data)
    scope = f_chunk isa Chunk ? f_chunk.scope : AnyScope()

    # Fetch inputs
    transfer_time = Threads.Atomic{UInt64}(0)
    transfer_size = Threads.Atomic{UInt64}(0)
    _data, _ids = if meta
        (Any[first(data)], Int[first(ids)]) # always fetch function
    else
        (data, ids)
    end
    fetch_tasks = map(Iterators.zip(_data,_ids)) do (x, id)
        @async begin
            timespan_start(ctx, :move, (;thunk_id, id), (;f, id, data=x))
            x = if x isa Chunk
                if haskey(CHUNK_CACHE, x)
                    get!(CHUNK_CACHE[x], to_proc) do
                        # TODO: Choose "closest" processor of same type first
                        some_proc = first(keys(CHUNK_CACHE[x]))
                        some_x = CHUNK_CACHE[x][some_proc]
                        move(some_proc, to_proc, some_x)
                    end
                else
                    time_start = time_ns()
                    _x = move(to_proc, x)
                    time_finish = time_ns()
                    if x.handle.size !== nothing
                        Threads.atomic_add!(transfer_time, time_finish - time_start)
                        Threads.atomic_add!(transfer_size, x.handle.size)
                    end
                    CHUNK_CACHE[x] = Dict{Processor,Any}()
                    CHUNK_CACHE[x][to_proc] = _x
                    _x
                end
            else
                move(to_proc, x)
            end
            timespan_finish(ctx, :move, (;thunk_id, id), (;f, id, data=x); tasks=[Base.current_task()])
            return x
        end
    end
    fetched = Any[]
    for task in fetch_tasks
        push!(fetched, fetch_report(task))
    end
    if meta
        append!(fetched, data[2:end])
    end
    f = popfirst!(fetched)
    @assert !(f isa Chunk) "Failed to unwrap thunk function"

    # Check if we'll go over capacity from running this thunk
    real_util = lock(TASK_SYNC) do
        AT = get!(()->Dict{Processor,Ref{UInt64}}(), PROC_UTILIZATION, uid)
        get!(()->Ref{UInt64}(UInt64(0)), AT, to_proc)
    end
    #cap = UInt64(capacity(OSProc(), typeof(to_proc))) * UInt64(1e9)
    cap = typemax(UInt64)

    # Wait for a free processor if necessary
    while true
        lock(TASK_SYNC)
        if ((extra_util isa MaxUtilization) && (real_util[] > 0)) ||
           ((extra_util isa Real) && (extra_util + real_util[] > cap))
            # Fully subscribed, wait and re-check
            @debug "($(myid())) $f ($thunk_id) Waiting for free $(typeof(to_proc)): $extra_util | $(real_util[])/$cap"
            wait(TASK_SYNC)
            unlock(TASK_SYNC)
        else
            # Under-subscribed, calculate extra utilization and execute thunk
            @debug "($(myid())) $f ($thunk_id) Using available $to_proc: $extra_util | $(real_util[])/$cap"
            extra_util = if extra_util isa MaxUtilization
                count(c->typeof(c)===typeof(to_proc), children(from_proc))
            else
                extra_util
            end
            real_util[] += extra_util
            unlock(TASK_SYNC)
            break
        end
    end

    timespan_start(ctx, :compute, thunk_id, (;f, to_proc))
    res = nothing
    threadtime_start = cputhreadtime()
    result_meta = try
        # Set TLS variables
        Dagger.set_tls!((
            sch_uid=uid,
            sch_handle=sch_handle,
            processor=to_proc,
            utilization=extra_util,
        ))

        res = Dagger.with_options(propagated) do
            # Execute
            execute!(to_proc, f, fetched...)
        end

        # Construct result
        send_result || meta ? res : tochunk(res, to_proc; persist=persist, cache=persist ? true : cache)
    catch ex
        bt = catch_backtrace()
        RemoteException(myid(), CapturedException(ex, bt))
    end
    threadtime = cputhreadtime() - threadtime_start
    timespan_finish(ctx, :compute, thunk_id, (;f, to_proc); tasks=Dagger.prof_tasks_take!(thunk_id))
    lock(TASK_SYNC) do
        real_util[] -= extra_util
        pop!(TASKS_RUNNING, thunk_id)
        notify(TASK_SYNC)
    end
    @debug "($(myid())) $f ($thunk_id) Releasing $(typeof(to_proc)): $extra_util | $(real_util[])/$cap"
    metadata = (
        pressure=real_util[],
        loadavg=((Sys.loadavg()...,) ./ Sys.CPU_THREADS),
        threadtime=threadtime,
        transfer_rate=(transfer_size[] > 0 && transfer_time[] > 0) ? round(UInt64, transfer_size[] / (transfer_time[] / 10^9)) : nothing,
    )
    (result_meta, metadata)
end

end # module Sch


text:
# # [Splitting a string into words and counting them in parallel](@id parallel-word-count)
#
# We start from the parallel algorithm presented in Guy Steele's 2009
# ICFP talk [(video)](https://vimeo.com/6624203).  It splits a
# space-separated string into list of strings (words).  The repeating
# theme in the talk was to build "singleton solutions" and then merge
# them together using an associative function.  We will follow this
# guideline and slightly extend the algorithm.
#
# It is highly recommended to just watch the talk for understanding
# the algorithm.  However, we briefly describe how it works.
#
# When a certain contiguous region of a string is processed, we either
# already have seen at least one space or not.  These two states are
# tracked using following two types.  If there is no space so far, we
# only have a chunk of a possibly larger word (see example below):

struct Chunk
    s::String
end

# If there are one or more spaces, (possibly zero) words that are
# already determined and left/right "chunks" have to be tracked
# separately:

struct Segment
    l::String
    A::Vector{String}
    r::String
end

# Here is an example taken from the talk:
#
# ```
# Segment("Here", ["is", "a"], "")
#    |
#    |       Segment("lian", [], "string")
#  __|_____            _|______
# |        |          |        |
# Here is a sesquipedalian string of words
#           |________|          |________|
#    Chunk("sesquipeda")        Segment("g", ["of"], "words")
# ```

# We then need a way to merge two results which can independently in
# one of the above two states.

⊕(x::Chunk, y::Chunk) = Chunk(x.s * y.s)
⊕(x::Chunk, y::Segment) = Segment(x.s * y.l, y.A, y.r)
⊕(x::Segment, y::Chunk) = Segment(x.l, x.A, x.r * y.s)
⊕(x::Segment, y::Segment) =
    Segment(x.l,
            append!(append!(x.A, maybewordv(x.r * y.l)), y.A),
            y.r)

maybewordv(s::String) = isempty(s) ? String[] : [s]
nothing  # hide

# Input is a sequence of `Char`s.  Each of them has to be converted
# into a "singleton solution" which can be merged with already
# aggregated (or another singleton) solution with `⊕`:

segment_or_chunk(c::Char) = c == ' ' ? Segment("", [], "") : Chunk(string(c))
nothing  # hide

# Putting them together, we get:

function collectwords(s::String)
    g = mapfoldl(segment_or_chunk, ⊕, s; init=Segment("", [], ""))
    if g isa Char
        return maybewordv(g.s)
    else
        return append!(append!(maybewordv(g.l), g.A), maybewordv(g.r))
    end
end
nothing  # hide

# Let's run a few tests covering some edge cases:

using Test
@testset begin
    @test collectwords("This is a sample") == ["This", "is", "a", "sample"]
    @test collectwords(" Here is another sample ") == ["Here", "is", "another", "sample"]
    @test collectwords("JustOneWord") == ["JustOneWord"]
    @test collectwords(" ") == []
    @test collectwords("") == []
end
nothing  # hide

# ## String-splitting transducer
#
# Let's make it re-usable by packaging it into a transducer.

# Rather than accumulating words into a vector, we are going to write
# a transducer that "emits" a word as soon as it is ready.  The
# downstream transducer may choose to record everything or only
# aggregate, e.g., reduced statistics.  To this end, we replace
# `Segment` in the original algorithm with

struct Vacant
    l::String
    r::String
end

# and output the words in the "middle" without accumulating it.  So,
# instead of `segment_or_chunk`, we now have:

vacant_or_chunk(c::Char) = c == ' ' ? Vacant("", "") : Chunk(string(c))
nothing  # hide

# The idea is to create a custom transducer `WordsXF` that is used as
# in
#
# ```julia
# ... |> Map(vacant_or_chunk) |> WordsXF() |> Filter(!isnothing) |> ...
# ```
#
# so that the whole transducer streams non-empty words to the
# downstream.  That is to say, the input stream is first processed by
# `vacant_or_chunk` which returns either a `Vacant` or a `Chunk`.
# This is processed by `WordsXF()` which outputs either a word (a
# `String`) or `nothing`.  We are using `Filter(!isnothing)` in the
# downstream to simplify the definition of `WordsXF`.
#
# We define a function `extract(x::Union{Chunk,Vacant},
# y::Union{Chunk,Vacant}) -> (output, state)`.  It is something like
# `⊕` but works with `Chunk` and `Vacant`:

extract(x::Chunk, y::Chunk) = nothing, Chunk(x.s * y.s)
extract(x::Chunk, y::Vacant) = nothing, Vacant(x.s * y.l, y.r)
extract(x::Vacant, y::Chunk) = nothing, Vacant(x.l, x.r * y.s)
extract(x::Vacant, y::Vacant) = maybeword(x.r * y.l), Vacant(x.l, y.r)

maybeword(s) = isempty(s) ? nothing : s
nothing  # hide

# Let's wrap this in a [`Transducer`](@ref).

using Transducers
using Transducers:
    @next, R_, Transducer, combine, complete, inner, next, start, unwrap, wrap, wrapping

# First, we declare a transducer type:

struct WordsXF <: Transducer end
nothing  # hide

# Since this transducer has to keep "unfinished" words as its own
# private state, we use [`wrap`](@ref Transducers.wrap) inside
# [`start`](@ref Transducers.start) to prepare the state for it:

Transducers.start(rf::R_{WordsXF}, init) = wrap(rf, Chunk(""), start(inner(rf), init))
nothing  # hide

# Inside of [`next`](@ref Transducers.complete) (i.e., "loop body") we
# call `extract` defined above to combine the input
# `x::Union{Chunk,Vacant}` into `state::Union{Chunk,Vacant}`.  If
# `extract` returns a word, it is passed to the inner reducing
# function:

function Transducers.next(rf::R_{WordsXF}, acc, x)
    wrapping(rf, acc) do state, iacc
        word, state = extract(state, x)
        iacc = next(inner(rf), iacc, word)
        return state, iacc
    end
end
nothing  # hide

# At the end of a fold, [`complete`](@ref Transducers.complete) is
# called.  We can process unfinished words at this stage.  Note that
# we need to use [`combine`](@ref Transducers.combine) of the inner
# reducing function (assuming it is associative) to "prepend" a word
# to the accumulated state of the inner reducing function.

function Transducers.complete(rf::R_{WordsXF}, acc)
    state, iacc = unwrap(rf, acc)
    if state isa Vacant
        pre = @next(inner(rf), start(inner(rf), Init), maybeword(state.l))
        iacc = combine(inner(rf), pre, iacc)  # prepending `state.l`
        iacc = @next(inner(rf), iacc, maybeword(state.r))  # appending `state.r`
    else
        @assert state isa Chunk
        iacc = @next(inner(rf), iacc, maybeword(state.s))
    end
    return complete(inner(rf), iacc)
end
nothing  # hide

# That's all we need for using this transducer with sequential folds.
# For parallel reduce we need [`combine`](@ref Transducers.combine).
# It is more or less identical to `next`:

function Transducers.combine(rf::R_{WordsXF}, a, b)
    ua, ira = unwrap(rf, a)
    ub, irb = unwrap(rf, b)
    word, uc = extract(ua, ub)
    ira = @next(inner(rf), ira, word)
    irc = combine(inner(rf), ira, irb)
    return wrap(rf, uc, irc)
end
nothing  # hide

using Compat: isnothing  #src
wordsxf = opcompose(Map(vacant_or_chunk), WordsXF(), Filter(!isnothing))
nothing  # hide

# Test:

@testset begin
    @test collect(wordsxf, "This is a sample") == ["This", "is", "a", "sample"]
    @test collect(wordsxf, " Here is another sample ") == ["Here", "is", "another", "sample"]
    @test collect(wordsxf, "JustOneWord") == ["JustOneWord"]
    @test collect(wordsxf, " ") == []
    @test collect(wordsxf, "") == []
end
nothing  # hide

# ## Word-counting transducer
#
# We can pipe the resulting words into various transducers.

using MicroCollections: SingletonDict

processcount(word) = SingletonDict(word => 1)
countxf = opcompose(wordsxf, Map(processcount))

# Transducer `countxf` constructs a "singleton solution" as a
# dictionary which then accumulated with the associative reducing step
# function `mergewith!!(+)` from BangBang.jl:

using BangBang: mergewith!!
nothing  # hide

# Putting the transducer and reducing function together, we get:

countwords(s; kwargs...) =
    foldxt(mergewith!!(+), countxf, s; init = CopyInit(Dict{String,Int}()), kwargs...)
nothing  # hide

# Side note: We use [`CopyInit`](@ref) to create a fresh initial
# state for each sub-reduce to avoid overwriting mutable data between
# threads.

# Let's run some tests with different `basesize` (`length(s) /
# basesize` corresponds to number of tasks to be used):

@testset for basesize in [1, 2, 4]
    @test countwords("This is a sample", basesize=basesize) ==
        Dict("This" => 1, "is" => 1, "a" => 1, "sample" => 1)
    @test countwords(" Here is another sample ", basesize=basesize) ==
        Dict("Here" => 1, "is" => 1, "another" => 1, "sample" => 1)
    @test countwords("JustOneWord", basesize=basesize) ==
        Dict("JustOneWord" => 1)
    @test countwords(" ", basesize=basesize) == Dict()
    @test countwords("", basesize=basesize) == Dict()
    @test countwords("aaa bb aaa aaa bb bb aaa", basesize=basesize) ==
        Dict("aaa" => 4, "bb" => 3)
    @test countwords("あああ いい あああ あああ いい いい あああ", basesize=basesize) ==
        Dict("あああ" => 4, "いい" => 3)
end
nothing  # hide


text:
# Simply supported square plate with center concentrated load
module bent_strip_examples

using FinEtools
using FinEtoolsDeforLinear
using FinEtoolsFlexStructures.FESetShellT3Module: FESetShellT3
using FinEtoolsFlexStructures.FESetShellQ4Module: FESetShellQ4
using FinEtoolsFlexStructures.FEMMShellT3FFModule
using FinEtoolsFlexStructures.RotUtilModule: initial_Rfield, update_rotation_field!
using VisualStructures: plot_nodes, plot_midline, render, plot_space_box, plot_midsurface, space_aspectratio, save_to_json
using FinEtools.MeshExportModule.VTKWrite: vtkwrite

    # Generate a graphical display of resultants
function   cartesian!(csmatout::FFltMat, XYZ::FFltMat, tangents::FFltMat, fe_label::FInt)
    csmatout[:, 1] .= (1.0, 0.0, 0.0)
    csmatout[:, 2] .= (0.0, 1.0, 0.0)
    csmatout[:, 3] .= (0.0, 0.0, 1.0)using .Main.tippling_examples; Main.tippling_examples.allrun()
f
    return csmatout
end

function _execute_dsg_model(formul, n = 2, visualize = true)
    E = 1e3;
    nu = 0.3;
    thickness = 0.1
    W = 1.0
    L = 10.0;
    q = 1.0 * thickness^3;
    # analytical solution for the vertical deflection under the load
    analyt_sol=-1.5875;

    tolerance = L/n/1000
    fens, fes = T3block(L, W, 5*n, n);
    fens.xyz = xyz3(fens)

    mater = MatDeforElastIso(DeforModelRed3D, E, nu)

    sfes = FESetShellT3()
    accepttodelegate(fes, sfes)
    femm = formul.make(IntegDomain(fes, TriRule(1), thickness), mater)
    stiffness = formul.stiffness
    associategeometry! = formul.associategeometry!

    # Construct the requisite fields, geometry and displacement
    # Initialize configuration variables
    geom0 = NodalField(fens.xyz)
    u0 = NodalField(zeros(size(fens.xyz,1), 3))
    Rfield0 = initial_Rfield(fens)
    dchi = NodalField(zeros(size(fens.xyz,1), 6))

    # Apply EBC's
    l1 = selectnode(fens; box = Float64[0 0 -Inf Inf -Inf Inf], inflate = tolerance)
    for i in [1, 2, 3, 6]
        setebc!(dchi, l1, true, i)
    end
    l1 = selectnode(fens; box = Float64[L L -Inf Inf -Inf Inf], inflate = tolerance)
    for i in [1, 2, 3, 6]
        setebc!(dchi, l1, true, i)
    end
    applyebc!(dchi)
    numberdofs!(dchi);

    # Assemble the system matrix
    associategeometry!(femm, geom0)
    K = stiffness(femm, geom0, u0, Rfield0, dchi);

    # Load
    lfemm = FEMMBase(IntegDomain(fes, TriRule(3)))
    fi = ForceIntensity(FFlt[0, 0, -q, 0, 0, 0]);
    F = distribloads(lfemm, geom0, dchi, fi, 2);

    # Solve
    U = K\F
    scattersysvec!(dchi, U[:])
    @show minimum(dchi.values[:, 3])/analyt_sol*100

    # Visualization
    if visualize

    ocsys = CSys(3, 3, cartesian!)
    scalars = []
    for nc in 1:3
        fld = fieldfromintegpoints(femm, geom0, dchi, :moment, nc, outputcsys = ocsys)
        # fld = elemfieldfromintegpoints(femm, geom0, dchi, :moment, nc)
        push!(scalars, ("m$nc", fld.values))
    end
    vtkwrite("bent_strip_examples-m.vtu", fens, fes; scalars = scalars)
        # Generate a graphical display of resultants
    scalars = []
    for nc in 1:2
        fld = fieldfromintegpoints(femm, geom0, dchi, :shear, nc, outputcsys = ocsys)
        push!(scalars, ("q$nc", fld.values))
        fld = elemfieldfromintegpoints(femm, geom0, dchi, :shear, nc, outputcsys = ocsys)
        push!(scalars, ("eq$nc", fld.values))
    end
    vtkwrite("bent_strip_examples-q.vtu", fens, fes; scalars = scalars)

    scattersysvec!(dchi, (L/4)/maximum(abs.(U)).*U)
    update_rotation_field!(Rfield0, dchi)
    plots = cat(plot_space_box([[0 0 -L/2]; [L/2 L/2 L/2]]),
        #plot_nodes(fens),
        plot_midsurface(fens, fes; x = geom0.values, u = dchi.values[:, 1:3], R = Rfield0.values);
        dims = 1)
    pl = render(plots)
    return true
end
end

function test_convergence()
    formul = FEMMShellT3FFModule
    @info "Simply supported square plated with concentrated force,"
    @info " formulation=$(formul)"

    for n in [2, 4, 8, ]
        _execute_dsg_model(formul, n, true)
    end
    return true
end

function allrun()
    println("#####################################################")
    println("# test_convergence ")
    test_convergence()
    return true
end # function allrun

@info "All examples may be executed with "
println("using .$(@__MODULE__); $(@__MODULE__).allrun()")

end # module
nothing


text:
using SparseArrays
using LinearAlgebra
using JuMP

include("proof_helper.jl")
gurobi_env = Gurobi.Env()
k = 5
r = round(Int64,(k+1)/2)
alph = 0.57
standard = false

if standard
    n = 200
    n1 = round(Int64,floor(n*alph))
    n2 = n - n1
    b1, b2 = n1_n2_baselines(k,n1,n2)
else
    b1 = arbitrary_baselines(k,alph)
    b2 = arbitrary_baselines(k,1-alph)
end

X,eps = monotonic_lp_jump(b1,b2,k)

Dk, Db, R, B = Bmat(b1,b2)

Xs = Xstanard(k,n1,n2)
Xa = Xasym(k,alph)

# Try another way
B = monotonic_LP_B(b1,b2,k)
Xopt, Yopt, eps = PrimalDualXY(gurobi_env,B,k)

if standard
    type = "Standard"
else
    type = "Asymptotic"
end
println("")

if eps> 0
    println("\tMonotonic homophily is possible!")
else
    println("\n\tMonotonic homophily is impossible.")
end

## Check alternate dual variables
x = Xa
x = X
D1 = sum(i*x[i+1] for i = 1:k)
D2 = sum(i*x[k-i+1] for i = 1:k)
r = round(Int64,(k+1)/2)
for t = r:k
    @assert((t)*x[t+1]/b1[t] - (t-1)*x[t]/b1[t-1] >= 0 )
    @assert((t)*x[k-t+1]/b2[t] - (t-1)*x[k-t+2]/b2[t-1] >= 0)
end

## Check more complicated baselines (from other hypergraph)

include("../src/helper.jl")
data = matread("../data/Polblogs_H.mat")

H = data["H"]; classes = vec(round.(Int64,data["labels"]))

## Get scores
r = 13                    # maximum k size we care about
Hw, Hm = relative_affinities(H,classes,r)
N = get_hyperedge_counts(H,classes)
Nmat = N[1:r,1:r+1]
# Get baselines
b1c = Hw[k,1:k]
b2c = Hm[k,1:k]

## Find solution to the LP

k = 11
N = Nmat[k,1:k+1]
N = N.+1
r = round(Int64,(k+1)/2)

b1 = zeros(k)
b2 = zeros(k)
D1 = sum(i*N[i+1] for i = 1:k)
D2 = sum(i*N[k-i+1] for i = 1:k)
for i = 1:k
    b1[i] = i*N[i+1]/D1
    b2[i] = i*N[k-i+1]/D2
end

B = monotonic_LP_B(b1,b2,k)
Xopt, Yopt, eps = PrimalDualXY(gurobi_env,B,k)

## Check dual variables

S = [r/b2[r] -(r-1)/b1[r-1];
    -(r-1)/b2[r-1] r/b1[r]]

S*[D1;D2]

max_stars_repo_name: nveldt/HypergraphHomophily

max_issues_repo_name: nveldt/HypergraphHomophily

max_forks_repo_name: nveldt/HypergraphHomophily


text:
struct SpatialMassActionJump{A<:Union{AbstractVector,Nothing},B<:Union{AbstractMatrix,Nothing},S,U,V} <: AbstractMassActionJump
    uniform_rates::A # reactions that are uniform in space
    spatial_rates::B # reactions whose rate depends on the site
    reactant_stoch::S
    net_stoch::U
    param_mapper::V

    """
    uniform rates go first in ordering
    """
    function SpatialMassActionJump{A,B,S,U,V}(uniform_rates::A, spatial_rates::B, reactant_stoch::S, net_stoch::U, param_mapper::V, scale_rates::Bool, useiszero::Bool, nocopy::Bool) where {A<:Union{AbstractVector,Nothing},B<:Union{AbstractMatrix,Nothing}, S, U, V}
        uniform_rates = (nocopy || isnothing(uniform_rates)) ? uniform_rates : copy(uniform_rates)
        spatial_rates = (nocopy || isnothing(spatial_rates)) ? spatial_rates : copy(spatial_rates)
        reactant_stoch = nocopy ? reactant_stoch : copy(reactant_stoch)
        for i in eachindex(reactant_stoch)
            if useiszero && (length(reactant_stoch[i]) == 1) && iszero(reactant_stoch[i][1][1])
                reactant_stoch[i] = typeof(reactant_stoch[i])()
            end
        end
        num_unif_rates = isnothing(uniform_rates) ? 0 : length(uniform_rates)
        if scale_rates && num_unif_rates > 0
            scalerates!(uniform_rates, reactant_stoch)
        end
        if scale_rates && !isnothing(spatial_rates) && !isempty(spatial_rates)
            scalerates!(spatial_rates, reactant_stoch[num_unif_rates+1:end])
        end
        new(uniform_rates, spatial_rates, reactant_stoch, net_stoch, param_mapper)
    end

end

################ Constructors ##################

SpatialMassActionJump(urates::A, srates::B, rs::S, ns::U, pmapper::V; scale_rates = true, useiszero = true, nocopy=false) where {A<:Union{AbstractVector,Nothing},B<:Union{AbstractMatrix,Nothing},S,U,V} = SpatialMassActionJump{A,B,S,U,V}(urates, srates, rs, ns, pmapper, scale_rates, useiszero, nocopy)
SpatialMassActionJump(urates::A, srates::B, rs, ns; scale_rates = true, useiszero = true, nocopy=false) where {A<:Union{AbstractVector,Nothing},B<:Union{AbstractMatrix,Nothing}} = SpatialMassActionJump(urates, srates, rs, ns, nothing; scale_rates=scale_rates, useiszero=useiszero, nocopy=nocopy)

SpatialMassActionJump(srates::B, rs, ns, pmapper; scale_rates = true, useiszero = true, nocopy=false) where {B<:Union{AbstractMatrix,Nothing}} = SpatialMassActionJump(nothing, srates, rs, ns, pmapper; scale_rates = scale_rates, useiszero = useiszero, nocopy=nocopy)
SpatialMassActionJump(srates::B, rs, ns; scale_rates = true, useiszero = true, nocopy=false) where {B<:Union{AbstractMatrix,Nothing}} = SpatialMassActionJump(nothing, srates, rs, ns, nothing; scale_rates = scale_rates, useiszero = useiszero, nocopy=nocopy)

SpatialMassActionJump(urates::A, rs, ns, pmapper; scale_rates = true, useiszero = true, nocopy=false) where {A<:Union{AbstractVector,Nothing}} = SpatialMassActionJump(urates, nothing, rs, ns, pmapper; scale_rates = scale_rates, useiszero = useiszero, nocopy=nocopy)
SpatialMassActionJump(urates::A, rs, ns; scale_rates = true, useiszero = true, nocopy=false) where {A<:Union{AbstractVector,Nothing}} = SpatialMassActionJump(urates, nothing, rs, ns, nothing; scale_rates = scale_rates, useiszero = useiszero, nocopy=nocopy)

SpatialMassActionJump(ma_jumps::MassActionJump{T,S,U,V}; scale_rates = true, useiszero = true, nocopy=false) where {T,S,U,V} = SpatialMassActionJump(ma_jumps.scaled_rates, ma_jumps.reactant_stoch, ma_jumps.net_stoch, ma_jumps.param_mapper; scale_rates = scale_rates, useiszero = useiszero, nocopy=nocopy)

##############################################

get_num_majumps(spatial_majump::SpatialMassActionJump{Nothing,Nothing,S,U,V}) where {S,U,V} = 0
get_num_majumps(spatial_majump::SpatialMassActionJump{Nothing,B,S,U,V}) where {B,S,U,V} = size(spatial_majump.spatial_rates, 1)
get_num_majumps(spatial_majump::SpatialMassActionJump{A,Nothing,S,U,V}) where {A,S,U,V} = length(spatial_majump.uniform_rates)
get_num_majumps(spatial_majump::SpatialMassActionJump{A,B,S,U,V}) where {A<:AbstractVector,B<:AbstractMatrix,S,U,V} = length(spatial_majump.uniform_rates) + size(spatial_majump.spatial_rates, 1)
using_params(spatial_majump::SpatialMassActionJump) = false

rate_at_site(rx, site, spatial_majump::SpatialMassActionJump{Nothing,B,S,U,V}) where {B,S,U,V} = spatial_majump.spatial_rates[rx, site]
rate_at_site(rx, site, spatial_majump::SpatialMassActionJump{A,Nothing,S,U,V}) where {A,S,U,V} = spatial_majump.uniform_rates[rx]
function rate_at_site(rx, site, spatial_majump::SpatialMassActionJump{A,B,S,U,V}) where {A<:AbstractVector,B<:AbstractMatrix,S,U,V}
    num_unif_rxs = length(spatial_majump.uniform_rates)
    rx <= num_unif_rxs ? spatial_majump.uniform_rates[rx] : spatial_majump.spatial_rates[rx-num_unif_rxs, site]
end

function evalrxrate(speciesmat::AbstractMatrix{T}, rxidx::S, majump::SpatialMassActionJump, site::Int) where {T,S}
    val = one(T)
    @inbounds for specstoch in majump.reactant_stoch[rxidx]
        specpop = speciesmat[specstoch[1], site]
        val    *= specpop
        @inbounds for k = 2:specstoch[2]
            specpop -= one(specpop)
            val     *= specpop
        end
    end
    @inbounds return val * rate_at_site(rxidx, site, majump)
end


text:
using CairoMakie, Random
Random.seed!(14)
n, m = 25, 50
xPos = LinRange(-π/3,2π,n)
startPos = sin.(xPos)/3
widths = 0.1 .+ 0.2rand(n)
colors = rand(n)
cmap1 = to_colormap(:starrynight)[5:end-5]
cmap2 = to_colormap(:vangogh)[5:end-5]
with_theme(theme_dark()) do
    fig = Figure(resolution = (600,800))
    ax1 = Axis(fig[1,1])
    ax2 = Axis(fig[2,1])
    for i in 0:3.5:m
        indxs = shuffle(1:n)
        barplot!(ax1, xPos[indxs], rand(1:3,n) .+ i; fillto = startPos .+ i,
        color = colors, colormap = cmap1, width = widths, strokewidth = 0.5,
        strokecolor = colors[indxs], direction = :y)
        barplot!(ax2, xPos[indxs], rand(1:3,n) .+ i; fillto = startPos .+ i,
        color = colors, colormap = cmap2, width = widths, strokewidth = 0.5,
        strokecolor = colors[indxs], direction = :x)
    end
    hidedecorations!(ax1)
    hidedecorations!(ax2)
    save("barplot2.svg", fig)
    display(fig)
end


text:
#####
##### Testing distributed self-play
#####

ENV["JULIA_CUDA_MEMORY_POOL"] = "split"

using AlphaZero
using ProgressMeter

include("../../games/connect-four/main.jl")
using .ConnectFour: Game, Training

params = Training.params
network = Training.Network{Game}(Training.netparams)
env = AlphaZero.Env{Game}(params, network)
duel = Training.benchmark[1]

const sp_progress = Progress(params.self_play.num_games)
const arena_progress = Progress(params.arena.num_games)
const bench_progress = Progress(duel.num_games)

struct Handler end
AlphaZero.Handlers.game_played(::Handler) = next!(sp_progress)
AlphaZero.Handlers.checkpoint_game_played(::Handler) = next!(arena_progress)

println("Running on $(Threads.nthreads()) threads.")

report, t, mem, gct = @timed AlphaZero.self_play_step!(env, Handler())

# (avgr, redundancy), t, mem, gct = @timed AlphaZero.evaluate_network(
#   env.curnn, env.bestnn, env.params, Handler())

# report, t, mem, gct = @timed AlphaZero.Benchmark.run(env, duel, bench_progress)

println("Total time: $t")
println("Spent in GC: $gct")


text:
module D91PVInversion

using OffsetArrays
using SparseArrays
using LinearMaps
using LinearAlgebra
using IterativeSolvers
using JLD2

# structs.jl
export Params, float_type
export Domain, new_field, new_rhs, dx, dy, dz, halo,
       field_from_rhs, field_from_rhs!,
       fields_from_linearized_rhs, fields_from_linearized_rhs!,
       rhs_from_field, rhs_from_field!
export Solver, is_converged
# profiles.jl
export background_ϕ, background_ψ, background_θ
# initialization.jl
export allocate_fields, allocate_rhs, allocate_linearized_rhs,
       set_background_ψ!, set_background_ϕ!, set_q!, set_q′!
# operators.jl
export generate_∇², generate_linearized_L, generate_Lψ, generate_Lϕ,
       ∂x, ∂y, ∂z, ∂x², ∂y², ∂z², ∂xy, ∂xz, ∂yz
# sparse_matrix_operators.jl
export generate_sparse_linearized_L
# halos.jl
export fill_ψ′_halos!, fill_ϕ′_halos!, fill_ψ_halos!, fill_ϕ_halos!
# boundaries.jl
export set_∂ψ!, set_∂ϕ!
# rhsides.jl
export set_linearized_b!, set_linear_bϕ!, set_bψ!, set_bϕ!
# diagnostics.jl
export diagnose_u, diagnose_v, diagnose_umag, diagnose_ζ, diagnose_θ, diagnose_penetration_depth
# residuals.jl
export compute_residual
# relaxation.jl
export relax!
# inversions.jl
export ColumnarVortex, LinearizedSparseInversion, LinearizedInversion, NLInversion,
       initialize!, solve!, iterate!, is_converged, save_inversion_results
# redimensionalize.jl
export dimensional_π, dimensional_p, dimensional_pseudoz, dimensional_r, dimensional_q, dimensional_u,
       redimensionalize_q!, redimensionalize_u!

include("sparse_utils.jl")
include("structs.jl")
include("profiles.jl")
include("initialization.jl")
include("operators.jl")
include("sparse_matrix_operators.jl")
include("halos.jl")
include("boundaries.jl")
include("rhsides.jl")
include("diagnostics.jl")
include("residuals.jl")
include("relaxation.jl")
include("inversions.jl")
include("redimensionalize.jl")

end # module


text:
module Solve

end


text:
# Plot a trajectory with a trail
@recipe function f(::Type{Val{:trail}}, x, y, z; trailcolor=nothing, symmetric=false)
    if length(x) > 1
        @series begin
            if !isnothing(trailcolor)
                seriescolor := trailcolor
            end
            seriescolor --> :black
            label := nothing
            seriestype := :path
            linealpha := (symmetric ? vcat(LinRange(0, 1, length(x)/2 |> floor |> Int), LinRange(1, 0, length(x)/2 |> ceil |> Int)) : LinRange(0, 1, length(x)))
            linewidth --> 2.5
            linewidth := LinRange(0.0, plotattributes[:linewidth], length(x))
            x := x
            y := y
            z := z
        end
    end
    seriestype := :scatter
    x := x[end]
    label --> nothing
    seriescolor --> :black
    markersize --> 5
    !isnothing(y) ? y := y[end] : nothing
    !isnothing(z) ? z := y[end] : nothing
    ()
end

function movie(P::Process; vars=1:length(P.X0), downsample=1, trail=10, seriestype=:trail, kwargs...)
    X = timeseries(P, vars)
    bounds = mapslices(extrema, X)
    𝐭 = times(P)
    I = length(𝐭)
    plot(;kwargs...)
    anim = @animate for i ∈ 1:downsample:I
        print("\u1b[1GFrame $(i÷downsample)/$(I÷downsample)")
        if seriestype == :trail
            trailIdxs = (i-trail:1:i) #i-(trail*size(X, 1)÷downsample):i
            trailIdxs = trailIdxs[trailIdxs .> 0]
            Xtrail = X[trailIdxs, 1]
            Ytrail = X[trailIdxs, 2]
            plot(Xtrail, Ytrail, seriestype=:trail, markersize=20,
            color=:black,
            size=(500, 500),
            xlims=bounds[1].*1.1,
            ylims=bounds[2].*1.1,
            border=:none,
            label=:none,
            kwargs...)
        else
            bounds = widen.(mapslices(extrema, X), (0.3,))
            x = [X[1:i, v] for v in 1:lastindex(X, 2)]
            if length(bounds) < 3
                bounds = hcat(bounds, [nothing])
            end
            plot!(x..., seriestype=seriestype, xlims=bounds[1], ylims=bounds[2], zlims=bounds[3])
        end
    end every 1
    anim
end
export movie

# Bring to the boil... Can extract trail plot function to remove
function doublePendulumMovie(P::Process; downsample=1, trail=10, barcolor=:gray, ballcolor=[:black, :black], trailcolor=[:cornflowerblue, :crimson], trailwidth=2, kwargs...)
    X = timeseries(P, 1:2)
    X2 = timeseries(P, 3:4)
    bound = 1.1*max(abs.([(mapslices(extrema, X2)...)...])...)
    𝐭 = times(P)
    I = length(𝐭)
    anim = @animate for i ∈ 1:downsample:I
        print("\u1b[1GFrame $(i÷downsample)/$(I÷downsample)")

        # Bars
        plot([0, X[i, 1]], [0, X[i, 2]], seriestype=:line, linewidth=5, color=barcolor, label=:none, linealpha=0.5, xlims=(-bound, bound), ylims=(-bound, bound))
        plot!([0], [0]; seriestype=:scatter,
            markersize=20,
            markercolor=:black,
            markeralpha=0.0,
            label=:none
        )
        plot!([X[i, 1], X2[i, 1]], [X[i, 2], X2[i, 2]], seriestype=:line, linewidth=5, color=barcolor, label=:none, linealpha=0.5)

        # Ball 1
        trailIdxs = i-(trail*size(X, 1)÷downsample):i
        trailIdxs = trailIdxs[trailIdxs .> 0]
        Xtrail = X[trailIdxs, 1]
        Ytrail = X[trailIdxs, 2]
        if length(trailIdxs) > 1
            plot!(Xtrail, Ytrail, seriestype=:path,
                    color=trailcolor[1],
                    linealpha=LinRange(0, 1, length(trailIdxs)),
                    label=:none,
                    linewidth=trailwidth,
            )
        end
        # Ball 2
        Xtrail = X2[trailIdxs, 1]
        Ytrail = X2[trailIdxs, 2]
        if length(trailIdxs) > 1
            plot!(Xtrail, Ytrail, seriestype=:path,
                    color=trailcolor[2],
                    linealpha=LinRange(0, 1, length(trailIdxs)),
                    label=:none,
                    linewidth=trailwidth,
            )
        end
        plot!([X[i, 1]], [X[i, 2]]; seriestype=:scatter,
            markersize=20,
            markercolor=ballcolor[1],
            size=(1000, 1000),
            border=:none,
            label=:none,
            kwargs...
        )
        plot!([X2[i, 1]], [X2[i, 2]]; seriestype=:scatter,
            markersize=20,
            markercolor=ballcolor[2],
            size=(1000, 1000),
            border=:none,
            label=:none,
            kwargs...
        )
    end every 1
    print("\n")
    anim

end
export doublePendulumMovie


text:
using Test
using Distributed
include(joinpath(Sys.BINDIR, "..", "share", "julia", "test", "testenv.jl"))
addprocs_with_testenv(3)
@test nworkers() == 3

@everywhere using Logging
for w in workers()
    # Do not log on worker nodes
    remotecall(()->global_logger(NullLogger()),w)
end

@everywhere using StochasticPrograms
@everywhere using JuMP
@everywhere using LinearAlgebra
@everywhere using GLPK
@everywhere using Ipopt

@everywhere using Distributions

@everywhere import StochasticPrograms: probability, expected

include("../decisions/decisions.jl")
TestDecisionVariable.run_dtests()
TestDecisionConstraint.run_dtests()
TestDecisionObjective.run_dtests()
include("../problems/problem_load.jl")
include("dfunctional_tests.jl")
include("dsolver_tests.jl")


text:
println("Testing modalparam: ")

m = simula01()
A = m["A"]
C = m["C"]
dt = m["dt"]
wm = m["wm"]
zm = m["zm"]
Vm = m["Vm"]
wm1,zm1,Vm1 = modalparam(A,C,dt)

if (sum((wm-wm1).^2) < 0.1)
    println(" wm OK")
end
if (sum((zm-zm1).^2) < 0.1) && (sum((Vm-Vm1).^2) < 0.1)
    println(" zm OK")
end
if (sum((Vm-Vm1).^2) < 0.1)
    println(" Vm OK")
end


text:
function tasksleep(;id=id::Int, workfor::Int=workfor, sleepfor::Int=sleepfor)

    stopcmd   = `kill -STOP $id`

    resumecmd = `kill -CONT $id`

    sleep(workfor)

    while true

        run(resumecmd)

        for i in 0:workfor-1
            local msg = @sprintf("\rProcess %d is working for %d seconds 🙂", id, workfor-i)
            print(Crayon(foreground = :green,  bold = true), @sprintf("%s",msg), Crayon(reset = true))
            sleep(1)
        end

        run(stopcmd)

        for i in 0:sleepfor-1
            local msg = @sprintf("\rProcess %d is sleeping for %d seconds 😴", id, sleepfor-i)
            print(Crayon(foreground = :light_red,  bold = true), @sprintf("%s",msg), Crayon(reset = true))
            sleep(1)

        end

    end

end


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

JLLWrappers.@generate_main_file_header("GnuPG")
JLLWrappers.@generate_main_file("GnuPG", UUID("1522389b-45f8-5faa-af4d-a301b79c50ac"))
end  # module GnuPG_jll


text:
include("testFirstOrderMethods.jl")
include("testNewtonMethods.jl")
include("testLineSearch.jl")
include("testTrustRegionMethods.jl")


text:
using GenomePermutations
using GenomicFeatures
using Documenter

push!(LOAD_PATH,"../src/")
DocMeta.setdocmeta!(GenomePermutations, :DocTestSetup, :(using GenomePermutations); recursive=true)

makedocs(;
    modules=[GenomePermutations],
    authors="Simone De Angelis <simone.deangelis@qmul.ac.uk> <side97.11+jl@gmail.com>",
    repo="https://github.com/sdangelis/GenomePermutations.jl/blob/{commit}{path}#{line}",
    sitename="GenomePermutations.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://sdangelis.github.io/GenomePermutations.jl",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
    ],
)

deploydocs(;
    repo="github.com/sdangelis/GenomePermutations.jl",
    devbranch="main",
)


text:
@testset "W361" begin
    s = """
    module Test361
    export foobar
    end
    """
    msgs = lintstr(s)
    @test msgs[1].code == :W361
    @test contains(msgs[1].message, "exporting undefined symbol")
end


text:
using DecidableDependentTypes

@macroexpand DecidableDependentTypes.@f 2x


text:
using Revise

using JLD
using DiffPrecTest
using DataFrames, CSV
using Random, LinearAlgebra, Distributions, SparseArrays, Statistics
using CovSel
using Plots

dir = "/home/mkolar/.julia/dev/DiffPrecTest/exper/sim5/res"

p = 20
n = 300

# generate model
Ωx = Matrix{Float64}(I, p, p)
mf = [1., 0.5, 0.4]
# block 1
α = 1.
bp = 1
ep = 10
for k=0:2
    v = α * mf[k+1]
    for l=bp:ep-k
        Ωx[l  , l+k] = v
        Ωx[l+k, l  ] = v
    end
end
# block 2
α = 2.
bp = 11
ep = 15
for k=0:2
    v = α * mf[k+1]
    for l=bp:ep-k
        Ωx[l  , l+k] = v
        Ωx[l+k, l  ] = v
    end
end
# block 3
α = 4.
bp = 16
ep = 20
for k=0:2
    v = α * mf[k+1]
    for l=bp:ep-k
        Ωx[l  , l+k] = v
        Ωx[l+k, l  ] = v
    end
end
Ωy = Matrix(Diagonal(Ωx))
Δ = Ωx - Ωy

Σx = inv(Symmetric(Ωx))
Σy = inv(Symmetric(Ωy))

dist_X = MvNormal(convert(Matrix, Σx))
dist_Y = MvNormal(convert(Matrix, Σy))

τArr = collect( range(10,0,length=500) )

NUM_REP=200
FPR_mine = zeros(NUM_REP, length(τArr))
TPR_mine = zeros(NUM_REP, length(τArr))
FPR_tr = zeros(NUM_REP, 50)
TPR_tr = zeros(NUM_REP, 50)

for rep=1:NUM_REP
    global FPR_mine, TPR_mine, FPR_tr, TPR_tr, Δ, τArr
    fname = "$(dir)/res_$(rep).jld"

    file = jldopen(fname, "r")
    eΔNormal = read(file, "eΔNormal")
    eΔDTr = read(file, "eΔDTr")
    close(file)

    for j=1:length(τArr)
        confusionMatrix = CovSel.getConfusionMatrix(Δ, eΔNormal[j])
        FPR_mine[rep, j] = CovSel.fpr(confusionMatrix)
        TPR_mine[rep, j] = CovSel.tpr(confusionMatrix)
    end

    for j=1:50
        confusionMatrix = CovSel.getConfusionMatrix(Δ, eΔDTr[j])
        FPR_tr[rep, j] = CovSel.fpr(confusionMatrix)
        TPR_tr[rep, j] = CovSel.tpr(confusionMatrix)
    end
end

pyplot()
plot(mean(FPR_mine, dims=1)', mean(TPR_mine, dims=1)')
plot!(mean(FPR_tr, dims=1)', mean(TPR_tr, dims=1)')

plot(FPR_mine[1, :], TPR_mine[1, :])
FPR_mine[1, :]
rep = 1
fname = "$(dir)/res_$(rep).jld"
file = jldopen(fname, "r")
eΔNormal = read(file, "eΔNormal")

# using JLD
# using DiffPrecTest
# using DataFrames, CSV
# using Random, LinearAlgebra, Distributions, SparseArrays
# using CovSel
#
# dir = ARGS[1]
#
# p = 100
# n = 300
#
# # generate model
# Ωx = Matrix{Float64}(I, p, p)
# ρ = 0.5
# for k=1:p-1
#     for l=1:p-k
#         Ωx[l  , l+k] = ρ^k
#         Ωx[l+k, l  ] = ρ^k
#     end
# end
#
# Ωy = copy(Ωx)
# k = div(p, 4)
# for l=1:p-k
#     Ωy[l  , l+k] = 0.2
#     Ωy[l+k, l  ] = 0.2
# end
# Δ = Ωx - Ωy
#
# NUM_REP = 100
# resN = Array{CovSel.ConfusionMatrix}(undef, NUM_REP)
# resBoot = Array{CovSel.ConfusionMatrix}(undef, NUM_REP)
# resDTr2 = Array{CovSel.ConfusionMatrix}(undef, NUM_REP)
# resDTrInf = Array{CovSel.ConfusionMatrix}(undef, NUM_REP)
# i = 0
# for rep=1:NUM_REP
#     global i, resN, resBoot, resDTr2, resDTrInf
#     fname = "$(dir)/res_$(rep).jld"
#
#     try
#       i += 1
#       file = jldopen(fname, "r")
#       eΔNormal = read(file, "eΔNormal")
#       eΔBoot = read(file, "eΔBoot")
#       eΔDTr = read(file, "eΔDTr")
#       i2 = read(file, "i2")
#       iInf = read(file, "iInf")
#       close(file)
#
#       resN[i] = CovSel.getConfusionMatrix(Δ, eΔNormal)
#       resBoot[i] = CovSel.getConfusionMatrix(Δ, eΔBoot)
#       resDTr2[i] = CovSel.getConfusionMatrix(Δ, eΔDTr[i2])
#       resDTrInf[i] = CovSel.getConfusionMatrix(Δ, eΔDTr[iInf])
#     catch
#          @show fname
#     end
# end
#
# @show mean([CovSel.tpr(resN[j]) for j=1:i])
# @show mean([CovSel.fpr(resN[j]) for j=1:i])
# @show mean([CovSel.precision(resN[j]) for j=1:i])
#
# @show mean([CovSel.tpr(resBoot[j]) for j=1:i])
# @show mean([CovSel.fpr(resBoot[j]) for j=1:i])
# @show mean([CovSel.precision(resBoot[j]) for j=1:i])
#
# @show mean([CovSel.tpr(resDTr2[j]) for j=1:i])
# @show mean([CovSel.fpr(resDTr2[j]) for j=1:i])
# @show mean([CovSel.precision(resDTr2[j]) for j=1:i])
#
# @show mean([CovSel.tpr(resDTrInf[j]) for j=1:i])
# @show mean([CovSel.fpr(resDTrInf[j]) for j=1:i])
# @show mean([CovSel.precision(resDTrInf[j]) for j=1:i])


text:
### A Nbdev script file ###
### Autogenerated file. Don't modify. ###

module Documenter
#export
using Base.Docs

#export
using TOML

#export
using Markdown

#export
using Distributed

#export
using ProgressMeter

#export
include("../src/Export.jl")

#export
include("../src/CodeRunner.jl")

#export
import Pluto: Notebook, Cell, load_notebook_nobackup

#export
begin
"""
> struct Section--> This is like the section of a page and is made up of one or multiple lines.
> * Fields:
>   * line--> String which makes up a section
"""
Base.@kwdef mutable struct Section
	line::String=""
end

"""
> Section(line)--> Helper function to call the constructor of the type Section
"""
Section(line)=Section(line=line)
"""
> line(section::Section)--> Getter for accessing the underlying field line.
"""
line(section::Section)=section.line
end

#export
begin
"""
> Page--> The type that represents a page in a document.
> * Fields:
>   * sections--> Array of Section type.
>   * name--> Name of the module being documented.
"""
struct Page
	sections::Array{Section,1}
    name::AbstractString
end

"""
> Page--> Helper function to create constructer for Page type.
"""
Page(sections, path)=Page(sections=sections, path=path)

"""
> sections--> Getter for accessing the underlying field sections of Page.
"""
sections(p::Page)=p.sections

"""
> name--> Getter for accessing the underlying field name of Page.
"""
name(p::Page)=p.name
end

#export
"""
> runandupdatenb(fn::AbstractString)--> Run the notebook in the supplied path and update the notebook with the output of each cell.
"""
function runandupdatenb(fn::AbstractString)
	nb=load_notebook_nobackup(fn)
	return CodeRunner.executecode(nb)
end

#export
begin
"""
> struct Functiondocs--> Stores the document of different objects.
> * funcdocs--> Array of strings.
"""
mutable struct Functiondocs
	funcdocs::Array{String, 1}
end

"""
> Functiondocs(funcdocs)--> Helper for accessing the FunctionDocs constructer.
"""
Functiondocs(funcdocs)=Functiondocs(funcdocs)
end

#export
begin

"""
> stitchcode(cell::Cell)--> Stitches the code in a Pluto notebook cell with the output of that code. The output is a code block.
"""
function stitchcode(cell::AbstractArray)
	op=' '
	if length(cell) > 1
	    op=values(cell[2])
	end
	string("```","\n$(cell[1])\n","------\nOutput\n------\n","$(op)\n", "```\n")
end

"""
> stitchcode(cellop::AbstractString)--> Removes the quotes from a string and creates a code block with that string inside the newely formed code block
"""
function stitchcode(cellop::AbstractString)
	cleanedop=Export.strip(Export.strip(cellop,"\""), "\"")
	string("```","\n$cleanedop\n","```\n")
end

"""
> stitchcode(fdocs::Functiondocs)--> When supplied with a FunctionDocs type, stitchCode appends together the object docstrings and generates documentation for that particular object
"""
function stitchcode(fdocs::Functiondocs)
		funcdocs=""

		for fdoc in fdocs.funcdocs
		    funcdocs=string(funcdocs, "$(fdoc)\n\n")
	    end

		funcdocs
end
end

#export
"""
> collectfuncdocs(obj)--> Collects objects (functions, methods, macro structs etc.) and creates an array of documents (generated from the docstrings). Creates aFunctionDocs type from these documents.
"""
function collectfuncdocs(obj)
	docs=doc(obj)
	searchurl = "https://github-link.vercel.app/api?ghUrl="
	giturl = getsettings(joinpath("..", "Project.toml"), ["settings", "github_url"])

	pattern = r"\\(.+\\)*(.+)(\.jl)"
	fn = uppercasefirst(Export.strip(basename(join(match(pattern, string(first(methods(obj)))).captures)), r"[0-9_]"))

	searchurl = "$searchurl$giturl/blob/master/src/$fn&q=$(string(nameof(obj)))"
    fdocs=["$(docs.meta[:results][i].object) [source]($searchurl)" for i=1:length(docs.meta[:results])]
	Functiondocs(fdocs)
end

#export
begin

"""
> showdoc(o)--> Takes an object and builds markdown documentation.
"""
function showdoc(o)
	docs=collectfuncdocs(o)
	stitchcode(docs)
end
end

#export
begin
function maintainheading(str:: AbstractString)
	res= nothing
	if startswith(str, """

""")
			res = Export.strip(Export.strip(Export.strip(str, """

"""), """
"""), """
""")
			res = string("## ", res)
	else
		res = str
	end
	res
end

"""
> CreatePage--> Creates the Page type from the markdown and example code cells of the supplied notebook. The filename is the name of the notebook which is parsed.
"""
function createpage(fn::AbstractString, nb::Notebook)
	sections=Section[]
	res=nothing

	for cell in nb.cells

		if cell.errored
			error("Build stopped. Seems like the code $cell.code has an error")
			break
	    end
	    if startswith(cell.code, "md")
			cleanop = maintainheading(cell.output.body)
			push!(sections, Section(cleanop))
		elseif !startswith(cell.code, "#export") && !startswith(cell.code, "#hide")
			if occursin( "showdoc", cell.code)
				cleanedop=Export.strip(cell.output.body, "\"")
				cleanedop=replace(cleanedop, "\\n"=>"\n")
				push!(sections, Section(cleanedop))
			elseif occursin( "#noop", cell.code)
				stitchedcode=stitchcode([cell.code])
			    push!(sections, Section(stitchedcode))
			else
				stitchedcode=stitchcode([cell.code, cell.output.body])
			    push!(sections, Section(stitchedcode))
			end
		end
	end
	if !isempty(sections)
	    res=Page(sections, basename(fn))
	end

	return res

end
end

#export
begin

"""
> savepage(io, page::Page)--> Take the contents from a Page type and write to the io
"""
function savepage(io, page::Page)

	pageheading=uppercasefirst(Export.strip(Export.strip(page.name, r"[0-9_]"), r".jl"))

    println(io, "
$pageheading
")

	for section in page.sections
			println(io, section.line)
    end

end

"""
> savepage(page::Page, path::String)--> Given a Page type and the required path, creates the related markdwon file in the specified path. The name of the resulting markdown file is same as the nameof the notebook for which the document is generated
"""
function savepage(page::Page, path::String)
	fn=lowercase(Export.strip(Export.strip(page.name, r"[0-9_]"), r".jl"))
	open(joinpath(path, fn*".md"), "w") do io
        savepage(io, page)
    end
end
end

#export
begin
"""
> export2md(fn::String, path::String)--> Generate document for a file in the given path
"""
function export2md(fn::String, path::String)
	nb=runandupdatenb(fn)
	page=createpage(fn, nb)
	if page != nothing
	    savepage(page, path)
	end
end

"""
> export2md(fns::AbstractVector, path::String)--> Map the `export2md(file, path)` to a given vector of file.
"""
function export2md(fns::AbstractVector, path::String)
	n = length(fns)
	p = Progress(n, 1)
	for fn in fns
	  next!(p)
      export2md(fn, path)
	end
end

"""
> export2md()--> Higher level API to generate documents for all the valid notebooks
"""
export2md(nbsdir)=export2md(Export.readfilenames(joinpath(pwd(), nbsdir)), "docs")
end

#export
export showdoc, export2md

#export
begin
"""
> export2readme()--> create readme from the contents of Index notebook
"""
function export2readme()
    cp(normpath(joinpath(".", "docs/index.md")), normpath(joinpath(".", "README.md")), force=true)
end
end

end