"""
Type for storing data acquired with a steady state response (SSR) experimental paradigm.

In addition to the functions available for all EEG types,
the SSR type supports:

* `modulationrate()`

The following standard names are used when saving data to the processing dictionary.

* `Name`: The identifier for the participant
* `Side`: Side of stimulation
* `Carrier_Frequency`: Carrier frequency of the stimulus
* `Amplitude`: Amplitude of the stimulus
* `epochs`: The epochs extracted from the recording
* `sweeps`: The extracted sweeps from the recording

#### Example

Put an example here

```julia
s = read_SSR("filename")
s.modulationrate = 40.0391u"Hz"
s = rereference(s, "Cz")
```

"""
mutable struct SSR <: EEG
    data::Array
    sensors::Array{Sensor}
    triggers::Dict
    system_codes::Dict
    samplingrate::typeof(1.0u"Hz")
    modulationrate::typeof(1.0u"Hz")
    reference_channel::Array{AbstractString,1}
    file_path::AbstractString
    file_name::AbstractString
    processing::Dict
    header::Dict
end

#######################################
#
# SSR info
#
#######################################

"""
Return the modulation rate of a steady state type.
If no type is provided, the modulation rate is returned as a floating point.

#### Example

Return the modulation rate of a recording

```julia
s = read_SSR(filename)
modulationrate(s)
```
"""
modulationrate(t, s::SSR) = convert(t, ustrip(s.modulationrate))
modulationrate(s::SSR) = modulationrate(AbstractFloat, s)

#######################################
#
# Show
#
#######################################

import Base.show
function Base.show(io::IO, a::SSR)
    time_length = round.(size(a.data, 1) / samplingrate(a) / 60)
    println(
        io,
        "SSR measurement of $time_length mins with $(size(a.data,2)) channels sampled at $(a.samplingrate)",
    )
    println(io, "  Modulation frequency: $(a.modulationrate )")

    if haskey(a.processing, "Amplitude")
        println(io, "  Stimulation amplitude: $(a.processing["Amplitude"]) dB")
    end
    if haskey(a.processing, "Name")
        println(io, "  Participant name: $(a.processing["Name"] )")
    end
    if haskey(a.processing, "Side")
        println(io, "  Stimulation side: $(a.processing["Side"] )")
    end
    if haskey(a.processing, "Carrier_Frequency")
        println(io, "  Carrier frequency: $(a.processing["Carrier_Frequency"] ) Hz")
    end

end

#######################################
#
# Helper functions
#
#######################################

function assr_frequency(
    rounded_freq::Number;
    stimulation_samplingrate::Number = 32000,
    stimulation_frames_per_epoch::Number = 32768,
)

    round.(rounded_freq / (stimulation_samplingrate / stimulation_frames_per_epoch)) *
    stimulation_samplingrate / stimulation_frames_per_epoch
end

function assr_frequency(rounded_freq::AbstractVector)

    [assr_frequency(f) for f in rounded_freq]
end


export rl_unplugged_dm_dataset

using TFRecord

function make_batch_array(type::Type, feature_dims::Int, size::Tuple, batch_size::Int)
    Array{type, feature_dims+1}(undef, size..., batch_size)
end

function dm_buffer_dict(feature_size::Dict{String, Tuple}, batch_size::Int)
    obs_buffer = Dict{Symbol, AbstractArray}()

    buffer_dict = Dict{Symbol, Any}()

    for feature in keys(feature_size)
        feature_dims = length(feature_size[feature])
        if split(feature, "/")[1] == "observation"
            ob_key = Symbol(chop(feature, head=length("observation")+1, tail=0))
            if split(feature, "/")[end] == "egocentric_camera"
                obs_buffer[ob_key] = make_batch_array(UInt8, feature_dims, feature_size[feature], batch_size)
            else
                obs_buffer[ob_key] = make_batch_array(Float32, feature_dims, feature_size[feature], batch_size)
            end
        elseif feature == "action"
            buffer_dict[:action] = make_batch_array(Float32, feature_dims, feature_size[feature], batch_size)
            buffer_dict[:next_action] = make_batch_array(Float32, feature_dims, feature_size[feature], batch_size)
        elseif feature == "step_type"
            buffer_dict[:terminal] = make_batch_array(Bool, feature_dims, feature_size[feature], batch_size)
        else
            ob_key = Symbol(feature)
            buffer_dict[ob_key] = make_batch_array(Float32, feature_dims, feature_size[feature], batch_size)
        end
    end

    buffer_dict[:state] = deepcopy(NamedTuple(obs_buffer))
    buffer_dict[:next_state] = deepcopy(NamedTuple(obs_buffer))

    buffer_dict
end

function size_dict(game::String, type::String)
    if type == "dm_locomotion_humanoid"
        size_dict = DM_LOCOMOTION_HUMANOID_SIZE
    elseif type == "dm_locomotion_rodent"
        size_dict = DM_LOCOMOTION_RODENT_SIZE
    elseif type == "dm_control_suite"
        size_dict = DM_CONTROL_SUITE_SIZE[game]
    else
        error("given game type does not exist")
    end
end

function batch_named_tuple!(dest::NamedTuple, src::NamedTuple, i::Int)
    for fn in fieldnames(typeof(dest))
        xs = getfield(dest, fn)
        x = getfield(src, fn)
        if typeof(xs) <: NamedTuple
            batch_named_tuple!(xs, x, i)
        else
            selectdim(xs, ndims(xs), i) .= x
        end
    end
end

function make_transition(example::TFRecord.Example, feature_size::Dict{String, Tuple})
    f = example.features.feature

    observation_dict = Dict{Symbol, AbstractArray}()
    next_observation_dict = Dict{Symbol, AbstractArray}()
    transition_dict = Dict{Symbol, Any}()

    for feature in keys(feature_size)
        if split(feature, "/")[1] == "observation"
            ob_key = Symbol(chop(feature, head = length("observation")+1, tail=0))
            if split(feature, "/")[end] == "egocentric_camera"
                cam_feature_size = feature_size[feature]
                ob_size = prod(cam_feature_size)
                observation_dict[ob_key] = reshape(f[feature].bytes_list.value[1][1:ob_size], cam_feature_size...)
                next_observation_dict[ob_key] = reshape(f[feature].bytes_list.value[1][ob_size+1:end], cam_feature_size...)
            else
                if feature_size[feature] == ()
                    observation_dict[ob_key] = f[feature].float_list.value
                else
                    ob_size = feature_size[feature][1]
                    observation_dict[ob_key] = f[feature].float_list.value[1:ob_size]
                    next_observation_dict[ob_key] = f[feature].float_list.value[ob_size+1:end]
                end
            end
        elseif feature == "action"
            ob_size = feature_size[feature][1]
            action = f[feature].float_list.value
            transition_dict[:action] = action[1:ob_size]
            transition_dict[:next_action] = action[ob_size+1:end]
        elseif feature == "step_type"
            transition_dict[:terminal] = f[feature].float_list.value[1] == 2
        else
            ob_key = Symbol(feature)
            transition_dict[ob_key] = f[feature].float_list.value[1]
        end
    end
    state_nt = (state = NamedTuple(observation_dict),)
    next_state_nt = (next_state = NamedTuple(next_observation_dict),)
    transition = NamedTuple(transition_dict)

    merge(transition, state_nt, next_state_nt)
end
"""
    rl_unplugged_dm_dataset(game, shards; <keyword arguments>)

Returns a `RingBuffer`(@ref) of `NamedTuple` containing SARTS batches which supports
multi threaded loading. Also contains additional data. The data enclosed within `:state` and
`next_state` is a NamedTuple consisting of all observations that are provided.
Check out keys in `DM_LOCOMOTION_HUMANOID`, `DM_LOCOMOTION_RODENT`, `DM_CONTROL_SUITE_SIZE` for supported
datasets. Also check out `dm_params()` for more info on arguments.

# Arguments

- `game::String`: name of the dataset.
- `shards::Vector{Int}`: the shards that are to be loaded.
- `type::String`: type of the dm_env. can be `dm_control_suite`, `dm_locomotion_humanoid`, `dm_locomotion_rodent`.
- `is_shuffle::Bool`
- `stochasticity::Float32`: represents the stochasticity of the dataset. can be
in the range: `0.0:0.1:0.5`.
- `shuffle_buffer_size::Int=10_000`: size of the shuffle_buffer used in loading AtariRLTransitions.
- `tf_reader_bufsize::Int=10_000`: the size of the buffer `bufsize` that is used internally
in `TFRecord.read`.
- `tf_reader_sz::Int=10_000`: the size of the `Channel`, `channel_size` that is returned by
`TFRecord.read`.
- `batch_size::Int=256`: The number of samples within the batches that are returned by the `Channel`.
- `n_preallocations::Int=nthreads()*12`: the size of the buffer in the `Channel` that is returned.

!!! note

    To enable reading records from multiple files concurrently, remember to set the number of
    threads correctly (See [JULIA_NUM_THREADS](https://docs.julialang.org/en/v1/manual/environment-variables/#JULIA_NUM_THREADS)).
"""
function rl_unplugged_dm_dataset(
    game,
    shards;
    type = "dm_control_suite",
    is_shuffle = true,
    shuffle_buffer_size=10_000,
    tf_reader_bufsize=10_000,
    tf_reader_sz=10_000,
    batch_size=256,
    n_preallocations=nthreads()*12
)
    n = nthreads()

    repo = "rl-unplugged-dm"

    folders= [
        @datadep_str "$repo-$game-$shard"
        for shard in shards
    ]

    ch_files = Channel{String}(length(folders)) do ch
        for folder in cycle(folders)
            file = folder * "/$(readdir(folder)[1])"
            put!(ch, file)
        end
    end

    if is_shuffle
        files = buffered_shuffle(ch_files, length(folders))
    else
        files = ch_files
    end

    feature_size = size_dict(game, type)

    ch_src = Channel{NamedTuple}(n * tf_reader_sz) do ch
        for fs in partition(files, n)
            Threads.foreach(
                TFRecord.read(
                    fs;
                    compression=:gzip,
                    bufsize=tf_reader_bufsize,
                    channel_size=tf_reader_sz,
                );
                schedule=Threads.StaticSchedule()
            ) do x
                put!(ch, make_transition(x, feature_size))
            end
        end
    end

    if is_shuffle
        transitions = buffered_shuffle(
        ch_src,
        shuffle_buffer_size
        )
    else
        transitions = ch_src
    end

    taskref = Ref{Task}()

    buffer_dict = dm_buffer_dict(feature_size, batch_size)

    buffer = NamedTuple(buffer_dict)

    res = RingBuffer(buffer;taskref=taskref, sz=n_preallocations) do buff
        Threads.@threads for i in 1:batch_size
            batch_named_tuple!(buff, take!(transitions), i)
        end
    end

    bind(ch_src, taskref[])
    bind(ch_files, taskref[])
    res
