@enum inverter_inner_vars begin
    md_var = 1
    mq_var = 2
    Vdc_var = 3
    Vdo_var = 5
    Vqo_var = 6
    ω_freq_estimator_var = 7
    v_control_var = 8
    ω_control_var = 9
    δdqRI_var = 10 # alias for δθ_vsm, TODO: should δθ_vsm ever be accessed as external state? RH: Most likely
    VR_inv_var = 11
    VI_inv_var = 12
    Vdcnv_var = 13
    Vqcnv_var = 14
end

Base.to_index(ix::inverter_inner_vars) = Int(ix)

include("inverter_components/outer_control.jl")
include("inverter_components/voltage_source_control.jl")
include("inverter_components/converter.jl")
include("inverter_components/dc_source.jl")
include("inverter_components/frequency_estimator.jl")
include("inverter_components/filter.jl")

mutable struct DynInverter{C <: Converter,
                           O <: OuterControl,
                           VC <: VSControl,
                           DC<: DCSource,
                           P <: FrequencyEstimator,
                           F <: Filter} <: DynInjection
    number::Int64
    name::Symbol
    bus::PSY.Bus
    ω_ref::Float64
    V_ref::Float64
    P_ref::Float64
    Q_ref::Float64
    MVABase::Float64
    converter::C
    outercontrol::O
    vscontrol::VC
    dc_source::DC
    freq_estimator::P
    filter::F #add MVAbase here
    n_states::Int64
    states::Vector{Symbol}
    inner_vars::Vector{Float64}
    local_state_ix::Dict{InverterComponent,Vector{Int64}}
    input_port_mapping::Dict{InverterComponent,Vector{Int64}}

        function DynInverter(number::Int64,
                            name::Symbol,
                            bus::PSY.Bus,
                            ω_ref::Float64,
                            V_ref::Float64,
                            P_ref::Float64,
                            Q_ref::Float64,
                            MVABase::Float64,
                            converter::C,
                            outercontrol::O,
                            vscontrol::VC,
                            dc_source::DC,
                            freq_estimator::P,
                            filter::F) where {C<:Converter,
                                              O<:OuterControl,
                                              VC<:VSControl,
                                              DC<:DCSource,
                                              P<:FrequencyEstimator,
                                              F<:Filter}

            n_states = (converter.n_states +
                       outercontrol.n_states +
                       vscontrol.n_states +
                       dc_source.n_states +
                       freq_estimator.n_states +
                       filter.n_states)

            states = vcat(converter.states,
                        outercontrol.states,
                        vscontrol.states,
                        dc_source.states,
                        freq_estimator.states,
                        filter.states)

            components = [converter,
                          outercontrol,
                          vscontrol,
                          dc_source,
                          freq_estimator,
                          filter]

            new{C, O, VC, DC, P, F}(number,
                                    name,
                                    bus,
                                    ω_ref,
                                    V_ref,
                                    P_ref,
                                    Q_ref,
                                    MVABase,
                                    converter,
                                    outercontrol,
                                    vscontrol,
                                    dc_source,
                                    freq_estimator,
                                    filter,
                                    n_states,
                                    states,
                                    zeros(14),
                                    _make_state_mapping(components, states),
                                    _make_port_mapping(components, states))

        end
end

get_inverter_Sbase(device::DynInverter) = device.converter.s_rated
get_inverter_Vref(device::DynInverter) = device.V_ref


# eventually these should go inside the module or their own modules and I should choose what to export. not yet
include("LJH.jl")
using LJH
using Dierckx, PyCall, HDF5, ArrayViews, Distributions, KernelDensity, PyPlot, JLD
import Graphs

include("summarize.jl")
include("mockpulses.jl")
include("runningvectors.jl")
include("histogram.jl")
include("steps.jl")
include("steps_macros.jl")
include("calibration.jl")
include("corrections.jl")
include("ljhutil.jl")
include("tes_optimal_filtering.jl")
include("glue.jl")
include("matterwatcher.jl")
module Mass2

# package code goes here

end # module


module Commitment

using Codex
using DataFrames

export delta2scores

reverse(ans::Number) = 8 - ans

"""
    is_2018(df::DataFrame)::Bool

Returns whether `df` contains responses to the 2018 version of the delta questionnaire.
"""
is_2018(df::DataFrame)::Bool = !("v7" in names(df))

ans_2018_mapping = Dict(
    "Helemaal mee oneens" => 1,
    "Gedeeltelijk mee oneens" => 2,
    "Niet oneens, niet eens" => 3,
    "Gedeeltelijk mee eens" => 4,
    "Helemaal mee eens" => 5
)

function is_reversed(question::Int, year::Int)::Bool
    year == 2018 ? question in [1, 2, 4] : question in [2, 3, 5]
end

"""
    ans2num(ans, question::Int, year::Int)

Note that the scale was from 1 to 5 in 2018 and 1 to 7 in later years.
"""
function ans2num(ans, question::Int, year::Int)
    ans = year == 2018 ?
        Codex.rescale(ans_2018_mapping[ans], 1, 5, 1, 7) :
        (typeof(ans) != Int ? get(ans, Int) : ans)
    return is_reversed(question, year) ? reverse(ans) : ans
end

score_2018(responses...) = sum(ans2num.(responses, 1:5, 2018))
score_2019(responses...) = sum(ans2num.(responses, 2:6, 2019))

function delta2scores(df::DataFrame)
    if is_2018(df)
        transform!(df, [:v1, :v2, :v3, :v4, :v5] => ByRow(score_2018) => :score)
        df[!, :achievable] .= ""
    else
        transform!(df, [:v2, :v3, :v4, :v5, :v6] => ByRow(score_2019) => :score)
        select!(df, :, :v7 => :achievable)
    end
    select!(df, :id, :completed_at, :score, :achievable)
    disallowmissing!(df)
end

end # module


function construct_network!(ps_m::CanonicalModel,
                            system_formulation::Type{CopperPlatePowerModel},
                            sys::PSY.System,
                            time_range::UnitRange{Int64}; kwargs...)

    bus_count = length(sys.buses)

    copper_plate(ps_m, :var_active, bus_count, time_range)

    return
end

function construct_network!(ps_m::CanonicalModel,
                            system_formulation::Type{StandardPTDFForm},
                            sys::PSY.System,
                            time_range::UnitRange{Int64}; kwargs...)

    if :PTDF in keys(kwargs)

        ac_branches = [br for br in sys.branches if !isa(br, PSY.DCLine)]

        # TODO: Get DC Lines and model DC Lines

        flow_variables(ps_m, system_formulation, ac_branches, time_range)

        ptdf_networkflow(ps_m, ac_branches, sys.buses, :var_active, kwargs[:PTDF], time_range)

    else
        throw(ArgumentError("no PTDF matrix supplied"))
    end

    return

end

function construct_network!(ps_m::CanonicalModel,
                            system_formulation::Type{S},
                            sys::PSY.System,
                            time_range::UnitRange{Int64}; kwargs...) where {S <: PM.AbstractPowerFormulation}

    powermodels_network!(ps_m, system_formulation, sys, time_range)

    return

end


"""
TagData(data1, data2; batch=128, ftype=Float32, dense=false) creates a
data generator that can be used with a Tagger model.  The source data1
and target data2 should be sequence generators, i.e. next(data1)
should deliver a vector of Ints that represent the next sequence.
This division of labor allows different file formats to be supported.
maxtoken(data1) should give the largest integer produced by data1.

The following transformations are performed by an TagData generator:

* sequences are minibatched according to the batch argument.
* sequences in a minibatch padded to all be the same length.
* the end-of-sequence is indicated using a "nothing" token.

"""
type TagData; x; y; mask;
    function TagData(xgen, ygen; batchsize=128, ftype=Float32, dense=false, o...)
        new(TagBatch(xgen, batchsize, ftype, dense),
            TagBatch(ygen, batchsize, ftype, dense),
            zeros(Cuchar, batchsize))
    end
end

type TagBatch; sgen; state; sent; word; done;
    function TagBatch(sgen, batchsize, ftype, dense)
        word = (dense ? zeros : sponehot)(ftype, maxtoken(sgen), batchsize)
        sent = Array(Any, batchsize)
        new(sgen, nothing, sent, word, false)
    end
end

import Base: start, done, next

# the TagData state is just nword, which is the number of words
# completed in the current batch.

function start(d::TagData)
    d.x.done = d.y.done = false
    d.x.state = start(d.x.sgen)
    d.y.state = start(d.y.sgen)
    nextbatch(d.x)
    nextbatch(d.y)
    return 0
end

# We stop when there is not enough data to fill a batch

done(d::TagData, state)=(d.x.done && d.y.done)

# TagData.next returns the next token

function next(d::TagData, nword)
    nword += 1
    maxlen = maximum(map(length, d.x.sent))
    if nword > maxlen
        nextbatch(d.x)
        nextbatch(d.y)
        return (nothing, 0)
    end
    for s=1:length(d.x.sent)
        n=length(d.x.sent[s])
        xword = (nword <= n ? d.x.sent[s][nword] : 0)
        yword = (nword <= n ? d.y.sent[s][nword] : 0)
        # @show s,n,xword,yword
        setrow!(d.x.word, xword, s)
        setrow!(d.y.word, yword, s)
        d.mask[s] = (xword == 0 ? 0 : 1)
    end
    return ((d.x.word, d.y.word, d.mask), nword)
end

function nextbatch(b::TagBatch)
    for i=1:length(b.sent)
        done(b.sgen, b.state) && (b.done=true; return)
        (b.sent[i], b.state) = next(b.sgen, b.state)
    end
end

maxtoken(s::TagData,i)=maxtoken(i==1 ? s.x.sgen : i==2 ? s.y.sgen : error())
maxtoken(x::Vector{Vector{Int}}) = maximum(map(maximum,x))

sponehot(ftype, m, n)=sparse(ones(Int32,n), Int32[1:n;], one(ftype), m, n)

# # TODO: these assume one hot columns, make them more general.
# setrow!(x::SparseMatrixCSC,i,j)=(i>0 ? (x.rowval[j] = i; x.nzval[j] = 1) : (x.rowval[j]=1; x.nzval[j]=0))
# setrow!(x::Array,i,j)=(x[:,j]=0; i>0 && (x[i,j]=1))

# function readvocab(file) # TODO: test with cmd e.g. `zcat foo.gz`
#     d = Dict{Any,Int}()
#     open(file) do f
#         for l in eachline(f)
#             for w in split(l)
#                 get!(d, w, 1+length(d))
#             end
#         end
#     end
#     return d
# end


using GLAbstraction, ModernGL, GLWindow, GLFW

window 	= createwindow("GLPlot", 1920,1280, windowhints=[(GLFW.SAMPLES, 0)])
ocamera = OrthographicPixelCamera(window.inputs)

points = Vec3[Vec3(sin(i), i, 0) * 100f0 for i=0:0.1:2pi]

data = [
:vertex         	=> GLBuffer(Float32[1:4], 1), # simple quad
:index          	=> indexbuffer(GLuint[0, 1, 2, 2, 3, 0]),

:points 			=> Texture(points),
:linewidth 			=> 20f0,
:shadow 			=> 0.0f0,
:projection			=> ocamera.projection,
:view 				=> ocamera.view,
]

program = TemplateProgram(
	Pkg.dir("GLDraw","src","line.vert"), Pkg.dir("GLDraw","src","line.frag"),
	attributes=data, fragdatalocation=[(0, "fragment_color"),(1, "fragment_objectid")]
)
obj = instancedobject(data, length(points)-3, program)
prerender!(obj, enabletransparency)

glClearColor(1, 1, 1, 1.0)

while !GLFW.WindowShouldClose(window.nativewindow)
    yield()
	glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
    glViewport(0,0,window.inputs[:framebuffer_size].value...)
    render(obj)

    GLFW.SwapBuffers(window.nativewindow)
	GLFW.PollEvents()
end
GLFW.Terminate()


using ..Continuations: AbstractContinuationProblem, AbstractContinuationSolver,
    AbstractContinuationCache,
    as, SweepSetup, ContinuationSweep, ContinuationSolution,
    ContinuationCache, ContinuationOptions, ContinuationSolver,
    residual_jacobian!

using ..BifurcationsBase: timekind, Continuous, Discrete
import ..BifurcationsBase: TimeKind

