text:
export FlowModel, CompiledFlowModel
export compile, nr_params
export getlayers

using TupleTools: flatten

import Base: length, eltype

@doc raw"""
The FlowModel structure is the most generic type of Flow model, in which the layers are not constrained to be of a specific type. The FlowModel structure contains the input dimensionality and a tuple of layers and can be constructed as `FlowModel( dim, (layer1, layer2, ...) )`.

Note: this model can be specialized by constraining the types of layers. This potentially allows for more efficient specialized methods that can deal with specifics of these layers, such as triangular jacobian matrices.
"""
struct FlowModel{N,T<:NTuple{N,AbstractLayer}} <: AbstractFlowModel
    dim     :: Int
    layers  :: T
end

@doc raw"""
Creates a flow model with uninitialized parameters.

Input arguments:
- `dim::Int` - input dimensionality.
- `layers<:NTuple{N,AbstractLayer}` - arbitrarily sized tuple containing abstractlayers.

Return arguments:
- `::Flowmodel` - model containing layers of which are appropriately sized according to the input dimensionality.
"""
function FlowModel(dim::Int, layers::T) where { T <: NTuple{N,AbstractLayerPlaceholder} where { N } }
    return FlowModel(dim, flatten(prepare.(dim, layers)))
end
function FlowModel(layers::T) where { T <: NTuple{N, Union{AbstractLayer, AbstractLayerPlaceholder}} where { N } }
    @assert typeof(first(layers)) <: InputLayer "The FlowModel requires an input dimension to be specified. This can be achieved, either by preceding the layers tuple with an integer as `FlowModel(dim, layers)`, or by starting the tuple of layers with an `InputLayer(dim)` as `FlowModel((InputLayer(dim), layers...))`."
    return FlowModel(getdim(first(layers)), flatten(prepare.(getdim(first(layers)), Base.tail(layers))))
end

# prepare function for setting correct sizes in the layers (without assigning the parameters yet!)
prepare(dim::Int, layers::T) where { T <: NTuple{N,Union{AbstractLayer, AbstractLayerPlaceholder}} where { N }} = _prepare(dim, layers)
Broadcast.broadcasted(::typeof(prepare), dim::Int, layers::T) where { T <: NTuple{N,Union{AbstractLayer, AbstractLayerPlaceholder}} where { N }} = broadcast(_prepare, Ref(dim), layers)

@doc raw"""
The CompiledFlowModel structure is the most generic type of compiled Flow model, in which the layers are not constrained to be of a specific type. The FlowModel structure contains the input dimension and a tuple of compiled layers. Do not manually create a CompiledFlowModel! Instead create a FlowModel first and compile it with `compile(model::FlowModel)`. This will make sure that all layers/mappings are configured with the proper dimensionality and with randomly sampled parameters. Alternatively, if you would like to pass your own parameters, call `compile(model::FlowModel, params::Vector)`.

Note: this model can be specialized by constraining the types of layers. This potentially allows for more efficient specialized methods that can deal with specifics of these layers, such as triangular jacobian matrices.
"""
struct CompiledFlowModel{N,T<:NTuple{N,AbstractLayer}} <: AbstractCompiledFlowModel
    dim     :: Int
    layers  :: T
end

@doc raw"""
`compile()` compiles a model by setting its parameters. It randomly sets parameter values in the layers and flows such that inference in the model can be obtained.

Input arguments
- `model::FlowModel` - a model of which the dimensionality of its layers/flows has been initialized, but its parameters have not been set.

Return arguments
- `::CompiledFlowModel` - a compiled model with set parameters, such that it can be used for processing data.
"""
function compile(model::FlowModel)
    # do not create parameters here for decentralized initializers (TODO)
    return CompiledFlowModel(model.dim, compile.(model.layers))
end

@doc raw"""
`compile(model::FlowModel, params::Vector)` lets you initialize a model `model` with a vector of parameters `params`.

Input arguments
- `model::FlowModel` - a model of which the dimensionality of its layers/flows has been initialized, but its parameters have not been set.
- `params::Vector`   - a vector of parameters with which the model should be compiled.

Return arguments
- `::CompiledFlowModel` - a compiled model with set parameters, such that it can be used for processing data.
"""
function compile(model::FlowModel, params::Vector)

    # assert whether the parameter vector complies with the model
    @assert nr_params(model) == length(params) "The number of parameters in the model does not match the passed number of parameters."

    # compile layers
    return CompiledFlowModel(getdim(model), compile(getlayers(model), params))

end
function compile(x::Tuple, params::Vector)

    # fetch number of parameters in layer
    params_in_layer = nr_params(first(x))

    if params_in_layer > 0
        # if layer has parameters, compile that layer and the other layers with parameters
        return (compile(first(x), params[1:params_in_layer]), compile(Base.tail(x), params[params_in_layer+1:end])...)
    else
        # if layer has no parameters, compile that layer on its own and the other layers with parameters
        return (compile(first(x)), compile(Base.tail(x), params)...)
    end

end
function compile(x::Tuple)
    return (compile(first(x)), compile(Base.tail(x))...)
end
# compile functions to stop the recursion
compile(x::Tuple{}, params::Vector) = ()
compile(x::Tuple{})         = ()

# get-functions for the FlowModel structure
getdim(model::FlowModel)                    = model.dim
getdim(model::CompiledFlowModel)            = model.dim
getlayers(model::FlowModel)                 = model.layers
getlayers(model::CompiledFlowModel)         = model.layers
getforward(model::CompiledFlowModel)        = (x) -> forward(model, x)
getbackward(model::CompiledFlowModel)       = (x) -> backward(model, x)
getjacobian(model::CompiledFlowModel)       = (x) -> jacobian(model, x)
getinv_jacobian(model::CompiledFlowModel)   = (x) -> inv_jacobian(model, x)

# custom Base function for the FlowModel structure
eltype(model::CompiledFlowModel) = promote_type(map(eltype, model.layers)...)
length(model::CompiledFlowModel) = length(model.layers)

# fetch number of parameters of model
nr_params(model::FlowModel)         = nr_params(getlayers(model))
nr_params(model::CompiledFlowModel) = nr_params(getlayers(model))
nr_params(layers::T) where { T <: NTuple{N,AbstractLayer} where N} = nr_params(first(layers)) + nr_params(Base.tail(layers))
nr_params(layers::Tuple{}) = return 0

# forward pass through the Flow model
function _forward(model::CompiledFlowModel, input::AbstractVector{T}) where { T <: Real }

    # promote type for allocating output
    Ti = promote_type(eltype(model), T)

    # allocate space for result
    output = zeros(Ti, size(input))

    # perform calculations
    forward!(output, model, input)

    # return result
    return output

end

# when calling forward, redirect to _forward
forward(model::CompiledFlowModel, input::AbstractVector{ <: Real }) = _forward(model, input)

# for broadcasting over forward, fix the model for multiple inputs
Broadcast.broadcasted(::typeof(forward), model::CompiledFlowModel, input::AbstractVector{ <: AbstractVector{ <: Real } }) = broadcast(_forward, Ref(model), input)

# inplace forward pass through the Flow model
function forward!(output::AbstractVector{ <: Real }, model::CompiledFlowModel, input::AbstractVector{T}) where { T <: Real }

    # fetch layers
    layers = getlayers(model)

    # promote type for allocating output
    Ti = promote_type(eltype(model), T)

    # decouple changing input from actual input
    input_new = zeros(Ti, size(input))
    input_new .= input

    # perform forward pass over all layers
    forward!(output, layers, input_new)

end

# inplace forward pass through a tuple of layers
function forward!(output::AbstractVector{ <: Real }, layers::T, input::AbstractVector{ <: Real }) where { T <: NTuple{N,AbstractLayer} where N}

    # perform pass over first layer
    forward!(output, first(layers), input)

    # update intermediate input
    input .= output

    # perform pass over all remaining layers
    forward!(output, Base.tail(layers), input)

end

# when no layers are left, stop the inplace recursion
forward!(output::AbstractVector{ <: Real }, layers::Tuple{}, input::AbstractVector{ <: Real }) = return

# backward pass through the Flow model
function _backward(model::CompiledFlowModel, output::AbstractVector{T}) where { T <: Real }

    # promote type for allocating output
    Ti = promote_type(eltype(model), T)

    # allocate space for result
    input = zeros(Ti, size(output))

    # perform calculations
    backward!(input, model, output)

    # return result
    return input

end

# when calling backward, redirect to _backward
backward(model::CompiledFlowModel, output::AbstractVector{ <: Real }) = _backward(model, output)

# for broadcasting over backward, fix the model for multiple inputs
Broadcast.broadcasted(::typeof(backward), model::CompiledFlowModel, output::AbstractVector{ <: AbstractVector{ <: Real } }) = broadcast(_backward, Ref(model), output)

# inplace backward pass through the Flow model
function backward!(input::AbstractVector{ <: Real }, model::CompiledFlowModel, output::AbstractVector{T}) where { T <: Real }

    # fetch layers
    layers = getlayers(model)

    # promote type for allocating output
    Ti = promote_type(eltype(model), T)

    # decouple changing output from actual output
    output_new = zeros(Ti, size(output))
    output_new .= output

    # perform backward pass over all layers
    backward!(input, layers, output_new)

end

# inplace backward pass through a tuple of layers
function backward!(input::AbstractVector{ <: Real }, layers::T, output::AbstractVector{ <: Real }) where { T <: NTuple{N,AbstractLayer} where N}

    # perform pass over last layer
    backward!(input, last(layers), output)

    # update intermediate output
    output .= input

    # perform pass over all remaining layers
    backward!(input, Base.front(layers), output)

end

# when no layers are left, stop the inplace recursion
backward!(input::AbstractVector{ <: Real }, layers::Tuple{}, output::AbstractVector{ <: Real }) = return

# joint forward-jacobian of the Flow model
function _forward_jacobian(model::CompiledFlowModel, input::AbstractVector{T}) where { T <: Real }

    # fetch layers
    dim    = getdim(model)

    # promote type for allocating output
    Ti = promote_type(eltype(model), T)

    # allocate space for output and jacobian
    output = zeros(Ti, dim)
    J = zeros(Ti, dim, dim)
    for k = 1:dim
        J[k,k] = one(Ti)
    end

    # calculate jacobian
    forward_jacobian!(output, J, model, input)

    # return result
    return output, J