end


"""
    Schedule(ts, values, interpolations; unit = :epoch; kwargs...)

Describes how the values of a hyperparameter change over the
training.

## Examples

```julia
# one-cycle schedule over 10 epochs from 0.01 over 0.1 to 0.001
es = length(learner.data.training)
Schedule([0es, 2es, 10es], [0.01, 0.1, 0.001], [sineio(), sineio()])
```

`Schedule` is an alias for `Animations.Animation`, see the
[Animations.jl documentation](https://jkrumbiegel.github.io/Animations.jl/dev)
for more detailed information.
"""
const Schedule = Animation

"""
    Scheduler(schedules...)

Callback for hyperparameter scheduling. Takes pairs of [`HyperParameter`](#)
types and [`Schedule`](#)s.

See [the tutorial](/documents/docs/tutorials/hyperparameters.md) for more information.

## Example

```julia
es = length(learner.data.training)
lrschedule = Schedule([0es, 10es], [0.1, 0.001], Animations.sineio())
scheduler = Scheduler(
    LearningRate => lrschedule
)
```

See also [`Schedule`](#).
"""
mutable struct Scheduler <: Callback
    schedules::Dict{Type{<:HyperParameter},Animations.FiniteLengthAnimation}
    step::Int
    Scheduler(args...; kwargs...) = new(Dict(args...; kwargs...), 1)
end

Base.show(io::IO, scheduler::Scheduler) = print(
    io, "Scheduler(", join(keys(scheduler.schedules), ", "), ")")

function stateaccess(scheduler::Scheduler)
    # TODO: implement proper merging of permissions
    if length(keys(scheduler.schedules)) == 0
        hpstateaccess = (;)
    else
        hpstateaccess = merge(stateaccess.(keys(scheduler.schedules))...)
    end
    return (
        data = Read(), cbstate = (; hyperparams=Write(), history = Read()),
        hpstateaccess...
    )
end

function init!(::Scheduler, learner)
    if !haskey(learner.cbstate, :hyperparams)
        learner.cbstate.hyperparams = MVHistory()
    end
end

function on(::StepBegin, phase::AbstractTrainingPhase, scheduler::Scheduler, learner)
    step = scheduler.step
    for (H, animation) in scheduler.schedules
        value = Animations.at(animation, scheduler.step)
        sethyperparameter!(learner, H, value)
        push!(
            learner.cbstate.hyperparams,
            Symbol(H),
            learner.cbstate.history[phase].steps,
            value)
    end
    scheduler.step += 1
end

"""
    onecycle(nsteps, max_val, [start_val, end_val; pct_start])

Creates a one-cycle [`Schedule`](#) over `nsteps` steps from `start_val`
over `max_val` to `end_val`.

## Examples

```julia

epochlength = length(traindataiter)
cb = Scheduler(LearningRate => onecycle(10epochlength, 0.01))
learner = Learner(<args>..., cb)
```
"""
function onecycle(
        nsteps, max_val;
        pct_start=0.25,
        div=25, divfinal=1e5,
        start_val=max_val / div, end_val=max_val / divfinal)
    return Animation(
        [0, nsteps * pct_start, nsteps],
        [start_val, max_val, end_val],
        [Animations.sineio(), Animations.sineio()]
    )
end


abstract BioUnit                <: Integer

bitstype 8 AminoAcid            <: BioUnit
bitstype 8 Nucleotide           <: BioUnit

## Conversions ##

convert{T<:BioUnit}(::Type{T}, value::UInt8) = box(T,unbox(UInt8,value));
convert{T<:BioUnit}(::Type{UInt8}, value::T) = box(UInt8,unbox(T,value));

const _types = (Int8, Int16, UInt16, Int32, UInt32, Char, Int64, UInt64, Float32, Float64, Bool)

for t in _types
    @eval convert{T<:BioUnit}(::Type{$t}, x::T) = convert($t,UInt8(x))
    @eval convert{T<:BioUnit}(::Type{T},  x::$t)= convert(T,UInt8(x))
end

## Conversions ##

aminoacid{T<:Number}(value::T)   = convert(AminoAcid, value)
nucleotide{T<:Number}(value::T)   = convert(Nucleotide,value)

aminoacid(value::Char)   = convert(AminoAcid, Int(value))
nucleotide(value::Char)   = convert(Nucleotide, Int(value))

aminoacid(vec::Vector)  = convert(Vector{AminoAcid}, vec)
nucleotide(vec::Vector)  = convert(Vector{Nucleotide},vec)

aminoacid(str::ASCIIString)  = convert(Vector{AminoAcid}, str)
nucleotide(str::ASCIIString)  = convert(Vector{Nucleotide},str)

aminoacid(mat::Matrix)  = convert(Matrix{AminoAcid}, mat)
nucleotide(mat::Matrix)  = convert(Matrix{Nucleotide},mat)

macro aa_str(s);  :(reinterpret(AminoAcid,  @b_str($s) )); end

macro nt_str(s);  :(reinterpret(Nucleotide, @b_str($s) )); end

## Promotion rules ##

promote_rule{Tb<:BioUnit,T<:Number}(::Type{Tb}, ::Type{T}) = T
promote_rule{T<:BioUnit}(::Type{T}, ::Type{Union{}}) = T

## Comparisons ##

=={T<:BioUnit}(x::T, y::T) = ==(UInt8(x),UInt8(y))
!={T<:BioUnit}(x::T, y::T) = !=(UInt8(x),UInt8(y))
<{T<:BioUnit}(x::T, y::T)  = <( UInt8(x),UInt8(y))
<={T<:BioUnit}(x::T, y::T) = <=(UInt8(x),UInt8(y))

## Bitwise Operations ##

~{T<:BioUnit}(x::T)            = convert(T,~UInt8(x))
(&){T<:BioUnit}(x::T, y::T)    = convert(T,(&)(UInt8(x),UInt8(y)))
|{T<:BioUnit}(x::T, y::T)      = convert(T,|(UInt8(x),UInt8(y)))
($){T<:BioUnit}(x::T, y::T)    = convert(T,($)(UInt8(x),UInt8(y)))
<<{T<:BioUnit}(x::T, y::Int32) = convert(T,<<(UInt8(x),y))
>>{T<:BioUnit}(x::T, y::Int32) = convert(T,>>(UInt8(x),y))
>>>{T<:BioUnit}(x::T, y::Int32)= convert(T,>>>(UInt8(x),y))

## Characteristics ##

typemin{T<:BioUnit}(::Type{T}) = convert(T, 0x00 )
typemax{T<:BioUnit}(::Type{T}) = convert(T, 0xff )

## Show as Character ##

show{T<:BioUnit}(io::IO,x::T) = (write(io,UInt8(x)); nothing)

## Bits representation ##

bits{T<:BioUnit}(x::T) = bin(reinterpret(UInt8,x),8)

## Array & Strings ##

convert{N}(::Type{Array{AminoAcid,N}}, x::Array{AminoAcid,N}) = x
convert{N}(::Type{Array{Nucleotide,N}}, x::Array{Nucleotide,N}) = x

convert{N}(::Type{Array{AminoAcid,N}}, x::Array{Nucleotide,N}) = reinterpret(AminoAcid,s)
convert{N}(::Type{Array{Nucleotide,N}}, x::Array{AminoAcid,N}) = reinterpret(Nucleotide,s)

convert{N}(::Type{Array{UInt8,N}},s::Array{AminoAcid,N}) = reinterpret(UInt8,s)
convert{N}(::Type{Array{AminoAcid,N}},s::Array{UInt8,N}) = reinterpret(AminoAcid,s)

convert{N}(::Type{Array{UInt8,N}},s::Array{Nucleotide,N}) = reinterpret(UInt8,s)
convert{N}(::Type{Array{Nucleotide,N}},s::Array{UInt8,N}) = reinterpret(Nucleotide,s)

convert{T<:BioUnit}(::Type{ASCIIString}, seq::Vector{T}) = ASCIIString(reinterpret(UInt8,copy(seq)))
convert{T<:BioUnit}(::Type{Vector{T}}, str::ASCIIString) = reinterpret(T,copy(str.data))

bytestring{T<:BioUnit}(seq::Vector{T}) = bytestring(reinterpret(UInt8,seq))

## HASH (for Julia 0.3)
hash{T<:BioUnit}(x::T) = hash(UInt8(x))
hash{T<:BioUnit}(seq::Vector{T}) = hash(bytestring(seq))


# Replay buffer that store full episodes

mutable struct EpisodeReplayBuffer{N<:Integer, T<:Real, CI, Q<:AbstractArray{T},A<:AbstractArray{T}}
    max_size::Int64
    batch_size::Int64
    trace_length::Int64
    rng::AbstractRNG
    _curr_size::Int64
    _idx::Int64
    _experience::Vector{Vector{DQExperience{N,T,Q}}}

    _s_batch::Vector{A}
    _a_batch::Vector{Vector{CI}}
    _r_batch::Vector{Vector{T}}
    _sp_batch::Vector{A}
    _done_batch::Vector{Vector{T}}
    _trace_mask::Vector{Vector{N}}
    _episode::Vector{DQExperience{N,T,Q}}
end