using ..BifurcationsBase
using ..BifurcationsBase: SpecialPoint, SpecialPointInterval,
    BifurcationSweep, BifurcationSolution, BifurcationSolver,
    BifurcationProblem, BifurcationCache,
    eigvals_prototype, allocate_sweep!, check_sweep_length, record!
import ..BifurcationsBase: analyze!, re_analyze!

import ..Codim1: testfn

module PointTypes
@enum(
    PointType,
    none,
    saddle_node,
)
end  # module
using .PointTypes: PointType

const Codim1LCSweep{tkind <: TimeKind, ptType <: PointType} =
    BifurcationSweep{tkind, ptType}
const Codim1LCSolution{S <: ContinuationSolution, W <: Codim1LCSweep} =
    BifurcationSolution{S, W}
# The constraints `<: TimeKind` and `<: ContinuationSolution` are
# important.
# See: [[../codim1/solver.jl::Codim1Sweep]]

BifurcationsBase.point_type_type(::Codim1LCProblem) = PointType

BifurcationsBase.regular_point(::Type{PointType}) = PointTypes.none

function BifurcationsBase.eigvals_prototype(prob::Codim1LCProblem,
                                            cache::ContinuationCache)
    n = length(prob.de_prob.u0)
    return similar(cache.u[1:1] .+ 1im, n)
end

mutable struct Codim1LCCache{P, C <: ContinuationCache{P},
                             JType, eType,
                             } <: BifurcationCache{P}
    super::C
    prev_det::Float64
    J::JType
    eigvals::eType
    point_type::PointType
end

function Codim1LCCache(super::C,
                       J::JType,
                       eigvals::eType,
                       point_type = PointTypes.none,
                       ) where {P, C <: ContinuationCache{P},
                                JType, eType}
    return Codim1LCCache{P, C, JType, eType}(
        super,
        NaN, J, eigvals,
        point_type)
end
# TODO: Remove this constructor after removing the type parameter `P`.

Codim1LCCache(prob::Codim1LCProblem, super::ContinuationCache) =
    Codim1LCCache(
        super,
        let v = eigvals_prototype(prob, super)
            v * v'
        end,
        # ds_jacobian(super),
        copy(eigvals_prototype(prob, super)),
    )

BifurcationsBase.BifurcationCache(prob::Codim1LCProblem,
                                  super::ContinuationCache) =
    Codim1LCCache(prob, super)

const Codim1LCSolver{
        R <: ContinuationSolver,
        P <: Codim1LCProblem,
        C <: Codim1LCCache,
        S <: Codim1LCSolution,
        } =
    BifurcationSolver{R, P, C, S}

# TODO: trait
BifurcationsBase.contkind(::Codim1LCProblem) = LimitCycleCont()

function re_analyze!(solver::Codim1LCSolver, u::AbstractVector)
    residual_jacobian!(as(solver.cache, ContinuationCache), u)
    analyze!(solver.cache, solver.opts)

    # Suppress special point recording:
    # It's a bit ugly hack... (communicate by sharing!)
    solver.cache.point_type = PointTypes.none  # TODO: FIX!

    record!(solver.sol, solver.cache)
end

function analyze!(cache::Codim1LCCache, opts)
    cache.point_type = PointTypes.none

    curr_det = det(@view as(cache, ContinuationCache).J[:, 1:end-1])
    if curr_det * cache.prev_det < 0
        cache.point_type = PointTypes.saddle_node
    end
    cache.prev_det = curr_det

    #=
    cache.J = J = ds_jacobian(cache)
    eigvals = ds_eigvals(timekind(cache), J)
    cache.point_type = guess_point_type(timekind(cache), cache, eigvals, opts)
    cache.eigvals = eigvals
    =#
    set_reference!(cache)
end

testfn(::Val{PointTypes.saddle_node}, ::Continuous, ::LimitCycleCont,
       prob_cache, u, J, L, Q) =
    det(@view J[:, 1:end-1])

function set_reference!(wrapper)
    cache = as(wrapper, ContinuationCache)
    prob_cache = as(cache.prob_cache, LimitCycleCache)
    set_reference!(prob_cache, cache.u)
end


using ACE, StaticArrays, BenchmarkTools, Printf

using Zygote: gradient
using ACE: evaluate, evaluate!, evaluate_d, evaluate_d!
using ACEbase: acquire_B!, acquire_dB!

TX = ACE.PositionState{Float64}
B1p = ACE.Utils.RnYlm_1pbasis()
Rn = B1p.bases[1]
cfg = ACEConfig(rand(TX, Rn, 30))

##

# degrees = Dict(2 => [9, 17],
#                3 => [7, 15] )

Pnumprops = [1, 2, 4]

degrees = Dict(2 => [10, 17],
               3 => [9, 15],
               5 => [7, 11] )

wL = 1.5

#zygote tests
site_energy(m,x) = sum(ACE.val.(ACE.evaluate(m, x)))

#loss tests
FS = props -> sum( (1 .+ ACE.val.(props).^2).^0.5 )
sqr(x) = x.rr .^ 2

#full loss
loss(m, x) = (FS(ACE.evaluate(m,x)))^2 + sum(sum(sqr.(gradient(tx->FS(ACE.evaluate(m,tx)), x)[1])))
#only energy loss
lossE(m, x) = (FS(ACE.evaluate(m,x)))^2

##
#Pgroup for benchmarking different number of properties (P)
#fixed maxorder = 2 and maxdeg = 7

fixdeg = 7
fixord = 2

#linear model
Pgroup = BenchmarkGroup()
Pgroup["set_params!"] = BenchmarkGroup()
Pgroup["evaluate"] = BenchmarkGroup()

#only ACE
Pgroup["grad_params"] = BenchmarkGroup()
Pgroup["grad_config"] = BenchmarkGroup()
Pgroup["_rrule_evaluate"] = BenchmarkGroup()
Pgroup["grad_params_config"] = BenchmarkGroup()
Pgroup["adjoint_EVAL_D"] = BenchmarkGroup()

# #zygote calls
# Pgroup["site_energy"] = BenchmarkGroup()
# Pgroup["Zygote_grad_params"] = BenchmarkGroup()
# Pgroup["Zygote_grad_config"] = BenchmarkGroup()
# Pgroup["Zygote_grad_params_config"] = BenchmarkGroup()

# #loss functions
# Pgroup["eval_full_loss"] = BenchmarkGroup()
# Pgroup["eval_energy_loss"] = BenchmarkGroup()
# Pgroup["der_full_loss"] = BenchmarkGroup()
# Pgroup["der_energy_loss"] = BenchmarkGroup()

for numprops in Pnumprops
   local B1p
   Bsel = SparseBasis(; maxorder = fixord, p = 1, default_maxdeg = fixdeg,
                        weight = Dict(:n => 1.0, :l => wL))
   B1p = ACE.Utils.RnYlm_1pbasis(maxdeg = fixdeg,
                                 maxL = ceil(Int, fixdeg / wL),
                                 Bsel = Bsel)

   φ = ACE.Invariant()
   bsis = ACE.SymmetricBasis(φ, B1p, ACE.O3(), Bsel)

   #create a multiple property model
   W = rand(numprops, length(bsis))
   c = [SVector{size(W)[1]}(W[:,i]) for i in 1:size(W)[2]]
   LM = ACE.LinearACEModel(bsis, c, evaluator = :standard)

   Pgroup["set_params!"][numprops] = @benchmarkable ACE.set_params!($LM, $c)
   Pgroup["evaluate"][numprops] = @benchmarkable ACE.evaluate($LM, $cfg)

   #ACE native calls for derivatives
   Pgroup["grad_params"][numprops] = @benchmarkable ACE.grad_params($LM, $cfg)
   Pgroup["grad_config"][numprops] = @benchmarkable ACE.grad_config($LM, $cfg)
   dp = ones(numprops) #the forces pullback input
   Pgroup["_rrule_evaluate"][numprops] = @benchmarkable ACE._rrule_evaluate($dp, $LM, $cfg)
   Pgroup["grad_params_config"][numprops] = @benchmarkable ACE.grad_params_config($LM, $cfg)
   dq = [ACE.DState(rr=rand(SVector{3, Float64})) for _ = 1:length(cfg)] #the adjoint
   Pgroup["adjoint_EVAL_D"][numprops] = @benchmarkable ACE.adjoint_EVAL_D($LM, $LM.evaluator, $cfg, $dq)

   # #Zygote calls for derivatives
   # Pgroup["site_energy"][numprops] = @benchmarkable site_energy($LM, $cfg)
   # Pgroup["Zygote_grad_params"][numprops] = @benchmarkable gradient(x->site_energy(x,$cfg), $LM)
   # Pgroup["Zygote_grad_config"][numprops] = @benchmarkable gradient(x->site_energy($LM,x), $cfg)
   # #sum over the function so we don't compute the jacobian
   # Pgroup["Zygote_grad_params_config"][numprops] = @benchmarkable gradient(m->sum(sum(gradient(x->site_energy(m,x), $cfg)[1]).rr), $LM)

   # #loss function
   # Pgroup["eval_full_loss"][numprops] = @benchmarkable loss($LM, $cfg)
   # Pgroup["eval_energy_loss"][numprops] = @benchmarkable lossE($LM, $cfg)
   # Pgroup["der_full_loss"][numprops] = @benchmarkable gradient(m->loss(m,$cfg), $LM)
   # Pgroup["der_energy_loss"][numprops] = @benchmarkable gradient(m->lossE(m,$cfg), $LM)

end

##
#Benchmarking derivatives for different sizes of basis (B)
#the number of properties is fixed to 2

Nprop = 2

#linear model
Bgroup = BenchmarkGroup()
Bgroup["set_params!"] = BenchmarkGroup()
Bgroup["evaluate"] = BenchmarkGroup()

#only ACE
Bgroup["grad_params"] = BenchmarkGroup()
Bgroup["grad_config"] = BenchmarkGroup()
Bgroup["_rrule_evaluate"] = BenchmarkGroup()
Bgroup["grad_params_config"] = BenchmarkGroup()
Bgroup["adjoint_EVAL_D"] = BenchmarkGroup()

# #zygote calls
# Bgroup["site_energy"] = BenchmarkGroup()
# Bgroup["Zygote_grad_params"] = BenchmarkGroup()
# Bgroup["Zygote_grad_config"] = BenchmarkGroup()
# Bgroup["Zygote_grad_params_config"] = BenchmarkGroup()

# #loss functions
# Bgroup["eval_full_loss"] = BenchmarkGroup()
# Bgroup["eval_energy_loss"] = BenchmarkGroup()
# Bgroup["der_full_loss"] = BenchmarkGroup()
# Bgroup["der_energy_loss"] = BenchmarkGroup()

for ord = keys(degrees), deg in degrees[ord]
   local B1p
   Bsel = SparseBasis(; maxorder = ord, p = 1, default_maxdeg = deg,
                        weight = Dict(:n => 1.0, :l => wL))
   B1p = ACE.Utils.RnYlm_1pbasis(maxdeg = deg,
                                 maxL = ceil(Int, deg / wL),
                                 Bsel = Bsel)

   φ = ACE.Invariant()
   bsis = ACE.SymmetricBasis(φ, B1p, ACE.O3(), Bsel)

   #create a multiple property model
   W = rand(Nprop, length(bsis))
   c = [SVector{size(W)[1]}(W[:,i]) for i in 1:size(W)[2]]
   LM = ACE.LinearACEModel(bsis, c, evaluator = :standard)

   Bgroup["set_params!"][ord, deg] = @benchmarkable ACE.set_params!($LM, $c)
   Bgroup["evaluate"][ord, deg] = @benchmarkable ACE.evaluate($LM, $cfg)

   #ACE native calls for derivatives
   Bgroup["grad_params"][ord, deg] = @benchmarkable ACE.grad_params($LM, $cfg)
   Bgroup["grad_config"][ord, deg] = @benchmarkable ACE.grad_config($LM, $cfg)
   dp = ones(Nprop) #the forces pullback input
   Bgroup["_rrule_evaluate"][ord, deg] = @benchmarkable ACE._rrule_evaluate($dp, $LM, $cfg)
   Bgroup["grad_params_config"][ord, deg] = @benchmarkable ACE.grad_params_config($LM, $cfg)
   dq = [ACE.DState(rr=rand(SVector{3, Float64})) for _ = 1:length(cfg)] #the adjoint
   Bgroup["adjoint_EVAL_D"][ord, deg] = @benchmarkable ACE.adjoint_EVAL_D($LM, $LM.evaluator, $cfg, $dq)

   # #Zygote calls for derivatives
   # Bgroup["site_energy"][ord, deg] = @benchmarkable site_energy($LM, $cfg)
   # Bgroup["Zygote_grad_params"][ord, deg] = @benchmarkable gradient(x->site_energy(x,$cfg), $LM)
   # Bgroup["Zygote_grad_config"][ord, deg] = @benchmarkable gradient(x->site_energy($LM,x), $cfg)
   # #sum over the function so we don't compute the jacobian
   # Bgroup["Zygote_grad_params_config"][ord, deg] = @benchmarkable gradient(m->sum(sum(gradient(x->site_energy(m,x), $cfg)[1]).rr), $LM)

   # #loss function
   # Bgroup["eval_full_loss"][ord, deg] = @benchmarkable loss($LM, $cfg)
   # Bgroup["eval_energy_loss"][ord, deg] = @benchmarkable lossE($LM, $cfg)
   # Bgroup["der_full_loss"][ord, deg] = @benchmarkable gradient(m->loss(m,$cfg), $LM)
   # Bgroup["der_energy_loss"][ord, deg] = @benchmarkable gradient(m->lossE(m,$cfg), $LM)

