text:
# core

Base.length(dict::CardinalPairDict{V,K}) where {K,V} = sum(dict.present)
Base.endof(dict::CardinalPairDict{V,K}) where {K,V} = findlast(dict.present, true)
Base.isempty(dict::CardinalPairDict{V,K}) where {K,V} = !any(dict.present)
isfull(dict::CardinalPairDict{V,K}) where {K,V} = all(dict.present)

Base.eltype(dict::CardinalPairDict{V,K}) where {K,V} = Pair{K,V}

Base.:(==)(a_dict::D, b_dict::D) where D<:CardinalPairDict{V,K} where {K,V} =
    keymax(a_dict) == keymax(b_dict) && keys(a_dict) == keys(b_dict) && content(a_dict) == content(b_dict)
Base.:(==)(a_dict::CardinalPairDict{V,K}, b_dict::CardinalPairDict{J,W}) where {J,W,K,V} =
    false

Base.:(!=)(a_dict::D, b_dict::D) where D<:CardinalPairDict{V,K} where {K,V} =
    keymax(a_dict) != keymax(b_dict)  || keys(a_dict) != keys(b_dict) || content(a_dict) != content(b_dict)
Base.:(!=)(a_dict::CardinalPairDict{V,K}, b_dict::CardinalPairDict{J,W}) where {J,W,K,V} =
    true

# keys, content, in (∈, ∋, ∉, ∌)

function tally_active(dict::CardinalPairDict{V,K}) where {K,V}
    return sum(dict.present)
end
function tally_passive(dict::CardinalPairDict{V,K}) where {K,V}
    return length(dict.present) - sum(dict.present)
end

function internal_keys(dict::CardinalPairDict{V,K}) where {K,V}
    allkeys = one(K):keymax(dict)
    result = Vector{K}()
    for k in allkeys
        if unsafe_haskey(dict, k)
            push!(result,k)
        end
    end
    return result
end

function Base.keys(dict::CardinalPairDict{V,K}) where {K,V}
   allkeys = internal_keys(dict)
   result = Vector{Tuple{K,K}}()
   for k in allkeys
        push!(result, unpair(k))
   end
   return result
end

function Base.values(dict::CardinalPairDict{V,K}) where {K,V}
    allkeys = keys(dict)
    result = Vector{V}()
    for (x,y) in allkeys
        push!(result, getindex(dict, x, y))
    end
    return result
end

function Base.in(value::V, dict::CardinalPairDict{V,K}) where {K,V}
    return length(dict)!=0 && in(value, content(dict))
end

#@inline Base.:(∈)(value::V, dict::CardinalPairDict{V,K}) where {K,V} = in(value, dict) # not needed
@inline Base.:(∋)(dict::CardinalPairDict{V,K}, value::V) where {K,V} = in(value, dict)
@inline Base.:(∉)(value::V, dict::CardinalPairDict{V,K}) where {K,V} = !in(value, dict)
@inline Base.:(∌)(dict::CardinalPairDict{V,K}, value::V) where {K,V} = !in(value, dict)

# delete!, clearindex!, empty!

function Base.delete!(dict::CardinalPairDict{V,K}, key1::K, key2::K) where {K,V}
    check_index(dict, key1, key2)
    key = pair(key1, key2)
    @inbounds setindex!(dict.present, false, key)
    return dict
end
@inline function Base.delete!(dict::CardinalPairDict{V,K}, keys::Tuple{K,K}) where {K,V}
    return delete!(dict, keys...)
end
@inline Base.delete!(dict::CardinalPairDict{V,K}, key1::J, key2::J) where {J,K,V} =
    delete!(dict, key1%K, key2%K)
@inline Base.delete!(dict::CardinalPairDict{V,K}, keys::Tuple{J,J}) where {J,K,V} =
    delete!(dict, keys[1]%K, keys[2]%K)

function Base.delete!(dict::CardinalPairDict{V,K}, key::K) where {K,V}
    0 < key <= keymax(dict) || throw(ErrorException("Key (index) $(key) is outside of the domain 1:$(keymax(dict))."))
    @inbounds setindex!(dict.present, false, key)
    return dict
end
@inline Base.delete!(dict::CardinalPairDict{V,K}, key::J) where {J,K,V} =
    delete!(dict, key%K)

function Base.empty!(dict::CardinalPairDict{V,K}) where {K,V}
    0 < key <= keymax(dict) || throw(ErrorException("Key (index) $(key) is outside of the domain 1:$(keymax(dict))."))
    for k in keys(dict)
        @inbounds setindex!(dict.present, false, key)
    end
    return dict
end

# iteration

function Base.start(dict::CardinalPairDict{V,K}) where {K,V}
    (1, keys(dict))
end

function Base.next(dict::CardinalPairDict{V,K}, state) where {K,V}
    index, ks = state
    (ks[index], dict.present[index]), (index+1, ks)
end

function Base.done(dict::CardinalPairDict{V,K}, state) where {K,V}
    state[1] > length(dict)
end

# nonstandard

keymax(dict::CardinalPairDict{V,K}) where {K,V} = K(length(dict.present))
keyamax(dict::CardinalPairDict{V,K}) where {K,V} = unpair(K(length(dict.present)))

function clearindex!(dict::CardinalPairDict{V,K}, key1::K, key2::K) where {K,V}
    check_index(dict, key1, key2)
    key = pair(key1, key2)
    @inbounds setindex!(dict.present, false, key)
    return nothing
end
@inline function clearindex!(dict::CardinalPairDict{V,K}, keys::Tuple{K,K}) where {K,V}
    return clearindex!(dict, keys...)
end
@inline clearindex!(dict::CardinalPairDict{V,K}, key1::J, key2::J) where {J,K,V} =
    clearindex!(dict, key1%K, key2%K)
@inline clearindex!(dict::CardinalPairDict{V,K}, keys::Tuple{J,J}) where {J,K,V} =
    clearindex!(dict, keys[1]%K, keys[2]%K)


text:
function _pass_abstract_jump(optimizer::Union{Nothing, JuMP.OptimizerFactory},
                              parameters::Bool,
                              JuMPmodel::Union{JuMP.AbstractModel,Nothing})
    if isa(optimizer, Nothing)
        @info("The optimization model has no optimizer attached")
    end
    if !isnothing(JuMPmodel)
        if parameters
            if !haskey(JuMPmodel.ext, :params)
                @info("Model doesn't have Parameters enabled. Parameters will be enabled")
            end
            PJ.enable_parameters(JuMPmodel)
        end
        return JuMPmodel
    end
    JuMPmodel = JuMP.Model(optimizer)
    if parameters
        PJ.enable_parameters(JuMPmodel)
    end
    return JuMPmodel
end

function _make_container_array(V::DataType, parameters::Bool, ax...)

    if parameters
        return JuMP.Containers.DenseAxisArray{PGAE{V}}(undef, ax...)
    else
        return JuMP.Containers.DenseAxisArray{GAE{V}}(undef, ax...)
    end

    return

end

function _make_expressions_dict(transmission::Type{S},
                                V::DataType,
                                bus_numbers::Vector{Int64},
                                time_steps::UnitRange{Int64},
                                parameters::Bool) where {S<:PM.AbstractPowerModel}

    return DSDA(:nodal_balance_active =>  _make_container_array(V,
                                                                parameters,
                                                                bus_numbers,
                                                                time_steps),
                :nodal_balance_reactive => _make_container_array(V,
                                                                 parameters,
                                                                 bus_numbers,
                                                                 time_steps))

end

function _make_expressions_dict(transmission::Type{S},
                                V::DataType,
                                bus_numbers::Vector{Int64},
                                time_steps::UnitRange{Int64},
                                parameters::Bool) where {S<:PM.AbstractActivePowerModel}

    return DSDA(:nodal_balance_active =>  _make_container_array(V,
                                                                parameters,
                                                                bus_numbers,
                                                                time_steps))
end

function _canonical_init(bus_numbers::Vector{Int64},
                        jump_model::JuMP.AbstractModel,
                        optimizer::Union{Nothing,JuMP.OptimizerFactory},
                        transmission::Type{S},
                        time_steps::UnitRange{Int64},
                        resolution::Dates.Period,
                        use_forecast_data::Bool,
                        initial_time::Dates.DateTime,
                        make_parameters_container::Bool,
                        ini_con::DICKDA) where {S<:PM.AbstractPowerModel}

    V = JuMP.variable_type(jump_model)

    canonical = Canonical(jump_model,
                              optimizer,
                              time_steps,
                              resolution,
                              use_forecast_data,
                              initial_time,
                              DSDA(),
                              DSDA(),
                              zero(JuMP.GenericAffExpr{Float64, V}),
                              _make_expressions_dict(transmission,
                                                     V,
                                                     bus_numbers,
                                                     time_steps,
                                                     make_parameters_container),
                              make_parameters_container ? DRDA() : nothing,
                              ini_con,
                              nothing);

    return canonical

end

mutable struct Canonical
    JuMPmodel::JuMP.AbstractModel
    optimizer_factory::Union{Nothing, JuMP.OptimizerFactory}
    time_steps::UnitRange{Int64}
    resolution::Dates.Period
    use_forecast_data::Bool
    initial_time::Dates.DateTime
    variables::Dict{Symbol, JuMP.Containers.DenseAxisArray}
    constraints::Dict{Symbol, JuMP.Containers.DenseAxisArray}
    cost_function::JuMP.AbstractJuMPScalar
    expressions::Dict{Symbol, JuMP.Containers.DenseAxisArray}
    parameters::Union{Nothing, Dict{UpdateRef, JuMP.Containers.DenseAxisArray}}
    initial_conditions::DICKDA
    pm::Union{Nothing, PM.AbstractPowerModel}

    function Canonical(JuMPmodel::JuMP.AbstractModel,
                       optimizer_factory::Union{Nothing, JuMP.OptimizerFactory},
                       time_steps::UnitRange{Int64},
                       resolution::Dates.Period,
                       use_forecast_data::Bool,
                       initial_time::Dates.DateTime,
                       variables::Dict{Symbol, JuMP.Containers.DenseAxisArray},
                       constraints::Dict{Symbol, JuMP.Containers.DenseAxisArray},
                       cost_function::JuMP.AbstractJuMPScalar,
                       expressions::Dict{Symbol, JuMP.Containers.DenseAxisArray},
                       parameters::Union{Nothing, Dict{UpdateRef, JuMP.Containers.DenseAxisArray}},
                       initial_conditions::DICKDA,
                       pm::Union{Nothing, PM.AbstractPowerModel})

        new(JuMPmodel,
            optimizer_factory,
            time_steps,
            resolution,
            use_forecast_data,
            initial_time,
            variables,
            constraints,
            cost_function,
            expressions,
            parameters,
            initial_conditions,
            pm)

    end

end