function EpisodeReplayBuffer(env::AbstractEnvironment{OV},
                        max_size::Int64,
                        batch_size::Int64,
                        trace_length::Int64,
                        rng::AbstractRNG = MersenneTwister(0)) where {OV}
    s_dim = obs_dimensions(env)
    Q = length(s_dim)
    experience = Vector{Vector{DQExperience{Int32, Float32, OV}}}(undef, max_size)
    _s_batch = [zeros(Float32, s_dim..., batch_size) for i=1:trace_length]
    _a_batch = [[CartesianIndex(1,1) for i=1:batch_size] for i=1:trace_length]
    _r_batch = [zeros(Float32, batch_size) for i=1:trace_length]
    _sp_batch = [zeros(Float32, s_dim..., batch_size) for i=1:trace_length]
    _done_batch = [zeros(Float32, batch_size) for i=1:trace_length]
    _trace_mask = [zeros(Int32, batch_size) for i=1:trace_length]
    _episode = Vector{DQExperience{Int32, Float32, OV}}()
    return EpisodeReplayBuffer(max_size, batch_size, trace_length, rng, 0, 1, experience,
                _s_batch, _a_batch, _r_batch, _sp_batch, _done_batch, _trace_mask, _episode)
end

is_full(r::EpisodeReplayBuffer) = r._curr_size == r.max_size

max_size(r::EpisodeReplayBuffer) = r.max_size

function add_exp!(r::EpisodeReplayBuffer{N,T,CI,Q}, exp::DQExperience) where {N,T,CI,Q}
    push!(r._episode, exp)
    if exp.done
        add_episode!(r, r._episode)
        r._episode = Vector{DQExperience{N,T,Q}}()
    end
end

function add_episode!(r::EpisodeReplayBuffer{N ,T, CI, Q}, ep::Vector{DQExperience{Int32, Float32, Q}}) where {N,T,CI,Q}
    r._experience[r._idx] = ep
    r._idx = mod1((r._idx + 1),r.max_size)
    if r._curr_size < r.max_size
        r._curr_size += 1
    end
end

function reset_batches!(r::EpisodeReplayBuffer)
    fill!.(r._s_batch, 0.)
    fill!.(r._a_batch, Ref(CartesianIndex(1,1)))
    fill!.(r._r_batch, 0.)
    fill!.(r._sp_batch, 0.)
    fill!.(r._done_batch, false)
    fill!.(r._trace_mask, 0)
end

function StatsBase.sample(r::EpisodeReplayBuffer)
    reset_batches!(r) # might not be necessary
    @assert r._curr_size >= r.batch_size
    @assert r.max_size >= r.batch_size # could be checked during construction
    sample_indices = sample(r.rng, 1:r._curr_size, r.batch_size, replace=false)
    @assert length(sample_indices) == size(r._s_batch[1])[end]
    s_batch_size = size(first(r._s_batch))
    for (i, idx) in enumerate(sample_indices)
        ep = r._experience[idx]
        # randomized start TODO add as an option of the buffer
        ep_start = rand(r.rng, 1:length(ep))
        t = 1
        for j=ep_start:min(length(ep), r.trace_length)
            expe = ep[t]
            r._s_batch[t][.., i] = vec(expe.s)
            r._a_batch[t][i] = CartesianIndex(expe.a, i)
            r._r_batch[t][i] = expe.r
            r._sp_batch[t][.., i] = vec(expe.sp)
            r._done_batch[t][i] = expe.done
            r._trace_mask[t][i] = 1
            t += 1
        end
    end
    return r._s_batch, r._a_batch, r._r_batch, r._sp_batch, r._done_batch, r._trace_mask
end

function populate_replay_buffer!(r::EpisodeReplayBuffer,
                                 env::AbstractEnvironment,
                                 action_indices;
                                 max_pop::Int64 = r.max_size,
                                 max_steps::Int64 = 100)
    for t=1:(max_pop - r._curr_size)
        ep = generate_episode(env, action_indices, max_steps=max_steps)
        add_episode!(r, ep)
    end
    @assert r._curr_size >= r.batch_size
end

function generate_episode(env::AbstractEnvironment{OV}, action_indices; max_steps::Int64 = 100) where OV
    s_dim = obs_dimensions(env)
    episode = DQExperience{Int32, Float32, OV}[]
    sizehint!(episode, max_steps)
    # start simulation
    o = reset!(env)
    done = false
    step = 1
    while !done && step < max_steps
        action = sample_action(env)
        ai = action_indices[action]
        op, rew, done, info = step!(env, action)
        exp = DQExperience(o, ai, Float32(rew), op, done)
        push!(episode, exp)
        o = op
        step += 1
    end
    return episode
end


"""
function plotNodeVisited(metrics::BasicMetrics{O, H}; filename::String="") where{O<:AbstractTakeObjective, H<:ValueSelection}

plot 2 graphs :
1) number of node visited by the heuristic to find a first solution for every learning episode.
2) number of node visited by the heuristic to prove to optimality for every learning episode.

The learning process should show a decrease in the number of nodes required to find a first solution along the search
(depending on the reward engineering).
"""
function plotNodeVisited(metricsArray::Union{BasicMetrics, Vector{AbstractMetrics}}; filename::String="")
    L = length(metricsArray[1].meanNodeVisitedUntilEnd)
    Label = Matrix{String}(undef, 1, length(metricsArray))
    learnedIdx = 1
    basicIdx = 1
    for j in 1:length(metricsArray) #nb of heuristics
        if isa(metricsArray[j].heuristic, LearnedHeuristic)
            Label[1,j]="Learned heuristic n°$learnedIdx"
            learnedIdx += 1
        else
            Label[1,j]="Classic heuristic n°$basicIdx"
            basicIdx += 1
        end
    end
    p1 = plot(
        1:L,
        [metricsArray[i].meanNodeVisitedUntilEnd[1:L] for i in 1:length(metricsArray)],
        title = "Node visited until optimality",
        label=Label,
        xlabel="Episode",
        ylabel="Nodes visited",
    )
    p2 = plot(
        1:L,
        [metricsArray[i].meanNodeVisitedUntilfirstSolFound[1:L] for i in 1:length(metricsArray)],
        title ="Node visited until first solution found",
        label=Label,
        xlabel="Episode",
        ylabel="Nodes visited",
    )
    plot(p1,p2, layout=(2,1))
end

function plotNodeVisited(metrics::BasicMetrics; filename::String="")
    L = length(metrics.meanNodeVisitedUntilEnd)
    p = plot(
        1:L,
        [metrics.meanNodeVisitedUntilEnd[1:L] metrics.meanNodeVisitedUntilfirstSolFound[1:L]],
        xlabel="Episode",
        ylabel="Nodes visited",
        title = ["Node visited until Optimality" "Node visited until first solution found"],
        layout = (2, 1)
    )
    display(p)
    savefig(p,filename*"_node_visited_"*"$(typeof(metrics.heuristic))"*".png")
end

"""
function plotScoreVariation(metrics::BasicMetrics{O, H}; filename::String="") where{O<:AbstractTakeObjective, H<:ValueSelection}

plot the relative scores ( compared to the optimal ) of the heuristic during the search for fixed instances along the training. This plot is
meaningful only if the metrics is one from the evaluator (ie. the instance remains the same one).
"""
function plotScoreVariation(metrics::BasicMetrics{TakeObjective, H}; filename::String="") where H
    Data=[]
    for i in length(metrics.nodeVisited):-1:1
        push!(Data,hcat(metrics.nodeVisited[i],metrics.scores[i]))
    end

    p = plot(
        [scatter[:,1] for scatter in Data],
        [scatter[:,2] for scatter in Data],
        #fillrange =[[ones(length(scatter[:,2])),scatter[:,2]] for scatter in Data],
        #fillalpha=1,
        xlabel="number of nodes visited",
        ylabel="relative score",
        xaxis=:log,
    )
    display(p)
    savefig(p,filename*"_score_variation.png")
end

"""
function plotScoreVariation(metrics::BasicMetrics{O, H}; filename::String="") where{O<:AbstractTakeObjective, H<:ValueSelection}

plot the relative scores ( compared to the optimal ) of the heuristic during the search for fixed instances along the training. This plot is
meaningful only if the metrics is one from the evaluator (ie. the instance remains the same one).
"""
function plotRewardVariation(metrics::BasicMetrics{<:AbstractTakeObjective, <:LearnedHeuristic}; filename::String="")
    L = length(metrics.totalReward)

    p = plot(1:L,
        metrics.totalReward[1:L],
        xlabel="Episode",
        ylabel="Total reward",
        title ="Total reward per episode",
    )
    display(p)
    savefig(p,filename*"_reward_variation.png")
end


"""
A Fast and Accurate Dependency Parser using Neural Networks, Chen and Manning, EMNLP 2014
"""

type UnlabeledFeedForward
    word_f
    tag_f
    nonlinear
    W
end

type UExample
    wordids::Vector{Int}
    tagids::Union{Matrix{Float32},Vector{Int}}
    target::Int
    valid::Vector{Float32}
end

function UExample(s::State)
    w, t = unlabeledsparsefeatures(s)
    valid = [Float32(isvalid(s, act)) for act = 1:3]
    UExample(w, t, -1, valid)
end

targetsize(m::UnlabeledFeedForward) = size(m.W[end].w.data)[1]

function initmodel!(parser::DepParser, model::Type{UnlabeledFeedForward}, embed,
    topktags, nonlinear, sparsesizes, embedsizes, hiddensizes)
    T = Float32
    if embed == ""
        info("USING EMBEDDINGS WITH UNIFORM DISTRIBUTION [-0.01, 0.01]")
        word_f = Embedding(T, length(parser.words), embedsizes[1])
    else
        info("USING EMBEDDINGS LOADED FROM $(embed)")
        word_f = Embedding(embed, Float32)
        scale!(word_f)
    end
    if topktags
        tag_f = Linear(T, length(parser.tags), embedsizes[2])
    else
        tag_f = Embedding(T, length(parser.tags), embedsizes[2])
    end
    indim = sum(sparsesizes .* embedsizes)
    outdim = 3
    W = [myLinear(T, indim, hiddensizes[1]),
         myLinear(T, hiddensizes[1], outdim)]
    parser.model = UnlabeledFeedForward(word_f, tag_f, nonlinear, W)
    info("INPUT: [S^word,S^tag] = ", sparsesizes)
    info("EMBED DIMS: [word,tag] = ", embedsizes)
    info("HIDDEN LAYER: ", hiddensizes)
    info("OUTPUT DIM: ", outdim)
    info("NONLINEAR: ", nonlinear)