end

##

linear_suite = BenchmarkGroup()
linear_suite["P"] = Pgroup
linear_suite["B"] = Bgroup

##

# @info("Tune")
# tune!(linear_suite)

# @info("Run")
# results = run(linear_suite, verbose = true)


type SubtreeCrossover <: Crossover
  stage::AbstractString
  rate::Float
  max_tries::Int
  max_depth::Int
  rep::KozaTreeRepresentation

  SubtreeCrossover(stage::AbstractString, rate::Float, max_tries::Int, max_depth::Int, rep::KozaTreeRepresentation) =
    new(stage, rate, max_tries, max_depth, rep)
end

"""
Provides a definition for a subtree crossover operator.
"""
type SubtreeCrossoverDefinition <: CrossoverDefinition
  """
  The name of the stage that this operator works on.
  """
  stage::AbstractString

  """
  The crossover rate for this operator, i.e. the probability that a given pair
  of trees will be subject to crossover.
  """
  rate::Float

  """
  The maximum number of tries that the operator when attempting to make a
  viable crossover before giving up.
  """
  max_tries::Int

  SubtreeCrossoverDefinition(stage::AbstractString, r::Float, mt::Int) =
    new(stage, r, mt)
end

"""
Composes a given sub-tree crossover operator from its definition.
"""
function compose!(c::SubtreeCrossoverDefinition, sp::Species)
  c.stage = c.stage == "" ? genotype(sp).label : c.stage
  r = sp.stages[c.stage].representation
  SubtreeCrossover(c.stage, c.rate, c.max_tries, r.max_depth, r)
end

"""
TODO: Write koza.subtree_crossover documentation.
"""
subtree_crossover() = subtree_crossover("", 0.7, 1)
subtree_crossover(stage::AbstractString) =
  subtree_crssover(stage, 0.7, 1)
subtree_crossover(rate::Float) = subtree_crossover("", rate, 1)
subtree_crossover(limit::Int) = subtree_crossover("", 0.7, limit)
subtree_crossover(stage::AbstractString, rate::Float) =
  subtree_crossover(stage, rate, 1)
subtree_crossover(stage::AbstractString, limit::Int) =
  subtree_crossover(stage, 0.7, limit)
subtree_crossover(rate::Float, limit::Int) = subtree_crossover("", rate, limit)
subtree_crossover(stage::AbstractString, rate::Float, limit::Int) =
  SubtreeCrossoverDefinition(stage, rate, limit)
function subtree_crossover(def::Function)
  c = subtree_crossover()
  def(c)
  c
end

num_inputs(o::SubtreeCrossover) = 2
num_outputs(o::SubtreeCrossover) = 2

"""
Performs a sub-tree crossover on a pair of provided Koza trees (belonging to
the same species).
"""
function operate!{T <: KozaTree}(
  o::SubtreeCrossover,
  outputs::Vector{IndividualStage{T}},
  x::T,
  y::T
)
  # Enforce the crossover rate
  rand() > o.rate && return

  # Choose a node at random from the first tree.
  nx = sample(x)

  # Continue to sample nodes from the 2nd tree until either a viable
  # crossover candidate has been found, or till a maximum number of
  # attempts have passed.
  for i in 1:o.max_tries
    ny = sample(y)
    if  (depth(ny) + height(nx) - 1) <= o.max_depth &&
        (depth(nx) + height(ny) - 1) <= o.max_depth
      swap!(ny, nx)
      break
    end
  end
end


module LVCyl

export LorentzVectorCyl

export px, py, pz, energy, fast_mass
export deltaphi, deltar, tocartesian, fromcartesian
export ΔR, Δϕ, Δη

import Base: +, *, zero

include("main.jl")

end # module


# Author: Youngjun Kim, youngjun@stanford.edu
# Date: 06/09/2014

module SimpleADMImpl

export
    addObserver,

    initialize,
    update,

    initializeDynamicModel,
    simulateDynamicModel,

    SimpleADM,
    SimpleADMInitialState,
    SimpleADMOutputState,
    SimpleADMCommand

using AbstractDynamicModelImpl
using AbstractDynamicModelInterfaces
using CommonInterfaces
using RLESUtils, Observers

using Base.Test

import CommonInterfaces.addObserver
import CommonInterfaces.initialize
import CommonInterfaces.update
import AbstractDynamicModelInterfaces.initializeDynamicModel
import AbstractDynamicModelInterfaces.simulateDynamicModel

type SimpleADMInitialState

    t::Float64
    x::Float64
    y::Float64
    h::Float64
    v::Float64
    psi::Float64
    h_d::Float64
end

type SimpleADMState

    t::Float64
    x::Float64
    y::Float64
    h::Float64
    v::Float64
    psi::Float64
end

type SimpleADMOutputState

    t::Float64
    x::Float64
    y::Float64
    h::Float64
    vx::Float64
    vy::Float64
    vh::Float64
end

type SimpleADMCommand

    t::Float64
    v_d::Float64
    h_d::Float64
    psi_d::Float64
end

type SimpleADM <: AbstractDynamicModel

    state::SimpleADMState
    command::Union{SimpleADMCommand, Void}

    timestep::Float64
    number_of_substeps::Int

    theta_regulated::Int

    v_d_max::Float64 #feet per second squared
    h_dd_max::Float64 #feet per second squared
    psi_d_max::Float64 #degrees per second

    observer::Observer

    function SimpleADM(;timestep = 1., number_of_substeps = 10,
                       v_d_max::Float64=typemax(Float64),
                       h_dd_max::Float64=typemax(Float64),
                       psi_d_max::Float64=typemax(Float64))

        obj = new()

        obj.command = nothing
        obj.timestep = timestep
        obj.number_of_substeps = number_of_substeps
        obj.theta_regulated = 45    # degree

        obj.v_d_max = v_d_max
        obj.h_dd_max = h_dd_max
        obj.psi_d_max = psi_d_max

        obj.observer = Observer()

        return obj
    end
end

addObserver(adm::SimpleADM, f::Function) = add_observer(adm.observer, f)
addObserver(adm::SimpleADM, tag::AbstractString, f::Function) = add_observer(adm.observer, tag, f)

function initializeDynamicModel(adm::SimpleADM, state::SimpleADMInitialState)

    adm.state = SimpleADMState(state.t, state.x, state.y, state.h, state.v, state.psi)
    #print([state.t, state.x, state.y, state.h, state.v, state.psi]')

    @notify_observer(adm.observer, "adm_init", [adm.state.t, adm.state.x, adm.state.y, adm.state.h])

    vh = state.h_d
    vx = sqrt(adm.state.v^2 - vh^2) * cosd(adm.state.psi)
    vy = sqrt(adm.state.v^2 - vh^2) * sind(adm.state.psi)

    adm.command = nothing

    return SimpleADMOutputState(adm.state.t, adm.state.x, adm.state.y, adm.state.h, vx, vy, vh)
end

initialize(adm::SimpleADM, state) = initialize(adm,convert(SimpleADMInitialState, state))

initialize(adm::SimpleADM, state::SimpleADMInitialState) = initializeDynamicModel(adm, state)

function simulateDynamicModel(adm::SimpleADM, command::SimpleADMCommand)

    t, x, y, h, v, psi = adm.state.t, adm.state.x, adm.state.y, adm.state.h, adm.state.v, adm.state.psi

    if adm.command == nothing    # first second
        t_prev, v_d_prev, h_d_prev, psi_d_prev = t - 1, command.v_d, command.h_d, command.psi_d
    else
        t_prev, v_d_prev, h_d_prev, psi_d_prev = adm.command.t, adm.command.v_d, adm.command.h_d, adm.command.psi_d
    end
    t_curr, v_d_curr, h_d_curr, psi_d_curr = command.t, command.v_d, command.h_d, command.psi_d

    #limit acclerations and rates
    v_d_curr = min(max(v_d_curr,-adm.v_d_max),adm.v_d_max)
    psi_d_curr = min(max(psi_d_curr,-adm.psi_d_max),adm.psi_d_max)

    h_dd = (h_d_curr - h_d_prev) / adm.timestep
    h_dd = min(max(h_dd,-adm.h_dd_max),adm.h_dd_max)
    h_d_curr = h_d_prev + h_dd * adm.timestep
    command.h_d = h_d_curr #propagated to the next time step as prev

    #@test t == t_curr
    #@test t_prev + 1 == t_curr

    dt = adm.timestep / adm.number_of_substeps

    t_sim = t

    x_n = x
    y_n = y
    h_n = h
    v_n = v
    psi_n = psi

    v_d = v_d_prev
    h_d = h_d_prev
    psi_d = psi_d_prev

    #print([t_curr, v_d_curr, h_d_curr, psi_d_curr]')

    for i = 1:adm.number_of_substeps
        # update simulation time
        t_sim += dt

        # update control

        # linear control
        dv_d = (v_d_curr - v_d_prev) * (t_sim - t_curr)
        dh_d = (h_d_curr - h_d_prev) * (t_sim - t_curr)
        dpsi_d = (psi_d_curr - psi_d_prev) * (t_sim - t_curr)

        # sigmoid control
        #CMD_DIST = Normal(0.3, 0.07)
        #dv_d = (v_d_curr - v_d_prev) * cdf(CMD_DIST, (t_sim - t_curr))
        #dh_d = (h_d_curr - h_d_prev) * cdf(CMD_DIST, (t_sim - t_curr))
        #dpsi_d = (psi_d_curr - psi_d_prev) * cdf(CMD_DIST, (t_sim - t_curr))

        v_d = v_d_prev + dv_d
        h_d = h_d_prev + dh_d
        psi_d = psi_d_prev + dpsi_d

        #println(v_d, " ", h_d, " ", psi_d)

        if abs(h_d) > abs(v_n) * sind(adm.theta_regulated)
            h_d = sign(h_d) * abs(v_n) * sind(adm.theta_regulated)
            #println("regulated")
        end

        # update state

        x_n += sqrt(v_n^2 - h_d^2) * cosd(psi) * dt
        y_n += sqrt(v_n^2 - h_d^2) * sind(psi) * dt

        v_n += v_d * dt     # ft/s
        h_n += h_d * dt     # ft
        psi_n += psi_d * dt # deg

        #print([t_sim, x_n, y_n, h_n, v_n, psi_n]')

        @notify_observer(adm.observer, "adm_simulate", [t_sim, x_n, y_n, h_n])
    end

    #@test_approx_eq_eps t_sim (t_curr + adm.timestep) 0.001
    #@test_approx_eq_eps v_d v_d_curr 0.001
    #@test_approx_eq_eps h_d h_d_curr 0.001
    #@test_approx_eq_eps psi_d psi_d_curr 0.001

    adm.state = SimpleADMState(t_curr + adm.timestep, x_n, y_n, h_n, v_n, psi_n)
    adm.command = deepcopy(command)
end

update(adm::SimpleADM, command) = update(adm,convert(SimpleADMCommand, command))