end

# when calling forward_jacobian, redirect to _forward_jacobian
forward_jacobian(model::CompiledFlowModel, input::AbstractVector{ <: Real }) = _forward_jacobian(model, input)

# for broadcasting over forward_jacobian, fix the model for multiple inputs
Broadcast.broadcasted(::typeof(forward_jacobian), model::CompiledFlowModel, input::AbstractVector{ <: AbstractVector{ <: Real } }) = broadcast(_jacobian, Ref(model), input)

# inplace forward_jacobian of the Flow model
function forward_jacobian!(output::AbstractVector{ <: Real }, J::AbstractMatrix{T}, model::CompiledFlowModel, input::AbstractVector{ <: Real }) where { T <: Real }

    # fetch layers
    layers = getlayers(model)
    dim = getdim(model)

    # promote type for allocating output
    Ti = promote_type(eltype(model), T)

    # allocate space for intermediate output and jacobian
    input_new = zeros(Ti, dim)
    input_new .= input
    J_new = zeros(Ti, dim, dim)
    J_old = zeros(Ti, dim, dim)
    for k = 1:dim
        J_old[k,k] = one(Ti)
    end

    # calculate forward_jacobian over all layers
    forward_jacobian!(J, J_new, J_old, output, input_new, layers)

end

# inplace forward_jacobian calculation for a tuple of layers
function forward_jacobian!(J::AbstractMatrix{ <: Real }, J_new::AbstractMatrix{ <: Real }, J_old::AbstractMatrix{ <: Real }, output::AbstractVector{ <: Real }, input_new::AbstractVector{ <: Real }, layers::T) where { T <: NTuple{N,AbstractLayer} where N}

    # perform pass through first layers
    forward_jacobian!(J, J_new, J_old, output, input_new, first(layers))

    # update unless we are on the last layer
    if length(layers) > 1

        # update intermediate input
        input_new .= output

        # update old jacobian
        J_old .= J

    end

    # calculate forward_jacobian of remaining layers
    forward_jacobian!(J, J_new, J_old, output, input_new, Base.tail(layers))

end

# specialized methods
function forward_jacobian!(J::AbstractMatrix{ <: Real }, J_new::AbstractMatrix{ <: Real }, J_old::AbstractMatrix{ <: Real }, output::AbstractVector{ <: Real }, input_new::AbstractVector{ <: Real }, layer::PermutationLayer)
    # calculate new output
    forward!(output, layer, input_new)

    # calculate total jacobian
    mul!(J, jacobian(layer, input_new), J_old)
end

# standard method
function forward_jacobian!(J::AbstractMatrix{ <: Real }, J_new::AbstractMatrix{ <: Real }, J_old::AbstractMatrix{ <: Real }, output::AbstractVector{ <: Real }, input_new::AbstractVector{ <: Real }, layer::AbstractLayer)
    # calculate new jacobian
    forward_jacobian!(output, J_new, layer, input_new)

    # calculate total jacobian
    mul!(J, J_new, J_old)
end

# when no layers are left, stop the inplace recursion
forward_jacobian!(J::AbstractMatrix{ <: Real }, J_new::AbstractMatrix{ <: Real }, J_old::AbstractMatrix{ <: Real }, output::AbstractVector{ <: Real }, input_new::AbstractVector{ <: Real }, layers::Tuple{}) = return

# joing backward inverse jacobian of the Flow model
function _backward_inv_jacobian(model::CompiledFlowModel, output::AbstractVector{T}) where { T <: Real }

    # fetch layers
    dim    = getdim(model)

    # promote type for allocating output
    Ti = promote_type(eltype(model), T)

    # allocate space for jacobian
    J = zeros(Ti, dim, dim)
    input = zeros(Ti, dim)
    for k = 1:dim
        J[k,k] = one(Ti)
    end

    # calculate jacobian
    backward_inv_jacobian!(input, J, model, output)

    # return result
    return input, J

end

# when calling backward inverse jacobian, redirect to _backward_inv_jacobian
backward_inv_jacobian(model::CompiledFlowModel, output::AbstractVector{ <: Real}) = _backward_inv_jacobian(model, output)

# for broadcasting over backward inverse jacobian, fix the model for multiple inputs
Broadcast.broadcasted(::typeof(backward_inv_jacobian), model::CompiledFlowModel, output::AbstractVector{ <: AbstractVector{ <: Real } }) = broadcast(_backward_inv_jacobian, Ref(model), output)

# inplace inverse backward jacobian of the Flow model
function backward_inv_jacobian!(input::AbstractVector{ <: Real }, J::AbstractMatrix{T}, model::CompiledFlowModel, output::AbstractVector{ <: Real }) where { T <: Real }

    # fetch layers
    layers = getlayers(model)
    dim = getdim(model)

    # promote type for allocating output
    Ti = promote_type(eltype(model), T)

    # allocate space for intermediate output and jacobian
    output_new = zeros(Ti, dim)
    output_new .= output
    J_new = zeros(Ti, dim, dim)
    J_old = zeros(Ti, dim, dim)
    for k = 1:dim
        J_old[k,k] = one(Ti)
    end

    # calculate backward inverse jacobian over all layers
    backward_inv_jacobian!(J, J_new, J_old, input, output_new, layers)

end

# inplace backward inverse jacobian calculation for a tuple of layers
function backward_inv_jacobian!(J::AbstractMatrix{ <: Real }, J_new::AbstractMatrix{ <: Real }, J_old::AbstractMatrix{ <: Real }, input::AbstractVector{ <: Real }, output_new::AbstractVector{ <: Real }, layers::T) where { T <: NTuple{N,AbstractLayer} where N}

    # perform backward pass through layers
    backward_inv_jacobian!(J, J_new, J_old, input, output_new, last(layers))

    # update unless we are on the last/first layer
    if length(layers) > 1

        # update intermediate output
        output_new .= input

        # update old jacobian
        J_old .= J

    end

    # calculate backward inv jacobian of remaining layers
    backward_inv_jacobian!(J, J_new, J_old, input, output_new, Base.front(layers))

end

# specialized methods
function backward_inv_jacobian!(J::AbstractMatrix{ <: Real }, J_new::AbstractMatrix{ <: Real }, J_old::AbstractMatrix{ <: Real }, input::AbstractVector{ <: Real }, output_new::AbstractVector{ <: Real }, layer::PermutationLayer)
    # perform forward pass over last layer
    backward!(input, layer, output_new)

    # calculate total jacobian
    mul!(J, inv_jacobian(layer, output_new), J_old)
end

# standard method
function backward_inv_jacobian!(J::AbstractMatrix{ <: Real }, J_new::AbstractMatrix{ <: Real }, J_old::AbstractMatrix{ <: Real }, input::AbstractVector{ <: Real }, output_new::AbstractVector{ <: Real }, layer::AbstractLayer)
    # perform forward pass over last layer
    backward_inv_jacobian!(input, J_new, layer, output_new)

    # calculate total jacobian
    mul!(J, J_new, J_old)
end

# when no layers are left, stop the inplace recursion
backward_inv_jacobian!(J::AbstractMatrix{ <: Real }, J_new::AbstractMatrix{ <: Real }, J_old::AbstractMatrix{ <: Real }, input::AbstractVector{ <: Real }, output_new::AbstractVector{ <: Real }, layers::Tuple{}) = return

# specify jacobian and inv_jacobian functions based on joint functions
_jacobian(model::CompiledFlowModel, input::AbstractVector{ <: Real })      = forward_jacobian(model, input)[2]
jacobian(model::CompiledFlowModel, input::AbstractVector{ <: Real })       = _jacobian(model, input)
Broadcast.broadcasted(::typeof(jacobian), model::CompiledFlowModel, input::AbstractVector{ <: AbstractVector{ <: Real } }) = broadcast(_jacobian, Ref(model), input)
_inv_jacobian(model::CompiledFlowModel, output::AbstractVector{ <: Real }) = backward_inv_jacobian(model, output)[2]
inv_jacobian(model::CompiledFlowModel, output::AbstractVector{ <: Real })  = _inv_jacobian(model, output)
Broadcast.broadcasted(::typeof(inv_jacobian), model::CompiledFlowModel, output::AbstractVector{ <: AbstractVector{ <: Real } }) = broadcast(_inv_jacobian, Ref(model), output)
function jacobian!(J::AbstractMatrix{T1}, model::CompiledFlowModel, input::AbstractVector{T2}) where { T1 <: Real, T2 <: Real }

    # fetch dimension
    dim    = getdim(model)

    # promote type for allocating output
    T = promote_type(eltype(model), T1, T2)

    # allocate space for output and jacobian
    output = zeros(T, dim)
    J .= zero(T1)
    for k = 1:dim
        J[k,k] = one(T1)
    end

    # calculate jacobian
    forward_jacobian!(output, J, model, input)

end
function inv_jacobian!(J::AbstractMatrix{T1}, model::CompiledFlowModel, output::AbstractVector{T2}) where { T1 <: Real, T2 <: Real }

    # fetch dimension
    dim    = getdim(model)

    # promote type for allocating output
    T = promote_type(eltype(model), T1, T2)

    # allocate space for input and jacobian
    input = zeros(T, dim)
    J .= zero(T1)
    for k = 1:dim
        J[k,k] = one(T1)
    end

    # calculate jacobian
    backward_inv_jacobian!(input, J, model, output)

end

# fallback joint functions over layers
function forward_jacobian!(output::AbstractVector{ <: Real }, J_new::AbstractMatrix{ <: Real }, layer::AbstractLayer, input::AbstractVector{ <: Real })
    forward!(output, layer, input)
    jacobian!(J_new, layer, input)
end
function backward_inv_jacobian!(input::AbstractVector{ <: Real }, J_new::AbstractMatrix{ <: Real }, layer::AbstractLayer, output::AbstractVector{ <: Real })
    backward!(input, layer, output)
    inv_jacobian!(J_new, layer, output)
end