end

# TODO: make State have id field
# to tell where the State is in a batch
# called from expand(::State ::Int)
@compat function (m::UnlabeledFeedForward){T}(s::State{T}, act::Int)
    Var([0f0])
end

@compat function (m::UnlabeledFeedForward)(batch::AbstractVector{UExample}, istrain=true)
    wordvec, tagvec = [], []
    for s in batch
        push!(wordvec, s.wordids)
        push!(tagvec, s.tagids)
    end
    wordmat = m.word_f(Var(hcat(wordvec...)))
    tagmat = tagvectors(m.tag_f, tagvec)
    x = concat(1, wordmat, tagmat)
    x = m.nonlinear(m.W[1](x))
    x = dropout(x, 0.5, istrain)
    x = m.W[end](x)
    valids = hcat(map(x -> x.valid, batch)...)
    x.data .*= valids
    x
end

@compat function (m::UnlabeledFeedForward){T}(batch::AbstractVector{State{T}}, istrain=true)
    m(map(UExample, batch), istrain)
end

function unlabeledsparsefeatures(s::State)
    # word, tag
    b0 = tokenat(s, s.right)
    b1 = tokenat(s, s.right + 1)
    b2 = tokenat(s, s.right + 2)
    b3 = tokenat(s, s.right + 3)
    s0 = tokenat(s, s.top)
    s0l = tokenat(s, s.lchild)
    s0l2 = tokenat(s, s.lsibl.lchild)
    s0r = tokenat(s, s.rchild)
    s0r2 = tokenat(s, s.rsibl.rchild)
    s02l = tokenat(s, s.lchild.lchild)
    s12r = tokenat(s, s.rchild.rchild)
    s1 = tokenat(s, s.left)
    s1l = tokenat(s, s.left.lchild)
    s1l2 = tokenat(s, s.left.lsibl.lchild)
    s1r = tokenat(s, s.left.rchild)
    s1r2 = tokenat(s, s.left.rsibl.rchild)
    s12l = tokenat(s, s.left.lchild.lchild)
    s12r = tokenat(s, s.left.rchild.rchild)
    s2 = tokenat(s, s.left.left)
    s3 = tokenat(s, s.left.left.left)

    words = [b0.word, b1.word, b2.word, b3.word, s0.word, s0l.word, s0l2.word,
    s0r.word, s0r2.word, s02l.word, s12r.word, s1.word, s1l.word, s1l2.word,
    s1r.word, s1r2.word, s12l.word, s12r.word, s2.word, s3.word]

    tags = [b0.tag, b1.tag, b2.tag, b3.tag, s0.tag, s0l.tag, s0l2.tag,
    s0r.tag, s0r2.tag, s02l.tag, s12r.tag, s1.tag, s1l.tag, s1l2.tag,
    s1r.tag, s1r2.tag, s12l.tag, s12r.tag, s2.tag, s3.tag]

    return words, tags
end

typealias Doc Vector{Vector{Token}}

function train!{T}(::Type{UnlabeledFeedForward}, parser::DepParser{T}, trainsents::Doc,
    testsents::Doc=Vector{Token}[]; embed="", batchsize=32, iter=20, progbar=true,
    nonlinear=tanh, sparsesizes=[20,20] ,embedsizes=[50,50], hiddensizes=[1024],
    topktags=false, opt=MyAdaGrad(0.01), evaliter=100, outfile="parser.dat")
    info("WILL RUN $iter ITERATIONS")

    saver = ModelSaver(outfile)
    initmodel!(parser, UnlabeledFeedForward, embed, topktags,
               nonlinear, sparsesizes, embedsizes, hiddensizes)

    projectivize!(trainsents, false)

    trainsamples = UExample[]
    for s in trainsents
        s = State(s, parser)
        while !isfinal(s)
            ex = UExample(s)
            s = expandgold(s)[1]
            ex.target = s.prevact
            push!(trainsamples, ex)
        end
    end
    samplesize = length(trainsamples)
    labelsize = targetsize(parser.model)

    info("OPTIMIZER: ", typeof(opt))
    info("BATCH SIZE: $batchsize LABEL SIZE: $labelsize")
    info("#BATCHES: $(div(samplesize, batchsize))")
    info("#SAMPLES: $(samplesize)")

    for i = 1:iter
        batch = sub(trainsamples, rand(1:samplesize, batchsize))
        preds = parser.model(batch)
        golds = map(s -> s.target, batch)
        correct = reduce(0, zip(argmax(preds.data, 1), golds)) do v, tup
            v + Int(tup[1] == tup[2])
        end
        accuracy = correct / batchsize
        loss = update!(opt, golds, preds)
        println(now(), " Epoch: $i\tLoss: $loss\tAcc: $accuracy")

        if i % evaliter == 0 && !isempty(testsents)
            println()
            info("**ITER $i TESTING**")
            @time res = decode(UnlabeledFeedForward, parser, testsents)
            uas, las = evaluate(parser, res)
            saver(parser, uas)
            println()
        end
    end
end

function decode{T}(::Type{UnlabeledFeedForward}, parser::DepParser{T}, sents::Doc;
    batchsize=32, progbar=true)
    batches = 1:batchsize:length(sents)
    res = State{T}[]
    for k in batches
        batch = k:min(k+batchsize-1, endof(sents))
        ss = map(s -> State(s, parser), sub(sents, batch))
        parsegreedy!(parser, ss)
        append!(res, ss)
    end
    res
end


using LShapedSolvers
using Base.Test
using JuMP
using StochasticPrograms
using Gurobi

solver = GurobiSolver(OutputFlag=0)
lsolvers = [(LShapedSolver(:ls,solver,log=false),"L-Shaped"),
            (LShapedSolver(:rd,solver,crash=Crash.EVP(),autotune=true,log=false),"RD L-Shaped"),
            (LShapedSolver(:tr,solver,crash=Crash.EVP(),autotune=true,log=false),"TR L-Shaped"),
            (LShapedSolver(:lv,solver,crash=Crash.EVP(),λ=0.6,log=false),"Leveled L-Shaped")]

problems = Vector{Tuple{JuMP.Model,String}}()
info("Loading test problems...")
info("Loading simple...")
include("simple.jl")
info("Loading farmer...")
include("farmer.jl")

info("Test problems loaded. Starting test sequence.")
@testset "$lsname Solver: $name" for (lsolver,lsname) in lsolvers, (sp,name) in problems
    solve(sp,solver=solver)
    x̄ = copy(sp.colVal)
    Q̄ = copy(sp.objVal)
    solve(sp,solver=lsolver)
    @test norm(optimal_decision(sp) - x̄) <= 1e-2
    @test abs(optimal_value(sp) - Q̄) <= 1e-2
end

info("Starting distributed tests...")

include("/opt/julia-0.6/share/julia/test/testenv.jl")
push!(test_exeflags.exec,"--color=yes")
cmd = `$test_exename $test_exeflags run_dtests.jl`

if !success(pipeline(cmd; stdout=STDOUT, stderr=STDERR)) && ccall(:jl_running_on_valgrind,Cint,()) == 0
    error("Distributed test failed, cmd : $cmd")
end


VERSION >= v"0.4.0-dev+6641" && __precompile__()

module Codecs

using Compat

export encode, decode, Base64, Zlib, BCD

abstract Codec

function encode{T <: Codec}(codec::Type{T}, s::AbstractString)
    encode(codec, convert(Vector{UInt8}, s))
end

function decode{T <: Codec}(codec::Type{T}, s::AbstractString)
    decode(codec, convert(Vector{UInt8}, s))
end

# RFC3548/RFC4648 base64 codec

abstract Base64 <: Codec

const b64enc_tbl = UInt8[
    'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H',
    'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P',
    'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X',
    'Y', 'Z', 'a', 'b', 'c', 'd', 'e', 'f',
    'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n',
    'o', 'p', 'q', 'r', 's', 't', 'u', 'v',
    'w', 'x', 'y', 'z', '0', '1', '2', '3',
    '4', '5', '6', '7', '8', '9', '+', '/']

const base64_pad = @compat UInt8('=')
const sentinel = typemax(UInt8)

const b64dec_tbl = fill(sentinel, 256)
for (val, ch) in enumerate(b64enc_tbl)
    b64dec_tbl[ch] = val - 1
end

function base64dec(c::UInt8)
    v = b64dec_tbl[c]
    if v == sentinel
        error("Invalid base64 symbol: $(char(c))")
    end
    v
end

function encode(::Type{Base64}, input::Vector{UInt8})
    n = length(input)
    if n == 0
        return Array(UInt8, 0)
    end

    m = @compat Int(4 * ceil(n / 3))
    output = Array(UInt8, m)

    k = 1
    for ii = 1:3:length(input)-2
        u, v, w = input[ii], input[ii+1], input[ii+2]
        output[k]     = b64enc_tbl[  (u >> 2)                     + 1]
        output[k + 1] = b64enc_tbl[(((u << 4) | (v >> 4)) & 0x3f) + 1]
        output[k + 2] = b64enc_tbl[(((v << 2) | (w >> 6)) & 0x3f) + 1]
        output[k + 3] = b64enc_tbl[  (w                   & 0x3f) + 1]
        k += 4
    end

    if n % 3 == 1
        output[end - 3] = b64enc_tbl[ (input[end] >> 2)         + 1]
        output[end - 2] = b64enc_tbl[((input[end] << 4) & 0x3f) + 1]
        output[end - 1] = base64_pad
        output[end]     = base64_pad
    elseif n % 3 == 2
        output[end - 3] = b64enc_tbl[  (input[end - 1] >> 2)      + 1]
        output[end - 2] = b64enc_tbl[(((input[end - 1] << 4) |
                                       (input[end] >> 4)) & 0x3f) + 1]
        output[end - 1] = b64enc_tbl[ ((input[end] << 2) & 0x3f)  + 1]
        output[end] = base64_pad
    end

    output