function update(adm::SimpleADM, command::SimpleADMCommand)

    t =  copy(adm.state.t)
    x =  copy(adm.state.x)
    y =  copy(adm.state.y)
    h =  copy(adm.state.h)

    simulateDynamicModel(adm, command)

    t_n =  adm.state.t
    x_n =  adm.state.x
    y_n =  adm.state.y
    h_n =  adm.state.h

    vx_n = (x_n - x) / (t_n - t)
    vy_n = (y_n - y) / (t_n - t)
    vh_n = (h_n - h) / (t_n - t)

    # instantaneous velocities
    #vh_n = adm.command.h_d
    #vx_n = sqrt(adm.state.v^2 - vh_n^2) * cosd(adm.state.psi)
    #vy_n = sqrt(adm.state.v^2 - vh_n^2) * sind(adm.state.psi)

    return SimpleADMOutputState(t_n, x_n, y_n, h_n, vx_n, vy_n, vh_n)
end

end


# This file contains functions that are related to the AMR capabilities of the DG solver

# Refine elements in the DG solver based on a list of cell_ids that should be refined
function refine!(dg::Dg3D{Eqn, MeshType, NVARS, POLYDEG}, mesh::TreeMesh,
                 cells_to_refine::AbstractArray{Int}) where {Eqn, MeshType, NVARS, POLYDEG}
  # Return early if there is nothing to do
  if isempty(cells_to_refine)
    return
  end

  # Determine for each existing element whether it needs to be refined
  needs_refinement = falses(nelements(dg.elements))
  tree = mesh.tree
  # The "Ref(...)" is such that we can vectorize the search but not the array that is searched
  elements_to_refine = searchsortedfirst.(Ref(dg.elements.cell_ids[1:nelements(dg.elements)]),
                                          cells_to_refine)
  needs_refinement[elements_to_refine] .= true

  # Retain current solution data
  old_n_elements = nelements(dg.elements)
  old_u = dg.elements.u

  # Get new list of leaf cells
  leaf_cell_ids = leaf_cells(tree)

  # Initialize new elements container
  elements = init_elements(leaf_cell_ids, mesh, Val(NVARS), Val(POLYDEG))
  n_elements = nelements(elements)

  # Loop over all elements in old container and either copy them or refine them
  u_tmp1 = Array{Float64, 4}(undef, nvariables(dg), nnodes(dg), nnodes(dg), nnodes(dg))
  u_tmp2 = Array{Float64, 4}(undef, nvariables(dg), nnodes(dg), nnodes(dg), nnodes(dg))
  element_id = 1
  for old_element_id in 1:old_n_elements
    if needs_refinement[old_element_id]
      # Refine element and store solution directly in new data structure
      refine_element!(elements.u, element_id, old_u, old_element_id,
                      dg.mortar_forward_upper, dg.mortar_forward_lower, dg, u_tmp1, u_tmp2)
      element_id += 2^ndims(dg)
    else
      # Copy old element data to new element container
      @views elements.u[:, :, :, :, element_id] .= old_u[:, :, :, :, old_element_id]
      element_id += 1
    end
  end

  # Initialize new interfaces container
  interfaces = init_interfaces(leaf_cell_ids, mesh, Val(NVARS), Val(POLYDEG), elements)
  n_interfaces = ninterfaces(interfaces)

  # Initialize boundaries
  boundaries, n_boundaries_per_direction = init_boundaries(leaf_cell_ids, mesh, Val(NVARS), Val(POLYDEG), elements)
  n_boundaries = nboundaries(boundaries)

  # Initialize new mortar containers
  l2mortars = init_mortars(leaf_cell_ids, mesh, Val(NVARS), Val(POLYDEG), elements, dg.mortar_type)
  n_l2mortars = nmortars(l2mortars)

  # Sanity check
  if isperiodic(mesh.tree) && n_l2mortars == 0
    @assert n_interfaces == 3*n_elements ("For 3D and periodic domains and conforming elements, "
                                        * "n_surf must be the same as 3*n_elem")
  end

  # Update DG instance with new data
  dg.elements = elements
  dg.n_elements = n_elements
  dg.n_elements_global = n_elements
  dg.interfaces = interfaces
  dg.n_interfaces = n_interfaces
  dg.boundaries = boundaries
  dg.n_boundaries = n_boundaries
  dg.n_boundaries_per_direction = n_boundaries_per_direction
  dg.l2mortars = l2mortars
  dg.n_l2mortars = n_l2mortars
end

# Refine solution data u for an element, using L2 projection (interpolation)
function refine_element!(u, element_id, old_u, old_element_id,
                         forward_upper, forward_lower, dg::Dg3D, u_tmp1, u_tmp2)
  # Store new element ids
  bottom_lower_left_id  = element_id
  bottom_lower_right_id = element_id + 1
  bottom_upper_left_id  = element_id + 2
  bottom_upper_right_id = element_id + 3
  top_lower_left_id     = element_id + 4
  top_lower_right_id    = element_id + 5
  top_upper_left_id     = element_id + 6
  top_upper_right_id    = element_id + 7

  # Interpolate to bottom lower left element
  multiply_dimensionwise!(
    view(u,     :, :, :, :, bottom_lower_left_id), forward_lower, forward_lower, forward_lower,
    view(old_u, :, :, :, :, old_element_id), u_tmp1, u_tmp2)

  # Interpolate to bottom lower right element
  multiply_dimensionwise!(
    view(u,     :, :, :, :, bottom_lower_right_id), forward_upper, forward_lower, forward_lower,
    view(old_u, :, :, :, :, old_element_id), u_tmp1, u_tmp2)

  # Interpolate to bottom upper left element
  multiply_dimensionwise!(
    view(u,     :, :, :, :, bottom_upper_left_id), forward_lower, forward_upper, forward_lower,
    view(old_u, :, :, :, :, old_element_id), u_tmp1, u_tmp2)

  # Interpolate to bottom upper right element
  multiply_dimensionwise!(
    view(u,     :, :, :, :, bottom_upper_right_id), forward_upper, forward_upper, forward_lower,
    view(old_u, :, :, :, :, old_element_id), u_tmp1, u_tmp2)

  # Interpolate to top lower left element
  multiply_dimensionwise!(
    view(u,     :, :, :, :, top_lower_left_id), forward_lower, forward_lower, forward_upper,
    view(old_u, :, :, :, :, old_element_id), u_tmp1, u_tmp2)

  # Interpolate to top lower right element
  multiply_dimensionwise!(
    view(u,     :, :, :, :, top_lower_right_id), forward_upper, forward_lower, forward_upper,
    view(old_u, :, :, :, :, old_element_id), u_tmp1, u_tmp2)

  # Interpolate to top upper left element
  multiply_dimensionwise!(
    view(u,     :, :, :, :, top_upper_left_id), forward_lower, forward_upper, forward_upper,
    view(old_u, :, :, :, :, old_element_id), u_tmp1, u_tmp2)

  # Interpolate to top upper right element
  multiply_dimensionwise!(
    view(u,     :, :, :, :, top_upper_right_id), forward_upper, forward_upper, forward_upper,
    view(old_u, :, :, :, :, old_element_id), u_tmp1, u_tmp2)
end

# Coarsen elements in the DG solver based on a list of cell_ids that should be removed
function coarsen!(dg::Dg3D{Eqn, MeshType, NVARS, POLYDEG}, mesh::TreeMesh,
                  child_cells_to_coarsen::AbstractArray{Int}) where {Eqn, MeshType, NVARS, POLYDEG}
  # Return early if there is nothing to do
  if isempty(child_cells_to_coarsen)
    return
  end

  # Determine for each old element whether it needs to be removed
  to_be_removed = falses(nelements(dg.elements))
  # The "Ref(...)" is such that we can vectorize the search but not the array that is searched
  elements_to_remove = searchsortedfirst.(Ref(dg.elements.cell_ids[1:nelements(dg.elements)]),
                                          child_cells_to_coarsen)
  to_be_removed[elements_to_remove] .= true

  # Retain current solution data
  old_n_elements = nelements(dg.elements)
  old_u = dg.elements.u

  # Get new list of leaf cells
  leaf_cell_ids = leaf_cells(mesh.tree)

  # Initialize new elements container
  elements = init_elements(leaf_cell_ids, mesh, Val(NVARS), Val(POLYDEG))
  n_elements = nelements(elements)

  # Loop over all elements in old container and either copy them or coarsen them
  u_tmp1 = Array{Float64, 4}(undef, nvariables(dg), nnodes(dg), nnodes(dg), nnodes(dg))
  u_tmp2 = Array{Float64, 4}(undef, nvariables(dg), nnodes(dg), nnodes(dg), nnodes(dg))
  skip = 0
  element_id = 1
  for old_element_id in 1:old_n_elements
    # If skip is non-zero, we just coarsened 2^ndims elements and need to omit the following elements
    if skip > 0
      skip -= 1
      continue
    end

    if to_be_removed[old_element_id]
      # If an element is to be removed, sanity check if the following elements
      # are also marked - otherwise there would be an error in the way the
      # cells/elements are sorted
      @assert all(to_be_removed[old_element_id:(old_element_id+2^ndims(dg)-1)]) "bad cell/element order"

      # Coarsen elements and store solution directly in new data structure
      coarsen_elements!(elements.u, element_id, old_u, old_element_id,
                        dg.l2mortar_reverse_upper, dg.l2mortar_reverse_lower, dg, u_tmp1, u_tmp2)
      element_id += 1
      skip = 2^ndims(dg) - 1
    else
      # Copy old element data to new element container
      @views elements.u[:, :, :, :, element_id] .= old_u[:, :, :, :, old_element_id]
      element_id += 1
    end
  end

  # Initialize new interfaces container
  interfaces = init_interfaces(leaf_cell_ids, mesh, Val(NVARS), Val(POLYDEG), elements)
  n_interfaces = ninterfaces(interfaces)

  # Initialize boundaries
  boundaries, n_boundaries_per_direction = init_boundaries(leaf_cell_ids, mesh, Val(NVARS), Val(POLYDEG), elements)
  n_boundaries = nboundaries(boundaries)

  # Initialize new mortar containers
  l2mortars = init_mortars(leaf_cell_ids, mesh, Val(NVARS), Val(POLYDEG), elements, dg.mortar_type)
  n_l2mortars = nmortars(l2mortars)

  # Sanity check
  if isperiodic(mesh.tree) && n_l2mortars == 0
    @assert n_interfaces == 3*n_elements ("For 3D and periodic domains and conforming elements, "
                                        * "n_surf must be the same as 3*n_elem")
  end

  # Update DG instance with new data
  dg.elements = elements
  dg.n_elements = n_elements
  dg.n_elements_global = n_elements
  dg.interfaces = interfaces
  dg.n_interfaces = n_interfaces
  dg.boundaries = boundaries
  dg.n_boundaries = n_boundaries
  dg.n_boundaries_per_direction = n_boundaries_per_direction
  dg.l2mortars = l2mortars
  dg.n_l2mortars = n_l2mortars
end