# extra utility functions
det_jacobian(model::CompiledFlowModel, input::AbstractVector{<: Real})           = det(jacobian(model, input))
absdet_jacobian(model::CompiledFlowModel, input::AbstractVector{<: Real})        = abs(det_jacobian(model, input))
logdet_jacobian(model::CompiledFlowModel, input::AbstractVector{<: Real})        = logdet(jacobian(model, input))
logabsdet_jacobian(model::CompiledFlowModel, input::AbstractVector{<: Real})     = logabsdet(jacobian(model, input))

detinv_jacobian(model::CompiledFlowModel, output::AbstractVector{<: Real})       = det(inv_jacobian(model, output))
absdetinv_jacobian(model::CompiledFlowModel, output::AbstractVector{<: Real})    = abs(detinv_jacobian(model, output))
logdetinv_jacobian(model::CompiledFlowModel, output::AbstractVector{<: Real})    = logdet(inv_jacobian(model, output))
logabsdetinv_jacobian(model::CompiledFlowModel, output::AbstractVector{<: Real}) = logabsdet(inv_jacobian(model, output))

# throw an error when the model has not yet been compiled
_forward(model::FlowModel, input::AbstractVector{ <: Real })                                                        = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
forward(model::FlowModel, input::AbstractVector{ <: Real })                                                         = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
Broadcast.broadcasted(::typeof(forward), model::FlowModel, input::AbstractVector{ <: AbstractVector{ <: Real } })   = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
forward!(output::AbstractVector{ <: Real }, model::FlowModel, input::AbstractVector{ <: Real })                     = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
_backward(model::FlowModel, output::AbstractVector{ <: Real })                                                      = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
backward(model::FlowModel, output::AbstractVector{ <: Real })                                                       = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
Broadcast.broadcasted(::typeof(backward), model::FlowModel, output::AbstractVector{ <: AbstractVector{ <: Real } }) = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
backward!(input::AbstractVector{ <: Real }, model::FlowModel, output::AbstractVector{ <: Real })                    = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
jacobian!(J_new::AbstractMatrix{ <: Real }, model::FlowModel, input::AbstractVector{ <: Real })                     = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
jacobian(model::FlowModel, input::AbstractVector{ <: Real })                                                        = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
Broadcast.broadcasted(::typeof(jacobian), model::FlowModel, input::AbstractVector{ <: AbstractVector{ <: Real } })  = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
_jacobian(model::FlowModel, input::AbstractVector{ <: Real })                                                       = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
_inv_jacobian(model::FlowModel, output::AbstractVector{ <: Real })                                                  = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
inv_jacobian!(J_new::AbstractMatrix{ <: Real }, model::FlowModel, output::AbstractVector{ <: Real })                = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))

det_jacobian(model::FlowModel, input::AbstractVector{ <: Real })           = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
absdet_jacobian(model::FlowModel, input::AbstractVector{ <: Real })        = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
logdet_jacobian(model::FlowModel, input::AbstractVector{ <: Real })        = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
logabsdet_jacobian(model::FlowModel, input::AbstractVector{ <: Real })     = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))

detinv_jacobian(model::FlowModel, output::AbstractVector{ <: Real })       = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
absdetinv_jacobian(model::FlowModel, output::AbstractVector{ <: Real })    = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
logdetinv_jacobian(model::FlowModel, output::AbstractVector{ <: Real })    = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
logabsdetinv_jacobian(model::FlowModel, output::AbstractVector{ <: Real }) = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))


text:
export Guild

"""
A [`Guild`](@ref)'s verification level.
More details [here](https://discordapp.com/developers/docs/resources/guild#guild-object-verification-level).
"""
@enum VerificationLevel VL_NONE VL_LOW VL_MEDIUM VL_HIGH VL_VERY_HIGH
@boilerplate VerificationLevel :lower

"""
A [`Guild`](@ref)'s default message notification level.
More details [here](https://discordapp.com/developers/docs/resources/guild#guild-object-default-message-notification-level).
"""
@enum MessageNotificationLevel MNL_ALL_MESSAGES MNL_ONLY_MENTIONS
@boilerplate MessageNotificationLevel :lower

"""
A [`Guild`](@ref)'s explicit content filter level.
More details [here](https://discordapp.com/developers/docs/resources/guild#guild-object-explicit-content-filter-level).
"""
@enum ExplicitContentFilterLevel ECFL_DISABLED ECFL_MEMBERS_WITHOUT_ROLES ECFL_ALL_MEMBERS
@boilerplate ExplicitContentFilterLevel :lower

"""
A [`Guild`](@ref)'s MFA level.
More details [here](https://discordapp.com/developers/docs/resources/guild#guild-object-mfa-level).
"""
@enum MFALevel ML_NONE ML_ELEVATED
@boilerplate MFALevel :lower

"""
A Discord guild (server).
Can either be an [`UnavailableGuild`](@ref) or a [`Guild`](@ref).
"""
abstract type AbstractGuild end

function AbstractGuild(d::Dict{String, Any})
    return get(d, "unavailable", length(d) <= 2) === true ? UnavailableGuild(d) : Guild(d)
end

"""
An unavailable Discord guild (server).
More details [here](https://discordapp.com/developers/docs/resources/guild#unavailable-guild-object).
"""
struct UnavailableGuild <: AbstractGuild
    id::Snowflake
    unavailable::Union{Bool, Missing}
end
@boilerplate UnavailableGuild :dict :docs :lower :merge

"""
A Discord guild (server).
More details [here](https://discordapp.com/developers/docs/resources/guild#guild-object).

The `djl_*` fields are internal fields used for cache performance.
"""
struct Guild <: AbstractGuild
    id::Snowflake
    name::String
    icon::Union{String, Nothing}
    splash::Union{String, Nothing}
    owner::Union{Bool, Missing}
    owner_id::Union{Snowflake, Missing}  # Missing in Invite.
    permissions::Union{Int, Missing}
    region::Union{String, Missing}  # Invite
    afk_channel_id::Union{Snowflake, Missing, Nothing}  # Invite
    afk_timeout::Union{Int, Missing}  # Invite
    embed_enabled::Union{Bool, Missing}
    embed_channel_id::Union{Snowflake, Missing, Nothing}  # Not supposed to be nullable.
    verification_level::VerificationLevel
    default_message_notifications::Union{MessageNotificationLevel, Missing}  # Invite
    explicit_content_filter::Union{ExplicitContentFilterLevel, Missing}  # Invite
    roles::Union{Vector{Role}, Missing}  # Invite
    emojis::Union{Vector{Emoji}, Missing}  # Invite
    features::Vector{String}
    mfa_level::Union{MFALevel, Missing}  # Invite
    application_id::Union{Snowflake, Missing, Nothing}  # Invite
    widget_enabled::Union{Bool, Missing}
    widget_channel_id::Union{Snowflake, Missing, Nothing}  # Not supposed to be nullable.
    system_channel_id::Union{Snowflake, Missing, Nothing}  # Invite
    joined_at::Union{DateTime, Missing}
    large::Union{Bool, Missing}
    unavailable::Union{Bool, Missing}
    member_count::Union{Int, Missing}
    voice_states::Union{Vector{VoiceState}, Missing}
    members::Union{Vector{Member}, Missing}
    channels::Union{Vector{DiscordChannel}, Missing}
    presences::Union{Vector{Presence}, Missing}
    djl_users::Union{Set{Snowflake}, Missing}
    djl_channels::Union{Set{Snowflake}, Missing}
end
@boilerplate Guild :dict :docs :lower :merge

Base.merge(x::UnavailableGuild, y::Guild) = y
Base.merge(x::Guild, y::UnavailableGuild) = x

ext: jl


text:
using TreeverseAlgorithm
using Test

@testset "TreeverseAlgorithm.jl" begin
    include("treeverse.jl")
end


text:
#!/usr/bin/env julia
#

using Printf

e_str1 = "2.718"
e = parse(Float64, e_str1)
println(5e)
#> 13.5914
num_15 = parse(Int, "15")
println(3num_15)
#> 45

@printf "e = %0.2f\n" e
#> 2.718
e_str2 = @sprintf("%0.3f", e)

# show that the 2 strings are the same
println("e_str1 == e_str2: $(e_str1 == e_str2)")
#> e_str1 == e_str2: true

# available number format characters are f, e, g, c, s, p, d
# (pi is a predefined constant; however, since its type is
# "MathConst" it has to be converted to a float to be formatted)
@printf "fix trailing precision: %0.3f\n" float(pi)
#> fix trailing precision: 3.142
@printf "scientific form: %0.6e\n" 1000pi
#> scientific form: 3.141593e+03
# g is not implemented yet
@printf "a character: %c\n" 'α'
#> a character: α
@printf "a string: %s\n" "look I'm a string!"
#> a string: look I'm a string!
@printf "right justify a string: %50s\n" "width 50, text right justified!"
#> right justify a string:                    width 50, text right justified!
@printf "a pointer: %p\n" 100000000
#> a pointer: 0x0000000005f5e100
@printf "print a integer: %d\n" 1e10
#> print an integer: 10000000000

max_stars_repo_name: wbeebe/julia

max_issues_repo_name: wbeebe/julia

max_forks_repo_name: wbeebe/julia


text:
"""
  derivative_series_from_initial!(; initial_series::ComplexOffsetVector,
    derivatives_series ::ComplexOffsetMatrix)

Transform initial series into array of derivatives' and integrals' series

Given long Taylor series
``f(z) = \\sum_{n=0}^N a_n z^n``
generates array of series
``\\frac{1}{k!}f^{(k)}(z) = \\sum_{n=0}^M b_{k,n} z^n, 0 < z ≤ K_2``,
``f^{(k)}(z) = \\sum_{n=0}^M b_{k,n} z^n, K_1 ≤ z < 0``,
where K1 and K2 are first and last indices if *derivative_series*.

Assuming  N ≥ M + (K_2 - K_1 + 1).
"""
function derivative_series_from_initial!(;
        initial_series     ::ComplexOffsetVector,
        derivatives_series ::ComplexOffsetMatrix
        )
    #####
    upper_term       = lastindex(derivatives_series, 1)
    upper_derivative = lastindex(derivatives_series, 2)
    lower_derivative = firstindex(derivatives_series, 2)

    #derivatives_series[0 : upper_term, 0] .= initial_series[0 : upper_term]
    for n in 0 : upper_term
        derivatives_series[n, 0] = initial_series[n]
    end

    for k in 1 : upper_derivative
        for n in 1 : (upper_term + upper_derivative - k + 1)
            initial_series[n-1] = initial_series[n] * n / (k+1)
        end
        initial_series[upper_term + upper_derivative - k + 1] = .0im

        #@views derivatives_series[:, k] .= initial_series[0 : upper_term]
        for n in 0 : upper_term
            derivatives_series[n, k] = initial_series[n]
        end
    end

    for k in -1 : -1 : lower_derivative
        derivatives_series[0, k] = initial_series[k]
        for n in 1 : upper_term
            derivatives_series[n, k] = derivatives_series[n-1, k+1] / n
        end
    end