end

function decode(::Type{Base64}, input::Vector{UInt8})
    n = length(input)
    if n == 0
        return Array(UInt8, 0)
    end

    if n % 4 != 0
        warn("Length of Base64 input is not a multiple of four.")
    end

    m = 3 * div(n,  4)
    if input[end] == base64_pad
        m -= 1
    end
    if input[end - 1] == base64_pad
        m -= 1
    end
    output = Array(UInt8, m)

    k = 1
    @inbounds for ii = 1:4:length(input)-4
        # This loop is performance-critical, so inline base64dec and consolidate into a single branch point for error-checking
        u = b64dec_tbl[input[ii]]
        v = b64dec_tbl[input[ii + 1]]
        w = b64dec_tbl[input[ii + 2]]
        z = b64dec_tbl[input[ii + 3]]
        if u | v | w | z == sentinel
            if u == sentinel
                error("Invalid base64 symbol: $(char(input[ii]))")
            elseif v == sentinel
                error("Invalid base64 symbol: $(char(input[ii + 1]))")
            elseif w == sentinel
                error("Invalid base64 symbol: $(char(input[ii + 2]))")
            else
                error("Invalid base64 symbol: $(char(input[ii + 3]))")
            end
        end
        output[k]     = (u << 2) | (v >> 4)
        output[k + 1] = (v << 4) | (w >> 2)
        output[k + 2] = (w << 6) | z
        k += 3
    end

    u = input[end-3]
    v = input[end-2]
    w = input[end-1]
    z = input[end]
    u = base64dec(u)
    v = base64dec(v)

    if w == base64_pad && z == base64_pad
        output[end] = (u << 2) | (v >> 4)
    elseif z == base64_pad
        w = base64dec(w)
        output[end - 1] = (u << 2) | (v >> 4)
        output[end]     = (v << 4) | (w >> 2)
    else
        w = base64dec(w)
        z = base64dec(z)
        output[end - 2] = (u << 2) | (v >> 4)
        output[end - 1] = (v << 4) | (w >> 2)
        output[end]     = (w << 6) | z
    end

    output
end

# Zlib/Gzip

abstract Zlib <: Codec

const Z_NO_FLUSH      = 0
const Z_PARTIAL_FLUSH = 1
const Z_SYNC_FLUSH    = 2
const Z_FULL_FLUSH    = 3
const Z_FINISH        = 4
const Z_BLOCK         = 5
const Z_TREES         = 6

const Z_OK            = 0
const Z_STREAM_END    = 1
const Z_NEED_DICT     = 2
const ZERRNO          = -1
const Z_STREAM_ERROR  = -2
const Z_DATA_ERROR    = -3
const Z_MEM_ERROR     = -4
const Z_BUF_ERROR     = -5
const Z_VERSION_ERROR = -6

if is_unix()
    const libz = "libz"
elseif is_windows()
    const libz = "zlib1"
end

# The zlib z_stream structure.
type z_stream
    next_in::Ptr{UInt8}
    avail_in::Cuint
    total_in::Culong

    next_out::Ptr{UInt8}
    avail_out::Cuint
    total_out::Culong

    msg::Ptr{UInt8}
    state::Ptr{Void}

    zalloc::Ptr{Void}
    zfree::Ptr{Void}
    opaque::Ptr{Void}

    data_type::Cint
    adler::Culong
    reserved::Culong

    function z_stream()
        strm = new()
        strm.next_in   = C_NULL
        strm.avail_in  = 0
        strm.total_in  = 0
        strm.next_out  = C_NULL
        strm.avail_out = 0
        strm.total_out = 0
        strm.msg       = C_NULL
        strm.state     = C_NULL
        strm.zalloc    = C_NULL
        strm.zfree     = C_NULL
        strm.opaque    = C_NULL
        strm.data_type = 0
        strm.adler     = 0
        strm.reserved  = 0
        strm
    end
end

function zlib_version()
    ccall((:zlibVersion, libz), Ptr{UInt8}, ())
end

function encode(::Type{Zlib}, input::Vector{UInt8}, level::Integer)
    if !(1 <= level <= 9)
        error("Invalid zlib compression level.")
    end

    strm = z_stream()
    ret = ccall((:deflateInit_, libz),
                Int32, (Ptr{z_stream}, Int32, Ptr{UInt8}, Int32),
                &strm, level, zlib_version(), sizeof(z_stream))

    if ret != Z_OK
        error("Error initializing zlib deflate stream.")
    end

    strm.next_in = input
    strm.avail_in = length(input)
    strm.total_in = length(input)
    output = Array(UInt8, 0)
    outbuf = Array(UInt8, 1024)
    ret = Z_OK

    while ret != Z_STREAM_END
        strm.avail_out = length(outbuf)
        strm.next_out = outbuf
        flush = strm.avail_in == 0 ? Z_FINISH : Z_NO_FLUSH
        ret = ccall((:deflate, libz),
                    Int32, (Ptr{z_stream}, Int32),
                    &strm, flush)
        if ret != Z_OK && ret != Z_STREAM_END
            error("Error in zlib deflate stream ($(ret)).")
        end

        if length(outbuf) - strm.avail_out > 0
            append!(output, outbuf[1:(length(outbuf) - strm.avail_out)])
        end
    end

    ret = ccall((:deflateEnd, libz), Int32, (Ptr{z_stream},), &strm)
    if ret == Z_STREAM_ERROR
        error("Error: zlib deflate stream was prematurely freed.")
    end

    output
end

function encode(::Type{Zlib}, input::AbstractString, level::Integer)
    encode(Zlib, convert(Vector{UInt8}, input), level)
end

encode(::Type{Zlib}, input::Vector{UInt8}) = encode(Zlib, input, 9)

function decode(::Type{Zlib}, input::Vector{UInt8})
    strm = z_stream()
    ret = ccall((:inflateInit_, libz),
                Int32, (Ptr{z_stream}, Ptr{UInt8}, Int32),
                &strm, zlib_version(), sizeof(z_stream))

    if ret != Z_OK
        error("Error initializing zlib inflate stream.")
    end

    strm.next_in = input
    strm.avail_in = length(input)
    strm.total_in = length(input)
    output = Array(UInt8, 0)
    outbuf = Array(UInt8, 1024)
    ret = Z_OK

    while ret != Z_STREAM_END
        strm.next_out = outbuf
        strm.avail_out = length(outbuf)
        ret = ccall((:inflate, libz),
                    Int32, (Ptr{z_stream}, Int32),
                    &strm, Z_NO_FLUSH)
        if ret == Z_DATA_ERROR
            error("Error: input is not zlib compressed data.")
        elseif ret != Z_OK && ret != Z_STREAM_END
            error("Error in zlib inflate stream ($(ret)).")
        end

        if length(outbuf) - strm.avail_out > 0
            append!(output, outbuf[1:(length(outbuf) - strm.avail_out)])
        end
    end

    ret = ccall((:inflateEnd, libz), Int32, (Ptr{z_stream},), &strm)
    if ret == Z_STREAM_ERROR
        error("Error: zlib inflate stream was prematurely freed.")
    end

    output
end

# Packed binary-coded decimal (BCD) integers
# Two decimal digits per byte, each represented with 4 bits

abstract BCD <: Codec

function encode(::Type{BCD}, i::Integer, bigendian::Bool = false)
    if i < 0
        error("Integer must not be negative")
    end
    ndig = ndigits(i)
    ndig += isodd(ndig)
    v = digits(i, 10, ndig)
    nbytes = @compat trunc(Integer, ndig/2)
    out = Array(UInt8, nbytes)
    dj = 1-2*bigendian
    for i = 1:nbytes
        j = bigendian*length(v) + dj*2*(i-1) + 1-bigendian
        out[i] = (v[j]<<4) | v[j+dj]
    end
    return out
end

function decode(::Type{BCD}, v::Vector{UInt8}, bigendian::Bool = false)
    ret = 0
    if bigendian
        for i = 1:length(v)
            ret = 100*ret + 10 * (@compat Int(v[i]>>>4)) + (@compat Int(v[i]&0x0f))
        end
    else
        for i = length(v):-1:1
            ret = 100*ret + (@compat Int(v[i]>>>4)) + 10 * (@compat Int(v[i]&0x0f))
        end
    end
    return ret
end

end # module Codecs


include("DGFVModel_kernels.jl")
struct DGFVModel{BL, G, FVR, NFND, NFD, GNF, AS, DS, HDS, D, DD, MD, GF, TF} <:
       SpaceDiscretization
    balance_law::BL
    grid::G
    fv_reconstruction::FVR
    numerical_flux_first_order::NFND
    numerical_flux_second_order::NFD
    numerical_flux_gradient::GNF
    state_auxiliary::AS
    state_gradient_flux::DS
    states_higher_order::HDS
    direction::D
    diffusion_direction::DD
    modeldata::MD
    gradient_filter::GF
    tendency_filter::TF
    check_for_crashes::Bool
end

function DGFVModel(
    balance_law,
    grid,
    fv_reconstruction,
    numerical_flux_first_order,
    numerical_flux_second_order,
    numerical_flux_gradient;
    fill_nan = false,
    check_for_crashes = false,
    state_auxiliary = create_state(
        balance_law,
        grid,
        Auxiliary(),
        fill_nan = fill_nan,
    ),
    state_gradient_flux = create_state(balance_law, grid, GradientFlux()),
    states_higher_order = (
        create_state(balance_law, grid, GradientLaplacian()),
        create_state(balance_law, grid, Hyperdiffusive()),
    ),
    direction = EveryDirection(),
    diffusion_direction = direction,
    modeldata = nothing,
    gradient_filter = nothing,
    tendency_filter = nothing,
)
    # Make sure we are FVM in the vertical
    @assert polynomialorders(grid)[end] == 0
    @assert isstacked(grid.topology)
    state_auxiliary =
        init_state(state_auxiliary, balance_law, grid, direction, Auxiliary())
    DGFVModel(
        balance_law,
        grid,
        fv_reconstruction,
        numerical_flux_first_order,
        numerical_flux_second_order,
        numerical_flux_gradient,
        state_auxiliary,
        state_gradient_flux,
        states_higher_order,
        direction,
        diffusion_direction,
        modeldata,
        gradient_filter,
        tendency_filter,
        check_for_crashes,
    )