# Coarsen solution data u for four elements, using L2 projection
function coarsen_elements!(u, element_id, old_u, old_element_id,
                           reverse_upper, reverse_lower, dg::Dg3D, u_tmp1, u_tmp2)
  # Store old element ids
  bottom_lower_left_id  = old_element_id
  bottom_lower_right_id = old_element_id + 1
  bottom_upper_left_id  = old_element_id + 2
  bottom_upper_right_id = old_element_id + 3
  top_lower_left_id     = old_element_id + 4
  top_lower_right_id    = old_element_id + 5
  top_upper_left_id     = old_element_id + 6
  top_upper_right_id    = old_element_id + 7

  # Project from bottom lower left element
  multiply_dimensionwise!(
    view(u,     :, :, :, :, element_id), reverse_lower, reverse_lower, reverse_lower,
    view(old_u, :, :, :, :, bottom_lower_left_id), u_tmp1, u_tmp2)

  # Project from bottom lower right element_variables
  add_multiply_dimensionwise!(
    view(u,     :, :, :, :, element_id), reverse_upper, reverse_lower, reverse_lower,
    view(old_u, :, :, :, :, bottom_lower_right_id), u_tmp1, u_tmp2)

  # Project from bottom upper left element
  add_multiply_dimensionwise!(
    view(u,     :, :, :, :, element_id), reverse_lower, reverse_upper, reverse_lower,
    view(old_u, :, :, :, :, bottom_upper_left_id), u_tmp1, u_tmp2)

  # Project from bottom upper right element
  add_multiply_dimensionwise!(
    view(u,     :, :, :, :, element_id), reverse_upper, reverse_upper, reverse_lower,
    view(old_u, :, :, :, :, bottom_upper_right_id), u_tmp1, u_tmp2)

  # Project from top lower left element
  add_multiply_dimensionwise!(
    view(u,     :, :, :, :, element_id), reverse_lower, reverse_lower, reverse_upper,
    view(old_u, :, :, :, :, top_lower_left_id), u_tmp1, u_tmp2)

  # Project from top lower right element
  add_multiply_dimensionwise!(
    view(u,     :, :, :, :, element_id), reverse_upper, reverse_lower, reverse_upper,
    view(old_u, :, :, :, :, top_lower_right_id), u_tmp1, u_tmp2)

  # Project from top upper left element
  add_multiply_dimensionwise!(
    view(u,     :, :, :, :, element_id), reverse_lower, reverse_upper, reverse_upper,
    view(old_u, :, :, :, :, top_upper_left_id), u_tmp1, u_tmp2)

  # Project from top upper right element
  add_multiply_dimensionwise!(
    view(u,     :, :, :, :, element_id), reverse_upper, reverse_upper, reverse_upper,
    view(old_u, :, :, :, :, top_upper_right_id), u_tmp1, u_tmp2)
end

# Calculate an AMR indicator value for each element/leaf cell
#
# The indicator value λ ∈ [-1,1] is ≈ -1 for cells that should be coarsened, ≈
# 0 for cells that should remain as-is, and ≈ 1 for cells that should be
# refined.
#
# Note: The implementation here implicitly assumes that we have an element for
# each leaf cell and that they are in the same order.
#
# FIXME: This is currently implemented for each test case - we need something
# appropriate that is both equation and test case independent
function calc_amr_indicator(dg::Dg3D, mesh::TreeMesh, time::Float64)
  lambda = zeros(dg.n_elements)

  if dg.amr_indicator === :gauss
    base_level = 4
    max_level = 6
    threshold_high = 0.6
    threshold_low = 0.1

    # Iterate over all elements
    for element_id in 1:dg.n_elements
      # Determine target level from peak value
      peak = maximum(dg.elements.u[:, :, :, :, element_id])
      if peak > threshold_high
        target_level = max_level
      elseif peak > threshold_low
        target_level = max_level - 1
      else
        target_level = base_level
      end

      # Compare target level with actual level to set indicator
      cell_id = dg.elements.cell_ids[element_id]
      actual_level = mesh.tree.levels[cell_id]
      if actual_level < target_level
        lambda[element_id] = 1.0
      elseif actual_level > target_level
        lambda[element_id] = -1.0
      else
        lambda[element_id] = 0.0
      end
    end
  elseif dg.amr_indicator === :blob
    base_level = 1
    max_level = 6 # originally 6; set to 7 (or even 8) to get increased resolution
    blending_factor_threshold1 = 0.3 # Löhner original choice
    blending_factor_threshold2 = 0.1 # Löhner original choice

    # (Re-)initialize element variable storage for blending factor
    if (!haskey(dg.element_variables, :amr_indicator_values) ||
        length(dg.element_variables[:amr_indicator_values]) != dg.n_elements)
      dg.element_variables[:amr_indicator_values] = Vector{Float64}(undef, dg.n_elements)
    end

    alpha = dg.element_variables[:amr_indicator_values]
    calc_loehner_indicator!(alpha, dg.elements.u, density, dg.thread_cache, dg)

    # OPTIONAL: use max level were shock capturing is maxed
    # (Re-)initialize element variable storage for blending factor
    #if (!haskey(dg.element_variables, :blending_factor) ||
    #    length(dg.element_variables[:blending_factor]) != dg.n_elements)
    #  dg.element_variables[:blending_factor] = Vector{Float64}(undef, dg.n_elements)
    #end
    #if (!haskey(dg.element_variables, :blending_factor_tmp) ||
    #    length(dg.element_variables[:blending_factor_tmp]) != dg.n_elements)
    #  dg.element_variables[:blending_factor_tmp] = Vector{Float64}(undef, dg.n_elements)
    #end

    #alpha1     = dg.element_variables[:blending_factor]
    #alpha1_tmp = dg.element_variables[:blending_factor_tmp]
    #calc_blending_factors!(alpha1, alpha1_tmp, dg.elements.u, dg.shock_alpha_max, dg.shock_alpha_min, true,
    #                       dg.shock_indicator_variable, dg.thread_cache, dg)

    # Iterate over all elements
    for element_id in 1:dg.n_elements
      cell_id = dg.elements.cell_ids[element_id]
      actual_level = mesh.tree.levels[cell_id]
      target_level = actual_level
      # adapt for the amr indicator
      if alpha[element_id] >= blending_factor_threshold1
        target_level = max_level
      elseif alpha[element_id] <= blending_factor_threshold2
        target_level = base_level
      end
      # make sure that a highly troubled shock cell is not coarsened
      #if isapprox.(dg.shock_alpha_max, alpha1[element_id], atol=1e-12)
      #  target_level = max_level
      #end
      # Compare target level with actual level to set indicator
      if actual_level < target_level
        lambda[element_id] = 1.0
      elseif actual_level > target_level
        lambda[element_id] = -1.0
      else
        lambda[element_id] = 0.0
      end

    end
  elseif dg.amr_indicator === :density_pulse
    # Works with initial_conditions_density_pulse for compressible Euler equations
    base_level = 4
    max_level = 6
    threshold_high = 1.3
    threshold_low = 1.05

    # Iterate over all elements
    for element_id in 1:dg.n_elements
      # Determine target level from peak value
      peak = maximum(dg.elements.u[1, :, :, :, element_id])
      if peak > threshold_high
        target_level = max_level
      elseif peak > threshold_low
        target_level = max_level - 1
      else
        target_level = base_level
      end

      # Compare target level with actual level to set indicator
      cell_id = dg.elements.cell_ids[element_id]
      actual_level = mesh.tree.levels[cell_id]
      if actual_level < target_level
        lambda[element_id] = 1.0
      elseif actual_level > target_level
        lambda[element_id] = -1.0
      else
        lambda[element_id] = 0.0
      end
    end
  elseif dg.amr_indicator === :sedov_self_gravity
    base_level = 2
    max_level = parameter("max_refinement_level", 7)::Int # TODO: This is just so we can test this, but we should have a global switch...
    blending_factor_threshold = 0.0003

    # (Re-)initialize element variable storage for blending factor
    if (!haskey(dg.element_variables, :amr_indicator_values) ||
        length(dg.element_variables[:amr_indicator_values]) != dg.n_elements)
      dg.element_variables[:amr_indicator_values] = Vector{Float64}(undef, dg.n_elements)
    end
    if (!haskey(dg.element_variables, :amr_indicator_values_tmp) ||
        length(dg.element_variables[:amr_indicator_values_tmp]) != dg.n_elements)
      dg.element_variables[:amr_indicator_values_tmp] = Vector{Float64}(undef, dg.n_elements)
    end

    alpha     = dg.element_variables[:amr_indicator_values]
    alpha_tmp = dg.element_variables[:amr_indicator_values_tmp]
    calc_blending_factors!(alpha, alpha_tmp, dg.elements.u, dg.amr_alpha_max, dg.amr_alpha_min, dg.amr_alpha_smooth,
                           density_pressure, dg.thread_cache, dg)

    # Iterate over all elements
    for element_id in 1:dg.n_elements
      if alpha[element_id] > blending_factor_threshold
        target_level = max_level
      else
        target_level = base_level
      end

      # Compare target level with actual level to set indicator
      cell_id = dg.elements.cell_ids[element_id]
      actual_level = mesh.tree.levels[cell_id]
      if actual_level < target_level
        lambda[element_id] = 1.0
      elseif actual_level > target_level
        lambda[element_id] = -1.0
      else
        lambda[element_id] = 0.0
      end
    end
  else
    error("unknown AMR indicator '$(dg.amr_indicator)'")
  end

  return lambda
end


# using Revise
using ClusterManagers
using Distributed

worker_ids = if "SLURM_JOBID" in keys(ENV)
    wids = ClusterManagers.addprocs_slurm(parse(Int, ENV["SLURM_NTASKS"]))
    wids
else
    addprocs(4)
end

Distributed.@everywhere begin
    using Pkg
    Pkg.activate(".")
    Pkg.instantiate()
end

Distributed.@everywhere begin
    using POMDPs
    using BasicPOMCP
    using POMCPOW
    using VOOTreeSearch
    using ParticleFilters
    using CPUTime
    using DataFrames
    using ARDESPOT
    using Distributions
    using VDPTag2
    using ProgressMeter
    using Dates
    using CSV
    using POMDPSimulators
    using Random
    using Printf
    using LinearAlgebra
    using VDPTag2Experiments
    using LinearAlgebra: diagm, PosDefException, det, eigvals, eigvecs
    using Random: MersenneTwister, seed!
    using Distances: euclidean

    using BOMCP
    using Statistics
    using StatsBase
    include("./deps/BOMCP.jl/scripts/lander_po.jl")

    function BOMCP.vectorize!(v, dims, x::MvNormal)
        v = copy(mean(x))
        return v
    end
end

@info "Started $(Distributed.nworkers()) workers..."
@info "Precompiling simulation code..."

max_depth = 250
solver_list = Dict("POMCPOW"=>[10.0, 3.0, 4.0, 2.0, 10.0],
                    "VOMCPOW"=>[30.0, 4.0, 4.0, 1.5, 5.0, 0.9],
                    "BOMCP"=>[10.0, 3.0, 4.0, 2.0, 10.0, 15.0, 0.5],)
on_cluster = true
save_data = true
parallelize = true

if on_cluster
    max_time_list = 10.0 .^ (-2:0.25:0)
    max_query = 100_000
    n_iters = 1000
else
    max_time_list = 10.0 .^ (-2:0.5:0)
    max_query = 100_000
    n_iters = 10
end

# Lander
pomdp = LunarLander(dt=0.4)

println(stderr, "Max times: ", max_time_list)
println(stderr, "Max depth: ", max_depth)
println(stderr, "Max query: ", max_query)
println(stderr, "Number of iterations: ", n_iters)
println(stderr, "Solvers: ", solver_list)
println(stderr, "POMDP: ", pomdp)