end

function construct_dict_of_wei_series(;
    delta_to_nterms ::RComplex2IntDict,
    wei             ::Weierstrass
    )
#####
    max_n_terms = 0
    for v in values(delta_to_nterms)
        if v[1] > max_n_terms
            max_n_terms = v[1]
        end
    end

    wei_sequence = OffsetArray{ComplexF64}(undef, -2 : 2*max_n_terms)

    wei_series_dict = Dict{RationalComplex, ComplexOffsetMatrix}()

    for delta in keys(delta_to_nterms)
        upper_term = delta_to_nterms[delta]

        # Для вычисления разложений в ряд производных нужно разложить исходную
        # функцию с запасом членов.
        # Всего нужно n производных по n членов в каждой.
        diff_upper_term = 2 * upper_term

        if delta != 0
            z = real(delta)*wei.ω1 + imag(delta)*wei.ω3
            norm_factor = abs(z)

            weierstrass_normalized!(z=z,
                                norm_factor=norm_factor,
                                w=wei,
                                upper_term=diff_upper_term,
                                output=wei_sequence)
        else
            expand_at_pole!(output=wei_sequence,
                            weierstrass_data=wei,
                            upper_term=diff_upper_term)
        end

        series_array = OffsetArray{ComplexF64}(undef, 0:upper_term, -1:upper_term-2)

        derivative_series_from_initial!(initial_series=wei_sequence,
                                        derivatives_series=series_array)

        wei_series_dict[delta] = series_array
    end

    return wei_series_dict
end


text:
@testset "Dynamic Weights" begin
    @testset "Known V, Unknown W" begin
        Y = [[-0.112, 1.98, -2.28, 0.771, 21.2, 21.5],
             [-5.18, -2.44, -8.11, 0.165, 15.6, 18.6],
             [-19.3, -21.7, -22.0, -6.82, 13.6, 28.7],
             [-45.5, -32.2, -41.6, -35.0, 12.0, 27.3],
             [-64.4, -52.6, -70.3, -56.5, 2.44, 9.62],
             [-93.9, -77.0, -92.8, -75.8, -2.38, -1.88]]
        Y = to_matrix(Y)
        η = hcat([0.45, 0.45, 0.40, 0.30, 0.17, 0.05],
                 [0.45, 0.45, 0.40, 0.30, 0.17, 0.05],
                 [0.10, 0.10, 0.20, 0.40, 0.66, 0.90])

        F = [1. 0. 0.; 0. 1. 0.]
        G = [1. 0. 1.; 0. 1. 1.; 0. 0. 1.]
        V = Symmetric(diagm([12., 12.]))

        # Filter
        a, R, m, C = kfilter(Y, F, G, V, η, 0.8)

        @test a[end,:] ≈ [-27.190627146772364, -18.995342748051463, -4.916889343012473]
        @test R[end][1,3] ≈ R[end][2,3]
        @test R[end][3,3] ≈ 1.169478699932751
        @test m[end,:] ≈ [-18.453716253522465, -11.895375806382145, -2.5258277006501766]
        @test C[end][3,3] ≈ 0.6027170489902908

        # Smoother
        s, S = ksmoother(G, a, R, m, C)

        @test s[1,:] ≈ [-2.030157490822254, 3.400235808546456, -2.892548710163002]
        @test S[1][1,3] ≈ S[1][2,3]
        @test S[1][1,2] ≈ 1.9434098177645183

        # Fitted
        f, Q = fitted(F, V, s, S)

        @test f[1,:] ≈ [-2.030157490822254, 3.400235808546456]
        @test Q[1][1,1] ≈ Q[1][2,2]
        @test Q[end][1,2] ≈ 2.121050612190616

        # Forecast
        fh, Qh = forecast(F, G, V, 0.8, m[end,:], C[end], 6)

        @test fh[end,:] ≈ [-33.60868245742352, -27.0503420102832]
        @test Qh[end][1,1] ≈ Qh[end][2,2]
        @test Qh[end][1,2] ≈ 137.1312402019029
    end

    @testset "Unknown V, Unknown W" begin
        Y = [[-1.07, -4.79, 8.75, -3.38, 101.0, 95.3],
             [1.77, -4.83, -0.701, 8.53, 95.6, 91.6],
             [-2.44, 1.05, -1.99, -10.2, 110.0, 98.4],
             [-22.3, -25.6, -18.9, -27.4, 109.0, 116.0],
             [-66.0, -71.0, -64.1, -76.8, 142.0, 131.0],
             [-137.0, -138.0, -129.0, -138.0, 162.0, 163.0]]
        Y = to_matrix(Y)
        η = hcat([0.45, 0.475, 0.35, 0.20, 0.10, 0.025],
                 [0.45, 0.475, 0.35, 0.20, 0.10, 0.025],
                 [0.10, 0.150, 0.30, 0.60, 0.80, 0.950])

        F = [1. 0. 0.; 0. 1. 0.]
        G = [1. 0. 1.; 0. 1. 1.; 0. 0. 1.]

        # Estimate
        θ, Σ = estimate(Y, F, G, η, 0.8)

        @test θ[end,:] ≈ [128.5975140749007, 125.88516545696008, 25.250530953592737]
        @test diag(Σ) ≈ [17002.330502123168, 17452.4131082849]

        # Evolutional Covariances
        Σₑ = evolutional_covariances(Y, F, G, Σ, η, 0.8)

        @test Σₑ[end][:,3] ≈ [422.8068930004123, 424.1022162558545, 127.25615793643713]

        # Fitted
        f, Q = fitted(F, Σ, θ, Σₑ)

        @test f[1,:] ≈ [6.910174516807798, 4.197813682749192]
        @test diag(Q[1]) ≈ [23077.42980723912, 23582.512590829927]
        @test Q[1][1,2] ≈ 4270.747555928984

        # Forecast
        fh, Qh = forecast(F, G, Σ, 0.8, θ[end,:], Σₑ[end], 10)

        @test fh[end,:] ≈ [381.102823610828, 378.3904749928874]
        @test Qh[end][1,2] / (sqrt(Qh[end][1,1]) * sqrt(Qh[end][2,2])) ≈ 0.8104078439438087
        @test Qh[end][1,1] ≈ 163819.5096757467
    end
end


text:
import Pkg
import UUIDs

@inline function add_caps(strategy::CapStrategy,
                          option::LatestVersionOption,
                          registry_path::AbstractString)
    _registry_paths = Any[registry_path]
    add_caps(strategy, option, _registry_paths)
    return nothing
end

@inline function add_caps(strategy::CapStrategy,
                          option::LatestVersionOption,
                          registry_paths::AbstractVector)
    _registry_paths = Vector{String}(undef, 0)
    for path in registry_paths
        push!(_registry_paths, strip(path))
    end
    add_caps(strategy, option, _registry_paths)
    return nothing
end

@inline function add_caps(strategy::CapStrategy,
                          option::LatestVersionOption,
                          registry_paths::Vector{String})
    pkg_to_path,
        pkg_to_num_versions,
        pkg_to_latest_version,
        pkg_to_latest_zero_version = parse_registry(registry_paths)
    add_caps(strategy,
             option,
             registry_paths,
             pkg_to_path,
             pkg_to_num_versions,
             pkg_to_latest_version,
             pkg_to_latest_zero_version)
    return nothing
end

@inline function add_caps(strategy::CapStrategy,
                          option::LatestVersionOption,
                          registry_paths::Vector{String},
                          pkg_to_path::AbstractDict{Package, String},
                          pkg_to_num_versions::AbstractDict{Package, Int},
                          pkg_to_latest_version::AbstractDict{Package, VersionNumber},
                          pkg_to_latest_zero_version::AbstractDict{Package, <:Union{VersionNumber, Nothing}})
    all_pkgs = collect(keys(pkg_to_path))
    n = length(all_pkgs)
    for i = 1:n
        pkg = all_pkgs[i]
        @debug("Package $(i) of $(n)", pkg)
        pkg_path = pkg_to_path[pkg]::String
        add_caps(strategy,
                 option,
                 registry_paths,
                 pkg_to_latest_version,
                 pkg_to_latest_zero_version,
                 pkg,
                 pkg_path)
    end
    return nothing
end

@inline function add_caps(strategy::CapStrategy,
                          option::LatestVersionOption,
                          registry_paths::Vector{String},
                          pkg_to_latest_version::AbstractDict{Package, VersionNumber},
                          pkg_to_latest_zero_version::AbstractDict{Package, <:Union{VersionNumber, Nothing}},
                          pkg::Package,
                          pkg_path::String)
    all_versions = get_all_versions(pkg_path)
    latest_version = _get_latest_version(all_versions)
    compat_toml = joinpath(pkg_path, "Compat.toml")
    deps_toml = joinpath(pkg_path, "Deps.toml")
    if isfile(compat_toml) && isfile(deps_toml)
        compat = Compress.load(compat_toml)
        deps = Compress.load(deps_toml)
        m = length(all_versions)
        for j = 1:m
            version = all_versions[j]
            if (version != latest_version) | (option isa CapLatestVersion)
                add_caps!(compat,
                          deps,
                          strategy,
                          registry_paths,
                          pkg_to_latest_version,
                          pkg_to_latest_zero_version,
                          pkg,
                          pkg_path,
                          version)
            end
        end
        Compress.save(compat_toml, compat)
    end
    return nothing
end