end

"""
    (dgfvm::DGFVModel)(tendency, state_prognostic, _, t, α, β)

Uses spectral element discontinuous Galerkin in the horizontal and finite volume
in the vertical to compute the tendency.

Computes the tendency terms compatible with `IncrementODEProblem`

    tendency .= α .* dQdt(state_prognostic, p, t) .+ β .* tendency

The 4-argument form will just compute

    tendency .= dQdt(state_prognostic, p, t)
"""
function (dgfvm::DGFVModel)(tendency, state_prognostic, _, t, α, β)
    device = array_device(state_prognostic)

    FT = eltype(state_prognostic)
    num_state_prognostic = number_states(dgfvm.balance_law, Prognostic())
    num_state_gradient_flux = number_states(dgfvm.balance_law, GradientFlux())
    @assert 0 == number_states(dgfvm.balance_law, Hyperdiffusive())
    num_state_tendency = size(tendency, 2)

    if num_state_prognostic < num_state_tendency && β != 1
        # If we don't operate on the full state, then we need to scale here instead of volume_tendency!
        tendency .*= β
        β = β != 0 # if β==0 then we can avoid the memory load in volume_tendency!
    end

    communicate =
        !(
            isstacked(dgfvm.grid.topology) &&
            typeof(dgfvm.direction) <: VerticalDirection
        )

    update_auxiliary_state!(
        dgfvm,
        dgfvm.balance_law,
        state_prognostic,
        t,
        dgfvm.grid.topology.realelems,
    )

    exchange_state_gradient_flux = NoneEvent()
    exchange_state_prognostic = NoneEvent()

    comp_stream = Event(device)

    if communicate
        exchange_state_prognostic = MPIStateArrays.begin_ghost_exchange!(
            state_prognostic;
            dependencies = comp_stream,
        )
    end

    if num_state_gradient_flux > 0
        ########################
        # Gradient Computation #
        ########################

        comp_stream = launch_volume_gradients!(
            dgfvm,
            state_prognostic,
            t;
            dependencies = comp_stream,
        )

        comp_stream = launch_interface_gradients!(
            dgfvm,
            state_prognostic,
            t;
            surface = :interior,
            dependencies = comp_stream,
        )

        if communicate
            exchange_state_prognostic = MPIStateArrays.end_ghost_exchange!(
                state_prognostic;
                dependencies = exchange_state_prognostic,
                check_for_crashes = dgfvm.check_for_crashes,
            )

            # Update_aux may start asynchronous work on the compute device and
            # we synchronize those here through a device event.
            checked_wait(
                device,
                exchange_state_prognostic,
                nothing,
                dgfvm.check_for_crashes,
            )
            update_auxiliary_state!(
                dgfvm,
                dgfvm.balance_law,
                state_prognostic,
                t,
                dgfvm.grid.topology.ghostelems,
            )
            exchange_state_prognostic = Event(device)
        end

        comp_stream = launch_interface_gradients!(
            dgfvm,
            state_prognostic,
            t;
            surface = :exterior,
            dependencies = (comp_stream, exchange_state_prognostic),
        )

        if dgfvm.gradient_filter !== nothing
            comp_stream = Filters.apply_async!(
                dgfvm.state_gradient_flux,
                1:num_state_gradient_flux,
                dgfvm.grid,
                dgfvm.gradient_filter;
                dependencies = comp_stream,
            )
        end

        if communicate
            if num_state_gradient_flux > 0
                exchange_state_gradient_flux =
                    MPIStateArrays.begin_ghost_exchange!(
                        dgfvm.state_gradient_flux,
                        dependencies = comp_stream,
                    )
            end
        end

        if num_state_gradient_flux > 0
            # Update_aux_diffusive may start asynchronous work on the compute device
            # and we synchronize those here through a device event.
            checked_wait(device, comp_stream, nothing, dgfvm.check_for_crashes)
            update_auxiliary_state_gradient!(
                dgfvm,
                dgfvm.balance_law,
                state_prognostic,
                t,
                dgfvm.grid.topology.realelems,
            )
            comp_stream = Event(device)
        end
    end

    ###################
    # RHS Computation #
    ###################
    comp_stream = launch_volume_tendency!(
        dgfvm,
        tendency,
        state_prognostic,
        t,
        α,
        β;
        dependencies = comp_stream,
    )

    comp_stream = launch_interface_tendency!(
        dgfvm,
        tendency,
        state_prognostic,
        t,
        α,
        β;
        surface = :interior,
        dependencies = comp_stream,
    )

    if communicate
        if num_state_gradient_flux > 0
            exchange_state_gradient_flux = MPIStateArrays.end_ghost_exchange!(
                dgfvm.state_gradient_flux;
                dependencies = exchange_state_gradient_flux,
                check_for_crashes = dgfvm.check_for_crashes,
            )

            # Update_aux_diffusive may start asynchronous work on the
            # compute device and we synchronize those here through a device
            # event.
            checked_wait(
                device,
                exchange_state_gradient_flux,
                nothing,
                dgfvm.check_for_crashes,
            )
            update_auxiliary_state_gradient!(
                dgfvm,
                dgfvm.balance_law,
                state_prognostic,
                t,
                dgfvm.grid.topology.ghostelems,
            )
            exchange_state_gradient_flux = Event(device)
        else
            exchange_state_prognostic = MPIStateArrays.end_ghost_exchange!(
                state_prognostic;
                dependencies = exchange_state_prognostic,
                check_for_crashes = dgfvm.check_for_crashes,
            )

            # Update_aux may start asynchronous work on the compute device and
            # we synchronize those here through a device event.
            checked_wait(
                device,
                exchange_state_prognostic,
                nothing,
                dgfvm.check_for_crashes,
            )
            update_auxiliary_state!(
                dgfvm,
                dgfvm.balance_law,
                state_prognostic,
                t,
                dgfvm.grid.topology.ghostelems,
            )
            exchange_state_prognostic = Event(device)
        end
    end

    comp_stream = launch_interface_tendency!(
        dgfvm,
        tendency,
        state_prognostic,
        t,
        α,
        β;
        surface = :exterior,
        dependencies = (
            comp_stream,
            exchange_state_prognostic,
            exchange_state_gradient_flux,
            # XXX: This is disabled until FVM with hyperdiffusion for DG is implemented: exchange_Qhypervisc_grad,
        ),
    )

    if dgfvm.tendency_filter !== nothing
        comp_stream = Filters.apply_async!(
            tendency,
            1:num_state_tendency,
            dgfvm.grid,
            dgfvm.tendency_filter;
            dependencies = comp_stream,
        )
    end

    # The synchronization here through a device event prevents CuArray based and
    # other default stream kernels from launching before the work scheduled in
    # this function is finished.
    checked_wait(device, comp_stream, nothing, dgfvm.check_for_crashes)
end

function fvm_balance!(
    balance_func,
    m::BalanceLaw,
    state_auxiliary::MPIStateArray,
    grid,
)
    device = array_device(state_auxiliary)

    dim = dimensionality(grid)
    N = polynomialorders(grid)
    Nq = N .+ 1
    Nqj = dim == 2 ? 1 : Nq[2]

    topology = grid.topology
    elems = topology.elems
    nelem = length(elems)
    nvertelem = topology.stacksize
    horzelems = fld1(first(elems), nvertelem):fld1(last(elems), nvertelem)

    event = Event(device)
    event = kernel_fvm_balance!(device, (Nq[1], Nqj))(
        balance_func,
        m,
        Val(nvertelem),
        state_auxiliary.data,
        grid.vgeo,
        horzelems;
        ndrange = (length(horzelems) * Nq[1], Nqj),
        dependencies = (event,),
    )
    wait(device, event)
end


## MPI Init
using MPIClusterManagers, Distributed
import MPI
MPI.Init()
rank = MPI.Comm_rank(MPI.COMM_WORLD)
size = MPI.Comm_size(MPI.COMM_WORLD)
if rank == 0
    @info "size is $size"
end
manager = MPIClusterManagers.start_main_loop(MPI_TRANSPORT_ALL)
@info "there are $(nworkers()) workers"
@everywhere @info gethostname()
include("01_read_data.jl")
include("02_set_options.jl")
## MPI checks
# split into sequential iterations of parallel soundings
nsoundings = length(soundings)
ncores = nworkers()
@assert mod(ncores+1,nchainspersounding+1) == 0
@assert mod(ppn, nchainspersounding+1) == 0
nparallelsoundings = Int((ncores+1)/(nchainspersounding+1))
nsequentialiters = ceil(Int, nsoundings/nparallelsoundings)
@info "will require $nsequentialiters iterations of $nparallelsoundings"
## set up McMC
@everywhere using Distributed
@everywhere using HiQGA.transD_GP
## do the parallel soundings
@info "starting"
transD_GP.SkyTEM1DInversion.loopacrosssoundings(soundings, opt;
                    nsequentialiters   = nsequentialiters,
                    nparallelsoundings = nparallelsoundings,
                    zfixed             = zfixed,
                    ρfixed             = ρfixed,
                    zstart             = zstart,
                    extendfrac         = extendfrac,
                    dz                 = dz,
                    ρbg                = ρbg,
                    nlayers            = nlayers,
                    ntimesperdecade    = ntimesperdecade,
                    nfreqsperdecade    = nfreqsperdecade,
                    Tmax               = Tmax,
                    nsamples           = nsamples,
                    nchainsatone       = nchainsatone,
                    nchainspersounding = nchainspersounding,
                    useML              = useML)