function Canonical(::Type{T},
                   sys::PSY.System,
                   optimizer::Union{Nothing,JuMP.OptimizerFactory};
                   kwargs...) where {T<:PM.AbstractPowerModel}

    PSY.check_forecast_consistency(sys)
    user_defined_model = get(kwargs, :JuMPmodel, nothing)
    ini_con = get(kwargs, :initial_conditions, DICKDA())
    make_parameters_container = get(kwargs, :use_parameters, false)
    use_forecast_data = get(kwargs, :use_forecast_data, true)
    jump_model = _pass_abstract_jump(optimizer, make_parameters_container, user_defined_model)
    initial_time = get(kwargs, :initial_time, PSY.get_forecasts_initial_time(sys))

    if use_forecast_data
        horizon = get(kwargs, :horizon, PSY.get_forecasts_horizon(sys))
        time_steps = 1:horizon
        if length(time_steps) > 100
            @warn("The number of time steps in the model is over 100. This will result in
                  large multiperiod optimization problem")
        end
        resolution = PSY.get_forecasts_resolution(sys)
    else
        resolution = PSY.get_forecasts_resolution(sys)
        time_steps = 1:1
    end

    bus_numbers = sort([PSY.get_number(b) for b in PSY.get_components(PSY.Bus, sys)])

    return _canonical_init(bus_numbers,
                           jump_model,
                           optimizer,
                           T,
                           time_steps,
                           resolution,
                           use_forecast_data,
                           initial_time,
                           make_parameters_container,
                           ini_con)

end

function InitialCondition(canonical::Canonical,
                          device::T,
                          access_ref::Symbol,
                          value::Float64,
                          cache::Union{Nothing, Type{<:AbstractCache}}=nothing) where T <: PSY.Device

    if model_has_parameters(canonical)
        return InitialCondition(device,
                                UpdateRef{JuMP.VariableRef}(access_ref),
                                PJ.add_parameter(canonical.JuMPmodel, value),
                                cache)
    else
        !hasfield(T, access_ref) && error("Device of of type $(T) doesn't contain
                                            the field $(access_ref)")
        return InitialCondition(device,
                                UpdateRef{T}(access_ref),
                                value,
                                cache)
    end

end

function get_initial_conditions(canonical::Canonical, key::ICKey)
    return get(canonical.initial_conditions, key, Vector{InitialCondition}())
end

# Var_ref
function get_value(canonical::Canonical, ref::UpdateRef{JuMP.VariableRef})
    return get_variable(canonical, ref.access_ref)
end

# param_ref
function get_value(canonical::Canonical, ref::UpdateRef{PJ.ParameterRef})
    for (k, v) in canonical.parameters
        if k.access_ref == ref.access_ref
            return v
        end
    end
    return
end

_variable_type(cm::Canonical) = JuMP.variable_type(cm.JuMPmodel)
model_time_steps(canonical::Canonical) = canonical.time_steps
model_resolution(canonical::Canonical) = canonical.resolution
model_has_parameters(canonical::Canonical) = !isnothing(canonical.parameters)
model_uses_forecasts(canonical::Canonical) = canonical.use_forecast_data
model_initial_time(canonical::Canonical) = canonical.initial_time
#Internal Variables, Constraints and Parameters accessors
get_variables(canonical::Canonical) = canonical.variables
get_constraints(canonical::Canonical) = canonical.constraints
get_variable(canonical::Canonical, name::Symbol) = canonical.variables[name]
get_constraint(canonical::Canonical, name::Symbol) = canonical.constraints[name]
get_parameters(canonical::Canonical, param_reference::UpdateRef) = canonical.parameters[param_reference]
get_expression(canonical::Canonical, name::Symbol) = canonical.expressions[name]
get_initial_conditions(canonical::Canonical) = canonical.initial_conditions

ext: jl


text:
"""
    VerticallyStretchedRectilinearGrid{FT, TX, TY, TZ, R, A} <: AbstractRectilinearGrid{FT, TX, TY, TZ}

A rectilinear grid with with constant horizontal grid spacings `Δx` and `Δy`, and
non-uniform or stretched vertical grid spacing `Δz` between cell centers and cell faces,
topology `{TX, TY, TZ}`, and coordinate ranges of type `R` (where a range can be used) and
`A` (where an array is needed).
"""
struct VerticallyStretchedRectilinearGrid{FT, TX, TY, TZ, R, A} <: AbstractRectilinearGrid{FT, TX, TY, TZ}

    # Number of grid points in (x,y,z).
    Nx :: Int
    Ny :: Int
    Nz :: Int

    # Halo size in (x,y,z).
    Hx :: Int
    Hy :: Int
    Hz :: Int

    # Domain size [m].
    Lx :: FT
    Ly :: FT
    Lz :: FT

    # Grid spacing [m].
       Δx :: FT
       Δy :: FT
    Δzᵃᵃᶜ :: A
    Δzᵃᵃᶠ :: A

    # Range of coordinates at the centers of the cells.
    xᶜᵃᵃ :: R
    yᵃᶜᵃ :: R
    zᵃᵃᶜ :: A

    # Range of grid coordinates at the faces of the cells.
    # Note: there are Nx+1 faces in the x-dimension, Ny+1 in the y, and Nz+1 in the z.
    xᶠᵃᵃ :: R
    yᵃᶠᵃ :: R
    zᵃᵃᶠ :: A
end

function VerticallyStretchedRectilinearGrid(FT=Float64; architecture = CPU(),
                                              size, x, y, zF,
                                              halo = (1, 1, 1),
                                          topology = (Periodic, Periodic, Bounded))

    TX, TY, TZ = validate_topology(topology)
    size = validate_size(TX, TY, TZ, size)
    halo = validate_halo(TX, TY, TZ, halo)
    Lx, Ly, x, y = validate_vertically_stretched_grid_xy(TX, TY, FT, x, y)

    Nx, Ny, Nz = size
    Hx, Hy, Hz = halo

    # Initialize vertically-stretched arrays on CPU
    Lz, zᵃᵃᶠ, zᵃᵃᶜ, Δzᵃᵃᶜ, Δzᵃᵃᶠ = generate_stretched_vertical_grid(FT, topology[3], Nz, Hz, zF)

    # Construct uniform horizontal grid
    Lh, Nh, Hh, X₁ = (Lx, Ly), size[1:2], halo[1:2], (x[1], y[1])
    Δx, Δy = Δh = Lh ./ Nh

    # Face-node limits in x, y, z
    xF₋, yF₋ = XF₋ = @. X₁ - Hh * Δh
    xF₊, yF₊ = XF₊ = @. XF₋ + total_extent(topology[1:2], Hh, Δh, Lh)

    # Center-node limits in x, y, z
    xC₋, yC₋ = XC₋ = @. XF₋ + Δh / 2
    xC₊, yC₊ = XC₊ = @. XC₋ + Lh + Δh * (2Hh - 1)

    # Total length of Center and Face quantities
    TFx, TFy, TFz = total_length.(Face, topology, size, halo)
    TCx, TCy, TCz = total_length.(Center, topology, size, halo)

    # Include halo points in coordinate arrays
    xᶠᵃᵃ = range(xF₋, xF₊; length = TFx)
    yᵃᶠᵃ = range(yF₋, yF₊; length = TFy)

    xᶜᵃᵃ = range(xC₋, xC₊; length = TCx)
    yᵃᶜᵃ = range(yC₋, yC₊; length = TCy)

    xᶜᵃᵃ = OffsetArray(xᶜᵃᵃ,  -Hx)
    yᵃᶜᵃ = OffsetArray(yᵃᶜᵃ,  -Hy)
    zᵃᵃᶜ = OffsetArray(zᵃᵃᶜ,  -Hz)

    xᶠᵃᵃ = OffsetArray(xᶠᵃᵃ,  -Hx)
    yᵃᶠᵃ = OffsetArray(yᵃᶠᵃ,  -Hy)
    zᵃᵃᶠ = OffsetArray(zᵃᵃᶠ,  -Hz)

    Δzᵃᵃᶠ = OffsetArray(Δzᵃᵃᶠ, -Hz)
    Δzᵃᵃᶜ = OffsetArray(Δzᵃᵃᶜ, -Hz)

    # Needed for pressure solver solution to be divergence-free.
    # Will figure out why later...
    Δzᵃᵃᶠ[Nz] = Δzᵃᵃᶠ[Nz-1]

    # Seems needed to avoid out-of-bounds error in viscous dissipation
    # operators wanting to access Δzᵃᵃᶠ[Nz+2].
    Δzᵃᵃᶠ = OffsetArray(cat(Δzᵃᵃᶠ[0], Δzᵃᵃᶠ..., Δzᵃᵃᶠ[Nz], dims=1), -Hz-1)

    # Convert to appropriate array type for arch
    zᵃᵃᶠ  = OffsetArray(arch_array(architecture,  zᵃᵃᶠ.parent),  zᵃᵃᶠ.offsets...)
    zᵃᵃᶜ  = OffsetArray(arch_array(architecture,  zᵃᵃᶜ.parent),  zᵃᵃᶜ.offsets...)
    Δzᵃᵃᶜ = OffsetArray(arch_array(architecture, Δzᵃᵃᶜ.parent), Δzᵃᵃᶜ.offsets...)
    Δzᵃᵃᶠ = OffsetArray(arch_array(architecture, Δzᵃᵃᶠ.parent), Δzᵃᵃᶠ.offsets...)

    return VerticallyStretchedRectilinearGrid{FT, TX, TY, TZ, typeof(xᶠᵃᵃ), typeof(zᵃᵃᶠ)}(
        Nx, Ny, Nz, Hx, Hy, Hz, Lx, Ly, Lz, Δx, Δy, Δzᵃᵃᶜ, Δzᵃᵃᶠ, xᶜᵃᵃ, yᵃᶜᵃ, zᵃᵃᶜ, xᶠᵃᵃ, yᵃᶠᵃ, zᵃᵃᶠ)
end

#####
##### Vertically stretched grid utilities
#####

get_z_face(z::Function, k) = z(k)
get_z_face(z::AbstractVector, k) = z[k]

lower_exterior_Δzᵃᵃᶜ(z_topo,          zFi, Hz) = [zFi[end - Hz + k] - zFi[end - Hz + k - 1] for k = 1:Hz]
lower_exterior_Δzᵃᵃᶜ(::Type{Bounded}, zFi, Hz) = [zFi[2]  - zFi[1] for k = 1:Hz]

upper_exterior_Δzᵃᵃᶜ(z_topo,          zFi, Hz) = [zFi[k + 1] - zFi[k] for k = 1:Hz]
upper_exterior_Δzᵃᵃᶜ(::Type{Bounded}, zFi, Hz) = [zFi[end]   - zFi[end - 1] for k = 1:Hz]

function generate_stretched_vertical_grid(FT, z_topo, Nz, Hz, zF_generator)

    # Ensure correct type for zF and derived quantities
    interior_zF = zeros(FT, Nz+1)

    for k = 1:Nz+1
        interior_zF[k] = get_z_face(zF_generator, k)
    end

    Lz = interior_zF[Nz+1] - interior_zF[1]

    # Build halo regions
    ΔzF₋ = lower_exterior_Δzᵃᵃᶜ(z_topo, interior_zF, Hz)
    ΔzF₊ = lower_exterior_Δzᵃᵃᶜ(z_topo, interior_zF, Hz)

    z¹, zᴺ⁺¹ = interior_zF[1], interior_zF[Nz+1]

    zF₋ = [z¹   - sum(ΔzF₋[k:Hz]) for k = 1:Hz] # locations of faces in lower halo
    zF₊ = [zᴺ⁺¹ + ΔzF₊[k]         for k = 1:Hz] # locations of faces in width of top halo region

    zF = vcat(zF₋, interior_zF, zF₊)

    # Build cell centers, cell center spacings, and cell interface spacings
    TCz = total_length(Center, z_topo, Nz, Hz)
     zC = [ (zF[k + 1] + zF[k]) / 2 for k = 1:TCz ]
    ΔzC = [  zC[k] - zC[k - 1]      for k = 2:TCz ]

    # Trim face locations for periodic domains
    TFz = total_length(Face, z_topo, Nz, Hz)
    zF = zF[1:TFz]

    ΔzF = [zF[k + 1] - zF[k] for k = 1:TFz-1]

    return Lz, zF, zC, ΔzF, ΔzC
end

# We cannot reconstruct a VerticallyStretchedRectilinearGrid without the zF_generator.
# So the best we can do is tell the user what they should have done.
function with_halo(new_halo, old_grid::VerticallyStretchedRectilinearGrid)
    new_halo != halo_size(old_grid) &&
        @error "You need to construct your VerticallyStretchedRectilinearGrid with the keyword argument halo=$new_halo"
    return old_grid
end

@inline x_domain(grid::VerticallyStretchedRectilinearGrid{FT, TX, TY, TZ}) where {FT, TX, TY, TZ} = domain(TX, grid.Nx, grid.xᶠᵃᵃ)
@inline y_domain(grid::VerticallyStretchedRectilinearGrid{FT, TX, TY, TZ}) where {FT, TX, TY, TZ} = domain(TY, grid.Ny, grid.yᵃᶠᵃ)
@inline z_domain(grid::VerticallyStretchedRectilinearGrid{FT, TX, TY, TZ}) where {FT, TX, TY, TZ} = domain(TZ, grid.Nz, grid.zᵃᵃᶠ)

short_show(grid::VerticallyStretchedRectilinearGrid{FT, TX, TY, TZ}) where {FT, TX, TY, TZ} =
    "VerticallyStretchedRectilinearGrid{$FT, $TX, $TY, $TZ}(Nx=$(grid.Nx), Ny=$(grid.Ny), Nz=$(grid.Nz))"

function show(io::IO, g::VerticallyStretchedRectilinearGrid{FT, TX, TY, TZ}) where {FT, TX, TY, TZ}
    print(io, "VerticallyStretchedRectilinearGrid{$FT, $TX, $TY, $TZ}\n",
              "                   domain: $(domain_string(g))\n",
              "                 topology: ", (TX, TY, TZ), '\n',
              "  resolution (Nx, Ny, Nz): ", (g.Nx, g.Ny, g.Nz), '\n',
              "   halo size (Hx, Hy, Hz): ", (g.Hx, g.Hy, g.Hz))
end

Adapt.adapt_structure(to, grid::VerticallyStretchedRectilinearGrid{FT, TX, TY, TZ}) where {FT, TX, TY, TZ} =
    VerticallyStretchedRectilinearGrid{FT, TX, TY, TZ, typeof(grid.xᶠᵃᵃ), typeof(Adapt.adapt(to, grid.zᵃᵃᶠ))}(
        grid.Nx, grid.Ny, grid.Nz,
        grid.Hx, grid.Hy, grid.Hz,
        grid.Lx, grid.Ly, grid.Lz,
        grid.Δx, grid.Δy,
        Adapt.adapt(to, grid.Δzᵃᵃᶜ),
        Adapt.adapt(to, grid.Δzᵃᵃᶠ),
        grid.xᶜᵃᵃ, grid.yᵃᶜᵃ,
        Adapt.adapt(to, grid.zᵃᵃᶜ),
        grid.xᶠᵃᵃ, grid.yᵃᶠᵃ,
        Adapt.adapt(to, grid.zᵃᵃᶠ))

#####
##### Should merge with grid_utils.jl at some point
#####

@inline xnode(::Type{Center}, i, grid::VerticallyStretchedRectilinearGrid) = @inbounds grid.xᶜᵃᵃ[i]
@inline xnode(::Type{Face}, i, grid::VerticallyStretchedRectilinearGrid) = @inbounds grid.xᶠᵃᵃ[i]

@inline ynode(::Type{Center}, j, grid::VerticallyStretchedRectilinearGrid) = @inbounds grid.yᵃᶜᵃ[j]
@inline ynode(::Type{Face}, j, grid::VerticallyStretchedRectilinearGrid) = @inbounds grid.yᵃᵃᵃ[j]

@inline znode(::Type{Center}, k, grid::VerticallyStretchedRectilinearGrid) = @inbounds grid.zᵃᵃᶜ[k]
@inline znode(::Type{Face}, k, grid::VerticallyStretchedRectilinearGrid) = @inbounds grid.zᵃᵃᶠ[k]


text:
module JSServe

import AssetRegistry, Sockets
using UUIDs, Hyperscript, Hyperscript, JSON3, Observables
import Sockets: send, TCPServer
using Hyperscript: Node, children, tag
using HTTP, Markdown
using HTTP: Response, Request
using HTTP.Streams: Stream
using WebSockets
using Base64

include("compat.jl")
include("types.jl")
include("serialization.jl")
include("js_source.jl")
include("session.jl")
include("observables.jl")
include("dependencies.jl")
include("http.jl")
include("util.jl")
include("widgets.jl")
include("hyperscript_integration.jl")
include("display.jl")
include("jscall.jl")
include("markdown_integration.jl")

function __init__()
    url = get(ENV, "JULIA_WEBIO_BASEURL") do
        base = get(ENV, "WEBIO_SERVER_HOST_URL", "127.0.0.1")
        port = get(ENV, "WEBIO_HTTP_PORT", "8081")
        return "http://" * base * ":" * port
    end
    if endswith(url, "/")
        url = url[1:end-1]
    end
    server_proxy_url[] = url
end

end # module


text:
using NumCME
using Test
using BenchmarkTools
using SparseArrays
using LinearAlgebra: mul!, norm

𝕊 = [[-1, 1, 0] [1, -1, 0] [0, 0, 1] [0, 0, -1]]

x₀ = [1, 0, 0]
k₀₁ = 0.05
k₁₀ = 0.1
λ = 5.0
γ = 0.5
L = 20.0
θ = [k₀₁, k₁₀, λ, γ, L]
a1 = propensity() do x, p
    p[1] * x[1]
end
a2 = propensity() do x, p
    p[2] * x[2]
end
a3 = propensity() do x, p
    p[3] * x[2]
end
a4 = propensity((t, p) -> max(0.0, 1.0 - sin(π * t / p[5]))) do x, p
    p[4] * x[3]
end
propensities = [a1, a2, a3, a4]
RNACOUNT_MAX = 500
sensmodel = CmeModelWithSensitivity(CmeModel(𝕊, propensities, θ))
test_space = StateSpaceSparse(get_stoich_matrix(sensmodel), [[[1, 0, i] for i in 0:RNACOUNT_MAX]; [[0, 1, i] for i in 0:RNACOUNT_MAX]])

## Test that output vectors of mavec sum to ≈ 0
function test_zero_sum()
    pcount = get_parameter_count(sensmodel)
    sensmat = ForwardSensFspMatrixSparse{Float64}(sensmodel, test_space);
    n = get_rowcount(sensmat.fspmatrix)
    v = ones(n + pcount * n)
    out = similar(v)
    matvec!(out, 20.0, sensmat, v)
    return isapprox(sum(out) / sum(v), 0.0, atol = n*(pcount+1)*eps())
end
@test test_zero_sum()

## Test against analytic matrices
function getA(nmax::Integer, t::AbstractFloat, parameters::AbstractVector{<:AbstractFloat})
    rowids = Vector{UInt32}()
    colids = Vector{UInt32}()
    nzvals = Vector{Float64}()
    for gene in 0:1
        for rnacount in 0:nmax
            x = [1-gene, gene, rnacount]
            idx = gene*(nmax+1) + rnacount+1
            # Diagonal
            push!(rowids, idx)
            push!(colids, idx)
            push!(nzvals, -1.0*(a1(x, parameters) + a2(x, parameters) + a3(x, parameters) + a4(t, x, parameters)))
            # Offdiagonal for goff->gon
            if gene == 0
                push!(colids, idx)
                push!(rowids, nmax+1 + rnacount+1)
                push!(nzvals, a1(x, parameters))
            end
            # Offdiagonal for gon->goff
            if gene == 1
                push!(colids, idx)
                push!(rowids, rnacount+1)
                push!(nzvals, a2(x, parameters))
            end
            # Offdiagonal for gon->gon+rna
            if gene == 1
                push!(colids, idx)
                push!(rowids, (rnacount < nmax) ? nmax+1+rnacount+2 : 2*(nmax+1)+3)
                push!(nzvals, a3(x, parameters))
            end
            # Offdiagonal for rna->0
            if rnacount > 0
                push!(colids, idx)
                push!(rowids, gene*(nmax+1) + rnacount)
                push!(nzvals, a4(t, x, parameters))
            end
        end
    end
    return sparse(rowids, colids, nzvals, 2*(nmax+1)+4, 2*(nmax+1)+4)
end
# Generate the partial derivative matrix wrt k10 corresponding to a rectangular state space
function ∂A∂k01(nmax::Integer)
    sink_count = 4
    state_count = 2 * (nmax + 1)
    rowids = Vector{UInt32}()
    colids = Vector{UInt32}()
    nzvals = Vector{Float64}()
    for rnacount in 0:nmax
        # Fill column corresponding to state (1, 0, rnacount)
        push!(rowids, rnacount + 1)
        push!(colids, rnacount + 1)
        push!(nzvals, -1.0)
        push!(rowids, nmax + 1 + rnacount + 1)
        push!(colids, rnacount + 1)
        push!(nzvals, 1.0)
    end
    return sparse(rowids, colids, nzvals, state_count + sink_count, state_count + sink_count)
end
function ∂A∂k10(nmax::Integer)
    sink_count = 4
    state_count = 2 * (nmax + 1)
    rowids = Vector{UInt32}()
    colids = Vector{UInt32}()
    nzvals = Vector{Float64}()
    for rnacount in 0:nmax
        # Fill column corresponding to state (0, 1, rnacount)
        push!(colids, nmax + 1 + rnacount + 1)
        push!(rowids, nmax + 1 + rnacount + 1)
        push!(nzvals, -1.0)
        push!(colids, nmax + 1 + rnacount + 1)
        push!(rowids, rnacount + 1)
        push!(nzvals, 1.0)
    end
    return sparse(rowids, colids, nzvals, state_count + sink_count, state_count + sink_count)
end
function ∂A∂λ(nmax::Integer)
    sink_count = 4
    state_count = 2 * (nmax + 1)
    rowids = Vector{UInt32}()
    colids = Vector{UInt32}()
    nzvals = Vector{Float64}()
    for rnacount in 0:nmax
        # Fill column corresponding to state (0, 1, rnacount)
        push!(colids, nmax + 1 + rnacount + 1)
        push!(rowids, nmax + 1 + rnacount + 1)
        push!(nzvals, -1.0)
        push!(colids, nmax + 1 + rnacount + 1)
        push!(rowids, (rnacount < nmax) ? nmax + 1 + rnacount + 2 : state_count + 3)
        push!(nzvals, 1.0)
    end
    return sparse(rowids, colids, nzvals, state_count + sink_count, state_count + sink_count)
end
function ∂A∂γ(nmax::Integer, t, p)
    sink_count = 4
    state_count = 2 * (nmax + 1)
    rowids = Vector{UInt32}()
    colids = Vector{UInt32}()
    nzvals = Vector{Float64}()
    for rnacount in 1:nmax
        # Fill column corresponding to state (1, 0, rnacount)
        push!(rowids, rnacount + 1)
        push!(colids, rnacount + 1)
        push!(nzvals, -max(0.0, 1.0 - sin(π * t / p[5])) * rnacount)
        push!(rowids, rnacount)
        push!(colids, rnacount + 1)
        push!(nzvals, max(0.0, 1.0 - sin(π * t / p[5])) * rnacount)

        # Fill column corresponding to state (0, 1, rnacount)
        push!(colids, nmax + 1 + rnacount + 1)
        push!(rowids, nmax + 1 + rnacount + 1)
        push!(nzvals, -max(0.0, 1.0 - sin(π * t / p[5])) * rnacount)
        push!(rowids, nmax + 1 + rnacount)
        push!(colids, nmax + 1 + rnacount + 1)
        push!(nzvals, max(0.0, 1.0 - sin(π * t / p[5])) * rnacount)

    end
    return sparse(rowids, colids, nzvals, state_count + sink_count, state_count + sink_count)
end
function ∂A∂L(nmax::Integer, t, p)
    sink_count = 4
    state_count = 2 * (nmax + 1)
    rowids = Vector{UInt32}()
    colids = Vector{UInt32}()
    nzvals = Vector{Float64}()
    for rnacount in 1:nmax
        # Fill column corresponding to state (1, 0, rnacount)
        push!(rowids, rnacount + 1)
        push!(colids, rnacount + 1)
        push!(nzvals, (1.0-sin(π*t/p[5]) >= 0) ? -1.0*p[4]*rnacount*cos(π*t/p[5])*π*t/(p[5]^2) : 0.0)
        push!(rowids, rnacount)
        push!(colids, rnacount + 1)
        push!(nzvals, (1.0-sin(π*t/p[5]) >= 0) ? 1.0*p[4]*rnacount*cos(π*t/p[5])*π*t/(p[5]^2) : 0.0)

        # Fill column corresponding to state (0, 1, rnacount)
        push!(colids, nmax + 1 + rnacount + 1)
        push!(rowids, nmax + 1 + rnacount + 1)
        push!(nzvals, (1.0-sin(π*t/p[5]) >= 0) ? -1.0*p[4]*rnacount*cos(π*t/p[5])*π*t/(p[5]^2) : 0.0)
        push!(rowids, nmax + 1 + rnacount)
        push!(colids, nmax + 1 + rnacount + 1)
        push!(nzvals, (1.0-sin(π*t/p[5]) >= 0) ? 1.0*p[4]*rnacount*cos(π*t/p[5])*π*t/(p[5]^2) : 0.0)
    end
    return sparse(rowids, colids, nzvals, state_count + sink_count, state_count + sink_count)
end

function test_against_analytic_matrix(t_test::AbstractFloat, rnacount_max::Integer)
    sensmat = ForwardSensFspMatrixSparse{Float64}(sensmodel, test_space)
    pcount = get_parameter_count(sensmodel)
    n = get_rowcount(sensmat.fspmatrix)
    v = ones(n + pcount * n)
    v ./= sum(v)
    out_sensmat = similar(v)
    matvec!(out_sensmat, t_test, sensmat, v)

    A = getA(rnacount_max, t_test, θ)
    dAs = []
    push!(dAs, ∂A∂k01(rnacount_max))
    push!(dAs, ∂A∂k10(rnacount_max))
    push!(dAs, ∂A∂λ(rnacount_max))
    push!(dAs, ∂A∂γ(rnacount_max, t_test, θ))
    push!(dAs, ∂A∂L(rnacount_max, t_test, θ))
    out_ref = similar(v)
    mul!(view(out_ref, 1:n), A, view(v, 1:n))
    for ip in 1:pcount
        mul!(view(out_ref, ip*n+1:(ip+1)*n), dAs[ip], view(v, 1:n))
        mul!(view(out_ref, ip*n+1:(ip+1)*n), A, view(v, ip*n+1:(ip+1)*n), 1, 1)
    end
    return isapprox(norm(out_sensmat- out_ref, Inf), 0.0, atol=n*eps())
end

for t_test in [10.0, 20.0, 30.0, 100.0]
    @test test_against_analytic_matrix(t_test, RNACOUNT_MAX)
end


text:
__precompile__(true)

module ColorVectorSpace

using Colors, FixedPointNumbers, SpecialFunctions

import Base: ==, +, -, *, /, ^, <, ~
import Base: abs, abs2, clamp, convert, copy, div, eps, isfinite, isinf,
    isnan, isless, length, mapreduce, oneunit,
    promote_op, promote_rule, zero, trunc, floor, round, ceil, bswap,
    mod, rem, atan, hypot, max, min, real, typemin, typemax
import LinearAlgebra: norm
import StatsBase: histrange, varm
import SpecialFunctions: gamma, lgamma, lfact
import Statistics: middle

export nan

# The unaryOps
import Base:      conj, sin, cos, tan, sinh, cosh, tanh,
                  asin, acos, atan, asinh, acosh, atanh,
                  sec, csc, cot, asec, acsc, acot,
                  sech, csch, coth, asech, acsch, acoth,
                  sinc, cosc, cosd, cotd, cscd, secd,
                  sind, tand, acosd, acotd, acscd, asecd,
                  asind, atand, rad2deg, deg2rad,
                  log, log2, log10, log1p, exponent, exp,
                  exp2, exp10, expm1, cbrt, sqrt,
                  significand,
                  frexp, modf,
                  float

export dotc

# TODO: we get rid of these definitions or move them to ColorTypes.jl
TransparentRGBFloat{C<:AbstractRGB,T<:AbstractFloat} = TransparentColor{C,T,4}
TransparentGrayFloat{C<:AbstractGray,T<:AbstractFloat} = TransparentColor{C,T,2}
TransparentRGBNormed{C<:AbstractRGB,T<:Normed} = TransparentColor{C,T,4}
TransparentGrayNormed{C<:AbstractGray,T<:Normed} = TransparentColor{C,T,2}

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

# convert(RGB{Float32}, NaN) doesn't and shouldn't work, so we need to reintroduce nan
nan(::Type{T}) where {T<:AbstractFloat} = convert(T, NaN)
nan(::Type{C}) where {C<:MathTypes} = _nan(eltype(C), C)
_nan(::Type{T}, ::Type{C}) where {T<:AbstractFloat,C<:AbstractGray} = (x = convert(T, NaN); C(x))
_nan(::Type{T}, ::Type{C}) where {T<:AbstractFloat,C<:TransparentGray} = (x = convert(T, NaN); C(x,x))
_nan(::Type{T}, ::Type{C}) where {T<:AbstractFloat,C<:AbstractRGB} = (x = convert(T, NaN); C(x,x,x))
_nan(::Type{T}, ::Type{C}) where {T<:AbstractFloat,C<:TransparentRGB} = (x = convert(T, NaN); C(x,x,x,x))

## Generic algorithms
mapreduce(f, op::Union{typeof(&), typeof(|)}, a::MathTypes) = f(a)  # ambiguity
mapreduce(f, op, a::MathTypes) = f(a)
Base.add_sum(c1::MathTypes,c2::MathTypes) = mapc(Base.add_sum, c1, c2)
Base.reduce_first(::typeof(Base.add_sum), c::MathTypes) = mapc(x->Base.reduce_first(Base.add_sum, x), c)
function Base.reduce_empty(::typeof(Base.add_sum), ::Type{T}) where {T<:MathTypes}
    z = Base.reduce_empty(Base.add_sum, eltype(T))
    return zero(base_colorant_type(T){typeof(z)})
end

for f in (:trunc, :floor, :round, :ceil, :eps, :bswap)
    @eval $f(g::Gray{T}) where {T} = Gray{T}($f(gray(g)))
end
eps(::Type{Gray{T}}) where {T} = Gray(eps(T))

for f in (:trunc, :floor, :round, :ceil)
    @eval $f(::Type{T}, g::Gray) where {T<:Integer} = Gray{T}($f(T, gray(g)))
end

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

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

dotc(x::T, y::T) where {T<:Real} = acc(x)*acc(y)
dotc(x::Real, y::Real) = dotc(promote(x, y)...)

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

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

acctype(::Type{T}) where {T<:FixedPoint} = FixedPointNumbers.floattype(T)
acctype(::Type{T}) where {T<:Number} = T

acc(x::Number) = convert(acctype(typeof(x)), x)

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

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

arith_colorant_type(::C) where {C<:Colorant} = arith_colorant_type(C)
arith_colorant_type(::Type{C}) where {C<:Colorant} = base_colorant_type(C)
arith_colorant_type(::Type{Gray24}) = Gray
arith_colorant_type(::Type{AGray32}) = AGray
arith_colorant_type(::Type{RGB24}) = RGB
arith_colorant_type(::Type{ARGB32}) = ARGB

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

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

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

oneunit(::Type{C}) where {C<:AbstractRGB}     = C(1,1,1)
oneunit(::Type{C}) where {C<:TransparentRGB}  = C(1,1,1,1)

zero(::Type{C}) where {C<:AbstractRGB}    = C(0,0,0)
zero(::Type{C}) where {C<:TransparentRGB} = C(0,0,0,0)
zero(::Type{C}) where {C<:YCbCr} = C(0,0,0)
zero(::Type{C}) where {C<:HSV} = C(0,0,0)
oneunit(p::Colorant) = oneunit(typeof(p))
Base.one(c::Colorant) = Base.one(typeof(c))
zero(p::Colorant) = zero(typeof(p))

# These constants come from squaring the conversion to grayscale
# (rec601 luma), and normalizing
dotc(x::T, y::T) where {T<:AbstractRGB} = 0.200f0 * acc(red(x))*acc(red(y)) + 0.771f0 * acc(green(x))*acc(green(y)) + 0.029f0 * acc(blue(x))*acc(blue(y))
dotc(x::AbstractRGB, y::AbstractRGB) = dotc(promote(x, y)...)

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

middle(c::AbstractGray) = arith_colorant_type(c)(middle(gray(c)))
middle(x::C, y::C) where {C<:AbstractGray} = arith_colorant_type(C)(middle(gray(x), gray(y)))

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

norm(c::AbstractGray) = abs(gray(c))
abs(c::TransparentGray) = abs(gray(c))+abs(alpha(c)) # should this have a different name?
abs(c::TransparentGrayNormed) = Float32(gray(c)) + Float32(alpha(c)) # should this have a different name?
abs2(c::AbstractGray) = gray(c)^2
abs2(c::AbstractGray{T}) where {T<:Normed} = Float32(gray(c))^2
abs2(c::TransparentGray) = gray(c)^2+alpha(c)^2
abs2(c::TransparentGrayNormed) = Float32(gray(c))^2 + Float32(alpha(c))^2
atan(x::Gray, y::Gray) = atan(convert(Real, x), convert(Real, y))
hypot(x::Gray, y::Gray) = hypot(convert(Real, x), convert(Real, y))
norm(c::TransparentGray) = sqrt(abs2(c))

(<)(g1::AbstractGray, g2::AbstractGray) = gray(g1) < gray(g2)
(<)(c::AbstractGray, r::Real) = gray(c) < r
(<)(r::Real, c::AbstractGray) = r < gray(c)
isless(g1::AbstractGray, g2::AbstractGray) = isless(gray(g1), gray(g2))
isless(c::AbstractGray, r::Real) = isless(gray(c), r)
isless(r::Real, c::AbstractGray) = isless(r, gray(c))

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

zero(::Type{C}) where {C<:TransparentGray} = C(0,0)
oneunit(::Type{C}) where {C<:TransparentGray} = C(1,1)

dotc(x::T, y::T) where {T<:AbstractGray} = acc(gray(x))*acc(gray(y))
dotc(x::AbstractGray, y::AbstractGray) = dotc(promote(x, y)...)

float(::Type{T}) where {T<:Gray} = typeof(float(zero(T)))

# Mixed types
(+)(a::MathTypes, b::MathTypes) = (+)(promote(a, b)...)
(-)(a::MathTypes, b::MathTypes) = (-)(promote(a, b)...)

# Arrays---necessary methods
+(A::AbstractArray{C}) where {C<:MathTypes} = A
+(A::Array{C}) where {C<:MathTypes} = A

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

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

# To help type inference
promote_rule(::Type{T}, ::Type{C}) where {T<:Real,C<:AbstractGray} = promote_type(T, eltype(C))

typemin(::Type{T}) where {T<:ColorTypes.AbstractGray} = T(typemin(eltype(T)))
typemax(::Type{T}) where {T<:ColorTypes.AbstractGray} = T(typemax(eltype(T)))

typemin(::T) where {T<:ColorTypes.AbstractGray} = T(typemin(eltype(T)))
typemax(::T) where {T<:ColorTypes.AbstractGray} = T(typemax(eltype(T)))

## Deprecations

@deprecate (+)(A::AbstractArray{CV}, b::AbstractRGB) where {CV<:AbstractRGB} (.+)(A, b)
@deprecate (+)(b::AbstractRGB, A::AbstractArray{CV}) where {CV<:AbstractRGB} (.+)(b, A)
@deprecate (-)(A::AbstractArray{CV}, b::AbstractRGB) where {CV<:AbstractRGB} (.-)(A, b)
@deprecate (-)(b::AbstractRGB, A::AbstractArray{CV}) where {CV<:AbstractRGB} (.-)(b, A)
@deprecate (*)(A::AbstractArray{T}, b::AbstractRGB) where {T<:Number} A.*b
@deprecate (*)(b::AbstractRGB, A::AbstractArray{T}) where {T<:Number} A.*b

@deprecate (+)(A::AbstractArray{CV}, b::TransparentRGB) where {CV<:TransparentRGB} (.+)(A, b)
@deprecate (+)(b::TransparentRGB, A::AbstractArray{CV}) where {CV<:TransparentRGB} (.+)(b, A)
@deprecate (-)(A::AbstractArray{CV}, b::TransparentRGB) where {CV<:TransparentRGB} (.-)(A, b)
@deprecate (-)(b::TransparentRGB, A::AbstractArray{CV}) where {CV<:TransparentRGB} (.-)(b, A)
@deprecate (*)(A::AbstractArray{T}, b::TransparentRGB) where {T<:Number} A.*b
@deprecate (*)(b::TransparentRGB, A::AbstractArray{T}) where {T<:Number} A.*b

@deprecate (+)(A::AbstractArray{CV}, b::AbstractGray) where {CV<:AbstractGray} (.+)(A, b)
@deprecate (+)(b::AbstractGray, A::AbstractArray{CV}) where {CV<:AbstractGray} (.+)(b, A)
@deprecate (-)(A::AbstractArray{CV}, b::AbstractGray) where {CV<:AbstractGray} (.-)(A, b)
@deprecate (-)(b::AbstractGray, A::AbstractArray{CV}) where {CV<:AbstractGray} (.-)(b, A)
@deprecate (*)(A::AbstractArray{T}, b::AbstractGray) where {T<:Number} A.*b
@deprecate (*)(b::AbstractGray, A::AbstractArray{T}) where {T<:Number} A.*b
@deprecate (/)(A::AbstractArray{C}, b::AbstractGray) where {C<:AbstractGray} A./b

@deprecate (+)(A::AbstractArray{CV}, b::TransparentGray) where {CV<:TransparentGray} (.+)(A, b)
@deprecate (+)(b::TransparentGray, A::AbstractArray{CV}) where {CV<:TransparentGray} (.+)(b, A)
@deprecate (-)(A::AbstractArray{CV}, b::TransparentGray) where {CV<:TransparentGray} (.-)(A, b)
@deprecate (-)(b::TransparentGray, A::AbstractArray{CV}) where {CV<:TransparentGray} (.-)(b, A)
@deprecate (*)(A::AbstractArray{T}, b::TransparentGray) where {T<:Number} A.*b
@deprecate (*)(b::TransparentGray, A::AbstractArray{T}) where {T<:Number} A.*b

end


text:
a = [1,2,3+1im,8,5]
V=Vandermonde(a)

# Test element
@test V[4,5] == a[4]^4

# Test full matrix conversions
M = Matrix(V)
@test Matrix(V') == M'
@test Matrix(transpose(V)) == transpose(M)

# Test solving with vector and matrix rhs
y = [1im,1,5,0,2]
Y = [y 2*y]
for rhs=[y, Y]
    # Test solution
    @test isapprox(V\rhs, M\rhs)
    @test isapprox(V'\rhs, M'\rhs)
    @test isapprox(rhs'/V, rhs'/M)
    @test isapprox(transpose(V)\rhs, transpose(M)\rhs)
    # Check that overloading works
    x = zero(M\rhs)
    copyto!(x, rhs)
    SpecialMatrices.dvand!(a, x)
    @test V\rhs==x
    copyto!(x, rhs)
    SpecialMatrices.pvand!(a', x)
    @test V'\rhs==x
    @test rhs'/V==x'
    copyto!(x, rhs)
    SpecialMatrices.pvand!(a, x)
    @test transpose(V)\rhs==x
end

# Test dimension errors
rhs = zeros(2,2)
try
    V\rhs
catch e
    @test typeof(e)==DimensionMismatch
end
try
    V'\rhs
catch e
    @test typeof(e)==DimensionMismatch
end


text:
function convertInputsToFloats(U, g, lmd)
    U = convert(Array{Float64}, U)
    g = convert(Array{Float64}, g)
    lmd = convert(Float64, lmd)
    return U,g,lmd
end

function cancelZeroStates(U,g)
    ## cancel states with prior zero
    nonzeroStates = findall(a->a!=0, g)
    U = U[:,nonzeroStates]
    g = g[nonzeroStates]
    return U,g,nonzeroStates
end

function transform_utilities(U, lmd)
    # shiftterm = (((maximum(U)+minimum(U))/2))   ## reduce extremes of U for numerical stability
    # Z = exp.((U.-shiftterm)/lmd)    ## transform utilities
    shiftterm = 0
    Z = exp.(U/lmd)    ## transform utilities
    return Z, shiftterm
end

function checkForTriviality(Z,g,m)
    ## interupt if there is no RI problem
    if m==1  ## since only one action possible
        return ones(m), (log.(Z)*g)[1], 0, "Trivial: only one action is possible"
    end
    if typeof(Z)==Vector{Float64}  ## since only one state possible
        best_solution = findmax(Z[:,1])
        p = zeros(m)
        p[best_solution[2]]=1
        return p, log(best_solution[1]), 0, "Trivial: only one state is possible"
    end
    return "not trivial"
end

function checkForBoundarySolution(Z,g)
    ## check whether you have boundary solutions in case of two alternatives
    if (Z[1,:]./Z[2,:])'*g<1
        return [0, 1.0], log.(Z[2,:])'*g, 0,  "Found boundary solution (two alternatives only)"
    elseif (Z[2,:]./Z[1,:])'*g<1
        return [1.0, 0], log.(Z[1,:])'*g, 0,  "Found boundary solution (two alternatives only)"
    end
    return "no boundary solution"
end

function preprocessinput(U, g, lmd)
    m = size(U)[1]
    solution = (ones(m)/m,0,0,"Run Cover Algorithm")

    U,g = convertInputsToFloats(U, g, lmd)
    U,g,nonzeroStates = cancelZeroStates(U,g)
    Z,shiftterm = transform_utilities(U,lmd)
    intermediateSolution = checkForTriviality(Z,g,m)
    if intermediateSolution != "not trivial"
          solution = intermediateSolution
    end
    if m == 2
        intermediateSolution = checkForBoundarySolution(Z,g)
        if intermediateSolution != "no boundary solution"
              solution = intermediateSolution
        end
    end
    return Z,g,shiftterm,solution
end


text:
# 2D Curevelt operators: joCurvelet2D

## helper module
module joCurvelet2D_etc
    using JOLI: jo_convert, joLinearFunctionException
    function apply_fdct2Dwrap_real(v::Vector{vdt},n1::Integer,n2::Integer,m::Int128,rdt::DataType,nbs::Integer,nbac::Integer,actl::Integer,rctl::Integer,zfin::Integer) where vdt<:Union{AbstractFloat,Complex}
        C=Vector{Cdouble}(undef,m)
        eltype(v)<:Real || throw(joLinearFunctionException("joCurvelt2D: input vector must be real for real transform"))
        X=jo_convert(Cdouble,v,false)
        ccall((:jl_fdct_wrapping_real,:libdfdct_wrapping),Nothing,
            (Cint,Cint,Cint,Cint,Cint,Cint,Cint,Csize_t,Ptr{Array{Cdouble}},Ptr{Array{Cdouble}}),
            n1,n2,nbs,nbac,actl,rctl,zfin,m,X,C)
        C=jo_convert(rdt,C,false)
        return C
    end
    function apply_ifdct2Dwrap_real(v::Vector{vdt},n1::Integer,n2::Integer,m::Int128,rdt::DataType,nbs::Integer,nbac::Integer,actl::Integer,rctl::Integer,zfin::Integer) where vdt<:Union{AbstractFloat,Complex}
        X=Vector{Cdouble}(undef,n1*n2)
        eltype(v)<:Real || throw(joLinearFunctionException("joCurvelt2D: input vector must be real for real transform"))
        C=jo_convert(Cdouble,v,false)
        ccall((:jl_ifdct_wrapping_real,:libdfdct_wrapping),Nothing,
            (Cint,Cint,Cint,Cint,Cint,Cint,Cint,Csize_t,Ptr{Array{Cdouble}},Ptr{Array{Cdouble}}),
            n1,n2,nbs,nbac,actl,rctl,zfin,m,C,X)
        X=jo_convert(rdt,X,false)
        return X
    end
    function apply_fdct2Dwrap_cplx(v::Vector{vdt},n1::Integer,n2::Integer,m::Int128,rdt::DataType,nbs::Integer,nbac::Integer,actl::Integer,rctl::Integer,zfin::Integer) where vdt<:Union{AbstractFloat,Complex}
        C=Vector{Complex{Cdouble}}(undef,m)
        X=jo_convert(Complex{Cdouble},v,false)
        ccall((:jl_fdct_wrapping_cpx,:libdfdct_wrapping),Nothing,
            (Cint,Cint,Cint,Cint,Cint,Cint,Cint,Csize_t,Ptr{Array{Complex{Cdouble}}},Ptr{Array{Complex{Cdouble}}}),
            n1,n2,nbs,nbac,actl,rctl,zfin,m,X,C)
        C=jo_convert(rdt,C,false)
        return C
    end
    function apply_ifdct2Dwrap_cplx(v::Vector{vdt},n1::Integer,n2::Integer,m::Int128,rdt::DataType,nbs::Integer,nbac::Integer,actl::Integer,rctl::Integer,zfin::Integer) where vdt<:Union{AbstractFloat,Complex}
        X=Vector{Complex{Cdouble}}(undef,n1*n2)
        C=jo_convert(Complex{Cdouble},v,false)
        ccall((:jl_ifdct_wrapping_cpx,:libdfdct_wrapping),Nothing,
            (Cint,Cint,Cint,Cint,Cint,Cint,Cint,Csize_t,Ptr{Array{Complex{Cdouble}}},Ptr{Array{Complex{Cdouble}}}),
            n1,n2,nbs,nbac,actl,rctl,zfin,m,C,X)
        X=jo_convert(rdt,X,false)
        return X
    end
end
using .joCurvelet2D_etc

export joCurvelet2D
"""
    julia> op = joCurvelet2D(n1,n2;[DDT=joFloat,][RDT=DDT,]
                    [nbscales=...,][nbangles_coarse=...,][all_crvlts=...,]
                    [real_crvlts=...,][zero_finest=...,][name=...])

2D Curvelet transform (wrapping) over fast dimensions

# Signature

    function joCurvelet2D(n1::Integer,n2::Integer;DDT::DataType=joFloat,RDT::DataType=DDT,
            nbscales::Integer=0,
            nbangles_coarse::Integer=16,
            all_crvlts::Bool=false,
            real_crvlts::Bool=true,
            zero_finest::Bool=false,
            name::String="joCurvelt2D")

# Arguments

- `n1`,`n2`: image sizes
- keywords
    - `nbscales`: # of scales (requires #>=default; defaults to max(1,ceil(log2(min(n1,n2))-3)))
    - `nbangles_coarse`: # of angles at coarse scale (requires #%4==0, #>=8; defaults to 16)
    - `all_crvlts`: curvelets at finnest scales (defaults to false)
    - `real_crvlts`: real transform (defaults to true) and requires real input
    - `zero_finest`: zero out finnest scales (defaults to false)
    - `DDT`: domain data type
    - `RDT`: range data type
    - `name`: custom name

# Notes

- if DDT:<Real for complex transform then imaginary part will be neglected for transpose/adjoint
- isadjoint test at larger sizes (above 128) might require reseting tollerance to bigger number.

# Examples

real transform (64-bit)

    joCurvelet2D(32,32)

complex transform (64-bit)

    joCurvelet2D(32,32;real_crvlts=false)

real transform with curevelts at the finnest scales (64-bit)

    joCurvelet2D(32,32;all_crvlts=true)

real transform with zeros at the finnest scales (64-bit)

    joCurvelet2D(32,32;zero_finest=true)

complex transform with real 64-bit input for forward

    joCurvelet2D(32,32;DDT=Float64,real_crvlts=false)

complex transform with just precision specification for curvelets

    joCurvelet2D(32,32;DDT=Float32,RDT=Float64,real_crvlts=false)

complex transform with full type specification for curvelets (same as above)

    joCurvelet2D(32,32;DDT=Float32,RDT=ComplexF64,real_crvlts=false)

"""
function joCurvelet2D(n1::Integer,n2::Integer;DDT::DataType=joFloat,RDT::DataType=DDT,
            nbscales::Integer=0,
            nbangles_coarse::Integer=16,
            all_crvlts::Bool=false,
            real_crvlts::Bool=true,
            zero_finest::Bool=false,
            name::String="joCurvelt2D")

    nbs=convert(Cint,max(1,ceil(log2(min(n1,n2))-3)))
    nbs=convert(Cint,max(nbs,nbscales))
    nbs > 1 || throw(joLinearFunctionException("joCurvelt2D: not enough elements in one of dimensions"))
    if nbscales!=0 && nbs>nbscales
        @warn "Adjusted number of scales to required $nbs"
    end
    nbac=convert(Cint,max(8,nbangles_coarse))
    nbac%4==0 || throw(joLinearFunctionException("joCurvelt2D: nbangles_coarse must be multiple of 4"))
    if nbangles_coarse < 8
        @warn "Adjusted number of coarse angles to required $nbac"
    end
    actl=convert(Cint,all_crvlts)
    rctl=convert(Cint,real_crvlts)
    zfin=convert(Cint,zero_finest)
    if real_crvlts
        dtp=DDT
        rtp=RDT
        apply_fdct2Dwrap=joCurvelet2D_etc.apply_fdct2Dwrap_real
        apply_ifdct2Dwrap=joCurvelet2D_etc.apply_ifdct2Dwrap_real
        myname=name*"_wrapReal"
    else
        dtp=DDT
        rtp= RDT<:Complex ? RDT : Complex{RDT}
        apply_fdct2Dwrap=joCurvelet2D_etc.apply_fdct2Dwrap_cplx
        apply_ifdct2Dwrap=joCurvelet2D_etc.apply_ifdct2Dwrap_cplx
        myname=name*"_wrapCplx"
    end
    cfmap_size=ccall((:jl_fdct_sizes_map_size,:libdfdct_wrapping),Cint,
        (Cint,Cint,Cint),nbs,nbac,all_crvlts)
    cfmap=Vector{Cint}(undef,cfmap_size)
    m=Ref{Csize_t}(0)
    ccall((:jl_fdct_sizes,:libdfdct_wrapping),Nothing,
        (Cint,Cint,Cint,Cint,Cint,Ptr{Array{Cint}},Ref{Csize_t}),
        nbs,nbac,actl,n1,n2,cfmap,m)
    m=convert(Int128,m[])

    return joLinearFunction_A(m,n1*n2,
        v1->apply_fdct2Dwrap(v1,n1,n2,m,rtp,nbs,nbac,actl,rctl,zfin),
        v2->apply_ifdct2Dwrap(v2,n1,n2,m,dtp,nbs,nbac,actl,rctl,zfin),
        v3->apply_ifdct2Dwrap(v3,n1,n2,m,dtp,nbs,nbac,actl,rctl,zfin),
        v4->apply_fdct2Dwrap(v4,n1,n2,m,rtp,nbs,nbac,actl,rctl,zfin),
        dtp,rtp;
        name=myname
        )
end


text:
import ..Var

struct LinearΩ{TAGS, T, S, K}  <: AbstractΩ
  data::T
  tags::TAGS
  keymap::K
  subspace::S
end

function Var.recurse(exo::Var.PrimRandVar, ω::LinearΩ)
  newid = append(ω.subspace, exo.id)
  newexo = Var.Member(newid, exo.class)
  index = ω.keymap[newexo]
  ω.data[index]
end

## Tags
traits(::Type{LinearΩ{TAGS, T, S, K}}) where {TAGS, T, S, K} = traits(TAuGS)

max_stars_repo_name: zenna/expect

max_issues_repo_name: zenna/expect

max_forks_repo_name: zenna/expect


text:
module PolyPairPots

import JuLIP: energy, forces, virial, alloc_temp, alloc_temp_d, cutoff

import JuLIP.Potentials:  evaluate, evaluate_d, evaluate!, evaluate_d!
import Base: Dict, convert, ==

function alloc_B end
function alloc_dB end

include("jacobi.jl")

include("transforms.jl")

include("basis.jl")

include("calculator.jl")

include("repulsion.jl")

end # module


text:
#=
This file is auto-generated. Do not edit.
=#
"""
    mutable struct AVRFixed <: AVR
        Vf::Float64
        V_ref::Float64
        ext::Dict{String, Any}
        states::Vector{Symbol}
        n_states::Int
        states_types::Vector{StateTypes.StateType}
        internal::InfrastructureSystemsInternal
    end

Parameters of a AVR that returns a fixed voltage to the rotor winding

# Arguments
- `Vf::Float64`: Fixed voltage field applied to the rotor winding, validation range: `(0, nothing)`
- `V_ref::Float64`: Reference Voltage Set-point, validation range: `(0, nothing)`
- `ext::Dict{String, Any}`
- `states::Vector{Symbol}`: Fixed AVR has no states
- `n_states::Int`: Fixed AVR has no states
- `states_types::Vector{StateTypes.StateType}`: Fixed AVR has no states
- `internal::InfrastructureSystemsInternal`: power system internal reference, do not modify
"""
mutable struct AVRFixed <: AVR
    "Fixed voltage field applied to the rotor winding"
    Vf::Float64
    "Reference Voltage Set-point"
    V_ref::Float64
    ext::Dict{String, Any}
    "Fixed AVR has no states"
    states::Vector{Symbol}
    "Fixed AVR has no states"
    n_states::Int
    "Fixed AVR has no states"
    states_types::Vector{StateTypes.StateType}
    "power system internal reference, do not modify"
    internal::InfrastructureSystemsInternal
end

function AVRFixed(Vf, V_ref=1.0, ext=Dict{String, Any}(), )
    AVRFixed(Vf, V_ref, ext, Vector{Symbol}(), 0, Vector{StateTypes.StateType}(), InfrastructureSystemsInternal(), )
end

function AVRFixed(; Vf, V_ref=1.0, ext=Dict{String, Any}(), states=Vector{Symbol}(), n_states=0, states_types=Vector{StateTypes.StateType}(), internal=InfrastructureSystemsInternal(), )
    AVRFixed(Vf, V_ref, ext, states, n_states, states_types, internal, )
end

# Constructor for demo purposes; non-functional.
function AVRFixed(::Nothing)
    AVRFixed(;
        Vf=0,
        V_ref=0,
        ext=Dict{String, Any}(),
    )
end

"""Get [`AVRFixed`](@ref) `Vf`."""
get_Vf(value::AVRFixed) = value.Vf
"""Get [`AVRFixed`](@ref) `V_ref`."""
get_V_ref(value::AVRFixed) = value.V_ref
"""Get [`AVRFixed`](@ref) `ext`."""
get_ext(value::AVRFixed) = value.ext
"""Get [`AVRFixed`](@ref) `states`."""
get_states(value::AVRFixed) = value.states
"""Get [`AVRFixed`](@ref) `n_states`."""
get_n_states(value::AVRFixed) = value.n_states
"""Get [`AVRFixed`](@ref) `states_types`."""
get_states_types(value::AVRFixed) = value.states_types
"""Get [`AVRFixed`](@ref) `internal`."""
get_internal(value::AVRFixed) = value.internal

"""Set [`AVRFixed`](@ref) `Vf`."""
set_Vf!(value::AVRFixed, val) = value.Vf = val
"""Set [`AVRFixed`](@ref) `V_ref`."""
set_V_ref!(value::AVRFixed, val) = value.V_ref = val
"""Set [`AVRFixed`](@ref) `ext`."""
set_ext!(value::AVRFixed, val) = value.ext = val
"""Set [`AVRFixed`](@ref) `states_types`."""
set_states_types!(value::AVRFixed, val) = value.states_types = val


text:
function yeast(;
  R=50, L=2, RL=0, G=50, Ga=0, Gbg=0, Gd=0,
  k1=0.0038, k2=4e-4, k3=0.042, k4=0.01, k5=0.011, k6=0.1, k7=1.05e-3, k8=3.21)
  m = Network("Yeast polarization")

  m <= Species("R", R) # receptor
  m <= Species("L", L) # ligand
  m <= Species("RL", RL) # receptor-ligand
  m <= Species("G", G) # G protein
  m <= Species("Ga", Ga) # G alpha subunit
  m <= Species("Gbg", Gbg) # G beta-gamma subunit
  m <= Species("Gd", Gd) # G delta subunit

  m <= Reaction("receptor upregulation", k1, "0 --> R")
  m <= Reaction("receptor downregulation", k2, "R --> 0")
  m <= Reaction("ligand binding", k3, "L + R --> RL + L")
  m <= Reaction("ligand degradation", k4, "RL --> R")
  m <= Reaction("protein activation", k5, "RL + G --> Ga + Gbg + RL")
  m <= Reaction("dephosphorylation", k6, "Ga --> Gd")
  m <= Reaction("rebinding", k7, "Gd + Gbg --> G")
  m <= Reaction("bound receptor upregulation", k8, "0 --> RL")

  return m
end

model = yeast()


text:
using PackageCompiler, Pkg

dir(f...) = joinpath(@__DIR__, f...)
cd(@__DIR__)
build_executable(
    dir("makietest.jl"),
    "makie",
    dir("..", "examples", "program.c");
    snoopfile = dir("makiesnoop.jl"),
    builddir = dir("build"),
    verbose = true, quiet = false,
    cpu_target = "x86-64", optimize = "3"
)
PackageCompiler.build_object(
    dir("build", "julia_main.jl"), dir("build"), dir("build", "makie.o"), true,
    nothing, nothing, "x86-64", "3", nothing, nothing, nothing,
    nothing, nothing
)

packages = [
    "Quaternions",
    "GLVisualize",
    "StaticArrays",
    "GeometryTypes",
    "Reactive",
    "GLAbstraction",
    "GLWindow",
    "AbstractNumbers",
    "Contour",
    "FileIO",
    "Images",
    "UnicodeFun",
    "ColorBrewer",
    "Interact", # for displaying signals of Image - a bit unfortunate
    "Hiccup",
    "Media",
    "Juno",
    "ModernGL",
    "GLFW",
    "Fontconfig",
    "FreeType",
    "FreeTypeAbstraction",
    "ImageMagick",
]

for elem in packages
    cp(normpath(Base.find_package(elem), "..", ".."), dir("build", elem))
end


text:
"""
    AbstractCell

Inherit from this type to create a neuron cell (e.g. [`LIF`](@ref)).
"""
abstract type AbstractCell end

"""
    Soma{BT<:AbstractCell, TT<:AbstractThreshold}

A `Soma` is a cell body + a threshold.
"""
struct Soma{BT<:AbstractCell, TT<:AbstractThreshold}
    body::BT
    threshold::TT
end

getvoltage(soma::Soma) = getvoltage(soma.body)
isactive(soma::Soma, t::Integer; dt::Real = 1.0) = isactive(soma.body, t; dt = dt) || isactive(soma.threshold, t; dt = dt)

"""
    excite!(soma, current)

Inject `current` into `soma.body`.
"""
excite!(soma::T, current) where T<:Union{Soma, AbstractArray{<:Soma}} = excite!(soma.body, current)

"""
    evaluate!(soma::Soma, t::Integer; dt::Real = 1.0)
    (::Soma)(t::Integer; dt::Real = 1.0)
    evaluate!(somas::AbstractArray{<:Soma}, t::Integer; dt::Real = 1.0)

Evaluate the soma's cell body, decide whether to spike according to the
 threshold, then register the spike event with the cell body.
Return the spike event (0 for no spike or `t` for a spike).
"""
function evaluate!(soma::Soma, t::Integer; dt::Real = 1.0)
    spike = evaluate!(soma.threshold, t, soma.body(t; dt = dt); dt = dt)
    spike!(soma.body, spike; dt = dt)

    return spike
end
(soma::Soma)(t::Integer; dt::Real = 1.0) = evaluate!(soma, t; dt = dt)
function evaluate!(somas::T, t::Integer; dt::Real = 1.0) where T<:AbstractArray{<:Soma}
    voltage = evaluate!(somas.body, t; dt = dt)
    spikes = evaluate!(somas.threshold, t, voltage; dt = dt)
    spike!(somas.body, spikes; dt = dt)

    return spikes
end

"""
    reset!(soma::T) where T<:Union{Soma, AbstractArray{<:Soma}}

Reset `soma.body`.
"""
reset!(soma::T) where T<:Union{Soma, AbstractArray{<:Soma}} = reset!(soma.body)

"""
    Neuron{ST<:AbstractArray{<:AbstractSynapse}, CT<:Soma}

A `Neuron` is a vector of synapses feeding into a soma.
"""
struct Neuron{ST<:AbstractArray{<:AbstractSynapse}, CT<:Soma}
    synapses::ST
    soma::CT
end

"""
    Neuron(synapse::ST, body::BT, threshold::TT)
    Neuron{ST}(body::BT, threshold::TT) where ST<:AbstractSynapse

Create a `Neuron` specifying the cell `body`, `threshold`, and (optionally) `synapse`.
"""
Neuron(synapse::ST, body::BT, threshold::TT) where {ST<:AbstractSynapse, BT<:AbstractCell, TT<:AbstractThreshold} =
    Neuron(StructArray([synapse]; unwrap = t -> t <:AbstractSynapse), Soma(body, threshold))
Neuron{ST}(body::BT, threshold::TT) where {ST<:AbstractSynapse, BT<:AbstractCell, TT<:AbstractThreshold} =
    Neuron(StructArray{ST}(undef, 0), Soma(body, threshold))

"""
    connect!(neuron::Neuron, synapse::AbstractSynapse)

Connect `synapse` into `neuron`.
"""
function connect!(neuron::Neuron, synapse::AbstractSynapse)
    push!(neuron.synapses, synapse)

    return neuron
end

"""
    getvoltage(neuron::Neuron)

Get the current membrane potential of `neuron.soma`.
"""
getvoltage(neuron::Neuron) = getvoltage(neuron.soma)
isactive(neuron::Neuron, t::Integer; dt::Real = 1.0) = isactive(neuron.soma, t; dt = dt) || isactive(neuron.synapses, t; dt = dt)

"""
    excite!(neuron::Neuron, spike::Integer)
    excite!(neuron::Neuron, spikes::Array{<:Integer})
    excite!(neuron::Neuron, input, T::Integer; dt::Real = 1.0)

Excite a `neuron`'s synapses with spikes.
Or excite a `neuron`'s synapses with an arbitrary `input` function evaluated from `1:T`.
`input` must satisfy the following signature: `input(t; dt)`.
"""
excite!(neuron::Neuron, spike::Integer) = excite!(neuron.synapses, spike)
excite!(neuron::Neuron, spikes::Array{<:Integer}) = excite!(neuron.synapses, spikes)
function excite!(neuron::Neuron, input, T::Integer; dt::Real = 1.0)
    spikes = filter!(x -> x != 0, [input(t; dt = dt) for t = 1:T])
    excite!(neuron, spikes)

    return spikes
end
function excite!(neuron::Neuron, inputs::AbstractVector, T::Integer; dt::Real = 1.0)
    spikearray = Array[]
    for (i, input) in enumerate(inputs)
        spikes = filter!(x -> x != 0, [input(t; dt = dt) for t = 1:T])
        excite!(view(neuron.synapses, i), spikes)
        push!(spikearray, spikes)
    end

    return spikearray
end

"""
    evaluate!(neuron::Neuron, t::Integer; dt::Real = 1.0)
    (::Neuron)(t::Integer; dt::Real = 1.0)

Evaluate a neuron at time `t` by evaluating all its synapses,
 exciting the soma with current, then registering post-synaptic
 spikes with the synapses.
Return the spike event (0 for no spike or `t` for spike).
"""
function evaluate!(neuron::Neuron, t::Integer; dt::Real = 1.0)
    I = sum(evaluate!(neuron.synapses, t; dt = dt))
    excite!(neuron.soma, I)
    spike = evaluate!(neuron.soma, t; dt = dt)
    spike!(neuron.synapses, spike; dt = dt)

    return spike
end
(neuron::Neuron)(t::Integer; dt::Real = 1.0) = evaluate!(neuron, t; dt = dt)

"""
    reset!(neuron::T) where T<:Union{Neuron, AbstractArray{<:Neuron}}

Reset `neuron.synapses` and `neuron.soma`.
"""
function reset!(neuron::T) where T<:Union{Neuron, AbstractArray{<:Neuron}}
    reset!(neuron.synapses)
    reset!(neuron.soma)
end

"""
    simulate!(neuron::Neuron, T::Integer; dt::Real = 1.0, cb = () -> (), dense = false)

Fields:
- `neuron::Neuron`: the neuron to simulate
- `T::Integer`: number of time steps to simulate
- `dt::Real`: the length of simulation time step
- `cb::Function`: a callback function that is called at the start of each time step
- `dense::Bool`: set to `true` to evaluate every time step even in the absence of events
"""
function simulate!(neuron::Neuron, T::Integer; dt::Real = 1.0, cb = () -> (), dense = false)
    spikes = Int[]

    # step! neuron until queue is empty
    for t = 1:T
        if dense || isactive(neuron, t; dt = dt)
            cb()
            push!(spikes, evaluate!(neuron, t; dt = dt))
        end
    end

    return filter!(x -> x != 0, spikes)
end


text:
function autoscale{R <: Real}( img :: Array{R} )
  m = min(minimum(img), 0)
  d = max(maximum(img), 1) - m
  return clamp.((img - m) / d, 0, 1)
end

second{T, S}(t :: Tuple{T, S}) = t[2]

# --------------------------------------------------------------------------- #
# Import external functionality

#import Colors
#import ImageView

# --------------------------------------------------------------------------- #
# Utility function that shows a (greyscale) image as well as labels

function imshow( img :: GreyscaleImage, label :: Label = Label(zeros(2,0));
                 scale = 1., radius = 2, color=(0.1, 0.7, 0.2) )

  gdict = ImageView.imshow(scale * img.data)
  for (x,y) in label
    ImageView.annotate!(gdict, ImageView.AnnotationPoint(x, y,
                                                         shape='.',
                                                         size=radius,
                                                         color=Colors.RGB(color...)))
  end
  return gdict
end

function imshow( img :: GreyscaleImage, label;
                 scale = 1., radius = 2, color = nothing )

  if color == nothing
    color = repeat([ (0.1, 0.7, 0.2),
                     (0.1, 0.2, 0.7),
                     (0.7, 0.2, 0.1) ], outer=ceil(Int, length(label)/3))
  end

  gdict = ImageView.imshow(scale * img.data)
  for j in 1:length(label)
    for i in 1:length(label[j])
      ImageView.annotate!(gdict, ImageView.AnnotationPoint(label[j][i]...,
                                                           shape='.', size=radius,
                                                           color=Colors.RGB(color[j]...)))
    end
  end
  return gdict
end

# --------------------------------------------------------------------------- #
# Merging algorithms for labels

# adapted from
# https://stackoverflow.com/questions/19375675/simple-way-of-fusing-a-few-close-points
function declutter(lbl :: Label, dist :: Real)
  if dist <= 0
    return lbl
  end

  data = Tuple{Int, Int}[]
  n = length(lbl)
  merged = zeros(Bool, n)
  for i in 1:n
    if !merged[i]
      count = 1
      x0, y0 = lbl[i]
      x, y = x0, y0
      for j in (i+1):n
        v, w = lbl[j]
        if !merged[j] && (x0-v)^2 + (y0-w)^2 <= dist^2
          x += v
          y += w
          count += 1
          merged[j] = true
        end
      end
      merged[i] = true
      push!(data, round.(Int, (x / count, y / count)))
    end
  end
  return Label(data)
end

function merge(lbl :: Label, lbl2 :: Label, dist :: Real)
  if dist <= 0
    return join(lbl, lbl2)
  end

  data = copy(lbl2.data)
  for cm in lbl
    filter!(data) do ca
      norm([(cm .- ca)...]) > dist
    end
  end
  return join(lbl, Label(data))
end

max_stars_repo_name: tscode/DCellC

max_issues_repo_name: tscode/DCellC


text:
"""
    ce(m,q,z)

Compute the cosine-elliptic Mathieu function \$\\text{ce}_m(q,z)\$. `m` is the order of the solution and must be an integer \$m=0,1,2,3…\$.

Pass a vector, e.g. `m=0:4` and/or `z=LinRange(0,π,101)`, to compute the function for multiple orders and variables efficiently.
"""
function ce(m::Order,q::Number,z,N::Integer=matsize4(m,q))
    ie = map(iseven,m)
    all(ie)  && return cep(div.(m,2),q,z,N) # all even
    !any(ie) && return cea(div.(m,2),q,z,N) # all odd
    ce1 = cep(div.(filter(iseven,m),2),q,z,N)
    ce2 = cea(div.(filter(isodd, m),2),q,z,N)
    ce = similar(ce1,length(z),length(m))
    ce[:,ie] = ce1
    ce[:,.!ie] = ce2
    return ce
end

"""
    se(m,q,z)

Compute the sine-elliptic Mathieu function \$\\text{se}_m(q,z)\$. `m` is the order of the solution and must be an integer, excluding 0, \$m=1,2,3,4…\$.

Pass a vector, e.g. `m=1:5` and/or `z=LinRange(0,π,101)`, to compute the function for multiple orders and variables efficiently.
"""
function se(m::Order,q::Number,z,N::Integer=matsize4(m,q))
    ie = map(iseven,m)
    all(ie)  && return sep(div.(m.-2,2),q,z,N) # all even
    !any(ie) && return sea(div.(m.-1,2),q,z,N) # all odd
    se1 = sep(div.(filter(iseven,m).-2,2),q,z,N)
    se2 = sea(div.(filter(isodd, m).-1,2),q,z,N)
    se = similar(se1,length(z),length(m))
    se[:,ie] = se1
    se[:,.!ie] = se2
    return se
end

cep(n::Order,q::Number,z::Number,N::Integer=matsize3(n,q)) = cep(n,q,z:z,N)[1]
function cep(n::Order,q::Number,z::AbstractVector,N::Integer=matsize3(n,q))
    An = cep_coeffs(n,q,N) # Calculate coefficients
    # Evaluate cosine for all z values
    C = [cos(2k*z) for z in z, k in 0:N-1]
    # Calculate function at z, for all n
    return C * An
end

cea(n::Order,q::Number,z::Number,N::Integer=matsize3(n,q)) = cea(n,q,z:z,N)[1]
function cea(n::Order,q::Number,z::AbstractVector,N::Integer=matsize3(n,q))
    An = cea_coeffs(n,q,N) # Calculate coefficients
    # Evaluate cosine for all z values
    C = [cos((2k+1)*z) for z in z, k in 0:N-1]
    # Calculate function at z, for all n
    return C * An
end

sea(n::Order,q::Number,z::Number,N::Integer=matsize3(n,q)) = sea(n,q,z:z,N)[1]
function sea(n::Order,q::Number,z::AbstractVector,N::Integer=matsize3(n,q))
    Bn = sea_coeffs(n,q,N) # Calculate coefficients
    # Evaluate sine for all z values
    S = [sin((2k+1)*z) for z in z, k in 0:N-1]
    # Calculate function at z, for all n
    return S * Bn
end

sep(n::Order,q::Number,z::Number,N::Integer=matsize3(n,q)) = sep(n,q,z:z,N)[1]
function sep(n::Order,q::Number,z::AbstractVector,N::Integer=matsize3(n,q))
    Bn = sep_coeffs(n,q,N) # Calculate coefficients
    # Evaluate sine for all z values
    S = [sin((2k+2)*z) for z in z, k in 0:N-1]
    # Calculate function at z, for all n
    return S * Bn
end

fun_per(m::Order,q::Number,z::Number,N::Integer=matsize4(m,q)) = fun_per(m,q,z:z,N)[1]
function fun_per(m::Order,q::Number,z::AbstractVector,N::Integer=matsize4(m,q))
    ie = map(iseven,m)
    all(ie)  && return cep(div.(m,2),q,z,N)
    !any(ie) && return sep(div.(m,2),q,z,N)
    fp1 = cep(div.(filter(iseven,m),2),q,z,N)
    fp2 = sep(div.(filter(isodd, m),2),q,z,N)
    fp = similar(fp1,length(z),length(m))
    fp[:,ie] = fp1
    fp[:,.!ie] = fp2
    return fp
end


text:
"""
    stepforward!(prob)

Step forward `prob` one time step.
"""
stepforward!(prob::Problem) =
  stepforward!(prob.sol, prob.clock, prob.timestepper, prob.eqn, prob.vars, prob.params, prob.grid)

"""
    stepforward!(prob, nsteps::Int)

Step forward `prob` for `nsteps`.
"""
function stepforward!(prob::Problem, nsteps::Int)
  for _ in 1:nsteps
    stepforward!(prob)
  end

  return nothing
end

"""
    stepforward!(prob::Problem, diags, nsteps::Int)

Step forward `prob` for `nsteps`, incrementing `diags` along the way. `diags` may be a
single `Diagnostic` or a `Vector` of `Diagnostic`s.
"""
function stepforward!(prob::Problem, diags, nsteps::Int)
  for _ in 1:nsteps
    stepforward!(prob)
    increment!(diags)
  end

  return nothing
end

const fullyexplicitsteppers= [
  :ForwardEuler,
  :RK4,
  :AB3,
  :FilteredForwardEuler,
  :FilteredRK4,
  :FilteredAB3
]

isexplicit(stepper) = any(Symbol(stepper) .== fullyexplicitsteppers)

"""
    TimeStepper(stepper, equation, dt=nothing, dev=CPU(); kw...)

Instantiate the Time`stepper` for `equation` with timestep `dt` and
on the `dev`ice. The `kw` are passed to the timestepper constructor.
"""
function TimeStepper(stepper, equation, dt=nothing, dev::Device=CPU(); kw...)
  fullsteppername = Symbol(stepper, :TimeStepper)

  # Create expression that instantiates the time-stepper, depending on whether
  # timestepper is explicit or not.
  expr = isexplicit(stepper) ? Expr(:call, fullsteppername, equation, dev) :
                               Expr(:call, fullsteppername, equation, dt, dev)

  # Add keyword arguments
  length(kw) > 0 && push!(expr.args, Tuple(Expr(:kw, p.first, p.second) for p in kw)...)

  return eval(expr)
end

# The following time-steppers are implemented below
#
#   * Forward Euler
#   * Filtered Forward Euler
#   * RK4
#   * Filtered RK4
#   * ETDRK4
#   * Filtered ETDRK4
#   * AB3
#   * Filtered AB3
#
# Explicit time-steppers are constructed with the signature
#   ts = ExplicitTimeStepper(equation::Equation)
#
# Implicit time-steppers are constructed with the signature
#   ts = ImplicitTimeStepper(equation::Equation, dt)

# --
# Forward Euler
# --

"""
    ForwardEulerTimeStepper{T} <: AbstractTimeStepper{T}

A Forward Euler timestepper for time-stepping `∂u/∂t = RHS(u, t)` via:
```
uⁿ⁺¹ = uⁿ + dt * RHS(uⁿ, tⁿ)
```
"""
struct ForwardEulerTimeStepper{T} <: AbstractTimeStepper{T}
  N::T # Explicit linear and nonlinear terms
  ForwardEulerTimeStepper(N::T) where T = new{T}(0N)
end

"""
    ForwardEulerTimeStepper(equation::Equation, dev::Device=CPU())

Construct a Forward Euler timestepper for `equation` on device `dev`.
"""
ForwardEulerTimeStepper(equation::Equation, dev::Device=CPU()) =
  ForwardEulerTimeStepper(devzeros(dev, equation.T, equation.dims))

function stepforward!(sol, clock, ts::ForwardEulerTimeStepper, equation, vars, params, grid)
  equation.calcN!(ts.N, sol, clock.t, clock, vars, params, grid)
  @. sol += clock.dt * (equation.L * sol + ts.N)
  clock.t += clock.dt
  clock.step += 1

  return nothing
end

"""
    struct FilteredForwardEulerTimeStepper{T,Tf} <: AbstractTimeStepper{T}

A Forward Euler timestepper with spectral filtering. See [`ForwardEulerTimeStepper`](@ref).
"""
struct FilteredForwardEulerTimeStepper{T,Tf} <: AbstractTimeStepper{T}
       N :: T
  filter :: Tf
end

"""
    FilteredForwardEulerTimeStepper(equation, dev; filterkwargs...)

Construct a Forward Euler timestepper with spectral filtering for `equation` on device `dev`.
"""
function FilteredForwardEulerTimeStepper(equation::Equation, dev::Device=CPU(); filterkwargs...)
  filter = makefilter(equation; filterkwargs...)

  return FilteredForwardEulerTimeStepper(devzeros(dev, equation.T, equation.dims), filter)
end

function stepforward!(sol, clock, ts::FilteredForwardEulerTimeStepper, equation, vars, params, grid)
  equation.calcN!(ts.N, sol, clock.t, clock, vars, params, grid)
  @. sol = ts.filter * (sol + clock.dt * (ts.N + equation.L * sol))
  clock.t += clock.dt
  clock.step += 1

  return nothing
end

# --
# RK4
# --

"""
    RK4TimeStepper{T} <: AbstractTimeStepper{T}

A 4th-order Runge-Kutta timestepper for time-stepping `∂u/∂t = RHS(u, t)` via:
```
uⁿ⁺¹ = uⁿ + dt/6 * (k₁ + 2 * k₂ + 2 * k₃ + k₄)
```
where
```
k₁ = RHS(uⁿ, tⁿ)
k₂ = RHS(uⁿ + k₁ * dt/2, tⁿ + dt/2)
k₃ = RHS(uⁿ + k₂ * dt/2, tⁿ + dt/2)
k₄ = RHS(uⁿ + k₃ * dt, tⁿ + dt)
```
"""
struct RK4TimeStepper{T} <: AbstractTimeStepper{T}
  sol₁ :: T
  RHS₁ :: T
  RHS₂ :: T
  RHS₃ :: T
  RHS₄ :: T
end

"""
    RK4TimeStepper(equation::Equation, dev::Device=CPU())

Construct a 4th-order Runge-Kutta timestepper for `equation` on device `dev`.
"""
function RK4TimeStepper(equation::Equation, dev::Device=CPU())
  @devzeros typeof(dev) equation.T equation.dims N sol₁ RHS₁ RHS₂ RHS₃ RHS₄

  return RK4TimeStepper(sol₁, RHS₁, RHS₂, RHS₃, RHS₄)
end

"""
    FilteredRK4TimeStepper{T,Tf} <: AbstractTimeStepper{T}

A 4th-order Runge-Kutta timestepper with spectral filtering. See [`RK4TimeStepper`](@ref).
"""
struct FilteredRK4TimeStepper{T,Tf} <: AbstractTimeStepper{T}
    sol₁ :: T
    RHS₁ :: T
    RHS₂ :: T
    RHS₃ :: T
    RHS₄ :: T
  filter :: Tf
end

"""
    FilteredRK4TimeStepper(equation::Equation, dev::Device=CPU(); filterkwargs...)

Construct a 4th-order Runge-Kutta timestepper with spectral filtering for `equation` on device `dev`.
"""
function FilteredRK4TimeStepper(equation::Equation, dev::Device=CPU(); filterkwargs...)
  ts = RK4TimeStepper(equation, dev)
  filter = makefilter(equation; filterkwargs...)

  return FilteredRK4TimeStepper(getfield.(Ref(ts), fieldnames(typeof(ts)))..., filter)
end

function addlinearterm!(RHS, L, sol)
  @. RHS += L*sol

  return nothing
end

function substepsol!(newsol, sol, RHS, dt)
  @. newsol = sol + dt*RHS

  return nothing
end

function RK4substeps!(sol, clock, ts, equation, vars, params, grid, t, dt)
  # Substep 1
  equation.calcN!(ts.RHS₁, sol, t, clock, vars, params, grid)
  addlinearterm!(ts.RHS₁, equation.L, sol)

  # Substep 2
  substepsol!(ts.sol₁, sol, ts.RHS₁, dt/2)
  equation.calcN!(ts.RHS₂, ts.sol₁, t+dt/2, clock, vars, params, grid)
  addlinearterm!(ts.RHS₂, equation.L, ts.sol₁)

  # Substep 3
  substepsol!(ts.sol₁, sol, ts.RHS₂, dt/2)
  equation.calcN!(ts.RHS₃, ts.sol₁, t+dt/2, clock, vars, params, grid)
  addlinearterm!(ts.RHS₃, equation.L, ts.sol₁)

  # Substep 4
  substepsol!(ts.sol₁, sol, ts.RHS₃, dt)
  equation.calcN!(ts.RHS₄, ts.sol₁, t+dt, clock, vars, params, grid)
  addlinearterm!(ts.RHS₄, equation.L, ts.sol₁)

  return nothing
end

function RK4update!(sol, RHS₁, RHS₂, RHS₃, RHS₄, dt)
  @. sol += dt*(RHS₁ / 6 + RHS₂ / 3  + RHS₃ / 3 + RHS₄ / 6)

  return nothing
end

function RK4update!(sol, RHS₁, RHS₂, RHS₃, RHS₄, filter, dt)
  @. sol = filter * (sol + dt*(RHS₁ / 6 + RHS₂ / 3  + RHS₃ / 3 + RHS₄ / 6))

  return nothing
end

function stepforward!(sol, clock, ts::RK4TimeStepper, equation, vars, params, grid)
  RK4substeps!(sol, clock, ts, equation, vars, params, grid, clock.t, clock.dt)
  RK4update!(sol, ts.RHS₁, ts.RHS₂, ts.RHS₃, ts.RHS₄, clock.dt)
  clock.t += clock.dt
  clock.step += 1

  return nothing
end

function stepforward!(sol, clock, ts::FilteredRK4TimeStepper, equation, vars, params, grid)
  RK4substeps!(sol, clock, ts, equation, vars, params, grid, clock.t, clock.dt)
  RK4update!(sol, ts.RHS₁, ts.RHS₂, ts.RHS₃, ts.RHS₄, ts.filter, clock.dt)
  clock.t += clock.dt
  clock.step += 1

  return nothing
end

# ------
# ETDRK4
# ------

"""
    ETDRK4TimeStepper{T,TL} <: AbstractTimeStepper{T}

A 4th-order exponential-time-differencing Runge-Kutta timestepper for time-stepping
`∂u/∂t = L * u + N(u)`. The scheme treats the linear term `L` exact while for the
nonlinear terms `N(u)` it uses a 4th-order Runge-Kutta scheme. That is,
```
uⁿ⁺¹ = exp(L * dt) * uⁿ + RK4(N(uⁿ))
```
For more info refer to

> Kassam, A. K., & Trefethen, L. N. (2005). Fourth-order time-stepping for stiff PDEs. _SIAM Journal on Scientific Computing_, **26(4)**, 1214-1233.
"""
struct ETDRK4TimeStepper{T,TL} <: AbstractTimeStepper{T}
  # ETDRK4 coefficents
        ζ :: TL
        α :: TL
        β :: TL
        Γ :: TL
  expLdt  :: TL
  exp½Ldt :: TL
     sol₁ :: T
     sol₂ :: T
       N₁ :: T
       N₂ :: T
       N₃ :: T
       N₄ :: T
end

"""
    ETDRK4TimeStepper(equation::Equation, dt, dev::Device=CPU())

Construct a 4th-order exponential-time-differencing Runge-Kutta timestepper with timestep `dt`
for `equation` on device `dev`.
"""
function ETDRK4TimeStepper(equation::Equation, dt, dev::Device=CPU())
  dt = fltype(equation.T)(dt) # ensure dt is correct type.
  expLdt, exp½Ldt = getexpLs(dt, equation)
  ζ, α, β, Γ = getetdcoeffs(dt, equation.L)
  @devzeros typeof(dev) equation.T equation.dims sol₁ sol₂ N₁ N₂ N₃ N₄

  return ETDRK4TimeStepper(ζ, α, β, Γ, expLdt, exp½Ldt, sol₁, sol₂, N₁, N₂, N₃, N₄)
end

"""
    FilteredETDRK4TimeStepper{T,TL,Tf} <: AbstractTimeStepper{T}

A 4th-order exponential-time-differencing Runge-Kutta timestepper with spectral filtering.
See [`ETDRK4TimeStepper`](@ref).
"""
struct FilteredETDRK4TimeStepper{T,TL,Tf} <: AbstractTimeStepper{T}
  # ETDRK4 coefficents:
        ζ :: TL
        α :: TL
        β :: TL
        Γ :: TL
  expLdt  :: TL
  exp½Ldt :: TL
     sol₁ :: T
     sol₂ :: T
       N₁ :: T
       N₂ :: T
       N₃ :: T
       N₄ :: T
   filter :: Tf
end

"""
    FilteredETDRK4TimeStepper(equation, dt; filterkwargs...)

Construct a 4th-order exponential-time-differencing Runge-Kutta timestepper with timestep `dt` and with
spectral filtering for `equation` on device `dev`.
"""
function FilteredETDRK4TimeStepper(equation::Equation, dt, dev::Device=CPU(); filterkwargs...)
  timestepper = ETDRK4TimeStepper(equation, dt, dev)
  filter = makefilter(equation; filterkwargs...)

  return FilteredETDRK4TimeStepper(getfield.(Ref(timestepper), fieldnames(typeof(timestepper)))..., filter)
end

function ETDRK4update!(sol, expLdt, α, β, Γ, N₁, N₂, N₃, N₄)
  @. sol = (expLdt * sol +  α * N₁
                         + 2β * (N₂ + N₃)
                         +  Γ * N₄ )

  return nothing
end

function ETDRK4update!(sol, ts, filter)
  @. sol = filter * (ts.expLdt * sol +     ts.α * ts.N₁
                                     + 2 * ts.β * (ts.N₂ + ts.N₃)
                                     +     ts.Γ * ts.N₄ )

  return nothing
end

function ETDRK4substep12!(sol₁, exp½Ldt, sol, ζ, N)
  @. sol₁ = exp½Ldt * sol + ζ * N

  return nothing
end

function ETDRK4substep3!(sol₂, exp½Ldt, sol₁, ζ, N₁, N₃)
  @. sol₂ = exp½Ldt * sol₁ + ζ * (2N₃ - N₁)

  return nothing
end

function ETDRK4substeps!(sol, clock, ts, equation, vars, params, grid)
  # Substep 1
  equation.calcN!(ts.N₁, sol, clock.t, clock, vars, params, grid)
  ETDRK4substep12!(ts.sol₁, ts.exp½Ldt, sol, ts.ζ, ts.N₁)

  # Substep 2
  t2 = clock.t + clock.dt/2
  equation.calcN!(ts.N₂, ts.sol₁, t2, clock, vars, params, grid)
  ETDRK4substep12!(ts.sol₂, ts.exp½Ldt, sol, ts.ζ, ts.N₂)

  # Substep 3
  equation.calcN!(ts.N₃, ts.sol₂, t2, clock, vars, params, grid)
  ETDRK4substep3!(ts.sol₂, ts.exp½Ldt, ts.sol₁, ts.ζ, ts.N₁, ts.N₃)

  # Substep 4
  t3 = clock.t + clock.dt
  equation.calcN!(ts.N₄, ts.sol₂, t3, clock, vars, params, grid)

  return nothing
end

function stepforward!(sol, clock, ts::ETDRK4TimeStepper, equation, vars, params, grid)
  ETDRK4substeps!(sol, clock, ts, equation, vars, params, grid)
  ETDRK4update!(sol, ts.expLdt, ts.α, ts.β, ts.Γ, ts.N₁, ts.N₂, ts.N₃, ts.N₄)
  clock.t += clock.dt
  clock.step += 1

  return nothing
end

function stepforward!(sol, clock, ts::FilteredETDRK4TimeStepper, equation, vars, params, grid)
  ETDRK4substeps!(sol, clock, ts, equation, vars, params, grid)
  ETDRK4update!(sol, ts, ts.filter) # update
  clock.t += clock.dt
  clock.step += 1

  return nothing
end

# ---
# AB3
# ---

const ab3h1 = 23/12
const ab3h2 = 16/12
const ab3h3 = 5/12

"""
    AB3TimeStepper{T} <: AbstractTimeStepper{T}

A 3rd-order Adams-Bashforth timestepper for time-stepping `∂u/∂t = RHS(u, t)` via:
```
uⁿ⁺¹ = uⁿ + dt/12 * (23 * RHS(uⁿ, tⁿ) - 16 * RHS(uⁿ⁻¹, tⁿ⁻¹) + 5 * RHS(uⁿ⁻², tⁿ⁻²))
```

Adams-Bashforth is a multistep method, i.e., it not only requires information from the `n`-th time-step
(`uⁿ`) but also from the previous two timesteps (`uⁿ⁻¹` and `uⁿ⁻²`). For the first two timesteps, it
falls back to a forward Euler timestepping scheme:
```
uⁿ⁺¹ = uⁿ + dt * RHS(uⁿ, tⁿ)
```
"""
struct AB3TimeStepper{T} <: AbstractTimeStepper{T}
  RHS::T
  RHS₋₁::T
  RHS₋₂::T
end

"""
    AB3TimeStepper(equation::Equation, dev::Device=CPU())

Construct a 3rd order Adams-Bashforth timestepper for `equation` on device `dev`.
"""
function AB3TimeStepper(equation::Equation, dev::Device=CPU())
  @devzeros typeof(dev) equation.T equation.dims RHS RHS₋₁ RHS₋₂

  return AB3TimeStepper(RHS, RHS₋₁, RHS₋₂)
end

"""
    FilteredAB3TimeStepper{T} <: AbstractTimeStepper{T}

A 3rd order Adams-Bashforth timestepper with spectral filtering. See [`AB3TimeStepper`](@ref).
"""
struct FilteredAB3TimeStepper{T, Tf} <: AbstractTimeStepper{T}
   RHS   :: T
   RHS₋₁ :: T
   RHS₋₂ :: T
  filter :: Tf
end

"""
    FilteredAB3TimeStepper(equation::Equation, dev::Device=CPU(); filterkwargs...)

Construct a 3rd order Adams-Bashforth timestepper with spectral filtering for `equation` on device `dev`.
"""
function FilteredAB3TimeStepper(equation::Equation, dev::Device=CPU(); filterkwargs...)
  timestepper = AB3TimeStepper(equation, dev)
  filter = makefilter(equation; filterkwargs...)

  return FilteredAB3TimeStepper(getfield.(Ref(timestepper), fieldnames(typeof(timestepper)))..., filter)
end

function AB3update!(sol, ts, clock)
  if clock.step < 3  # forward Euler steps to initialize AB3
    @. sol += clock.dt * ts.RHS    # Update
  else   # Otherwise, stepforward with 3rd order Adams Bashforth:
    @. sol += clock.dt * (ab3h1 * ts.RHS - ab3h2 * ts.RHS₋₁ + ab3h3 * ts.RHS₋₂)
  end

  return nothing
end

function AB3update!(sol, ts::FilteredAB3TimeStepper, clock)
  if clock.step < 3  # forward Euler steps to initialize AB3
    @. sol = ts.filter * (sol + clock.dt * ts.RHS)    # Update
  else   # Otherwise, stepforward with 3rd order Adams Bashforth:
    @. sol = ts.filter * (sol + clock.dt * (ab3h1 * ts.RHS - ab3h2 * ts.RHS₋₁ + ab3h3 * ts.RHS₋₂))
  end

  return nothing
end

function stepforward!(sol, clock, ts::AB3TimeStepper, equation, vars, params, grid)
  equation.calcN!(ts.RHS, sol, clock.t, clock, vars, params, grid)
  addlinearterm!(ts.RHS, equation.L, sol)

  AB3update!(sol, ts, clock)
  clock.t += clock.dt
  clock.step += 1

  @. ts.RHS₋₂ = ts.RHS₋₁          # Store
  @. ts.RHS₋₁ = ts.RHS            # ... previous values of RHS

  return nothing
end

function stepforward!(sol, clock, ts::FilteredAB3TimeStepper, equation, vars, params, grid)
  equation.calcN!(ts.RHS, sol, clock.t, clock, vars, params, grid)
  addlinearterm!(ts.RHS, equation.L, sol)

  AB3update!(sol, ts, clock)
  clock.t += clock.dt
  clock.step += 1

  @. ts.RHS₋₂ = ts.RHS₋₁          # Store
  @. ts.RHS₋₁ = ts.RHS            # ... previous values of RHS

  return nothing
end

# --
# Timestepper utils
# --

function getexpLs(dt, equation)
  expLdt  = @. exp(dt * equation.L)
  exp½Ldt = @. exp(dt * equation.L/2)

  return expLdt, exp½Ldt
end

"""
    getetdcoeffs(dt, L; ncirc=32, rcirc=1)

Calculate ETDRK4 coefficients associated with the (diagonal) linear coefficient
`L` by integrating over a unit circle in the complex space.
"""
function getetdcoeffs(dt, L; ncirc=32, rcirc=1)
  shape = Tuple(cat(ncirc, ones(Int, ndims(L)), dims=1))

  circ = zeros(Complex{Float64}, shape) # use double precision for this calculation
  circ .= rcirc * exp.(2π * im/ncirc * (0.5:1:(ncirc-0.5)))
  circ = permutedims(circ, ndims(circ):-1:1)

  zc = dt * L .+ circ
  M = ndims(L) + 1

  # Four coefficients: ζ, α, β, Γ
  ζc = @.               ( exp(zc/2)-1 ) / zc
  αc = @. ( -4 - zc + exp(zc) * (4 - 3zc + zc^2) ) / zc^3
  βc = @.    ( 2  + zc + exp(zc) * (-2 + zc) ) / zc^3
  Γc = @. ( -4 - 3zc - zc^2 + exp(zc) * (4 - zc) ) / zc^3

  ζ = dt * dropdims(mean(ζc, dims=M), dims=M)
  α = dt * dropdims(mean(αc, dims=M), dims=M)
  β = dt * dropdims(mean(βc, dims=M), dims=M)
  Γ = dt * dropdims(mean(Γc, dims=M), dims=M)

  if eltype(L) <: Real # this is conservative, but unclear if necessary
    ζ = real.(ζ)
    α = real.(α)
    β = real.(β)
    Γ = real.(Γ)
  end

  return ζ, α, β, Γ
end

"""
    step_until!(prob, stop_time)

Step forward `prob` until `stop_time`. Cannot be used with ETDRK4 time steppers.

See also: [`stepforward!`](@ref)
"""
step_until!(prob, stop_time) = step_until!(prob, prob.timestepper, stop_time)

step_until!(prob, ::Union{ETDRK4TimeStepper, FilteredETDRK4TimeStepper}, stop_time) =
  error("step_until! requires fully explicit time stepper; does not work with ETDRK4")

function step_until!(prob, timestepper, stop_time)
  # Throw an error if stop_time is not greater than the current problem time
  stop_time > prob.clock.t || error("stop_time must be greater than prob.clock.t")

  # Extract current time step
  dt = prob.clock.dt

  # Step forward until just before stop_time
  time_interval = stop_time - prob.clock.t
  nsteps = floor(Int, time_interval / dt)
  stepforward!(prob, nsteps)

  # Take one final small step so that prob.clock.t = stop_time
  t_remaining = time_interval - prob.clock.t
  prob.clock.dt = t_remaining
  stepforward!(prob)

  # Restore previous time-step
  prob.clock.dt = dt

  return nothing
end