@inline function add_caps(strategy::MonotonicUpperBound,
                          option::LatestVersionOption,
                          registry_paths::Vector{String},
                          pkg_to_latest_version::AbstractDict{Package, VersionNumber},
                          pkg_to_latest_zero_version::AbstractDict{Package, <:Union{VersionNumber, Nothing}},
                          pkg::Package,
                          pkg_path::String)
    option isa CapLatestVersion || error("MonotonicUpperBound requires CapLatestVersion")
    all_versions = get_all_versions(pkg_path)
    latest_version = _get_latest_version(all_versions)
    compat_toml = joinpath(pkg_path, "Compat.toml")
    deps_toml = joinpath(pkg_path, "Deps.toml")
    if isfile(compat_toml) && isfile(deps_toml)
        compat = Compress.load(compat_toml)
        deps = Compress.load(deps_toml)
        m = length(all_versions)
        # Bound the latest version
        version = all_versions[end]
        add_caps!(compat,
                  deps,
                  UpperBound(),
                  registry_paths,
                  pkg_to_latest_version,
                  pkg_to_latest_zero_version,
                  pkg,
                  pkg_path,
                  version)
        # Propagate the bounds backwards
        for j = m-1:-1:1
            lastcompat = compat[version]
            version = all_versions[j]
            pkg_to_latest_version_tmp = copy(pkg_to_latest_version)
            pkg_to_latest_zero_version_tmp = copy(pkg_to_latest_zero_version)
            monotonize!(pkg_to_latest_version_tmp, lastcompat)
            monotonize!(pkg_to_latest_zero_version_tmp, lastcompat)
            add_caps!(compat,
                      deps,
                      UpperBound(),
                      registry_paths,
                      pkg_to_latest_version_tmp,
                      pkg_to_latest_zero_version_tmp,
                      pkg,
                      pkg_path,
                      version)
        end
        Compress.save(compat_toml, compat)
    end
    return nothing
end

@inline function add_caps!(compat::AbstractDict{VersionNumber, <:Any},
                           deps::AbstractDict{VersionNumber, <:Any},
                           strategy::CapStrategy,
                           registry_paths::Vector{String},
                           pkg_to_latest_version::AbstractDict{Package, VersionNumber},
                           pkg_to_latest_zero_version::AbstractDict{Package, <:Union{VersionNumber, Nothing}},
                           pkg::Package,
                           pkg_path::String,
                           version::VersionNumber)
    if haskey(deps, version)
        get!(() -> Dict{Any, Any}(), compat, version)
        always_assert(haskey(compat, version))
        for dep in keys(deps[version])
            if !is_stdlib(dep) && !is_jll(dep)
                latest_dep_version = pkg_to_latest_version[Package(dep)]
                latest_dep_zero_version = pkg_to_latest_zero_version[Package(dep)]
                current_compat_for_dep = _strip(get(compat[version], dep, ""))
                new_compat_for_dep = generate_compat_entry(pkg,
                                                           Package(dep),
                                                           strategy,
                                                           current_compat_for_dep,
                                                           latest_dep_version,
                                                           latest_dep_zero_version)
                compat[version][dep] = new_compat_for_dep
            end
        end
    end
    return nothing
end


text:
@testset "Basic Structural Model" begin
    # Univariate
    air_passengers = CSV.File(StateSpaceModels.AIR_PASSENGERS) |> DataFrame
    log_air_passengers = log.(air_passengers.passengers)

    @test has_fit_methods(BasicStructural)

    model = BasicStructural(log_air_passengers, 12)
    fit!(model)
    # Runned on Python statsmodels
    @test loglike(model) ≈ 234.33641 atol = 1e-5 rtol = 1e-5

    # forecasting
    forec = forecast(model, 10)
    @test monotone_forecast_variance(forec)
    # simualting
    scenarios = simulate_scenarios(model, 10, 10_000)
    test_scenarios_adequacy_with_forecast(forec, scenarios)

    log_ap32 = Float32.(log_air_passengers)
    model = BasicStructural(log_ap32, 12)
    @test model.system.Z == Float32[1.0; 0.0; 1.0; zeros(10)]
    @test_broken fit!(model)
    @test_broken loglike(model) ≈ 234.33641f0 atol = 1e-5 rtol = 1e-5
end


text:
using Documenter
using Impero

makedocs(
    pages = [
        "Impero" => "index.md",
        "Quick Start" => ["basics.md", "expression_rewriting.md"],
        "Slow Start" => ["equations.md"]
    ],
    sitename = "Impero",
    format = Documenter.HTML(collapselevel = 1),
    modules = [Impero],
)

deploydocs(
    repo = "github.com/CliMA/Impero.jl",
)


text:
import LinearAlgebra:normalize

function normalize(x::T, y::T, z::T) where {T<:Real}
  l::T = sqrt( *(x, x) + *(y, y) + *(z, z) )
  xn::T = x / l
  yn::T = y / l
  zn::T = z / l
  return xn, yn, zn
end

struct Point{T<:AbstractFloat}
    x::T
    y::T
    z::T
end

struct Line{T<:AbstractFloat}
    x1::T
    y1::T
    z1::T
    x2::T
    y2::T
    z2::T
end

struct Ray{T<:AbstractFloat}
    x::T
    y::T
    z::T
    dx::T
    dy::T
    dz::T
end

struct Coord{S<:Integer}
    x::S
    y::S
end

abstract type Asset{T<:AbstractFloat, S<:Integer} end

struct Hook{T<:AbstractFloat, S<:Integer} <: Asset{T, S}
    oidx::S
    pidx::S
    xoff::T
    yoff::T
    zoff::T
    beta::T
    gamma::T
    Hook{T,S}(oidx::S, pidx::S, x::T, y::T, z::T, beta::T, gamma::T) where {T,S} = new(oidx, pidx, x, y, z, beta, gamma)
end
Hook(oidx::Integer, pidx::Integer, x::AbstractFloat, y::AbstractFloat, z::AbstractFloat, beta::AbstractFloat, gamma::AbstractFloat) = Hook{typeof(x), typeof(oidx)}(oidx,pidx,x,y,z,beta,gamma)

struct Transform{T<:AbstractFloat, S<:Integer} <: Asset{T, S}
    oidx::S
    pidx::S
    matrix::Array{T, 2}
    inverse::Array{T, 2}
    function Transform{T,S}(oid::S, pid::S, m::AbstractArray{T, 2}) where {T,S}
        matrix = m
	inverse = inv(m)
        oid = oid
        pid = pid
        new(oid, pid, matrix, inverse)
    end
end
Transform(oid::S, pid::S, matrix::AbstractArray{T, 2}) where {T<:AbstractFloat, S<:Integer} = Transform{eltype(matrix), typeof(oid)}(oid, pid, matrix)

struct Bvh{T<:AbstractFloat, S<:Integer, A<:AbstractArray, E<:AbstractArray}
    xmin::A
    ymin::A
    zmin::A
    xmax::A
    ymax::A
    zmax::A
    bvh2bb::E
    left_child::E
    right_child::E
    Bvh{T,S,A,E}(a::AbstractArray{T,1}, b::AbstractArray{T,1}, c::AbstractArray{T,1}, d::AbstractArray{T,1}, e::AbstractArray{T,1}, f::AbstractArray{T,1}, p::AbstractArray{S}, q::AbstractArray{S}, r::AbstractArray{S}) where {T,S,A,E} = new(a,b,c,d,e,f,p,q,r)
end
Bvh(a::AbstractArray, b::AbstractArray, c::AbstractArray, d::AbstractArray, e::AbstractArray, f::AbstractArray, p::AbstractArray, q::AbstractArray, r::AbstractArray) = Bvh{eltype(a), eltype(p), typeof(a), typeof(p)}(a,b,c,d,e,f,p,q,r)

struct Camera{T<:AbstractFloat, S<:Integer, A<:AbstractArray{T,1}}
  eyePoint::A
  lookAtPoint::A
  fov::T
  xResolution::S
  yResolution::S
  nBands::S
  nRaysPerPixel::S
  maxNumberOfCycles::S
  viewDirection::A
  u::A
  w::A
  aspectRatio::T
  viewPlaneHalfHeight::T
  viewPlaneHalfWidth::T
  viewPlaneBottomLeftPoint::A
  xIncVector::A
  yIncVector::A
  function Camera{T,S,A}(eyePoint::A, lookAtPoint::A, fov::T, xResolution::S, yResolution::S, nBands::S, nRaysPerPixel::S, maxNumberOfCycles::S) where {T,S,A}
    viewDirection = LinearAlgebra.normalize(lookAtPoint .- eyePoint)
    u = LinearAlgebra.normalize(cross(LinearAlgebra.normalize([0.000000, 0.000001, 1.0]), viewDirection))
    w = LinearAlgebra.normalize(cross(u, viewDirection))
    viewPlaneHalfWidth = tan(fov/2.0)
    aspectRatio = float(yResolution) / float(xResolution)
    viewPlaneHalfHeight = aspectRatio .* viewPlaneHalfWidth
    viewPlaneBottomLeftPoint = (eyePoint .+ viewDirection) .- (w .* viewPlaneHalfHeight) .- (u .* viewPlaneHalfWidth)
    xIncVector = (u .* 2.0 .* viewPlaneHalfWidth) ./ float(xResolution)
    yIncVector = (w .* 2.0 .* viewPlaneHalfHeight) ./ float(yResolution)
    new(eyePoint, lookAtPoint, fov, xResolution, yResolution, nBands, nRaysPerPixel, maxNumberOfCycles, viewDirection, u, w, aspectRatio, viewPlaneHalfHeight, viewPlaneHalfWidth, viewPlaneBottomLeftPoint, xIncVector, yIncVector)
  end
end
Camera(v::A, w::A, p::T, q::S, r::S, s::S, t::S, u::S) where {T<:AbstractFloat, S<:Integer, A<:AbstractArray} = Camera{T,S,A}(v,w,p,q,r,s,t,u)

abstract type Geometry{T<:AbstractFloat, S<:Integer, A<:AbstractArray, E<:AbstractArray} end

struct Ball{T,S,A,E} <: Geometry{T,S,A,E}
    p1x::A
    p1y::A
    p1z::A
    r::A
    rsq::A
    mtl::E
    function Ball{T,S,A,E}(p::AbstractArray{T,2}, radius::AbstractArray{T,1}, mtl::AbstractArray{S,1}) where {T,S,A,E}
        n_balls = size(p, 1)
        p1x::Array{T, 1} = p[:,1]
        p1y::Array{T, 1} = p[:,2]
        p1z::Array{T, 1} = p[:,3]
        r::Array{T, 1} = radius
        rsq::Array{T, 1} = radius.^2
        mtl::Array{S, 1} = mtl
        new(p1x, p1y, p1z, r, rsq, mtl)
    end