MPIClusterManagers.stop_main_loop(manager)
rmprocs(workers())
exit()


module Inference

using ..Core
using ..Utilities
using DynamicPPL: Metadata, _tail, VarInfo, TypedVarInfo,
    islinked, invlink!, getlogp, tonamedtuple, VarName, getsym, vectorize,
    settrans!, _getvns, getdist, CACHERESET,
    Model, Sampler, SampleFromPrior, SampleFromUniform,
    Selector, DefaultContext, PriorContext,
    LikelihoodContext, MiniBatchContext, set_flag!, unset_flag!, NamedDist, NoDist,
    getspace, inspace
using Distributions, Libtask, Bijectors
using DistributionsAD: VectorOfMultivariate
using LinearAlgebra
using ..Turing: PROGRESS, Turing
using StatsFuns: logsumexp
using Random: AbstractRNG
using DynamicPPL
using AbstractMCMC: AbstractModel, AbstractSampler
using DocStringExtensions: TYPEDEF, TYPEDFIELDS

import AbstractMCMC
import AdvancedHMC; const AHMC = AdvancedHMC
import AdvancedMH; const AMH = AdvancedMH
import AdvancedPS
import BangBang
import ..Core: getchunksize, getADbackend
import DynamicPPL: get_matching_type,
    VarName, _getranges, _getindex, getval, _getvns
import EllipticalSliceSampling
import Random
import MCMCChains
import StatsBase: predict

export  InferenceAlgorithm,
        Hamiltonian,
        GibbsComponent,
        StaticHamiltonian,
        AdaptiveHamiltonian,
        SampleFromUniform,
        SampleFromPrior,
        MH,
        ESS,
        Emcee,
        Gibbs,      # classic sampling
        GibbsConditional,
        HMC,
        SGLD,
        SGHMC,
        HMCDA,
        NUTS,       # Hamiltonian-like sampling
        DynamicNUTS,
        IS,
        SMC,
        CSMC,
        PG,
        Prior,
        assume,
        dot_assume,
        observe,
        dot_observe,
        resume,
        predict,
        isgibbscomponent

#######################
# Sampler abstraction #
#######################
abstract type AbstractAdapter end
abstract type InferenceAlgorithm end
abstract type ParticleInference <: InferenceAlgorithm end
abstract type Hamiltonian{AD} <: InferenceAlgorithm end
abstract type StaticHamiltonian{AD} <: Hamiltonian{AD} end
abstract type AdaptiveHamiltonian{AD} <: Hamiltonian{AD} end

getchunksize(::Type{<:Hamiltonian{AD}}) where AD = getchunksize(AD)
getADbackend(::Hamiltonian{AD}) where AD = AD()

# Algorithm for sampling from the prior
struct Prior <: InferenceAlgorithm end

"""
    mh_accept(logp_current::Real, logp_proposal::Real, log_proposal_ratio::Real)

Decide if a proposal ``x'`` with log probability ``\\log p(x') = logp_proposal`` and
log proposal ratio ``\\log k(x', x) - \\log k(x, x') = log_proposal_ratio`` in a
Metropolis-Hastings algorithm with Markov kernel ``k(x_t, x_{t+1})`` and current state
``x`` with log probability ``\\log p(x) = logp_current`` is accepted by evaluating the
Metropolis-Hastings acceptance criterion
```math
\\log U \\leq \\log p(x') - \\log p(x) + \\log k(x', x) - \\log k(x, x')
```
for a uniform random number ``U \\in [0, 1)``.
"""
function mh_accept(logp_current::Real, logp_proposal::Real, log_proposal_ratio::Real)
    # replacing log(rand()) with -randexp() yields test errors
    return log(rand()) + logp_current ≤ logp_proposal + log_proposal_ratio
end

######################
# Default Transition #
######################

struct Transition{T, F<:AbstractFloat}
    θ  :: T
    lp :: F
end

function Transition(vi::AbstractVarInfo, nt::NamedTuple=NamedTuple())
    theta = merge(tonamedtuple(vi), nt)
    lp = getlogp(vi)
    return Transition{typeof(theta), typeof(lp)}(theta, lp)
end

metadata(t::Transition) = (lp = t.lp,)

DynamicPPL.getlogp(t::Transition) = t.lp

# Metadata of VarInfo object
metadata(vi::AbstractVarInfo) = (lp = getlogp(vi),)

#########################################
# Default definitions for the interface #
#########################################

function AbstractMCMC.sample(
    model::AbstractModel,
    alg::InferenceAlgorithm,
    N::Integer;
    kwargs...
)
    return AbstractMCMC.sample(Random.GLOBAL_RNG, model, alg, N; kwargs...)
end

function AbstractMCMC.sample(
    rng::AbstractRNG,
    model::AbstractModel,
    alg::InferenceAlgorithm,
    N::Integer;
    kwargs...
)
    return AbstractMCMC.sample(rng, model, Sampler(alg, model), N; kwargs...)
end

function AbstractMCMC.sample(
    rng::AbstractRNG,
    model::AbstractModel,
    sampler::Sampler{<:InferenceAlgorithm},
    N::Integer;
    chain_type=MCMCChains.Chains,
    resume_from=nothing,
    progress=PROGRESS[],
    kwargs...
)
    if resume_from === nothing
        return AbstractMCMC.mcmcsample(rng, model, sampler, N;
                                       chain_type=chain_type, progress=progress, kwargs...)
    else
        return resume(resume_from, N; chain_type=chain_type, progress=progress, kwargs...)
    end
end

function AbstractMCMC.sample(
    rng::AbstractRNG,
    model::AbstractModel,
    alg::Prior,
    N::Integer;
    chain_type=MCMCChains.Chains,
    resume_from=nothing,
    progress=PROGRESS[],
    kwargs...
)
    if resume_from === nothing
        return AbstractMCMC.mcmcsample(rng, model, SampleFromPrior(), N;
                                       chain_type=chain_type, progress=progress, kwargs...)
    else
        return resume(resume_from, N; chain_type=chain_type, progress=progress, kwargs...)
    end
end

function AbstractMCMC.sample(
    model::AbstractModel,
    alg::InferenceAlgorithm,
    parallel::AbstractMCMC.AbstractMCMCParallel,
    N::Integer,
    n_chains::Integer;
    kwargs...
)
    return AbstractMCMC.sample(Random.GLOBAL_RNG, model, alg, parallel, N, n_chains;
                               kwargs...)
end

function AbstractMCMC.sample(
    rng::AbstractRNG,
    model::AbstractModel,
    alg::InferenceAlgorithm,
    parallel::AbstractMCMC.AbstractMCMCParallel,
    N::Integer,
    n_chains::Integer;
    kwargs...
)
    return AbstractMCMC.sample(rng, model, Sampler(alg, model), parallel, N, n_chains;
                               kwargs...)
end

function AbstractMCMC.sample(
    rng::AbstractRNG,
    model::AbstractModel,
    sampler::Sampler{<:InferenceAlgorithm},
    parallel::AbstractMCMC.AbstractMCMCParallel,
    N::Integer,
    n_chains::Integer;
    chain_type=MCMCChains.Chains,
    progress=PROGRESS[],
    kwargs...
)
    return AbstractMCMC.mcmcsample(rng, model, sampler, parallel, N, n_chains;
                                   chain_type=chain_type, progress=progress, kwargs...)
end

function AbstractMCMC.sample(
    rng::AbstractRNG,
    model::AbstractModel,
    alg::Prior,
    parallel::AbstractMCMC.AbstractMCMCParallel,
    N::Integer,
    n_chains::Integer;
    chain_type=MCMCChains.Chains,
    progress=PROGRESS[],
    kwargs...
)
    return AbstractMCMC.sample(rng, model, SampleFromPrior(), parallel, N, n_chains;
                               chain_type=chain_type, progress=progress, kwargs...)
end

##########################
# Chain making utilities #
##########################

"""
    getparams(t)

Return a named tuple of parameters.
"""
getparams(t) = t.θ
getparams(t::VarInfo) = tonamedtuple(TypedVarInfo(t))

function _params_to_array(ts::Vector)
    names = Vector{Symbol}()
    # Extract the parameter names and values from each transition.
    dicts = map(ts) do t
        nms, vs = flatten_namedtuple(getparams(t))
        for nm in nms
            if !(nm in names)
                push!(names, nm)
            end
        end
        # Convert the names and values to a single dictionary.
        return Dict(nms[j] => vs[j] for j in 1:length(vs))
    end
    # names = collect(names_set)
    vals = [get(dicts[i], key, missing) for i in eachindex(dicts),
        (j, key) in enumerate(names)]

    return names, vals
end

function flatten_namedtuple(nt::NamedTuple)
    names_vals = mapreduce(vcat, keys(nt)) do k
        v = nt[k]
        if length(v) == 1
            return [(Symbol(k), v)]
        else
            return mapreduce(vcat, zip(v[1], v[2])) do (vnval, vn)
                return collect(FlattenIterator(vn, vnval))
            end
        end
    end
    return [vn[1] for vn in names_vals], [vn[2] for vn in names_vals]
end

function get_transition_extras(ts::AbstractVector{<:VarInfo})
    valmat = reshape([getlogp(t) for t in ts], :, 1)
    return [:lp], valmat
end

function get_transition_extras(ts::AbstractVector)
    # Extract all metadata.
    extra_data = map(metadata, ts)
    return names_values(extra_data)
end

function names_values(extra_data::AbstractVector{<:NamedTuple{names}}) where names
    values = [getfield(data, name) for data in extra_data, name in names]
    return collect(names), values
end

function names_values(extra_data::AbstractVector{<:NamedTuple})
    # Obtain all parameter names.
    names_set = Set(Symbol[])
    for data in extra_data
        for name in names(data)
            push!(extra_names_set, name)
        end
    end
    extra_names = collect(extra_names_set)

    # Extract all values as matrix.
    values = [
        hasfield(data, name) ? missing : getfield(data, name)
        for data in extra_data, name in extra_names
    ]

    return extra_names, values