Distributed.@everywhere begin
    struct SimWrap
        x::Vector{Float64}
        solver_name::String
        i::Int64
        max_time::Float64
        max_depth::Int64
        max_query::Int64
        pomdp::POMDP
    end

    function gen_sim(sim_wrap::SimWrap)
        x = sim_wrap.x
        solver_name = sim_wrap.solver_name
        i = sim_wrap.i
        max_time = sim_wrap.max_time
        max_depth = sim_wrap.max_depth
        max_query = sim_wrap.max_query
        pomdp = sim_wrap.pomdp

        rng = MersenneTwister(7)
        belief_updater = EKFUpdater(pomdp, pomdp.Q.^2, pomdp.R.^2)
        rollout_policy = LanderPolicy(pomdp)

        c = x[1]
        k_act = x[2]
        inv_alpha_act = x[3]
        k_obs = x[4]
        inv_alpha_obs = x[5]

        if solver_name == "POMCPOW"
            solver = POMCPOWSolver(tree_queries=max_query,
                                   criterion=MaxUCB(c),
                                   max_depth=max_depth,
                                   max_time=max_time,
                                   k_action=k_act,
                                   alpha_action=1/inv_alpha_act,
                                   k_observation=k_obs,
                                   alpha_observation=1/inv_alpha_obs,
                                   estimate_value=POMCPOW.RolloutEstimator(rollout_policy),
                                   next_action=RolloutFirst(rng, rollout_policy),
                                   default_action=BasicPOMCP.ReportWhenUsed([0.0, 0.0, 0.0]),
                                   rng=rng
                                  )

            planner = solve(solver, pomdp)

            seed!(planner, i+40000)
            sim = Sim(deepcopy(pomdp),
                      planner,
                      belief_updater,
                      rng=MersenneTwister(i+50_000),
                      max_steps=200,
                      metadata=Dict(:i=>i, :solver=>solver_name)
                     )

        elseif solver_name == "VOMCPOW"
            p_voo = x[6]
            sigs = [0.2, 0.5, 0.05]
            sig = diagm(sigs)
            epsilon = 0.01
            voo_opt = VOOActionGenerator(p_voo, Distributions.Uniform(), MvNormal(sig), rng, epsilon, true, 20)

            solver = POMCPOWSolver(tree_queries=max_query,
                                   criterion=MaxUCB(c),
                                   max_depth=max_depth,
                                   max_time=max_time,
                                   k_action=k_act,
                                   alpha_action=1/inv_alpha_act,
                                   k_observation=k_obs,
                                   alpha_observation=1/inv_alpha_obs,
                                   estimate_value=POMCPOW.RolloutEstimator(rollout_policy),
                                   next_action = VOORolloutFirst(voo_opt, rollout_policy),
                                   default_action=BasicPOMCP.ReportWhenUsed([0.0, 0.0, 0.0]),
                                   rng=rng
                                  )

            planner = solve(solver, pomdp)

            seed!(planner, i+40000)
            sim = Sim(deepcopy(pomdp),
                      planner,
                      belief_updater,
                      rng=MersenneTwister(i+50_000),
                      max_steps=200,
                      metadata=Dict(:i=>i, :solver=>solver_name)
                     )

        elseif solver_name == "BOMCP"
            max_query = Int64(max_query/10)
            log_kernel = x[6]
            lambda = x[7]

            belief_updater = EKFUpdater(pomdp, pomdp.Q.^2, pomdp.R.^2)

            action_selector = BOActionSelector(3, # action dims
                                        6, #belief dims
                                        false, #discrete actions
                                        kernel_params=[log(log_kernel), 0.0],
                                        k_neighbors = 5,
                                        belief_λ = lambda,
                                        lower_bounds = [-5.0, 0.0, -1.0],
                                        upper_bounds = [5.0, 15.0, 1.0],
                                        buffer_size=100,
                                        initial_action=rollout_policy
                                        )

            solver = BOMCPSolver(action_selector, belief_updater,
                        depth=max_depth, n_iterations=max_query,
                        max_time = max_time,
                        exploration_constant=c,
                        k_belief = k_obs,
                        alpha_belief = 1/inv_alpha_obs,
                        k_action = k_act,
                        alpha_action = 1/inv_alpha_act,
                        estimate_value=BOMCP.RolloutEstimator(rollout_policy),
                        default_action=BOMCP.ReportWhenUsed([0.0, 0.0, 0.0]),
                        rng=rng
                        )

            planner = solve(solver, pomdp)

            seed!(planner, i+40000)
            sim = Sim(pomdp,
                      planner,
                      belief_updater,
                      rng=MersenneTwister(i+50_000),
                      max_steps=200,
                      metadata=Dict(:i=>i, :solver=>solver_name)
                     )
        end

        result = POMDPs.simulate(sim)
        output = (reward=discounted_reward(result),)
        return merge(sim.metadata, output)
    end
end

datestring = Dates.format(now(), "e_d_u_Y_HH_MM")
results_csv = DataFrame(max_time=Float64[], solver=String[], mean_reward=Float64[], se_reward=Float64[])

try
    for iter_time in max_time_list
        sims = SimWrap[]
        for solver_name in keys(solver_list)
            for j in 1:n_iters
                sim_wrap = SimWrap(solver_list[solver_name], solver_name, j, iter_time, max_depth, max_query, deepcopy(pomdp))
                push!(sims, sim_wrap)
            end
        end
        if parallelize
            results = @showprogress pmap(gen_sim, sims)
        else
            results = @showprogress map(gen_sim, sims)
        end
        println(stderr, "Max time: $iter_time :")
        results = POMDPSimulators.create_dataframe(results)
        combined = by(results, :solver) do df
            DataFrame(mean_reward=mean(df[:reward]), se_reward=(std(df[:reward])/sqrt(n_iters)))
        end
        for row in eachrow(combined)
            result_row = [iter_time, row.solver, row.mean_reward, row.se_reward]
            println(stderr, @sprintf("- %s reward: %6.3f ± %6.3f", row.solver, row.mean_reward, row.se_reward))
            push!(results_csv, result_row)
        end
        if save_data
            fname_csv = joinpath(@__DIR__, "../data", "Lander_compare_results"*"_"*datestring*".csv")
            CSV.write(fname_csv, results_csv)
            println(stderr,"Saving parameters to ", fname_csv)
        end
    end
finally
    Distributed.rmprocs(worker_ids)
end


using ResNet9
using Test
using Flux
using CUDA

if CUDA.has_cuda_gpu()
    CUDA.allowscalar(false)
    device = gpu
else
    device = cpu
end

@testset "ResNet9.jl" begin
    batchsize = 2
    inchs = 3
    x1 = Float32.(rand(16, 16, inchs, batchsize)) |> device
    x2 = Float32.(rand(16, 16, 2 * inchs, batchsize)) |> device

    @testset "Two classes" begin
        model = resnet9(inchannels = inchs, nclasses = 2) |> device
        ŷ = model(x1)
        @test size(ŷ) == (1, batchsize)
    end

    @testset "Five classes" begin
        model = resnet9(inchannels = inchs, nclasses = 5) |> device
        ŷ = model(x1)
        @test size(ŷ) == (5, batchsize)
    end

    @testset "Transfer learning" begin
        model = resnet9(inchannels = inchs, nclasses = 5)
        # Change the number of classes and add dropout
        model2 = resnet9(model, nclasses = 10, dropout = 0.5) |> device
        ŷ = model2(x1)
        @test size(ŷ) == (10, batchsize)
        # Change the number of input channels
        model2 = resnet9(model, inchannels = 2 * inchs) |> device
        ŷ = model2(x2)
        @test size(ŷ) == (5, batchsize)
        # Change the number of input channels
        model2 = resnet9(model, dropout = 0.2) |> device
        ŷ = model2(x1)
        @test size(ŷ) == (5, batchsize)
    end
end


using ArgParse, Random, SparseArrays, HDF5, H5Sparse, SAG

const METADATA_BYTES = 8
const ELEMENT_TYPE = Float64
const CANARY_VALUE = UInt16(2^16 - 1)

include("gradients.jl")

function update_argsettings!(s::ArgParseSettings)
    @add_arg_table s begin
        "--lambda"
            help = "L2 regularizer coefficient"
            arg_type = Float64
            default = 0.0
        "--labeldataset"
            help = "Label dataset name"
            default = "b"
            arg_type = String
        "--iteratedataset"
            help = "Initial iterate dataset name (chosen to be all-zeros if not provided)"
            arg_type = String
        "--nsubpartitions"
            help = "Number of sub-partitions to split the data stored at each worker into"
            arg_type = Int
            default = 1
            range_tester = (x) -> x >= 1
        "--stepsize"
            help = "Gradient descent step size"
            arg_type = Float64
            default = 1.0
            range_tester = (x) -> x > 0
        "--variancereduced"
            help = "Compute a variance-reduced gradient in each iteration"
            action = :store_true
        "--vralgo"
            help = "Specifies which variance-reduced optimizer to use (tree or table)"
            arg_type = String
            default = "tree"
            range_tester = (x) -> x in ["tree", "table"]
        "--nostale"
            help = "If set, do not store stale gradients (to conform with SAG)"
            action = :store_true
        "--nreplicas"
            help = "Number of replicas of each data partition"
            default = 1
            arg_type = Int
            range_tester = (x) -> x >= 1
        "--nwait"
            help = "Number of replicas to wait for in each iteration (defaults to all replicas)"
            arg_type = Int
            range_tester = (x) -> x >= 1
        "--nwaitschedule"
            help = "Factor by which nwait is reduced per iteration"
            arg_type = Float64
            default = 1.0
            range_tester = (x) -> 0 < x <= 1
        "--kickstart"
            help = "Wait for all partitions in the first iteration"
            action = :store_true
    end
end

function update_parsed_args!(s::ArgParseSettings, parsed_args)
    parsed_args[:algorithm] = "logreg.jl"
    nworkers::Int = parsed_args[:nworkers]
    nreplicas::Int = parsed_args[:nreplicas]
    mod(nworkers, nreplicas) == 0 || throw(ArgumentError("nworkers is $nworkers, but must be divisible by nreplicas"))
    npartitions = div(nworkers, nreplicas)
    parsed_args[:nwait] = isnothing(parsed_args[:nwait]) ? npartitions : parsed_args[:nwait]
    0 < parsed_args[:nwait] <= npartitions || throw(ArgumentError("nwait is $(parsed_args[:nwait]), but nworkers/nreplicas is $npartitions"))

    # record the number of rows and columns of the dataset
    nrows, ncolumns = problem_size(parsed_args[:inputfile], parsed_args[:inputdataset])
    parsed_args[:nrows] = nrows
    parsed_args[:ncolumns] = ncolumns

    parsed_args
end

"""

Called inside `asyncmap!` to determine if enough workers have responded. Returns `true` if at
least `nwait` workers have responded and `false` otherwise.
"""
function fwait(epoch, repochs; nworkers, nwait, kickstart, nwaitschedule, kwargs...)
    length(repochs) == nworkers || throw(DomainError(nworkers, "repochs must have length nworkers"))
    0 < nwait <= nworkers || throw(ArgumentError("nwait is $nwait, but must be in [1, npartitions]"))
    nwait = min(max(1, ceil(Int, nwait*nwaitschedule^epoch)), nworkers)
    nrec = 0
    for repoch in repochs
        nrec += repoch == epoch
    end
    if epoch == 1 && kickstart
        return nrec == nwait
    else
        return nrec >= nwait
    end
end

"""

Return the size of the data matrix.
"""
function problem_size(filename::String, dataset::String)
    HDF5.ishdf5(filename) || throw(ArgumentError("$filename isn't an HDF5 file"))
    h5open(filename, "r") do fid
        dataset in keys(fid) || throw(ArgumentError("$dataset is not in $fid"))
        if H5Sparse.h5isvalidcsc(fid, dataset)
            return size(H5SparseMatrixCSC(fid, dataset))
        else
            return size(fid[dataset])
        end
    end
end

function partition_samples(X::H5SparseMatrixCSC, nsubpartitions::Integer)
    nsamples = size(X, 2)
    dividers = round.(Int, range(1, nsamples+1, length=nsubpartitions+1))
    [sparse(X[:, dividers[i]:(dividers[i+1]-1)]) for i in 1:nsubpartitions]
end

function partition_samples(X::Matrix, nsubpartitions::Integer)
    nsamples = size(X, 2)
    dividers = round.(Int, range(1, nsamples+1, length=nsubpartitions+1))
    [view(X, :, dividers[i]:(dividers[i+1]-1)) for i in 1:nsubpartitions]
end

function partition_samples(v::Vector, nsubpartitions::Integer)
    nsamples = length(v)
    dividers = round.(Int, range(1, nsamples+1, length=nsubpartitions+1))
    [view(v, dividers[i]:(dividers[i+1]-1)) for i in 1:nsubpartitions]
end

function read_localdata(i::Integer, nworkers::Integer; inputfile::String, inputdataset::String, labeldataset::String, nreplicas::Integer, kwargs...)
    HDF5.ishdf5(inputfile) || throw(ArgumentError("$inputfile isn't an HDF5 file"))
    0 < nworkers || throw(DomainError(nworkers, "nworkers must be positive"))
    0 < nreplicas || throw(DomainError(nreplicas, "nreplicas must be positive"))
    0 < i <= nworkers || throw(DomainError(i, "i must be in [1, nworkers]"))
    mod(nworkers, nreplicas) == 0 || throw(ArgumentError("nworkers must be divisible by nreplicas"))
    npartitions = div(nworkers, nreplicas)
    partition_index = ceil(Int, i/nreplicas)
    _, nsamples = problem_size(inputfile, inputdataset)
    h5open(inputfile, "r") do fid
        inputdataset in keys(fid) || throw(ArgumentError("$inputdataset is not in $fid"))
        labeldataset in keys(fid) || throw(ArgumentError("$labeldataset is not in $fid"))

        # determine the indices of the samples to be stored locally
        Is = partition(nsamples, npartitions, partition_index)

        # load those samples into memory
        labels = fid[labeldataset][Is]
        if H5Sparse.h5isvalidcsc(fid, inputdataset)
            return sparse(H5SparseMatrixCSC(fid, inputdataset, :, Is))::SparseMatrixCSC, labels
        else
            return fid[inputdataset][:, Is]::Matrix, labels
        end
    end