end
Ball(p::AbstractArray{T,2}, radius::AbstractArray{T,1}, mtl::AbstractArray{S,1}) where {T<:AbstractFloat, S<:Integer} = Ball{eltype(p), eltype(mtl), Array{T,1}, Array{S,1}}(p,radius,mtl)

struct Disk{T,S,A,E} <: Geometry{T,S,A,E}
    p1x::A
    p1y::A
    p1z::A
    nx::A
    ny::A
    nz::A
    r::A
    rsq::A
    mtl::E
    function Disk{T,S,A,E}(p::AbstractArray{T,2}, n::AbstractArray{T,2}, radius::AbstractArray{T,1}, mtl::AbstractArray{S,1}) where {T,S,A,E}
        n_disks = size(p, 1)
        p1x::Array{T, 1} = p[:,1]
        p1y::Array{T, 1} = p[:,2]
        p1z::Array{T, 1} = p[:,3]
        nx::Array{T, 1} = n[:,1]
        ny::Array{T, 1} = n[:,2]
        nz::Array{T, 1} = n[:,3]
        r::Array{T, 1} = radius
        rsq::Array{T, 1} = radius.^2
        mtl::Array{S, 1} = mtl
        new(p1x, p1y, p1z, nx, ny, nz, r, rsq, mtl)
    end
end
Disk(p::AbstractArray{T,2}, n::AbstractArray{T,2}, radius::AbstractArray{T,1}, mtl::AbstractArray{S,1}) where {T<:AbstractFloat, S<:Integer} = Disk{eltype(p), eltype(mtl), Array{T,1}, Array{S,1}}(p,n,radius,mtl)

struct Cylinder{T,S,A,E} <: Geometry{T,S,A,E}
    m1x::A
    m1y::A
    m1z::A
    m2x::A
    m2y::A
    m2z::A
    radius::A
    u::A
    v::A
    w::A
    length::A
    beta::A
    gamma::A
    rsq::A
    mtl::E
    function Cylinder{T,S,A,E}(m::AbstractArray{T,2}, l::AbstractArray{S,2}, radius::AbstractArray{T,1}, mtl::AbstractArray{S,1}) where {T,S,A,E}
        n_cylinders = size(l, 1)
        m1x::Array{T, 1} = zeros(n_cylinders)
        m1y::Array{T, 1} = zeros(n_cylinders)
        m1z::Array{T, 1} = zeros(n_cylinders)
        m2x::Array{T, 1} = zeros(n_cylinders)
        m2y::Array{T, 1} = zeros(n_cylinders)
        m2z::Array{T, 1} = zeros(n_cylinders)
        radius::Array{T, 1} = radius
        u::Array{T, 1} = zeros(n_cylinders)
        v::Array{T, 1} = zeros(n_cylinders)
        w::Array{T, 1} = zeros(n_cylinders)
        length::Array{T, 1} = zeros(n_cylinders)
        beta::Array{T, 1} = zeros(n_cylinders)
        gamma::Array{T, 1} = zeros(n_cylinders)
        rsq::Array{T, 1} = zeros(n_cylinders)
        mtl::Array{S, 1} = mtl
        for i = 1:n_cylinders
            idx1 = l[i,1]; idx2 = l[i,2]
            m1x[i] = m[idx1,1]; m1y[i] = m[idx1,2]; m1z[i] = m[idx1,3]
            m2x[i] = m[idx2,1]; m2y[i] = m[idx2,2]; m2z[i] = m[idx2,3]
            dx = m[idx2,1] - m[idx1,1]
            dy = m[idx2,2] - m[idx1,2]
            dz = m[idx2,3] - m[idx1,3]
            length_ = sqrt(dx^2 + dy^2 + dz^2)
            length[i] = length_
            beta[i] = acos(dz / length_)
            gamma[i] = atan(dy, dx) # formerly atan2
            rsq[i] = radius[i]^2
            # compute a 'normal' vector that is pointing up
            q = [dx, dy, dz]
            nv = normalize(cross([0.0034, 0.0071, 1.0], q))
            vq = cross(nv, q)
            u[i] = vq[1]
            v[i] = vq[2]
            w[i] = vq[3]
        end
        new(m1x, m1y, m1z, m2x, m2y, m2z, radius, u, v, w, length, beta, gamma, rsq, mtl)
    end
end
Cylinder(m::AbstractArray{T,2}, l::AbstractArray{S,2}, radius::AbstractArray{T,1}, mtl::AbstractArray{S,1}) where {T<:AbstractFloat, S<:Integer} = Cylinder{eltype(m), eltype(l), Array{T,1}, Array{S,1}}(m,l,radius,mtl)

struct Cone{T,S,A,E} <: Geometry{T,S,A,E}
    m1x::A
    m1y::A
    m1z::A
    m2x::A
    m2y::A
    m2z::A
    radius_base::A
    radius_tip::A
    u::A
    v::A
    w::A
    beta::A
    gamma::A
    length::A
    a::A
    aa::A
    r1a::A
    r1sq::A
    mtl::E
    function Cone{T,S,A,E}(m::AbstractArray{T,2}, l::AbstractArray{S,2}, radius_base::AbstractArray{T,1}, radius_tip::AbstractArray{T,1}, mtl::AbstractArray{S,1}) where {T,S,A,E}
        n_cones = size(l, 1)
        m1x::Array{T, 1} = zeros(n_cones)
        m1y::Array{T, 1} = zeros(n_cones)
        m1z::Array{T, 1} = zeros(n_cones)
        m2x::Array{T, 1} = zeros(n_cones)
        m2y::Array{T, 1} = zeros(n_cones)
        m2z::Array{T, 1} = zeros(n_cones)
        radius_base::Array{T, 1} = radius_base
        radius_tip::Array{T, 1} = radius_tip
        u::Array{T, 1} = zeros(n_cones)
        v::Array{T, 1} = zeros(n_cones)
        w::Array{T, 1} = zeros(n_cones)
        beta::Array{T, 1} = zeros(n_cones)
        gamma::Array{T, 1} = zeros(n_cones)
        length::Array{T, 1} = zeros(n_cones)
        a::Array{T, 1} = zeros(n_cones)
        aa::Array{T, 1} = zeros(n_cones)
        r1a::Array{T, 1} = zeros(n_cones)
        r1sq::Array{T, 1} = zeros(n_cones)
        mtl::Array{S, 1} = mtl
        for i = 1:n_cones
            idx1 = l[i,1]; idx2 = l[i,2]
            m1x[i] = m[idx1,1]; m1y[i] = m[idx1,2]; m1z[i] = m[idx1,3]
            m2x[i] = m[idx2,1]; m2y[i] = m[idx2,2]; m2z[i] = m[idx2,3]
            dx = m[idx2,1] - m[idx1,1]
            dy = m[idx2,2] - m[idx1,2]
            dz = m[idx2,3] - m[idx1,3]
            length_ = sqrt(dx^2 + dy^2 + dz^2)
            length[i] = length_
            beta[i] = acos(dz / length_)
            gamma[i] = atan(dy, dx) # formerly atan2
            a_ = (radius_tip[i] - radius_base[i]) / length_ # taper of the stem
            a[i] = a_
            aa[i] = a_^2
            r1a[i] = radius_base[i]*a_
            r1sq[i] = radius_base[i]^2
            # compute a 'normal' vector that is pointing up
            q = [dx, dy, dz]
            nv = normalize(cross([0.0034, 0.0071, 1.0], q))
            vq = cross(nv, q)
            u[i] = vq[1]
            v[i] = vq[2]
            w[i] = vq[3]
        end
        new(m1x, m1y, m1z, m2x, m2y, m2z, radius_base, radius_tip, u, v, w, beta, gamma, length, a, aa, r1a, r1sq, mtl)
    end
end
Cone(m::AbstractArray{T,2}, l::AbstractArray{S,2}, radius_base::AbstractArray{T,1}, radius_tip::AbstractArray{T,1}, mtl::AbstractArray{S,1}) where {T<:AbstractFloat, S<:Integer} = Cone{eltype(m), eltype(l), Array{T,1}, Array{S,1}}(m,l,radius_base,radius_tip,mtl)