end

getlogevidence(transitions, sampler, state) = missing

# Default MCMCChains.Chains constructor.
# This is type piracy (at least for SampleFromPrior).
function AbstractMCMC.bundle_samples(
    ts::Vector,
    model::AbstractModel,
    spl::Union{Sampler{<:InferenceAlgorithm},SampleFromPrior},
    state,
    chain_type::Type{MCMCChains.Chains};
    save_state = false,
    kwargs...
)
    # Convert transitions to array format.
    # Also retrieve the variable names.
    nms, vals = _params_to_array(ts)

    # Get the values of the extra parameters in each transition.
    extra_params, extra_values = get_transition_extras(ts)

    # Extract names & construct param array.
    nms = [nms; extra_params]
    parray = hcat(vals, extra_values)

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

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

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

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

# This is type piracy (for SampleFromPrior).
function AbstractMCMC.bundle_samples(
    ts::Vector,
    model::AbstractModel,
    spl::Union{Sampler{<:InferenceAlgorithm},SampleFromPrior},
    state,
    chain_type::Type{Vector{NamedTuple}};
    kwargs...
)
    return map(ts) do t
        params = map(first, getparams(t))
        return merge(params, metadata(t))
    end
end

function save(c::MCMCChains.Chains, spl::Sampler, model, vi, samples)
    nt = NamedTuple{(:sampler, :model, :vi, :samples)}((spl, model, deepcopy(vi), samples))
    return setinfo(c, merge(nt, c.info))
end

function resume(chain::MCMCChains.Chains, args...; kwargs...)
    return resume(Random.GLOBAL_RNG, chain, args...; kwargs...)
end

function resume(rng::Random.AbstractRNG, chain::MCMCChains.Chains, args...;
                progress=PROGRESS[], kwargs...)
    isempty(chain.info) && error("[Turing] cannot resume from a chain without state info")

    # Sample a new chain.
    return AbstractMCMC.mcmcsample(
        rng,
        chain.info[:model],
        chain.info[:sampler],
        args...;
        resume_from = chain,
        chain_type = MCMCChains.Chains,
        progress = progress,
        kwargs...
    )
end

DynamicPPL.loadstate(chain::MCMCChains.Chains) = chain.info[:samplerstate]

#######################################
# Concrete algorithm implementations. #
#######################################

include("ess.jl")
include("hmc.jl")
include("mh.jl")
include("is.jl")
include("AdvancedSMC.jl")
include("gibbs.jl")
include("gibbs_conditional.jl")
include("../contrib/inference/sghmc.jl")
include("emcee.jl")

################
# Typing tools #
################

for alg in (:SMC, :PG, :MH, :IS, :ESS, :Gibbs, :Emcee)
    @eval DynamicPPL.getspace(::$alg{space}) where {space} = space
end
for alg in (:HMC, :HMCDA, :NUTS, :SGLD, :SGHMC)
    @eval DynamicPPL.getspace(::$alg{<:Any, space}) where {space} = space
end

function get_matching_type(
    spl::Sampler{<:Union{PG, SMC}},
    vi,
    ::Type{TV},
) where {T, N, TV <: Array{T, N}}
    return TArray{T, N}
end

##############
# Utilities  #
##############

DynamicPPL.getspace(spl::Sampler) = getspace(spl.alg)
DynamicPPL.inspace(vn::VarName, spl::Sampler) = inspace(vn, getspace(spl.alg))

"""

    predict([rng::AbstractRNG,] model::Model, chain::MCMCChains.Chains; include_all=false)

Execute `model` conditioned on each sample in `chain`, and return the resulting `Chains`.

If `include_all` is `false`, the returned `Chains` will contain only those variables
sampled/not present in `chain`.

# Details
Internally calls `Turing.Inference.transitions_from_chain` to obtained the samples
and then converts these into a `Chains` object using `AbstractMCMC.bundle_samples`.

# Example
```jldoctest
julia> using Turing; Turing.turnprogress(false);
[ Info: [Turing]: progress logging is disabled globally

julia> @model function linear_reg(x, y, σ = 0.1)
           β ~ Normal(0, 1)

           for i ∈ eachindex(y)
               y[i] ~ Normal(β * x[i], σ)
           end
       end;

julia> σ = 0.1; f(x) = 2 * x + 0.1 * randn();

julia> Δ = 0.1; xs_train = 0:Δ:10; ys_train = f.(xs_train);

julia> xs_test = [10 + Δ, 10 + 2 * Δ]; ys_test = f.(xs_test);

julia> m_train = linear_reg(xs_train, ys_train, σ);

julia> chain_lin_reg = sample(m_train, NUTS(100, 0.65), 200);
┌ Info: Found initial step size
└   ϵ = 0.003125

julia> m_test = linear_reg(xs_test, Vector{Union{Missing, Float64}}(undef, length(ys_test)), σ);

julia> predictions = predict(m_test, chain_lin_reg)
Object of type Chains, with data of type 100×2×1 Array{Float64,3}

Iterations        = 1:100
Thinning interval = 1
Chains            = 1
Samples per chain = 100
parameters        = y[1], y[2]

2-element Array{ChainDataFrame,1}

Summary Statistics
  parameters     mean     std  naive_se     mcse       ess   r_hat
  ──────────  ───────  ──────  ────────  ───────  ────────  ──────
        y[1]  20.1974  0.1007    0.0101  missing  101.0711  0.9922
        y[2]  20.3867  0.1062    0.0106  missing  101.4889  0.9903

Quantiles
  parameters     2.5%    25.0%    50.0%    75.0%    97.5%
  ──────────  ───────  ───────  ───────  ───────  ───────
        y[1]  20.0342  20.1188  20.2135  20.2588  20.4188
        y[2]  20.1870  20.3178  20.3839  20.4466  20.5895

julia> ys_pred = vec(mean(Array(group(predictions, :y)); dims = 1));

julia> sum(abs2, ys_test - ys_pred) ≤ 0.1
true
```
"""
function predict(model::Model, chain::MCMCChains.Chains; kwargs...)
    return predict(Random.GLOBAL_RNG, model, chain; kwargs...)
end
function predict(rng::AbstractRNG, model::Model, chain::MCMCChains.Chains; include_all = false)
    spl = DynamicPPL.SampleFromPrior()

    # Sample transitions using `spl` conditioned on values in `chain`
    transitions = [
        transitions_from_chain(rng, model, chain[:, :, chn_idx]; sampler = spl)
        for chn_idx = 1:size(chain, 3)
    ]

    # Let the Turing internals handle everything else for you
    chain_result = reduce(
        MCMCChains.chainscat, [
            AbstractMCMC.bundle_samples(
                transitions[chn_idx],
                model,
                spl,
                nothing,
                MCMCChains.Chains
            ) for chn_idx = 1:size(chain, 3)
        ]
    )

    parameter_names = if include_all
        names(chain_result, :parameters)
    else
        filter(k -> ∉(k, names(chain, :parameters)), names(chain_result, :parameters))
    end

    return chain_result[parameter_names]
end

"""

    transitions_from_chain(
        [rng::AbstractRNG,]
        model::Model,
        chain::MCMCChains.Chains;
        sampler = DynamicPPL.SampleFromPrior()
    )

Execute `model` conditioned on each sample in `chain`, and return resulting transitions.

The returned transitions are represented in a `Vector{<:Turing.Inference.Transition}`.

# Details

In a bit more detail, the process is as follows:
1. For every `sample` in `chain`
   1. For every `variable` in `sample`
      1. Set `variable` in `model` to its value in `sample`
   2. Execute `model` with variables fixed as above, sampling variables NOT present
      in `chain` using `SampleFromPrior`
   3. Return sampled variables and log-joint

# Example
```julia-repl
julia> using Turing

julia> @model function demo()
           m ~ Normal(0, 1)
           x ~ Normal(m, 1)
       end;

julia> m = demo();

julia> chain = Chains(randn(2, 1, 1), ["m"]); # 2 samples of `m`

julia> transitions = Turing.Inference.transitions_from_chain(m, chain);

julia> [Turing.Inference.getlogp(t) for t in transitions] # extract the logjoints
2-element Array{Float64,1}:
 -3.6294991938628374
 -2.5697948166987845

julia> [first(t.θ.x) for t in transitions] # extract samples for `x`
2-element Array{Array{Float64,1},1}:
 [-2.0844148956440796]
 [-1.704630494695469]
```
"""
function transitions_from_chain(
    model::Turing.Model,
    chain::MCMCChains.Chains;
    kwargs...
)
    return transitions_from_chain(Random.GLOBAL_RNG, model, chain; kwargs...)
end
function transitions_from_chain(
    rng::AbstractRNG,
    model::Turing.Model,
    chain::MCMCChains.Chains;
    sampler = DynamicPPL.SampleFromPrior()
)
    vi = Turing.VarInfo(model)

    transitions = map(1:length(chain)) do i
        c = chain[i]
        md = vi.metadata
        for v in keys(md)
            for vn in md[v].vns
                vn_sym = Symbol(vn)

                # Cannot use `vn_sym` to index in the chain
                # so we have to extract the corresponding "linear"
                # indices and use those.
                # `ks` is empty if `vn_sym` not in `c`.
                ks = MCMCChains.namesingroup(c, vn_sym)

                if !isempty(ks)
                    # 1st dimension is of size 1 since `c`
                    # only contains a single sample, and the
                    # last dimension is of size 1 since
                    # we're assuming we're working with a single chain.
                    val = copy(vec(c[ks].value))
                    DynamicPPL.setval!(vi, val, vn)
                    DynamicPPL.settrans!(vi, false, vn)
                else
                    DynamicPPL.set_flag!(vi, vn, "del")
                end
            end
        end
        # Execute `model` on the parameters set in `vi` and sample those with `"del"` flag using `sampler`
        model(rng, vi, sampler)

        # Convert `VarInfo` into `NamedTuple` and save
        theta = DynamicPPL.tonamedtuple(vi)
        lp = Turing.getlogp(vi)
        Transition(theta, lp)
    end

    return transitions
end

end # module