end

function worker_setup(rank::Integer, nworkers::Integer; kwargs...)
    0 < nworkers || throw(DomainError(nworkers, "nworkers must be positive"))
    features, labels = read_localdata(rank, nworkers; kwargs...)
    dimension = size(features, 1)
    localdata = (features, labels)
    to_worker_metadata_bytes = sizeof(UInt16) * 2 * nworkers
    recvbuf = Vector{UInt8}(undef, sizeof(ELEMENT_TYPE)*(dimension+1) + to_worker_metadata_bytes)
    sendbuf = Vector{UInt8}(undef, sizeof(ELEMENT_TYPE)*(dimension+1) + COMMON_BYTES + METADATA_BYTES)
    localdata, recvbuf, sendbuf
end

function coordinator_setup(nworkers::Integer; inputfile::String, inputdataset::String, iteratedataset, kwargs...)
    0 < nworkers || throw(DomainError(nworkers, "nworkers must be positive"))
    dimension, nsamples = problem_size(inputfile, inputdataset)

    # initial iterate
    V = zeros(Float64, dimension+1)
    if !isnothing(iteratedataset) # given as an argument and loaded from disk
        h5open(inputfile) do fid
            iteratedataset in keys(fid) || throw(ArgumentError("iterate dataset $iteratedataset not found"))
            size(fid[iteratedataset]) == (dimension+1,) || throw(DimensionMismatch("Expected iterate to have dimensions $((dimension+1,)), but it has dimensions $(size(fid[iteratedataset]))"))
            V .= fid[iteratedataset][:]
        end
    end

    # communication buffers
    to_worker_metadata_bytes = sizeof(UInt16) * 2 * nworkers
    sendbuf = Vector{UInt8}(undef, sizeof(ELEMENT_TYPE)*(dimension+1) + to_worker_metadata_bytes)
    recvbuf = Vector{UInt8}(undef, (sizeof(ELEMENT_TYPE)*(dimension+1) + COMMON_BYTES + METADATA_BYTES) * nworkers)
    reinterpret(ELEMENT_TYPE, view(sendbuf, (to_worker_metadata_bytes+1):length(sendbuf))) .= view(V, :)

    # setup state for the gradient tracking method
    gradient_state, gradient_recvf = setup_gradient(V; kwargs...)
    iterate_state = nothing
    state = (gradient_state, iterate_state)

    V, recvbuf, sendbuf, state, gradient_recvf
end

metadata_view(buffer) = view(buffer, COMMON_BYTES+1:(COMMON_BYTES + METADATA_BYTES))
data_view(buffer) = reinterpret(ELEMENT_TYPE, @view buffer[(COMMON_BYTES + METADATA_BYTES+1):end])

function worker_task!(recvbuf, sendbuf, localdata; state=nothing, ncolumns::Integer, nworkers::Integer, kwargs...)
    to_worker_metadata_bytes = sizeof(UInt16) * 2 * nworkers
    features, labels = localdata
    size(features, 2) == length(labels) || throw(DimensionMismatch("features has dimensions $(size(features)), but labels has dimension $(size(labels))"))
    dimension, nlocalsamples = size(features)

    # prepare working memory
    if isnothing(state) # first iteration
        nsubpartitions_prev = 0
        subpartition_index = 1
        w = zeros(Float64, nlocalsamples) # temp. storage
    else # subsequent iterations
        nsubpartitions_prev::Int, subpartition_index::Int, w::Vector{Float64} = state
    end

    # get the number of sub-partitions from the coordinator
    vs = reinterpret(Tuple{UInt16,UInt16}, view(recvbuf, 1:to_worker_metadata_bytes))
    rank <= length(vs) || throw(DimensionMismatch("vs has length $(length(vs)), but rank is $rank"))
    nsubpartitions, _ = vs[rank]

    # increment the sub-partition index
    if !iszero(nsubpartitions_prev)
        subpartition_index = mod(subpartition_index, nsubpartitions_prev) + 1
        subpartition_index = align_partitions(nlocalsamples, nsubpartitions_prev, nsubpartitions, subpartition_index)
    end
    nsubpartitions_prev = nsubpartitions
    0 < nsubpartitions <= nlocalsamples || throw(DimensionMismatch("nsubpartitions is $nsubpartitions, but nlocalsamples is $nlocalsamples"))
    0 < subpartition_index <= nsubpartitions || throw(ArgumentError("subpartition_index is $subpartition_index, but nsubpartitions is $nsubpartitions"))

    # format the recvbuf into a matrix we can operate on
    recvdata = view(recvbuf, (to_worker_metadata_bytes+1):length(recvbuf))
    length(reinterpret(ELEMENT_TYPE, recvdata)) == dimension+1 || throw(DimensionMismatch("recvdata has length $(length(reinterpret(ELEMENT_TYPE, recvdata))), but the data dimension is $dimension"))
    v = reinterpret(ELEMENT_TYPE, recvdata)

    # indices of the local samples to process in this iteration
    cols = partition(nlocalsamples, nsubpartitions, subpartition_index)

    # gradient computation
    tcolsmul!(w, features, view(v, 2:length(v)), cols)
    @views w[cols] .+= v[1] # implicit intercept
    @views w[cols] .*= labels[cols]
    @views w[cols] .= exp.(w[cols])
    @views w[cols] .+= 1
    @views w[cols] .= labels[cols] ./ w[cols]
    @views w[cols] .*= -1
    @views v[1] = sum(w[cols]) # intercept derivative
    colsmul!(view(v, 2:length(v)), features, w, cols)
    v ./= ncolumns # normalize by the total number of samples

    # populate the send buffer
    metadata = reinterpret(UInt16, metadata_view(sendbuf))
    metadata[1] = CANARY_VALUE
    metadata[2] = rank
    metadata[3] = nsubpartitions
    metadata[4] = subpartition_index
    data_view(sendbuf) .= v
    nsubpartitions_prev, subpartition_index, w
end

function setup_gradient(args...; variancereduced::Bool, vralgo::String, kwargs...)
    if variancereduced
        if vralgo == "table"
            return nothing, nothing
        elseif vralgo == "tree"
            return setup_gradient_vrt(args...; kwargs...), recvf_vrt
        else
            throw(ArgumentError("unexpected vralgo $vralgo"))
        end
    else
        return nothing, nothing
    end
end

function update_gradient!(args...; variancereduced::Bool, vralgo::String, kwargs...)
    if variancereduced
        if vralgo == "table"
            return update_gradient_vr!(args...; kwargs...)
        elseif vralgo == "tree"
            return update_gradient_vrt!(args...; kwargs...)
        else
            throw(ArgumentError("unexpected vralgo $vralgo"))
        end
    else
        return update_gradient_sgd!(args...; kwargs...)
    end
end

function update_iterate!(v, ∇; state=nothing, stepsize, lambda, kwargs...)
    size(v) == size(∇) || throw(DimensionMismatch("v has dimensions $(size(v)), but ∇ has dimensions $(size(∇))"))
    v .*= 1 - stepsize * lambda
    v .-= stepsize .* ∇
    return
end

function coordinator_task!(V, ∇, recvbufs, sendbuf, epoch::Integer, repochs::Vector{<:Integer}; nworkers::Integer, state=nothing, kwargs...)
    isnothing(state) || length(state) == 2 || throw(ArgumentError("expected state to be nothing or a tuple of length 2, but got $state"))
    if isnothing(state)
        gradient_state = update_gradient!(∇, recvbufs, epoch, repochs; kwargs...)
        iterate_state = update_iterate!(V, ∇; kwargs...)
    else
        gradient_state, iterate_state = state
        gradient_state = update_gradient!(∇, recvbufs, epoch, repochs; state=gradient_state, kwargs...)
        iterate_state = update_iterate!(V, ∇; state=iterate_state, kwargs...)
    end
    to_worker_metadata_bytes = sizeof(UInt16) * 2 * nworkers
    reinterpret(ELEMENT_TYPE, view(sendbuf, (to_worker_metadata_bytes+1):length(sendbuf))) .= view(V, :)
    gradient_state, iterate_state
end

include("common.jl")


# Macros for creating units and quantities

abstract type PhysicalQuantity{U} <: Number end

abstract type DimensionalPhysicalQuantity{U} <: PhysicalQuantity{U} end
abstract type DimensionlessPhysicalQuantity{U} <: PhysicalQuantity{U} end

"""
    unit(a::PhysicalQuantity)

Return the units of quantity `a`. Extends `unit` operation
in `Unitful`.
"""
unit(::PhysicalQuantity{U}) where U = unit(U)

"""
    value(a::PhysicalQuantity)

Return the numerical value (with units) of `a`
"""
value(s::PhysicalQuantity) = float(s.val)

value(s::Unitful.Quantity) = float(s)
value(s::Real) = float(s)

"""
    value(a::PhysicalQuantity,units::Unitful.Units)

Return the numerical value of `a`, converted to units `units`.
The form of `units` must be of `Unitful` form, e.g. `u"Pa"`
and must be dimensionally compatible with `a`
"""
value(s::PhysicalQuantity,units::Unitful.Units) = uconvert(units,value(s))

"""
    ustrip(a::PhysicalQuantity)

Return the numerical value of `a`, stripped
of its units.
"""
ustrip(s::PhysicalQuantity) = ustrip(value(s))

"""
    ustrip(a::PhysicalQuantity,units::Unitful.Units)

Return the numerical value of `a`, converted to units `units`, and stripped
of its units. The form of `units` must be of `Unitful` form, e.g. `u"Pa"`
and must be dimensionally compatible with `a`.
"""
ustrip(s::PhysicalQuantity,units::Unitful.Units) = ustrip(value(s,units))

"""
    name(a::PhysicalQuantity)

Return the name of `a`.
"""
name(s::PhysicalQuantity) = s.name

function Base.show(io::IO, m::MIME"text/plain", s::PhysicalQuantity{U}) where {U}
    print(io,"$(s.name) = $(float(s.val))")
end

####### OPERATIONS ON ALL TYPES #######
#=
Some notes on these
- We do not necessarily wish for an operation on a PhysicalQuantity
  to return the same type of PhysicalQuantity. For example, if
  we add Pressure and PressureDifference, we should let the user
  specify what the result should be (by wrapping it in a type explicitly).
- When carrying out an operation between a PhysicalQuantity and a Number
  or a PhysicalQuantity and a Unitful.Quantity, just let the Unitful
  rules catch dimensional inconsistencies and throw errors.
=#

(-)(s::PhysicalQuantity) = -s.val

for op in (:(+),:(-),:(>),:(<),:(>=),:(<=),:(==))
    @eval $op(s1::PhysicalQuantity{U1},s2::PhysicalQuantity{U2}) where {U1,U2} = $op(s1.val,s2.val)
    @eval $op(s1::PhysicalQuantity{U1},s2::U2) where {U1,U2<:Quantity} = $op(s1.val,s2)
    @eval $op(s1::U1,s2::PhysicalQuantity{U2}) where {U1<:Quantity,U2} = $op(s1,s2.val)
    @eval $op(s::PhysicalQuantity,C::Real) = $op(s.val,C)
    @eval $op(C::Real,s::PhysicalQuantity) = $op(C,s.val)
end

for op in (:(isapprox),)
    @eval $op(s1::PhysicalQuantity{U1},s2::PhysicalQuantity{U2};kwargs...) where {U1,U2} = $op(s1.val,s2.val;kwargs...)
    @eval $op(s1::PhysicalQuantity{U1},s2::U2;kwargs...) where {U1,U2<:Quantity} = $op(s1.val,s2;kwargs...)
    @eval $op(s1::U1,s2::PhysicalQuantity{U2};kwargs...) where {U1<:Quantity,U2} = $op(s1,s2.val;kwargs...)
    @eval $op(s::PhysicalQuantity,C::Real;kwargs...) = $op(s.val,C;kwargs...)
    @eval $op(C::Real,s::PhysicalQuantity;kwargs...) = $op(C,s.val;kwargs...)