struct Mesh{T,S,A,E} <: Geometry{T,S,A,E}
    v1x::A
    v1y::A
    v1z::A
    v2x::A
    v2y::A
    v2z::A
    v3x::A
    v3y::A
    v3z::A
    mtl::E
    ux::A
    uy::A
    uz::A
    wx::A
    wy::A
    wz::A
    nx::A
    ny::A
    nz::A
    uu::A
    uw::A
    ww::A
    D::A
    function Mesh{T,S,A,E}(v::AbstractArray{T, 2}, t::AbstractArray{S, 2}, m::AbstractArray{S, 1}) where {T,S,A,E}
        n_triangles = size(t, 1)
        v1x::Array{T, 1} = zeros(n_triangles)
        v1y::Array{T, 1} = zeros(n_triangles)
        v1z::Array{T, 1} = zeros(n_triangles)
        v2x::Array{T, 1} = zeros(n_triangles)
        v2y::Array{T, 1} = zeros(n_triangles)
        v2z::Array{T, 1} = zeros(n_triangles)
        v3x::Array{T, 1} = zeros(n_triangles)
        v3y::Array{T, 1} = zeros(n_triangles)
        v3z::Array{T, 1} = zeros(n_triangles)
        mtl::Array{S, 1} = m
        ux::Array{T, 1} = zeros(n_triangles)
        uy::Array{T, 1} = zeros(n_triangles)
        uz::Array{T, 1} = zeros(n_triangles)
        wx::Array{T, 1} = zeros(n_triangles)
        wy::Array{T, 1} = zeros(n_triangles)
        wz::Array{T, 1} = zeros(n_triangles)
        nx::Array{T, 1} = zeros(n_triangles)
        ny::Array{T, 1} = zeros(n_triangles)
        nz::Array{T, 1} = zeros(n_triangles)
        uu::Array{T, 1} = zeros(n_triangles)
        uw::Array{T, 1} = zeros(n_triangles)
        ww::Array{T, 1} = zeros(n_triangles)
        D::Array{T, 1} = zeros(n_triangles)
        for i = 1:n_triangles
            idx1 = t[i,1]; idx2 = t[i,2]; idx3 = t[i,3]
            v1x[i] = v[idx1,1]; v1y[i] = v[idx1,2]; v1z[i] = v[idx1,3]
            v2x[i] = v[idx2,1]; v2y[i] = v[idx2,2]; v2z[i] = v[idx2,3]
            v3x[i] = v[idx3,1]; v3y[i] = v[idx3,2]; v3z[i] = v[idx3,3]
            # precompute some variables...
            ux_ = v2x[i] - v1x[i]
            uy_ = v2y[i] - v1y[i]
            uz_ = v2z[i] - v1z[i]
            wx_ = v3x[i] - v1x[i]
            wy_ = v3y[i] - v1y[i]
            wz_ = v3z[i] - v1z[i]
            u_ = [ux_, uy_, uz_]
            w_ = [wx_, wy_, wz_]
            n_ = cross(u_, w_)
            n2_ = n_.^2
            sm_ = sum(n2_)
            l_ = sqrt(sm_)
            n_ = n_./l_
            uu_ = dot(u_, u_)
            uw_ = dot(u_, w_)
            ww_ = dot(w_, w_)
            uw2_ = (uw_ * uw_)
            uuww_ = (uu_ * ww_)
            D_ = uw2_ - uuww_
            ux[i] = ux_
            uy[i] = uy_
            uz[i] = uz_
            wx[i] = wx_
            wy[i] = wy_
            wz[i] = wz_
            nx[i] = n_[1]
            ny[i] = n_[2]
            nz[i] = n_[3]
            uu[i] = uu_
            uw[i] = uw_
            ww[i] = ww_
            D[i] = D_
        end
        new(v1x, v1y, v1z, v2x, v2y, v2z, v3x, v3y, v3z, mtl, ux, uy, uz, wx, wy, wz, nx, ny, nz, uu, uw, ww, D)
    end
end
Mesh(v::AbstractArray{T,2}, t::AbstractArray{S,2}, m::AbstractArray{S,1}) where {T<:AbstractFloat, S<:Integer} = Mesh{eltype(v), eltype(t), Array{T,1}, Array{S,1}}(v,t,m)

struct Walls{T<:AbstractFloat}
    bottom::T
    top::T
    north::T
    south::T
    east::T
    west::T
    xsize::T
    ysize::T
    zsize::T
    function Walls{T}(BSW::Point{T}, TNE::Point{T}) where {T<:AbstractFloat}
        # BSW = Bottom South West corner; TNE = Top North East corner
        bottom = BSW.z
        top = TNE.z
        north = TNE.y
        south = BSW.y
        east = TNE.x
        west = BSW.x
        xsize = east - west
        ysize = north - south
        zsize = top - bottom
        new(bottom, top, north, south, east, west, xsize, ysize, zsize)
    end
end
Walls(p::Point{T}, q::Point{T}) where {T<:AbstractFloat} = Walls{T}(p,q)

max_stars_repo_name: martinvanleeuwen/RenderJay

max_issues_repo_name: martinvanleeuwen/RenderJay

max_forks_repo_name: martinvanleeuwen/RenderJay


text:
"""
    Quadrotor{R}

A standard quadrotor model, with simple aerodynamic forces. The orientation is represent by
a general rotation `R`. The body z-axis point is vertical, so positive controls cause acceleration
in the positive z direction.

# Constructor
    Quadrotor(; kwargs...)
    Quadrotor{R}(; kwargs...)

where `R <: Rotation{3}` and defaults to `UnitQuaternion{Float64}` if omitted. The keyword arguments are
* `mass` - mass of the quadrotor, in kg (default = 0.5)
* `J` - inertia of the quadrotor, in kg⋅m² (default = `Diagonal([0.0023, 0.0023, 0.004])`)
* `gravity` - gravity vector, in kg/m² (default = [0,0,-9.81])
* `motor_dist` - distane between the motors, in m (default = 0.1750)
* `km` - motor torque constant (default = 0.0245)
* `kf` - motor force constant (default = 1.0)
"""
@autodiff struct Quadrotor{R} <: RigidBody{R}
    mass::Float64
    J::SMatrix{3,3,Float64,9}
    Jinv::SMatrix{3,3,Float64,9}
    gravity::SVector{3,Float64}
    motor_dist::Float64
    kf::Float64
    km::Float64
    bodyframe::Bool  # velocity in body frame?
    ned::Bool
end
control_dim(::Quadrotor) = 4

function Quadrotor{R}(;
        mass=0.5,
        J=Diagonal(@SVector [0.0023, 0.0023, 0.004]),
        gravity=SVector(0,0,-9.81),
        motor_dist=0.1750,
        kf=1.0,
        km=0.0245,
        bodyframe=false,
        ned=false,
    ) where R
    @assert issymmetric(J)
    Quadrotor{R}(mass,J,inv(J),gravity,motor_dist,kf,km,bodyframe,ned)
end

(::Type{Quadrotor})(;kwargs...) = Quadrotor{UnitQuaternion{Float64}}(;kwargs...)

@inline RobotDynamics.velocity_frame(model::Quadrotor) = model.bodyframe ? :body : :world

function trim_controls(model::Quadrotor)
    @SVector fill(-model.gravity[3]*model.mass/4.0, size(model)[2])
end

function forces(model::Quadrotor, x, u)
    q = orientation(model, x)
    kf = model.kf
    g = model.gravity
    m = model.mass

    w1 = u[1]
    w2 = u[2]
    w3 = u[3]
    w4 = u[4]

    F1 = max(0,kf*w1);
    F2 = max(0,kf*w2);
    F3 = max(0,kf*w3);
    F4 = max(0,kf*w4);
    F = @SVector [0., 0., F1+F2+F3+F4] #total rotor force in body frame
    if model.ned
        F = SA[0,0,-F[3]]
        g = -g
    end

    f = m*g + q*F # forces in world frame
    return f
end

function moments(model::Quadrotor, x, u)

    kf, km = model.kf, model.km
    L = model.motor_dist

    w1 = u[1]
    w2 = u[2]
    w3 = u[3]
    w4 = u[4]

    F1 = max(0,kf*w1);
    F2 = max(0,kf*w2);
    F3 = max(0,kf*w3);
    F4 = max(0,kf*w4);

    M1 = km*w1;
    M2 = km*w2;
    M3 = km*w3;
    M4 = km*w4;
    tau = @SVector [L*(F2-F4), L*(F3-F1), (M1-M2+M3-M4)] #total rotor torque in body frame
    if model.ned
        tau = SA[tau[1], -tau[2], -tau[3]]
    end
    return tau
end

function wrenches(model::Quadrotor, x, u)
    F = forces(model, x, u)
    M = moments(model, x, u)
    return [F; M]

    q = orientation(model, x)
    C = forceMatrix(model)
    mass, g = model.mass, model.gravity

    # Calculate force and moments
    w = max.(u, 0)  # keep forces positive
    fM = forceMatrix(model)*w
    f = fM[1]
    M = @SVector [fM[2], fM[3], fM[4]]
    e3 = @SVector [0,0,1]
    F = mass*g - q*(f*e3)
    return F,M
end

function forceMatrix(model::Quadrotor)
    kf, km = model.kf, model.km
    L = model.motor_dist
    @SMatrix [
        kf   kf   kf   kf;
        0    L*kf 0   -L*kf;
       -L*kf 0    L*kf 0;
        km  -km   km  -km;
    ]
end

RobotDynamics.inertia(model::Quadrotor) = model.J
RobotDynamics.inertia_inv(model::Quadrotor) = model.Jinv
RobotDynamics.mass(model::Quadrotor) = model.mass

function Base.zeros(model::Quadrotor{R}) where R
    x = RobotDynamics.build_state(model, zero(RBState))
    u = @SVector fill(-model.mass*model.gravity[end]/4, 4)
    return x,u
end


text:
@testset "simple workflow" begin
    env = WrappedEnv(CloneStatePreprocessor(), CartPoleEnv{Float32}())
    agent = Agent(;
        policy = RandomPolicy(env),
        trajectory = VectorialCompactSARTSATrajectory(; state_type = Vector{Float32}),
    )
    N_EPISODE = 10000
    hook = TotalRewardPerEpisode()
    run(agent, env, StopAfterEpisode(N_EPISODE), hook)

    @test isapprox(sum(hook.rewards) / N_EPISODE, 21; atol = 2)
end


text:
# Convert a type into the SymbolicUtils type we'll use to represent it
# for example,
#     julia> SymbolicCodegen.sym(Int)
#     :(SymbolicCodegen.Sym{Int64})
#
#     julia> SymbolicCodegen.sym(Int, :n)
#     :(SymbolicCodegen.Sym{Int64}(:n))
#
sym(T::Type) = :(SymbolicCodegen.Sym{$T})

sym(T::Type, s::Symbol) = :($(sym(T))($(QuoteNode(s))))

export atoms

function atoms(t::Symbolic)
    f = operation(t)
    args = arguments(t)

    if hasproperty(f, :name) && f.name == :Sum
        return setdiff(atoms(args[1]), [args[2]])
    else
        return union(atoms(f), union(atoms.(args)...))
    end
end

atoms(s::Sym) = Set{Sym}([s])

atoms(x) = Set{Sym}()


text:
module NDIteration

export _Size, StaticSize, DynamicSize, get
export NDRange, blocks, workitems, expand

import Base.@pure

abstract type _Size end
struct DynamicSize <: _Size end
struct StaticSize{S} <: _Size
    function StaticSize{S}() where S
        new{S::Tuple{Vararg{Int}}}()
    end
end

@pure StaticSize(s::Tuple{Vararg{Int}}) = StaticSize{s}()
@pure StaticSize(s::Int...) = StaticSize{s}()
@pure StaticSize(s::Type{<:Tuple}) = StaticSize{tuple(s.parameters...)}()

# Some @pure convenience functions for `StaticSize`
@pure get(::Type{StaticSize{S}}) where {S} = S
@pure get(::StaticSize{S}) where {S} = S
@pure Base.getindex(::StaticSize{S}, i::Int) where {S} = i <= length(S) ? S[i] : 1
@pure Base.ndims(::StaticSize{S}) where {S} = length(S)
@pure Base.length(::StaticSize{S}) where {S} = prod(S)

"""
    NDRange

Encodes a blocked iteration space.

# Example
```
ndrange = NDRange{2, DynamicSize, DynamicSize}(CartesianIndices((256, 256)), CartesianIndices((32, 32)))
for block in ndrange
    for items in workitems(ndrange)
        I = expand(ndrange, block, items)
        checkbounds(Bool, A, I) || continue
        @inbounds A[I] = 2*A[I]
    end
end
```
"""
struct NDRange{N, StaticBlocks, StaticWorkitems, DynamicBlock, DynamicWorkitems}
    blocks::DynamicBlock
    workitems::DynamicWorkitems

    function NDRange{N, B, W}() where {N, B, W}
        new{N, B, W, Nothing, Nothing}(nothing, nothing)
    end

    function NDRange{N, B, W}(blocks, workitems) where {N, B, W}
        new{N, B, W, typeof(blocks), typeof(workitems)}(blocks, workitems)
    end
end

@inline workitems(range::NDRange{N, B, W}) where {N,B,W<:DynamicSize} = range.workitems::CartesianIndices{N}
@inline workitems(range::NDRange{N, B, W}) where {N,B,W<:StaticSize} = CartesianIndices(get(W))::CartesianIndices{N}
@inline blocks(range::NDRange{N, B}) where {N,B<:DynamicSize} = range.blocks::CartesianIndices{N}
@inline blocks(range::NDRange{N, B}) where {N,B<:StaticSize} = CartesianIndices(get(B))::CartesianIndices{N}

import Base.iterate
@inline iterate(range::NDRange) = iterate(blocks(range))
@inline iterate(range::NDRange, state) = iterate(blocks(range), state)

Base.length(range::NDRange) = length(blocks(range))

@inline function expand(ndrange::NDRange{N}, groupidx::CartesianIndex{N}, idx::CartesianIndex{N}) where N
    nI = ntuple(Val(N)) do I
        Base.@_inline_meta
        stride = size(workitems(ndrange), I)
        gidx = groupidx.I[I]
        (gidx-1)*stride + idx.I[I]
    end
    CartesianIndex(nI)
end

Base.@propagate_inbounds function expand(ndrange::NDRange, groupidx::Integer, idx::Integer)
    expand(ndrange, blocks(ndrange)[groupidx], workitems(ndrange)[idx])
end

Base.@propagate_inbounds function expand(ndrange::NDRange{N}, groupidx::CartesianIndex{N}, idx::Integer) where N
    expand(ndrange, groupidx, workitems(ndrange)[idx])
end

Base.@propagate_inbounds function expand(ndrange::NDRange{N}, groupidx::Integer, idx::CartesianIndex{N}) where N
    expand(ndrange, blocks(ndrange)[groupidx], idx)
end

"""
    partition(ndrange, workgroupsize)

Splits the maximum size of the iteration space by the workgroupsize.
Returns the number of workgroups necessary and whether the last workgroup
needs to perform dynamic bounds-checking.
"""
@inline function partition(ndrange, workgroupsize)
    @assert length(workgroupsize) <= length(ndrange)
    if length(workgroupsize) < length(ndrange)
        # pad workgroupsize with ones
        workgroupsize = ntuple(length(ndrange)) do I
            if I > length(workgroupsize)
                return 1
            else
                return workgroupsize[I]
            end
        end
    end

    dynamic = false
    blocks = ntuple(length(ndrange)) do I
        dynamic |= mod(ndrange[I], workgroupsize[I]) != 0
        return fld1(ndrange[I], workgroupsize[I])
    end

    return blocks, workgroupsize, dynamic
end

end #module


text:
"""
    ZmqSubscriber

A simple wrapper around a ZMQ subscriber, but only for protobuf messages.

# Construction

    Subscriber(context::ZMQ.Context, ipaddr, port; name)

To create a subscriber, pass in a `ZMQ.Context`, which allows all related
publisher / subscribers to be collected in a "group." The subscriber also
needs to be provided the IPv4 address (either as a string or as a `Sockets.IPv4`
object), and the port (either as an integer or a string).

A name can also be optionally provided via the `name` keyword, which can be used
to provide a helpful description about what the subscriber is subscribing to. It defaults
to "subscriber_#" where `#` is an increasing index.

# Usage
Use the blocking `subscribe` method to continually listen to the socket and
store data in a protobuf type:

    subscribe(sub::Subscriber, proto_msg::ProtoBuf.ProtoType)

Note that this function contains an infinite while loop so will block the calling
thread indefinately. It's usually best to assign the process to a separate thread / task:

```
sub_task = @task subscribe(sub, proto_msg)
schedule(sub_task)
```
"""
struct ZmqSubscriber <: Subscriber
    socket::ZMQ.Socket
    port::Int64
    ipaddr::Sockets.IPv4
    buffer::IOBuffer
    name::String
    socket_lock::ReentrantLock
    flags::SubscriberFlags
    zmsg::ZMQ.Message

    function ZmqSubscriber(
        ctx::ZMQ.Context,
        ipaddr::Sockets.IPv4,
        port::Integer;
        name = gensubscribername(),
        buffersize = 255,
    )
        local socket
        @catchzmq(
            socket = ZMQ.Socket(ctx, ZMQ.SUB),
            "Could not create socket for subscriber $name."
        )
        @catchzmq(
            ZMQ.subscribe(socket),
            "Could not set the socket as a subscriber for subscriber $name."
        )
        @catchzmq(
            set_conflate(socket, 1),
            "Could not set the conflate option for subscriber $name."
        )
        @catchzmq(
            ZMQ.connect(socket, "tcp://$ipaddr:$port"),
            "Could not connect subscriber $name to port $(tcpstring(ipaddr, port))."
        )

        @info "Subscribing $name to: tcp://$ipaddr:$port"
        new(
            socket,
            port,
            ipaddr,
            IOBuffer(zeros(UInt8, buffersize)),
            name,
            ReentrantLock(),
            SubscriberFlags(),
            ZMQ.Message(),
        )
    end
end

function ZmqSubscriber(ctx::ZMQ.Context, ipaddr, port::Integer; kwargs...)
    if !(ipaddr isa Sockets.IPv4)
        ipaddr = Sockets.IPv4(ipaddr)
    end
    ZmqSubscriber(ctx, ipaddr, port; kwargs...)
end

function ZmqSubscriber(ctx::ZMQ.Context, ipaddr, port::AbstractString; kwargs...)
    ZmqSubscriber(ctx, ipaddr, parse(Int, port); kwargs...)
end

getcomtype(::ZmqSubscriber) = :zmq
function Base.isopen(sub::ZmqSubscriber)
    return lock(() -> ZMQ.isopen(sub.socket), sub.socket_lock)
end

function Base.close(sub::ZmqSubscriber)
    lock(sub.socket_lock) do
        if isopen(sub.socket)
            @debug "Closing ZmqSubscriber: $(getname(sub))"
            ZMQ.close(sub.socket)
        end
    end
end

function forceclose(sub::ZmqSubscriber)
    @warn "Force closing ZmqSubscriber: $(getname(sub))"
    close(sub.socket)
end

function receive(sub::ZmqSubscriber, msg::MercuryMessage)
    did_receive = false
    getflags(sub).isreceiving = true
    bin_data = sub.zmsg

    bytes_read = Int32(0)
    if isopen(sub)
        bytes_read = ZMQ.msg_recv(sub.socket, bin_data, ZMQ.ZMQ_DONTWAIT)::Int32
        getflags(sub).bytesrecieved = bytes_read

        if bytes_read == -1
            ZMQ.zmq_errno() == ZMQ.EAGAIN || throw(ZMQ.StateError(ZMQ.jl_zmq_error_str()))
        else
            did_receive = got_new!(sub)
        end

        # Forces subscriber to conflate messages
        ZMQ.getproperty(sub.socket, :events)
    end

    # TODO: test this code to make sure it works in practice
    if bytes_read > length(sub.buffer.data)
        @warn "Increasing buffer size for subscriber $(getname(sub)) from $(length(sub.buffer.data)) to $bytes_read."
        sub.buffer.data = zeros(UInt8, bytes_read)
        sub.buffer.size = bytes_read
    end

    # Copy the data to the local buffer and decode

    if did_receive
        seek(sub.buffer, 0)
        copyto!(sub.buffer.data, 1, bin_data, 1, bytes_read)
        sub.buffer.size = bytes_read

        decode!(msg, sub.buffer)
    end

    return did_receive
end

portstring(sub::ZmqSubscriber) = tcpstring(sub.ipaddr, sub.port)

"""
    publish_until_receive(pub, sub, msg_out; [timeout])

Publish a message via the publisher `pub` until it's received by the subscriber `sub`.
Both `pub` and `sub` should have the same port and IP address.

The function returns `true` if a message was received before `timeout` seconds have passed,
    and `false` otherwise.
"""
function publish_until_receive(
    pub::ZmqPublisher,
    sub::ZmqSubscriber,
    msg_out::MercuryMessage,
    msg_in::MercuryMessage,
    timeout = 1.0,  # seconds
)
    @assert pub.ipaddr == sub.ipaddr && pub.port == sub.port "Publisher and subscriber must be on the same port!"
    @assert isopen(pub) "Publisher must be open"
    @assert isopen(sub) "Subscriber must be open"
    t_start = time()
    sub.flags.hasreceived = false
    while (time() - t_start < timeout)
        publish(pub, msg_out)
        sleep(0.001)
        receive(sub, msg_in)
        if sub.flags.hasreceived
            return true
        end
    end
    @warn "Publish until receive timed out"
    return false
end


text:
using WagnerModel
using Test

@testset "WagnerModel.jl" begin
    config = gen_config()
    @test config["test_param"]

    grn = [[1.0,0.0],[0.0,1.0]]
    grn = hcat(grn...)
    S_0 = [1.0,0.0]
    @test assess_stability(grn, S_0, config)[1]
    @test assess_stability(grn, S_0, config)[2] == [1.0,0.0]
    @test assess_stability(grn, S_0, config)[3] == 1
end