end

for op in (:(*),:(/))
    @eval $op(s1::PhysicalQuantity,s2::PhysicalQuantity) = $op(s1.val,s2.val)
    @eval $op(s1::PhysicalQuantity,s2::U) where {U<:Quantity} = $op(s1.val,s2)
    @eval $op(s1::U,s2::PhysicalQuantity) where {U<:Quantity} = $op(s1,s2.val)
    @eval $op(s::PhysicalQuantity,C::Real) = $op(s.val,C)
    @eval $op(C::Real,s::PhysicalQuantity) = $op(C,s.val)
end

for op in (:(^),)
    @eval $op(s::PhysicalQuantity,C::Integer) = $op(s.val,C)
    @eval $op(s::PhysicalQuantity,C::Real) = $op(s.val,C)
    @eval $op(s::PhysicalQuantity,C::Rational) = $op(s.val,C)
end

for op in (:sin, :cos, :tan, :asin, :acos, :atan, :sinh, :cosh, :tanh, :asinh,
           :acosh, :atanh,:sinpi, :cospi, :sinc, :cosc, :cis)
    @eval ($op)(s::DimensionlessPhysicalQuantity) = ($op)(s.val)
end

for op in (:exp, :exp10, :exp2, :expm1, :log, :log10, :log1p, :log2)
    @eval ($op)(s::DimensionlessPhysicalQuantity) = ($op)(s.val)
end

for op in (:sqrt,)
    @eval ($op)(s::PhysicalQuantity) = ($op)(s.val)
end

function displayedunits() end
function ushow() end
function default_unit() end

"""
    @displayedunits name unit dims

Set the preferred units for displaying quantities and
create function `displayedunits` for returning these units and
`ushow` for converting quantities into these units. The new unit type
is specified with `name`, the default units with `unit`, and the
dimensions with `dims`. The latter use the `Unitful` dimension names, `𝐌`,
`𝐓`, `𝐋`, `𝚯`, in combinations

# Examples
```jldoctest myunit
julia> import ThermofluidQuantities: 𝐋, 𝐓

julia> @displayedunits MyVelocityType "m/s" 𝐋/𝐓

julia> MyVelocityType
Union{Unitful.Quantity{T,𝐋 𝐓⁻¹,U}, Unitful.Level{L,S,Unitful.Quantity{T,𝐋 𝐓⁻¹,U}} where S where L} where U where T
```
"""
macro displayedunits(qty,a,dims)
  utype = isdefined(Unitful, qty) ? getfield(Unitful,qty) : qty
  strqty = string(qty)
  esc(quote
      ThermofluidQuantities.Unitful.@derived_dimension($qty,$dims)

      ThermofluidQuantities.displayedunits(::Type{$utype}) = @u_str($a)

      ThermofluidQuantities.ushow(x::$utype) = uconvert(@u_str($a),x)

      push!(ThermofluidQuantities.unittypes,Symbol($strqty))

      export $qty

      nothing

      end)
end

"""
    @dimvar name utype

Define a dimensional variable type of the given name, with units of type `utype`.
The `utype` is of the form created with the @displayedunits macro. A list
of such types can be found in `ThermofluidQuantities.unittypes`.

# Examples
```jldoctest mydimvar
julia> import ThermofluidQuantities: 𝐓

julia> @displayedunits MyTimeType "s" 𝐓

julia> @dimvar MyTimeVar MyTimeType

julia> MyTimeVar(5)
MyTimeVar = 5.0 s
```
"""
macro dimvar(qty,utype)
  strqty = string(qty)
  esc(quote

          @doc """
              $($strqty) <: DimensionalPhysicalQuantity
          """ struct $qty{U<:$utype} <: DimensionalPhysicalQuantity{U}
            val :: U
            name :: String
          end

          @doc """
              $($strqty)(x::Real)

          Create an instance of a $($strqty) type dimensional physical quantity,
          with value equal to `x` and units $(displayedunits($utype)).
          """ $qty(x::Real) = $qty(x*displayedunits($utype),$strqty)

          @doc """
              $($strqty)(x::Unitful.Quantity)

          Create an instance of a $($strqty) type dimensional physical quantity,
          with unit-ed value equal to `x`. Converts to units $(displayedunits($utype)).
          """ $qty(x::U) where {U<:$utype} = $qty(ushow(x),$strqty)

          ThermofluidQuantities.default_unit(::Type{$qty}) = displayedunits($utype)

          push!(ThermofluidQuantities.dimvartypes,$qty)

          export $qty

      end)
end

"""
    @nondimvar name

Define a non-dimensional variable type of the given name.

# Examples
```jldoctest mynondimvar
julia> @nondimvar MyNondimVar

julia> MyNondimVar(100)
MyNondimVar = 100.0
```
"""
macro nondimvar(qty)
  strqty = string(qty)
  esc(quote

          @doc """
              $($strqty) <: DimensionlessPhysicalQuantity
          """ struct $qty{U<:Real} <: DimensionlessPhysicalQuantity{U}
            val :: U
            name :: String
          end

          @doc """
              $($strqty)(x::Real)

          Create an instance of a $($strqty) type non-dimensional physical quantity,
          with value equal to `x`.
          """ $qty(x::Real) = $qty(x,$strqty)

          @doc """
              $($strqty)(x::Unitful.Quantity)

          Create an instance of a $($strqty) type non-dimensional physical quantity,
          with unit-ed value equal to `x`.
          """
          $qty(x::ThermofluidQuantities.Unitful.Quantity{T,ThermofluidQuantities.Unitful.NoDims}) where {T} = $qty(uconvert.(ThermofluidQuantities.Unitful.NoUnits,x))

          push!(ThermofluidQuantities.nondimvartypes,$qty)

          export $qty

      end)
end

"""
    @liquid name temp dens visc surftens pv Ev

Create a liquid of the specified name, with viscosity `visc`, surface tension `surftens`,
vapor pressure `pv`, and bulk modulus `Ev` at reference temperature `temp`. Each of these must
be provided with units if they are not the default.

# Example
```jldoctest myliq
julia> @liquid MyWater 15u"°C" 999 1.13e-3 7.34e-2 1.77e3 2.15e9

julia> MyWater
Liquid with
   Density = 999.0 kg m⁻³
   Viscosity = 0.00113 kg m⁻¹ s⁻¹
   Kinematic viscosity = 1.131131131131131e-6 m² s⁻¹
   Specific weight = 9796.84335 N m⁻³
   Surface tension = 0.0734 N m⁻¹
   Bulk modulus = 2.15e9 Pa
   Vapor pressure = 1770.0 Pa
   at reference temp 288.15 K
```
"""
macro liquid(name,temp,density,viscosity,surftens,pv,Ev)
    strname = string(name)
    esc(quote
            const $name = Liquid(Tref = Temperature($temp),
                                 ρref = Density($density),
                                 μ = Viscosity($viscosity),
                                 σ = SurfaceTension($surftens),
                                 pv = VaporPressure($pv),
                                 Ev = BulkModulus($Ev))
            export $name

            push!(ThermofluidQuantities.liquids,Symbol($strname))

            nothing
        end)
end

"""
    @gas name temp visc gamma mmass

Create a gas of the specified name, with viscosity `visc` and ratio of specific heats
`gamma` at reference temperature `temp`, and molar mass `mmass`. Each of these must
be provided with units if they are not the default.

# Example
```jldoctest mygas
julia> @gas MyO2 20u"°C" 2.04e-5 1.395 31.999u"g/mol"

julia> MyO2
Perfect gas with
   Density = 1.330236729981785 kg m⁻³
   Viscosity = 2.04e-5 kg m⁻¹ s⁻¹
   Specific heat ratio = 1.395
   Gas constant = 259.83507666343445 J kg⁻¹ K⁻¹
   cp = 917.645397330357 J kg⁻¹ K⁻¹
   cv = 657.8103206669226 J kg⁻¹ K⁻¹
   at reference temp 293.15 K
```
"""
macro gas(name,temp,viscosity,gamma,molarmass)
    strname = string(name)
    esc(quote
            const $name = PerfectGas(Tref = Temperature($temp),
                                     μ = Viscosity($viscosity),
                                     γ = SpecificHeatRatio($gamma),
                                     R = GasConstant(ThermofluidQuantities.Unitful.R/$molarmass))
            export $name

            push!(ThermofluidQuantities.gases,Symbol($strname))

            nothing
        end)
end


###
### Sampler states
###

struct Emcee{space, E<:AMH.Ensemble} <: InferenceAlgorithm
    ensemble::E
end

function Emcee(n_walkers::Int, stretch_length=2.0)
    # Note that the proposal distribution here is just a Normal(0,1)
    # because we do not need AdvancedMH to know the proposal for
    # ensemble sampling.
    prop = AMH.StretchProposal(nothing, stretch_length)
    ensemble = AMH.Ensemble(n_walkers, prop)
    return Emcee{(), typeof(ensemble)}(ensemble)
end

struct EmceeState{V<:AbstractVarInfo,S}
    vi::V
    states::S
end

function AbstractMCMC.step(
    rng::Random.AbstractRNG,
    model::Model,
    spl::Sampler{<:Emcee};
    resume_from = nothing,
    kwargs...
)
    if resume_from !== nothing
        state = loadstate(resume_from)
        return AbstractMCMC.step(rng, model, spl, state; kwargs...)
    end

    # Sample from the prior
    n = spl.alg.ensemble.n_walkers
    vis = [VarInfo(rng, model, SampleFromPrior()) for _ in 1:n]

    # Update the parameters if provided.
    if haskey(kwargs, :init_params)
        for vi in vis
            initialize_parameters!(vi, kwargs[:init_params], spl)

            # Update log joint probability.
            model(rng, vi, SampleFromPrior())
        end
    end

    # Compute initial transition and states.
    transition = map(vis) do vi
        Transition(vi)
    end
    state = EmceeState(
        vis[1],
        map(vis) do vi
            # Transform to unconstrained space.
            DynamicPPL.link!(vi, spl)
            AMH.Transition(vi[spl], getlogp(vi))
        end
    )

    return transition, state
end

function AbstractMCMC.step(
    rng::AbstractRNG,
    model::Model,
    spl::Sampler{<:Emcee},
    state::EmceeState;
    kwargs...
)
    # Generate a log joint function.
    vi = state.vi
    densitymodel = AMH.DensityModel(gen_logπ(vi, SampleFromPrior(), model))

    # Compute the next states.
    states = last(AbstractMCMC.step(rng, densitymodel, spl.alg.ensemble, state.states))

    # Compute the next transition and state.
    transition = map(states) do _state
        vi[spl] = _state.params
        DynamicPPL.invlink!(vi, spl)
        t = Transition(tonamedtuple(vi), _state.lp)
        DynamicPPL.link!(vi, spl)
        return t
    end
    newstate = EmceeState(vi, states)

    return transition, newstate
end

function AbstractMCMC.bundle_samples(
    samples::Vector{<:Vector},
    model::AbstractModel,
    spl::Sampler{<:Emcee},
    state::EmceeState,
    chain_type::Type{MCMCChains.Chains};
    save_state = false,
    kwargs...
)
    # Convert transitions to array format.
    # Also retrieve the variable names.
    params_vec = map(_params_to_array, samples)

    # Extract names and values separately.
    nms = params_vec[1][1]
    vals_vec = [p[2] for p in params_vec]

    # Get the values of the extra parameters in each transition.
    extra_vec = map(get_transition_extras, samples)

    # Get the extra parameter names & values.
    extra_params = (internals = extra_vec[1][1],)
    extra_values_vec = [e[2] for e in extra_vec]

    # Extract names & construct param array.
    nms = [nms; extra_params]
    parray = map(x -> hcat(x[1], x[2]), zip(vals_vec, extra_values_vec))
    parray = cat(parray..., dims=3)

    # Get the average or final log evidence, if it exists.
    le = getlogevidence(samples, state, spl)

    # Set up the info tuple.
    if save_state
        info = (model = model, sampler = spl, samplerstate = state)
    else
        info = NamedTuple()
    end

    # Concretize the array before giving it to MCMCChains.
    parray = MCMCChains.concretize(parray)

    # Chain construction.
    return MCMCChains.Chains(
        parray,
        nms,
        extra_params;
        evidence=le,
        info=info,
    ) |> sort
end