text:
using Test
using QuantumOptics
using LinearAlgebra

@testset "stochastic_master" begin

b_spin = SpinBasis(1//2)
sz = sigmaz(b_spin)
sm = sigmam(b_spin)
sp = sigmap(b_spin)
zero_op = 0*sz
γ = 0.1
noise_op = 0.5γ*sz

H = γ*(sp + sm)
Hs = [noise_op]

ψ0 = spindown(b_spin)

ρ0 = dm(ψ0)
rates = [0.1]
J = [sm]
C = [sm]
Jdagger = dagger.(J)
C .*= rates
Cdagger = dagger.(C)

dt = 1/30.0
T = [0:0.1:1;]
T_short = [0:dt:dt;]

# Function definitions for master_dynamic
function fdeterm_master(t, rho)
    H, J, Jdagger
end
function fstoch_master(t, rho)
    C, Cdagger
end

# Test master
tout, ρt_det = timeevolution.master(T, ψ0, H, J; rates=rates)
tout, ρt1 = stochastic.master(T, ψ0, H, J, 0 .*J; rates=rates, dt=dt)
tout, ρt2 = stochastic.master(T, ρ0, LazyProduct(H, one(H)), sqrt.(rates).*J, 0 .* J; dt=dt)
tout, ρt3 = stochastic.master_dynamic(T, ρ0, fdeterm_master, fstoch_master; rates=rates, dt=dt)
for i=1:length(tout)
    @test tracedistance(ρt1[i], ρt_det[i]) < dt
    @test tracedistance(ρt2[i], ρt_det[i]) < dt
    @test tracedistance(ρt3[i], ρt_det[i]) < dt
end

# Test remaining function calls for short times to test whether they work in principle
rates_mat = [0.1 0.05; 0.05 0.1]
tout, ρt = stochastic.master(T_short, ψ0, H, [sm, sm], [sm, sm]; rates=rates_mat, dt=dt)

# Test master dynamic
tout, ρt = stochastic.master_dynamic(T_short, ψ0, fdeterm_master, fstoch_master; noise_processes=1, dt=dt)
tout, ρt = stochastic.master_dynamic(T_short, ρ0, fdeterm_master, fstoch_master, dt=dt)

end # testset


text:
using AxisKeys
using PrettyTables

export AbstractCVMethod, AbstractCV

const POINTWISE_LABELS = (:cv_elpd, :naive_lpd, :p_eff, :ess, :pareto_k)
const CV_DESC = """
# Fields

  - `estimates::KeyedArray`: A KeyedArray with columns `:total, :se_total, :mean, :se_mean`,
    and rows `:cv_elpd, :naive_lpd, :p_eff`. See `# Extended help` for more.
      - `:cv_elpd` contains estimates for the out-of-sample prediction error, as
        estimated using leave-one-out cross validation.
      - `:naive_lpd` contains estimates of the in-sample prediction error.
      - `:p_eff` is the effective number of parameters -- a model with a `p_eff` of 2 is
        "about as overfit" as a model with 2 parameters and no regularization.
  - `pointwise::KeyedArray`: A `KeyedArray` of pointwise estimates with 5 columns --
      - `:cv_elpd` contains the estimated out-of-sample error for this point, as measured
      using leave-one-out cross validation.
      - `:naive_lpd` contains the in-sample estimate of error for this point.
      - `:p_eff` is the difference in the two previous estimates.
      - `:ess` is the L2 effective sample size, which estimates the simulation error caused
        by using Monte Carlo estimates. It does not measure model performance.
      - `:inf_ess` is the supremum-based effective sample size, which estimates the
        simulation error caused by using Monte Carlo estimates. It is more robust than
        `:ess` and should therefore be preferred. It does not measure model performance.
      - `:pareto_k` is the estimated value for the parameter `ξ` of the generalized Pareto
        distribution. Values above .7 indicate that PSIS has failed to approximate the true
        distribution.
  - `psis_object::Psis`: A `Psis` object containing the results of Pareto-smoothed
    importance sampling.
  - `gmpd`: The geometric mean of the predictive density. It is defined as the geometric
    mean of the probability assigned to each data point by the model, i.e. `exp(cv_avg)`.
    This measure is only interpretable for classifiers (variables with discrete outcomes).
    We can think of it as measuring how often the model was right: A model that always
    predicts incorrectly will have a GMPD of 0, while a model that always predicts
    correctly will have a GMPD of 1. However, the GMPD gives a model "Partial points"
    between 0 and 1 whenever the model assigns a probability other than 0 or 1 to the
    outcome that actually happened, making it a fully Bayesian measure of model quality.
  - `mcse`: A float containing the estimated Monte Carlo standard error for the total
    cross-validation estimate.

# Extended help

The total score depends on the sample size, and summarizes the weight of evidence for or
against a model. Total scores are on an interval scale, meaning that only differences of
scores are meaningful. *It is not possible to interpret a total score by looking at it.*
The total score is not a goodness-of-fit statistic (for this, see the average score).

The average score is the total score, divided by the sample size. It estimates the expected
log score, i.e. the expectation of the log probability density of observing the next point.
The average score is a relative goodness-of-fit statistic which does not depend on sample
size.

Unlike for chi-square goodness of fit tests, models do not have to be nested for model
comparison using cross-validation methods.
"""

##########################
#### CROSS VALIDATION ####
##########################

"""
    AbstractCV

An abstract type used in cross-validation.
"""
abstract type AbstractCV end

"""
    AbstractCVMethod

An abstract type used to dispatch the correct method for cross validation.
"""
abstract type AbstractCVMethod end

##########################
### BAYESIAN BOOTSTRAP ###
##########################

# """
#     BayesCV{
#         F <: Real,
#         AF <: AbstractArray{F},
#         VF <: AbstractVector{F},
#         I <: Integer,
#         VI <: AbstractVector{I},
#     } <: AbstractCV

# A struct containing the results of cross-validation using the Bayesian bootstrap.

# $CV_DESC

# See also: [`bayes_cv`]@ref, [`psis_loo`]@ref, [`psis`]@ref, [`Psis`]@ref
# """
# struct BayesCV{
#     F <: Real,
#     AF <: AbstractArray{F},
#     VF <: AbstractVector{F},
#     I <: Integer,
#     VI <: AbstractVector{I},
# } <: AbstractCV
#     estimates::KeyedArray
#     posteriors::KeyedArray
#     psis_object::Psis{F, AF, VF, I, VI}
# end

# function Base.show(io::IO, ::MIME"text/plain", cv_object::BayesCV)
#     table = cv_object.estimates
#     post_samples = cv_object.psis_object.posterior_sample_size
#     data_size = cv_object.psis_object.data_size
#     _throw_pareto_k_warning(cv_object.resamples(:pareto_k))
#     println(
#         "Results of Bayesian bootstrap CV with $post_samples Monte Carlo samples and " *
#         "$data_size data points.",
#     )
#     return pretty_table(
#         table;
#         compact_printing=false,
#         header=table.column,
#         row_names=table.statistic,
#         formatters=ft_printf("%5.2f"),
#         alignment=:r,
#     )
# end

ext: jl


text:
import Base.getindex, Base.setindex!, Base.length, Base.endof, Base.copy

type Path
  tg::TimeGrid
  values::Vector{Float64}
  # TODO ensure values size is the same as tg size
end

Path(tg::TimeGrid) = Path(tg, Vector{Float64}(length(tg.times)))

getindex(p::Path, i::Int) = p.values[i]
setindex!(p::Path, x::Float64, i::Int) = p.values[i] = x
length(p::Path) = length(p.tg.times)
endof(p::Path) = endof(p.values)

copy(p::Path) = Path(p.tg, copy(p.values))


text:
#### Cauchy matrix

export Cauchy
"""
[`Cauchy` matrix](http://en.wikipedia.org/wiki/Cauchy_matrix)
```julia
Cauchy(x,y)[i,j]=1/(x[i]+y[j])
Cauchy(x)=Cauchy(x,x)
cauchy(k::Number)=Cauchy(collect(1:k))

julia> Cauchy([1,2,3],[3,4,5])
3x3 Cauchy{Int64}:
 0.25      0.2       0.166667
 0.2       0.166667  0.142857
 0.166667  0.142857  0.125

julia> Cauchy([1,2,3])
3x3 Cauchy{Int64}:
 0.5       0.333333  0.25
 0.333333  0.25      0.2
 0.25      0.2       0.166667

julia> Cauchy(pi)
3x3 Cauchy{Float64}:
 0.5       0.333333  0.25
 0.333333  0.25      0.2
 0.25      0.2       0.166667
```
"""
struct Cauchy{T} <: AbstractMatrix{T}
    x::Vector{T} #
    y::Vector{T} #
end # immutable

function Cauchy(k::Number)
         Cauchy(collect(1:k),collect(1:k))
end

function Cauchy(x::Vector)
         Cauchy(x,x)
end

# Define its size

size(A::Cauchy, dim::Integer) = dim==1 ? length(A.x) : dim==2 ? length(A.y) : throw(ArgumentError("Invalid dimension $dim"))
size(A::Cauchy)= length(A.x), length(A.y)

# Index into a Cauchy
function getindex(A::Cauchy,i::Integer,j::Integer)
    return 1.0/(A.x[i]+A.y[j])
end # getindex

# Dense version of Cauchy
Matrix(A::Cauchy) = [A[i,j] for i=1:size(A,1), j=1:size(A,2)]


text:
module VertexStaticGraphs

import LightGraphs
const LG = LightGraphs

import LightGraphs.SimpleGraphs: badj
import LightGraphs.SimpleGraphs: fadj

using StaticArrays
using StaticArrays: SVector, @SVector

export VStaticGraph

include("vstaticgraph.jl")
include("lg_interface.jl")
include("simplegraph_interface.jl")

end # module


text:
module RestAPI

greet() = print("Hello World!")
export greet, Client

include("./models/datamodel.jl")
using .DataModel

include("./provider/connection.jl")
using .ConnectionPools

include("./provider/workers.jl")
using .Workers

include("./provider/auth.jl")
using .Auth

include("./provider/context.jl")
using .Contexts

include("./provider/schema.jl")
using .Schema

include("./blocs/service.jl")
using .Service

include("./router/router.jl")
using .Router

include("./client/clients.jl")
using .Client

# function init(dbfile, authkeysfile)
#     Workers.init()
#     Schema.init(dbfile)
#     Auth.init(authkeysfile)
#     Router.run()
# end

function init(dbfile)
    Workers.init()
    Schema.init(dbfile)
    # Auth.init(authkeysfile)
    Router.run()
end

end # module


text:
# used for debugging
function printtrace()
    show(stdout, MIME"text/plain"(), stacktrace()[2:5])
    println()
end

"""
Find the actual function call
    return example :(:call, :functionName, :arg1, :arg2, ...)
"""
function findFnCall(funExpr)
    funCall=funExpr
    if funExpr.head == :(=) || funExpr.head == :function
        funCall = funExpr.args[1]
    end
    if isa(funCall, Expr) && funCall.head == :where
        funCall = funCall.args[1]
    end
    if isa(funCall, Expr) && funCall.head == :call
        return funCall
    else
        dump(funCall)
        error("getFnCall: Could not detect function call")
    end
end

""" Find the symbol of the first argument to the function """
function findFnSelfArgNameSymbol(fun, ClsName=:Any)
    funCall = findFnCall(fun)
    self = funCall.args[2] # first arg
    if isa(self, Expr) && self.head == :parameters
        # if there is kw args, it goes first..
        self = funCall.args[3]
    end

    if isa(self, Expr)
        # check if first arg is of correct type
        if ClsName!=:Any && self.args[2] != :($ClsName)
            error("$(self.args[1]) can only be declared of type '$ClsName', '$(self.args[2])' found!")
        end
        self = self.args[1]
    end
    return self
end

"""
Overwrite the first arg with
selfArgExpr eg. :(self::MyClass)
"""
function setFnSelf!(funExpr, selfArgExpr)
    funCall = findFnCall(funExpr)
    if isa(funCall.args[2], Expr) && funCall.args[2].head == :parameters
        funCall.args[3] = selfArgExpr
    else
        funCall.args[2] = selfArgExpr
    end
end
"""
Delete the first argument
"""
function deleteFnSelf!(fun)
    funCall = findFnCall(fun)
    self = funCall.args[2]
    if isa(self, Expr) && self.head == :parameters
        deleteat!(funCall.args, 3)
    else
        deleteat!(funCall.args, 2)
    end
end

# No usages of this, so not bothering to support and thus add tests for this..
# function getFnParam(fun)
#     funCall = findFnCall(fun)
#     return funCall.args[2:end]
# end

""" Get the name of the function duh """
function getFnName(fun; withoutGeneric=false)
    funCall = findFnCall(fun)
    name = funCall.args[1]

    if withoutGeneric
        if isa(name, Symbol)
            return name
        else
            return name.args[1]
        end
    else
        if isa(name, Symbol)
            # maybe this should be :where
            return Expr(:curly, name)
        else
            return name
        end
    end
end

""" Set the name of the function duh """
function setFnName!(fun, name; withoutGeneric=false)
    if withoutGeneric
        funCall = findFnCall(fun)
        if isa(funCall.args[1], Symbol)
            funCall.args[1] = name
        else
            funCall.args[1].args[1] = name
        end
    else
        funCall = findFnCall(fun)
        funCall.args[1] = name
    end
end

"""
Set the type of the first argument
"""
function setFnSelfArgType!(fun, ClsName)
    selfArgNameSymbol = findFnSelfArgNameSymbol(fun)
    setFnSelf!(fun, :($selfArgNameSymbol::$ClsName))
end


text:
# Copyright © 2021 Alexander L. Hayes
# Copyright © 2019 John Tinnerhom

#=
MIT License

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
=#

module ExportPublic
    macro exportPublic()
        inI = []
        out = []
        res = []

        for name in Base.names(__module__, all = true)
            push!(out, :(export $(name)))
        end

        for name in Base.names(__module__, all = false, imported = true)
            push!(inI, :(export $(name)))
        end

        for name in Base.names(Main.Base, all = false, imported = true)
            push!(inI, :(export $(name)))
        end

        for name in Base.names(Main.Core, all = false, imported = true)
            push!(inI, :(export $(name)))
        end

        for e in out
            if e in inI || e == :(export include) || '#' in string(e) || "_" == SubString(string(e), 8, 8)
                continue
            end
            push!(res, :($(e)))
        end

        # Do not export things that we already have
        ret = Expr(:block, res...)
        return ret
    end

    @exportPublic()
end


text:
export pairwise_fst

# generate a wrapper struct so we can nicely print the results
struct PairwiseFST
    results::DataFrame
    method::String
end

# pretty-printing of FST results
function Base.show(io::IO, data::PairwiseFST)
    show(
        io,
        data.results,
        show_row_number=false,
        rowlabel = Symbol(" "),
        eltypes = false,
        row_names = names(data.results),
        title = "Pairwise FST: " * data.method
    )
end

function FST_global(data::AbstractDataFrame)
    # observed/expected het per locus per pop
    het_df = DataFrames.combine(
        groupby(data, [:locus, :population]),
        :genotype => nonmissing => :n,
        :genotype => hetero_o => :het_obs,
        :genotype => hetero_e => :het_exp,
        :genotype => allele_freq => :alleles
    )
    # collapse down to retrieve averages and counts
    n_df = DataFrames.combine(
        groupby(het_df, :locus),
        :n => count_nonzeros => :count,
        :n => (n -> count_nonzeros(n) / reciprocal_sum(n)) => :mn,
        [:het_obs, :het_exp, :n] => ((o,e,n) -> mean(skipmissing(gene_diversity_nei87.(e, o, count_nonzeros(n) / reciprocal_sum(n))))) => :HS,
        :het_obs => (o -> mean(skipmissing(o)))=> :Het_obs,
        :alleles => (alleles ->  sum(values(avg_allele_freq(alleles, 2))))=> :avg_freq
        )

    Ht = @inbounds 1.0 .- n_df.avg_freq .+ (n_df.HS ./ n_df.mn ./ n_df.count) - (n_df.Het_obs ./ 2.0 ./ n_df.mn ./ n_df.count)
    DST = @inbounds Ht .- n_df.HS
    DST′ = @inbounds n_df.count ./ (n_df.count .- 1) .* DST
    HT′ = @inbounds n_df.HS .+ DST′

    round(safemean(DST′) / safemean(HT′), digits = 4)
end

function _pairwise_Nei(data::PopData)
    idx_pdata = groupby(data.loci, :population)
    pops = getindex.(keys(idx_pdata), :population)
    npops = length(idx_pdata)
    results = zeros(npops, npops)
    for i in 2:npops
        for j in 1:(i-1)
            results[i,j] = FST_global(DataFrame(idx_pdata[[j,i]]))
        end
    end
    return PairwiseFST(DataFrame(results, Symbol.(pops)), "Nei 1987")
end

function _pairwise_WeirCockerham(data::PopData)
    idx_pdata = groupby(data.loci, :population)
    pops = getindex.(keys(idx_pdata), :population)
    npops = length(idx_pdata)
    results = zeros(npops, npops)
    for i in 2:npops
        for j in 1:(i-1)
            results[i,j] = weircockerham1984(DataFrame(idx_pdata[[j,i]]))
        end
    end
    return PairwiseFST(DataFrame(results, Symbol.(pops)), "Weir-Cockerham")
end

#TODO add methods and complete docstring
"""
    pairwise_fst(data::PopData; method::String)
Calculate pairwise FST between populations in a `PopData` object.
#### Methods:
- `"Nei87"`: Nei (1987) method
- `"WC84"` : Weir-Cockerham (1984) method (default)
"""
function pairwise_fst(data::PopData; method::String = "WC")
    occursin("nei", lowercase(method)) && return _pairwise_Nei(data)
    occursin("wc", lowercase(method)) && return _pairwise_WeirCockerham(data)
end

function FST_global(data::PopData)
    # observed/expected het per locus per pop
    het_df = DataFrames.combine(
        groupby(data, [:locus, :population]),
        :genotype => nonmissing => :n,
        :genotype => hetero_o => :het_obs,
        :genotype => (i -> hetero_e(i)) => :het_exp,
        :genotype => allele_freq => :alleles
    )
    # collapse down to retrieve averages and counts
    n_df = DataFrames.combine(
        groupby(het_df, :locus),
        :n => (n -> sum(.!iszero.(n)))=> :count,
        :n => (n -> sum(.!iszero.(n)) ./ sum(reciprocal.(n))) => :mn,
        [:het_obs, :het_exp, :n] => ((o,e,n) -> mean(skipmissing(gene_diversity_nei87.(e,o,sum(.!iszero.(n)) ./ sum(reciprocal.(n)))))) => :HS,
        :het_obs => (o -> mean(skipmissing(o)))=> :Het_obs,
        :alleles => (alleles ->  sum(values(avg_allele_freq(alleles, 2))))=> :avg_freq
        )

    Ht = @inbounds 1.0 .- n_df.avg_freq .+ (n_df.HS ./ n_df.mn ./ n_df.count) - (n_df.Het_obs ./ 2.0 ./ n_df.mn ./ n_df.count)
    DST = @inbounds Ht .- n_df.HS
    DST′ = @inbounds n_df.count ./ (n_df.count .- 1) .* DST
    HT′ = @inbounds n_df.HS .+ DST′

    HT = safemean(HT)
    DST = safemean(DST)
    DST′ = safemean(DST′)
    HT′ = safemean(HT′)

    return Dict{Symbol, Float64}(
        :FST => DST / HT,
        :FST′ => DST′ / HT′
    )
end

function FST_locus(data::PopData)
    # observed/expected het per locus per pop
    het_df = DataFrames.combine(
        groupby(data.loci, [:locus, :population]),
        :genotype => nonmissing => :n,
        :genotype => hetero_o => :het_obs,
        :genotype => (i -> hetero_e(i)) => :het_exp,
        :genotype => allele_freq => :alleles
    )
    # collapse down to retrieve averages and counts
    n_df = DataFrames.combine(
        [:het_obs, :het_exp, :n, :alleles] => (o,e,n,alleles) -> (
            count = sum(.!iszero.(n)),
            mn = sum(.!iszero.(n)) ./ sum(reciprocal.(n)),
            HS = mean(skipmissing(gene_diversity_nei87.(e,o,sum(.!iszero.(n)) ./ sum(reciprocal.(n))))),
            Het_obs = mean(skipmissing(o)),
            avg_freq = sum(values(avg_allele_freq(alleles)).^2)
            ),
        groupby(het_df, :locus)
    )

    HT = 1.0 .- n_df.avg_freq .+ (n_df.HS ./ n_df.mn ./ n_df.count) - (n_df.Het_obs ./ 2.0 ./ n_df.mn ./ n_df.count)
    DST = HT .- n_df.HS
    DST′ = n_df.count ./ (n_df.count .- 1) .* DST
    HT′ = n_df.HS .+ DST′

    return Dict{Symbol, Vector{Float64}}(
        :FST => DST ./ HT,
        :FST′ => DST′ ./ HT′
        )
end

function f_stat_p(data::PopData; nperm::Int = 1000)
    observed = FST_global(data)
    popnames = data.meta.population
    perm_dict = Dict{Symbol,Vector{Float64}}(
        :FST => Vector{Float64}(undef, nperm-1),
        :FST′ => Vector{Float64}(undef, nperm-1)
        )
    p = Progress(nperm-1, dt = 1, color = :blue)
    @inbounds @sync for n in 1:nperm-1
        Base.Threads.@spawn begin
            tmp = FST_global(permute_samples!(copy(data.loci), popnames))
            perm_dict[:FST][n] = tmp[:FST]
            perm_dict[:FST′][n] = tmp[:FST′]
            next!(p)
        end
    end
    @inbounds for (k,v) in perm_dict
        observed[k] = (1 + sum(observed[k] .<= v))/nperm
    end

    return observed
end


text:
@testset "jacobi correct values edge,origin" begin
    # this only has to go as high as 3 to show that a three term recurrence
    # relation behaves correctly, but as a run for record we go a bit higher
    for n in 1:10
        # design of jacobi polynomials is such that edge = 1
        @test jacobi(n, 0., 0., [1.]) ≈ [1.]
        correct = [-1.]
        if iseven(n)
            correct = [1.]
        end
        @test jacobi(n, 0., 0., [-1.]) ≈ correct
        if !iseven(n)
            @test jacobi(n, 0., 0., [0.]) ≈ [0.]
        end

        @test jacobi_weight.(0., 0., [-1., 0., 1.]) ≈ [1., 1., 1.]
        @test jacobi(0, 0., 0., [-1., 0., 1.]) ≈ [1., 1., 1.]
    end
end

@testset "jacobi sum correct values" begin
    terms = [1,3,5,7,9]
    weights = ones(size(terms)...)
    @test jacobi_sum(terms, weights, 0., 0., [0.]) ≈ [0.]
    terms = [1,2,3,4,5]
    weights = ones(size(terms)...)
    # -1 accumulates 3 times and +1 two times on left edge => expect -1.
    # 1 accumulates 5 times on right edge => expect 5.
    @test jacobi_sum(terms, weights, 0., 0., [-1., 1.]) ≈ [-1., 5.]
end

@testset "jacobi series correct values contiguous" begin
    terms = [1,2,3,4,5]
    ledge = [-1.,1.,-1.,1.,-1.]
    center = [0.,-0.5,0.,0.375,0.]
    redge = [1.,1.,1.,1.,1.]
    series = jacobi_series(terms, 0., 0., [-1., 0., 1.])
    @test series[1,:] ≈ ledge
    @test series[2,:] ≈ center
    @test series[3,:] ≈ redge
end

@testset "jacobi series correct values sparse" begin
    terms = [1,3,5]
    ledge = [-1.,-1.,-1.]
    center = [0.,0.,0.]
    redge = [1.,1.,1.]
    series = jacobi_series(terms, 0., 0., [-1., 0., 1.])
    @test series[1,:] ≈ ledge
    @test series[2,:] ≈ center
    @test series[3,:] ≈ redge
end

@testset "jacobi series doesn't break on edge cases" begin
    @test jacobi_series([0], 0., 0., [-1., 0., 1.]) ≈ [1., 1., 1.]
    @test jacobi_series([1, 2], 0., 0., [-1., 0., 1.]) ≈ [-1. 0. 1.; 1. -0.5 1.]'
    @test jacobi_series([0, 1], 0., 0., [-1., 0., 1.]) ≈ [1. 1. 1.; -1. 0. 1.]';
    @test jacobi_series([2], 0., 0., [0.]) ≈ [-0.5]
end


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

JLLWrappers.@generate_main_file_header("libgraphicsmagic")
JLLWrappers.@generate_main_file("libgraphicsmagic", UUID("3e975b7b-ab84-5bf4-b4b6-586754a53ef6"))
end  # module libgraphicsmagic_jll


text:
######################################################################
# Utility functions & macros used throughout the engine.
# -----
# Licensed under Apache License 2.0
export randport
randport(from::Integer, to::Integer = typemax(UInt16)) = randport(UInt16(from), UInt16(to))
randport(from::UInt16,  to::UInt16  = typemax(UInt16)) = floor(UInt16, rand() * (to - from) + from)

Base.isempty(buffer::IOBuffer) = buffer.ptr > buffer.size


text:
using Yota
using Zygote
using Zygote: @adjoint
# include("../src/core.jl")
using BenchmarkTools
using Distributions
import NNlib

logistic(x) = 1 / (1 + exp(-x))
softplus(x) = log(exp(x) + 1)

# Yota diff rules
@diffrule logistic(x::Number) x (logistic(x) * (1 - logistic(x)) * dy)
@diffrule softplus(x::Number) x logistic(x) * dy

# Zygote diff rules
@adjoint logistic(x) = logistic(x), Δ -> ((logistic(x) * (1 - logistic(x)) * Δ),)
@adjoint softplus(x) = softplus(x), Δ -> ((logistic(x) * Δ),)

function perf_test(f, args...)
    println("Compiling derivatives using Yota")
    @time Yota.grad(f, args...)
    r1 = @benchmark Yota.grad($f, $(args)...)
    show(stdout, MIME{Symbol("text/plain")}(), r1)
    println("\n")

    println("Compiling derivatives using Zygote")
    @time Zygote.gradient(f, args...)
    r2 = @benchmark Zygote.gradient($f, $(args)...)
    show(stdout, MIME{Symbol("text/plain")}(), r2)
    println("\n")
end

# VAE benchmark

function xavier_init(dim_in, dim_out; c=1)
    low = -c * sqrt(6.0 / (dim_in + dim_out))
    high = c * sqrt(6.0 / (dim_in + dim_out))
    return rand(Uniform(low, high), dim_in, dim_out)
end

# see examples/vae.jl for a more convenient API
function encode(We1, be1, We2, be2, We3, be3, We4, be4, x)
    he1 = softplus.(We1 * x .+ be1)
    he2 = softplus.(We2 * he1 .+ be2)
    mu = We3 * he2 .+ be3
    log_sigma2 = We4 * he2 .+ be4
    return mu, log_sigma2
end

function decode(Wd1, bd1, Wd2, bd2, Wd3, bd3, z)
    hd1 = softplus.(Wd1 * z .+ bd1)
    hd2 = softplus.(Wd2 * hd1 .+ bd2)
    x_rec = logistic.(Wd3 * hd2 .+ bd3)
    return x_rec
end

function vae_cost(We1, be1, We2, be2, We3, be3, We4, be4,
                  Wd1, bd1, Wd2, bd2, Wd3, bd3, eps, x)
    mu, log_sigma2 = encode(We1, be1, We2, be2, We3, be3, We4, be4, x)
    z = mu .+ sqrt.(exp.(log_sigma2)) .* eps
    x_rec = decode(Wd1, bd1, Wd2, bd2, Wd3, bd3, z)
    # loss
    rec_loss = -sum(x .* log.(1e-10 .+ x_rec) .+ (1 .- x) .* log.(1e-10 + 1.0 .- x_rec); dims=1)
    # KLD = -0.5 .* sum(1 .+ log_sigma2 .- mu .^ 2.0f0 - exp.(log_sigma2); dims=1)
    # cost = mean(rec_loss .+ KLD)
    cost = mean(rec_loss)
end

function benchmark_vae()
    n_inp, n_he1, n_he2, n_z, n_hd1, n_hd2, n_out = 784, 500, 500, 20, 500, 500, 784
    We1, be1, We2, be2, We3, be3, We4, be4 = (xavier_init(n_he1, n_inp),
                                              zeros(n_he1),
                                              xavier_init(n_he2, n_he1),
                                              zeros(n_he2),
                                              xavier_init(n_z, n_he2),
                                              zeros(n_z),
                                              xavier_init(n_z, n_he2),
                                              zeros(n_z))
    Wd1, bd1, Wd2, bd2, Wd3, bd3 = (xavier_init(n_hd1, n_z),
                                    zeros(n_hd1),
                                    xavier_init(n_hd2, n_hd1),
                                    zeros(n_hd2),
                                    xavier_init(n_out, n_hd2),
                                    zeros(n_out))
    eps = (rand(Normal(0, 1), size(We3, 1), 100))
    x = rand(784, 100)
    perf_test(vae_cost, We1, be1, We2, be2, We3, be3, We4, be4, Wd1, bd1, Wd2, bd2, Wd3, bd3, eps, x)
end


text:
module TestMeshInterop
using Test

using CombinatorialSpaces
using CombinatorialSpaces.MeshInterop

using GeometryBasics
using StaticArrays: SVector

const Point3D = SVector{3,Float64}

# Import Tooling
################

s = EmbeddedDeltaSet2D(joinpath(@__DIR__, "assets", "square.obj"))
@test s isa EmbeddedDeltaSet2D
@test triangle_vertices(s, 1) == [1,2,3]

sd = EmbeddedDeltaDualComplex2D{Bool, Float64, Point3D}(s)
subdivide_duals!(sd, Barycenter())

# Test consistency with conversion to/from mesh.
msh = Mesh(s)
s_msh = EmbeddedDeltaSet2D(msh)
@test point(s_msh) == point(s)
@test triangle_vertices(s_msh) == triangle_vertices(s)

end


text:
# MIT License
#
# Copyright (c) 2020 Joshua E Gentry

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
using Test

# -----------------------------------------------------------------------------------------
# Test the tess_set_param(inst, name, Int) method.
@testset "tess_set_param(inst, name, $type)" for type in [ Int8, Int16, Int32, Int64, Int]
    local inst = TessInst("eng", datadir)

    @test tess_set_param(inst, "editor_image_xpos", type(100)) == true
    @test tess_get_param(inst, "editor_image_xpos", type) == type(100)

    tess_delete!(inst)

    local err  = "Instance has been freed."
    @test (@test_logs (:error, err) tess_set_param(inst, "editor_image_xpos", type(100))) == false
end

# -----------------------------------------------------------------------------------------
# Test the tess_set_param(inst, name, Bool) method.
@testset "tess_set_param(inst, name, Bool)" begin
    local inst = TessInst("eng", datadir)

    @test tess_set_param(inst, "textord_noise_debug", true) == true
    @test tess_get_param(inst, "textord_noise_debug", Bool) == true

    tess_delete!(inst)

    local err  = "Instance has been freed."
    @test (@test_logs (:error, err) tess_set_param(inst, "textord_noise_debug", true)) == false
end

# -----------------------------------------------------------------------------------------
# Test the tess_set_param(inst, name, Float64) method.
@testset "tess_set_param(inst, name, Float64)" begin
    local inst = TessInst("eng", datadir)

    @test tess_set_param(inst, "classify_pico_feature_length", 0.025) == true
    @test tess_get_param(inst, "classify_pico_feature_length", Float64) == 0.025

    tess_delete!(inst)

    local err  = "Instance has been freed."
    @test (@test_logs (:error, err) tess_set_param(inst, "classify_pico_feature_length", 0.52)) == false
end

# -----------------------------------------------------------------------------------------
# Test the tess_set_param(inst, name, String) method.
@testset "tess_set_param(inst, name, String)" begin
    local inst = TessInst("eng", datadir)

    @test tess_set_param(inst, "page_separator", "\n* * *\n") == true
    @test tess_get_param(inst, "page_separator", String) == "\n* * *\n"

    tess_delete!(inst)

    local err  = "Instance has been freed."
    @test (@test_logs (:error, err) tess_set_param(inst, "tessedit_char_blacklist", "abc")) == false
end


text:
module MarkramTransmitter
export MarkramStates, Markramparams, fire

struct MarkramStates{T <: AbstractFloat}
    r::T
    w::T
end

struct Markramparams{T <: AbstractFloat}
    d::T
    f::T
    u::T
end

function fire(t, dt, states, params, updater, put_signal)
    rw = states.r * states.w
    updater(dt * (1 - states.r) / params.d - rw_exct,
            dt * (params.u - states.w) / params.f + params.u * (1 - states.w))
    put_signal(TimedDeltaCond(t, rw))
end

end # Module end


text:
"""
App level functionality -- loading and managing app-wide components like configs, models, initializers, etc.
"""
module App

import Revise
import Genie

const ASSET_FINGERPRINT = ""

### PRIVATE ###

"""
    bootstrap(context::Union{Module,Nothing} = nothing) :: Nothing

Kickstarts the loading of a Genie app by loading the environment settings.
(ja)
環境設定をロードすることで、Genieアプリのロードをキックスタートする。
"""
function bootstrap(context::Union{Module,Nothing} = Genie.default_context(context)) :: Nothing
  if haskey(ENV, "GENIE_ENV") && isfile(joinpath(Genie.config.path_env, ENV["GENIE_ENV"] * ".jl"))
    isfile(joinpath(Genie.config.path_env, Genie.GLOBAL_ENV_FILE_NAME)) && Base.include(context, joinpath(Genie.config.path_env, Genie.GLOBAL_ENV_FILE_NAME))
    isfile(joinpath(Genie.config.path_env, ENV["GENIE_ENV"] * ".jl")) && Base.include(context, joinpath(Genie.config.path_env, ENV["GENIE_ENV"] * ".jl"))
  else
    ENV["GENIE_ENV"] = Genie.Configuration.DEV
    Core.eval(context, Meta.parse("const config = Genie.Configuration.Settings(app_env = Configuration.DEV)"))
  end

  for f in fieldnames(typeof(context.config))
    setfield!(Genie.config, f, getfield(context.config, f))
  end

  printstyled("""

   _____         _
  |   __|___ ___|_|___
  |  |  | -_|   | | -_|
  |_____|___|_|_|_|___|

  """, color = :red, bold = true)

  printstyled("| Web: https://genieframework.com\n", color = :light_black, bold = true)
  printstyled("| GitHub: https://github.com/genieframework/Genie.jl\n", color = :light_black, bold = true)
  printstyled("| Docs: https://genieframework.github.io/Genie.jl/dev\n", color = :light_black, bold = true)
  printstyled("| Gitter: https://gitter.im/essenciary/Genie.jl\n", color = :light_black, bold = true)
  printstyled("| Twitter: https://twitter.com/GenieMVC\n\n", color = :light_black, bold = true)

  printstyled("Genie v$(Genie.Configuration.GENIE_VERSION)\n", color = :green, bold = true)
  printstyled("Active env: $(ENV["GENIE_ENV"] |> uppercase)\n\n", color = :light_blue, bold = true)

  nothing
end

end # module App


text:
# [test/runtests.jl]
using Nucleotide_Essentials
using Test

# Test scripts
include("readFastq_test.jl")
include("demultiplex_test.jl")


text:
# # Langmuir turbulence
#
# This script reproduces the Langmuir turbulence simulation reported in section 4 of
#
# [McWilliams, J. C. et al., "Langmuir Turbulence in the ocean," Journal of Fluid Mechanics (1997)](https://www.cambridge.org/core/journals/journal-of-fluid-mechanics/article/langmuir-turbulence-in-the-ocean/638FD0E368140E5972144348DB930A38).

using Pkg

Pkg.instantiate()

using Printf
using JLD2

## https://discourse.julialang.org/t/unable-to-display-plot-using-the-repl-gks-errors/12826/18
ENV["GKSwstype"] = "nul"
using Plots

using ArgParse

using Oceananigans
using Oceananigans.Units
using Oceananigans.BuoyancyModels: g_Earth

using LESbrary
using LESbrary.TurbulenceStatistics: TurbulentKineticEnergy, ViscousDissipation
using LESbrary.TurbulenceStatistics: first_through_second_order, turbulent_kinetic_energy_budget
using Oceanostics.TKEBudgetTerms: ZShearProduction

"Returns a dictionary of command line arguments."
function parse_command_line_arguments()
    settings = ArgParseSettings()

    @add_arg_table! settings begin
        "--Nh"
            help = "The number of grid points in x, y."
            default = 32
            arg_type = Int

        "--Nz"
            help = "The number of grid points in z."
            default = 32
            arg_type = Int

        "--Lh"
            help = "The physical extent of the domain in x, y."
            default = 128.0
            arg_type = Float64

        "--Lz"
            help = "The physical extent of the domain in z."
            default = 96.0
            arg_type = Float64

        "--stop-hours"
            help = "Number of hours to run the simulation for"
            default = 12.0
            arg_type = Float64

        "--advection-scheme"
            help = "Advection scheme: CenteredSecondOrder, WENO5, etc"
            default = :WENO5
            arg_type = Symbol

        "--pickup"
            help = "Whether or not to pick the simulation up from latest checkpoint"
            default = false
            arg_type = Bool
    end

    return parse_args(settings)
end

args = parse_command_line_arguments()

#####
##### Parameters
#####

Nh = args["Nh"]
Nz = args["Nz"]
Lh = args["Lh"]
Lz = args["Lz"]

f = 1e-4
Qᵘ = -3.72e-5 # m² s⁻²
Qᵇ = 2.307e-9 # m³ s⁻²
N² = 1.936e-5 # s⁻²

initial_mixed_layer_depth = 33 # m

inertial_period = 2π / f

stop_time = args["stop-hours"] * hours
snapshot_time_interval = 10minutes
averages_time_interval = 24hours
averages_time_window = inertial_period
averages_stride = 100
slice_depth = 8.0

amplitude = 0.8 # m

const wavenumber = 2π / 60 # m⁻¹
const Uˢ = amplitude^2 * wavenumber * sqrt(g_Earth * wavenumber) # m s⁻¹

uˢ(z) = Uˢ * exp(2wavenumber * z)
@inline ∂z_uˢ(z, t) = 2wavenumber * Uˢ * exp(2wavenumber * z)

#####
##### Grid
#####

grid = RectilinearGrid(CPU(), size=(Nh, Nh, Nz), extent=(Lh, Lh, Lz))

#####
##### Boundary conditions
#####

u_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(Qᵘ))

b_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(Qᵇ),
                                bottom = GradientBoundaryCondition(N²))

#####
##### Sponge layer
#####

gaussian_mask = GaussianMask{:z}(center=-grid.Lz, width=grid.Lz/10)

u_sponge = v_sponge = w_sponge = Relaxation(rate=1/hour, mask=gaussian_mask)

b_sponge = Relaxation(rate = 1/hour,
                      target = LinearTarget{:z}(intercept=0.0, gradient=N²),
                      mask = gaussian_mask)

#####
##### Model setup
#####

advection = eval(args["advection-scheme"])()

model = NonhydrostaticModel(; advection, grid,
                            timestepper = :RungeKutta3,
                            tracers = :b,
                            buoyancy = BuoyancyTracer(),
                            coriolis = FPlane(f=f),
                            closure = AnisotropicMinimumDissipation(),
                            stokes_drift = UniformStokesDrift(∂z_uˢ=∂z_uˢ),
                            boundary_conditions = (u=u_bcs, b=b_bcs),
                            forcing = (u=u_sponge, v=v_sponge, w=w_sponge, b=b_sponge))

#####
##### Initial conditions: Stokes drift + stratification + noise
#####

Ξ(z) = randn() * exp(z / 4)

stratification(z) = z < - initial_mixed_layer_depth ? N² * z : - N² * initial_mixed_layer_depth

bᵢ(x, y, z) = stratification(z)     + 1e-2 * Ξ(z) * N² * model.grid.Lz
uᵢ(x, y, z) = uˢ(z) + sqrt(abs(Qᵘ)) * 1e-2 * Ξ(z)
wᵢ(x, y, z) =         sqrt(abs(Qᵘ)) * 1e-2 * Ξ(z)

set!(model, u=uᵢ, w=wᵢ, b=bᵢ)

#####
##### Simulation setup
#####

simulation = Simulation(model; Δt=1.0, stop_time)

wizard = TimeStepWizard(cfl=1.0, max_change=1.1, max_Δt=30.0)
simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10))

wall_clock = [time_ns()]

function print_progress(simulation)
    model = simulation.model

    umax = maximum(abs, model.velocities.u)
    vmax = maximum(abs, model.velocities.v)
    wmax = maximum(abs, model.velocities.w)

    ## Print a progress message
    msg = @sprintf("i: %04d, t: %s, Δt: %s, umax = (%.1e, %.1e, %.1e) ms⁻¹, wall time: %s\n",
                   iteration(simulation),
                   prettytime(simulation),
                   prettytime(simulation.Δt),
                   umax, vmax, wmax,
                   prettytime(1e-9 * (time_ns() - wall_clock[1]))
                  )

    wall_clock[1] = time_ns()

    @info msg

    return nothing
end

simulation.callbacks[:progress] = Callback(print_progress, IterationInterval(10))

#####
##### Output setup
#####

prefix = @sprintf("langmuir_turbulence_Nx%d_Nz%d", grid.Nx, grid.Nz)
data_directory = joinpath(@__DIR__, "..", "data", prefix) # save data in /data/prefix

# "Primitive" statistics

b = BuoyancyField(model)
p = model.pressures.pHY′ + model.pressures.pNHS
w_scratch = ZFaceField(model.grid)
c_scratch = CenterField(model.grid)

primitive_statistics = first_through_second_order(model, b=b, p=p, w_scratch=w_scratch, c_scratch=c_scratch)

U = primitive_statistics[:u]
V = primitive_statistics[:v]

# Turbulent kinetic energy budget terms

e = TurbulentKineticEnergy(model, U=U, V=V)
shear_production = ZShearProduction(model, U=U, V=V)
dissipation = ViscousDissipation(model)

tke_budget_statistics = turbulent_kinetic_energy_budget(model, b=b, p=p, U=U, V=V, e=e,
                                                        shear_production=shear_production, dissipation=dissipation)

statistics_to_output = merge(primitive_statistics, tke_budget_statistics)

fields_to_output = merge(model.velocities, model.tracers,
                         (νₑ = model.diffusivity_fields.νₑ,
                           e = Field(e),
                          sp = Field(shear_production),
                           ϵ = Field(dissipation)))

# Output configured for pickup

pickup = args["pickup"]
force = pickup ? false : true

k_xy_slice = searchsortedfirst(znodes(Face, grid), -slice_depth)

simulation.output_writers[:checkpointer] =
    Checkpointer(model, schedule = TimeInterval(26hour), prefix = prefix * "_checkpointer", dir = data_directory)

simulation.output_writers[:xz] =
    JLD2OutputWriter(model, fields_to_output,
                         schedule = TimeInterval(snapshot_time_interval),
                           prefix = prefix * "_xz",
                     field_slicer = FieldSlicer(j=1),
                              dir = data_directory,
                            force = force)

simulation.output_writers[:yz] =
    JLD2OutputWriter(model, fields_to_output,
                         schedule = TimeInterval(snapshot_time_interval),
                           prefix = prefix * "_yz",
                     field_slicer = FieldSlicer(i=1),
                              dir = data_directory,
                            force = force)

simulation.output_writers[:xy] =
    JLD2OutputWriter(model, fields_to_output,
                         schedule = TimeInterval(snapshot_time_interval),
                           prefix = prefix * "_xy",
                     field_slicer = FieldSlicer(k=k_xy_slice),
                              dir = data_directory,
                            force = force)

simulation.output_writers[:statistics] =
    JLD2OutputWriter(model, statistics_to_output,
                     schedule = TimeInterval(snapshot_time_interval),
                       prefix = prefix * "_statistics",
                          dir = data_directory,
                        force = force)

simulation.output_writers[:averaged_statistics] =
    JLD2OutputWriter(model, statistics_to_output,
                     schedule = AveragedTimeInterval(averages_time_interval,
                                                     window = averages_time_window,
                                                     stride = averages_stride),
                       prefix = prefix * "_averaged_statistics",
                          dir = data_directory,
                        force = force)

#####
##### Run
#####

run!(simulation)

#####
##### Animations and analysis
#####

xw, yw, zw = nodes(model.velocities.w)
xu, yu, zu = nodes(model.velocities.u)
xc, yc, zc = nodes((Center, Center, Center), grid)

xyfile = jldopen(joinpath(data_directory, prefix * "_xy.jld2"))
xzfile = jldopen(joinpath(data_directory, prefix * "_xz.jld2"))

iterations = parse.(Int, keys(xyfile["timeseries/t"]))

function divergent_levels(c, clim, nlevels=31)
    levels = range(-clim, stop=clim, length=nlevels)
    cmax = maximum(abs, c)
    clim < cmax && (levels = vcat([-cmax], levels, [cmax]))
    return (-clim, clim), levels
end

function sequential_levels(c, clims, nlevels=31)
    levels = collect(range(clims[1], stop=clims[2], length=nlevels))
    cmin = minimum(c)
    cmax = maximum(c)
    cmin < clims[1] && pushfirst!(levels, cmin)
    cmax > clims[2] && push!(levels, cmax)
    return clims, levels
end

@info "Making an animation from the saved data..."

anim = @animate for (i, iter) in enumerate(iterations)

    @info "Drawing frame $i from iteration $iter \n"

    ## Load 3D fields from file
    t = xyfile["timeseries/t/$iter"]
    wxy = xyfile["timeseries/w/$iter"][:, :, 1]
    uxz = xzfile["timeseries/u/$iter"][:, 1, :]
    exz = xzfile["timeseries/e/$iter"][:, 1, :]

    wlims, wlevels = divergent_levels(wxy, 0.02)
    ulims, ulevels = divergent_levels(uxz, 0.1)
    elims, elevels = sequential_levels(exz, (0, 1e-4))

    wxy_plot = contourf(xw, yw, clamp.(wxy, wlims[1], wlims[2])';
                              color = :balance,
                          linewidth = 0,
                        aspectratio = :equal,
                              clims = wlims,
                             levels = wlevels,
                              xlims = (0, grid.Lx),
                              ylims = (0, grid.Ly),
                             xlabel = "x (m)",
                             ylabel = "y (m)")

    uxz_plot = contourf(xu, zu, clamp.(uxz, ulims[1], ulims[2])';
                              color = :balance,
                          linewidth = 0,
                        aspectratio = :equal,
                              clims = ulims,
                             levels = ulevels,
                              xlims = (0, grid.Lx),
                              ylims = (-grid.Lz, 0),
                             xlabel = "x (m)",
                             ylabel = "z (m)")

    exz_plot = contourf(xc, zc, clamp.(exz, elims[1], elims[2])';
                              color = :thermal,
                          linewidth = 0,
                        aspectratio = :equal,
                              clims = elims,
                             levels = elevels,
                              xlims = (0, grid.Lx),
                              ylims = (-grid.Lz, 0),
                             xlabel = "x (m)",
                             ylabel = "z (m)")

    w_title = @sprintf("w(z=0, t=%s) (m s⁻¹)", prettytime(t))
    u_title = @sprintf("u(y=0, t=%s) (m s⁻¹)", prettytime(t))
    e_title = @sprintf("e(y=0, t=%s) (m s⁻¹)", prettytime(t))

    plot(wxy_plot, uxz_plot, exz_plot, layout=(3, 1), size=(1200, 1000),
         title = [w_title u_title e_title])

    if iter == iterations[end]
        close(xyfile)
        close(xzfile)
    end
end

mp4(anim, prefix * ".mp4", fps = 8)

##### Turbulent kinetic energy budget analysis

xu, yu, zu = nodes(model.velocities.u)

statistics_file = jldopen(joinpath(data_directory, prefix * "_averaged_statistics.jld2"))
statistics_iterations = parse.(Int, keys(statistics_file["timeseries/t"]))

last_iteration = statistics_iterations[end]

U = statistics_file["timeseries/u/$last_iteration"][1, 1, :]
V = statistics_file["timeseries/v/$last_iteration"][1, 1, :]
E = statistics_file["timeseries/e/$last_iteration"][1, 1, :]

# Budget terms
shear_production = statistics_file["timeseries/tke_shear_production/$last_iteration"][1, 1, :]
dissipation = statistics_file["timeseries/tke_dissipation/$last_iteration"][1, 1, :]
buoyancy_flux = statistics_file["timeseries/tke_buoyancy_flux/$last_iteration"][1, 1, :]
advective_flux = statistics_file["timeseries/tke_advective_flux/$last_iteration"][1, 1, :]
pressure_flux = statistics_file["timeseries/tke_pressure_flux/$last_iteration"][1, 1, :]

close(statistics_file)

Δz = grid.Δzᵃᵃᶜ
pressure_flux_divergence = @. (pressure_flux[2:end] - pressure_flux[1:end-1]) ./ Δz
advective_flux_divergence = @. (advective_flux[2:end] - advective_flux[1:end-1]) ./ Δz

tendency = @. - pressure_flux_divergence - advective_flux_divergence + shear_production + buoyancy_flux - dissipation

kwargs = (linewidth=2, alpha=0.8, ylims=(-48, 0))

velocities_plot = plot(U, zu; label="\$ U \$", legend=:bottomleft, xlabel="Velocities (m s⁻¹)", ylabel="z (m)", kwargs...)
plot!(velocities_plot, V, zu; label="\$ V \$", kwargs...)
plot!(velocities_plot, sqrt.(E), zu; label="\$ \\sqrt{E} \$", kwargs...)

budget_plot = plot(-dissipation, zu; label="dissipation", legend=:bottomleft, xlabel="Turbulent kinetic energy budget", ylabel="z (m)", kwargs...)
plot!(budget_plot, shear_production, zu; label="shear production", kwargs...)
plot!(budget_plot, buoyancy_flux, zu; label="buoyancy flux", kwargs...)
plot!(budget_plot, -pressure_flux_divergence, zu; label="pressure flux divergence", kwargs...)
plot!(budget_plot, -advective_flux_divergence, zu; label="advective flux divergence", kwargs...)
plot!(budget_plot, tendency, zu; label="tendency", linestyle=:dash, kwargs...)

tendency_plot = plot(tendency, zu; xlabel="TKE tendency", ylabel="z (m)", legend=nothing, linestyle=:dash, kwargs...)

statistics_plot = plot(velocities_plot, budget_plot, tendency_plot,
                       layout = (1, 3), size = (1000, 600))

savefig(statistics_plot, prefix * "_statistics.png")

max_stars_repo_name: ali-ramadhan/LESbrary

max_issues_repo_name: ali-ramadhan/LESbrary

max_forks_repo_name: ali-ramadhan/LESbrary


text:
export unsafe_delete!,
       personality, personality!,
       callconv, callconv!,
       gc, gc!,
       entry

# forward declaration of Function in src/core/basicblock.jl
identify(::Type{Value}, ::Val{API.LLVMFunctionValueKind}) = Function

Function(mod::Module, name::String, ft::FunctionType) =
    Function(API.LLVMAddFunction(mod, name, ft))

unsafe_delete!(::Module, f::Function) = API.LLVMDeleteFunction(f)

function personality(f::Function)
    has_personality = convert(Core.Bool, API.LLVMHasPersonalityFn(f))
    return has_personality ? Function(API.LLVMGetPersonalityFn(f)) : nothing
end
personality!(f::Function, persfn::Union{Nothing,Function}) =
    API.LLVMSetPersonalityFn(f, persfn===nothing ? C_NULL : persfn)

callconv(f::Function) = API.LLVMGetFunctionCallConv(f)
callconv!(f::Function, cc) = API.LLVMSetFunctionCallConv(f, cc)

function gc(f::Function)
  ptr = API.LLVMGetGC(f)
  return ptr==C_NULL ? "" :  unsafe_string(ptr)
end
gc!(f::Function, name::String) = API.LLVMSetGC(f, name)

entry(f::Function) = BasicBlock(API.LLVMGetEntryBasicBlock(f))

# attributes

export function_attributes, parameter_attributes, return_attributes

struct FunctionAttrSet
    f::Function
    idx::API.LLVMAttributeIndex
end

function_attributes(f::Function) = FunctionAttrSet(f, API.LLVMAttributeFunctionIndex)
parameter_attributes(f::Function, idx::Integer) = FunctionAttrSet(f, API.LLVMAttributeIndex(idx))
return_attributes(f::Function) = FunctionAttrSet(f, API.LLVMAttributeReturnIndex)

Base.eltype(::FunctionAttrSet) = Attribute

function Base.collect(iter::FunctionAttrSet)
    elems = Vector{API.LLVMAttributeRef}(undef, length(iter))
    if length(iter) > 0
      # FIXME: this prevents a nullptr ref in LLVM similar to D26392
      API.LLVMGetAttributesAtIndex(iter.f, iter.idx, elems)
    end
    return Attribute[Attribute(elem) for elem in elems]
end

Base.push!(iter::FunctionAttrSet, attr::Attribute) =
    API.LLVMAddAttributeAtIndex(iter.f, iter.idx, attr)

Base.delete!(iter::FunctionAttrSet, attr::EnumAttribute) =
    API.LLVMRemoveEnumAttributeAtIndex(iter.f, iter.idx, kind(attr))

function Base.delete!(iter::FunctionAttrSet, attr::StringAttribute)
    k = kind(attr)
    API.LLVMRemoveStringAttributeAtIndex(iter.f, iter.idx, k, length(k))
end

function Base.length(iter::FunctionAttrSet)
    API.LLVMGetAttributeCountAtIndex(iter.f, iter.idx)
end

# parameter iteration

export Argument, parameters

@checked struct Argument <: Value
    ref::API.LLVMValueRef
end
identify(::Type{Value}, ::Val{API.LLVMArgumentValueKind}) = Argument

struct FunctionParameterSet
    f::Function
end

parameters(f::Function) = FunctionParameterSet(f)

Base.eltype(::FunctionParameterSet) = Argument

function Base.getindex(iter::FunctionParameterSet, i)
    @boundscheck 1 <= i <= length(iter) || throw(BoundsError(iter, i))
    return Argument(API.LLVMGetParam(iter.f, i-1))
end

function Base.iterate(iter::FunctionParameterSet, state=API.LLVMGetFirstParam(iter.f))
    state == C_NULL ? nothing : (Argument(state), API.LLVMGetNextParam(state))
end

Base.last(iter::FunctionParameterSet) =
    Argument(API.LLVMGetLastParam(iter.f))

Base.isempty(iter::FunctionParameterSet) =
    API.LLVMGetLastParam(iter.f) == C_NULL

Base.length(iter::FunctionParameterSet) = API.LLVMCountParams(iter.f)

# NOTE: optimized `collect`
function Base.collect(iter::FunctionParameterSet)
    elems = Vector{API.LLVMValueRef}(undef, length(iter))
    API.LLVMGetParams(iter.f, elems)
    return map(el->Argument(el), elems)
end

# basic block iteration

export blocks

struct FunctionBlockSet
    f::Function
end

blocks(f::Function) = FunctionBlockSet(f)

Base.eltype(::FunctionBlockSet) = BasicBlock

function Base.iterate(iter::FunctionBlockSet, state=API.LLVMGetFirstBasicBlock(iter.f))
    state == C_NULL ? nothing : (BasicBlock(state), API.LLVMGetNextBasicBlock(state))
end

Base.last(iter::FunctionBlockSet) =
    BasicBlock(API.LLVMGetLastBasicBlock(iter.f))

Base.isempty(iter::FunctionBlockSet) =
    API.LLVMGetLastBasicBlock(iter.f) == C_NULL

Base.length(iter::FunctionBlockSet) = API.LLVMCountBasicBlocks(iter.f)

# NOTE: optimized `collect`
function Base.collect(iter::FunctionBlockSet)
    elems = Vector{API.LLVMValueRef}(undef, length(iter))
    API.LLVMGetBasicBlocks(iter.f, elems)
    return BasicBlock[BasicBlock(elem) for elem in elems]
end

# intrinsics

export isintrinsic, Intrinsic, isoverloaded, declaration

isintrinsic(f::Function) = API.LLVMGetIntrinsicID(f) != 0

struct Intrinsic
    id::UInt32

    function Intrinsic(f::Function)
        id = API.LLVMGetIntrinsicID(f)
        id == 0 && throw(ArgumentError("Function is not an intrinsic"))
        new(id)
    end

    function Intrinsic(name::String)
        @assert version() >= v"9"
        new(API.LLVMLookupIntrinsicID(name, length(name)))
    end
end

Base.convert(::Type{UInt32}, intr::Intrinsic) = intr.id

function name(intr::Intrinsic)
    @assert version() >= v"8"
    len = Ref{Csize_t}()
    str = API.LLVMIntrinsicGetName(intr, len)
    unsafe_string(convert(Ptr{Cchar}, str), len[])
end

function name(intr::Intrinsic, params::Vector{<:LLVMType})
    @assert version() >= v"8"
    len = Ref{Csize_t}()
    str = API.LLVMIntrinsicCopyOverloadedName(intr, params, length(params), len)
    unsafe_message(convert(Ptr{Cchar}, str), len[])
end

function isoverloaded(intr::Intrinsic)
    @assert version() >= v"8"
    convert(Core.Bool, API.LLVMIntrinsicIsOverloaded(intr))
end

function Function(mod::Module, intr::Intrinsic, params::Vector{<:LLVMType}=LLVMType[])
    @assert version() >= v"8"
    Value(API.LLVMGetIntrinsicDeclaration(mod, intr, params, length(params)))
end

function FunctionType(intr::Intrinsic, params::Vector{<:LLVMType}=LLVMType[];
                  ctx::Context=GlobalContext())
    @assert version() >= v"8"
    LLVMType(API.LLVMIntrinsicGetType(ctx, intr, params, length(params)))
end

function Base.show(io::IO, intr::Intrinsic)
    print(io, "Intrinsic($(intr.id))")
    if version() >= v"8"
        if isoverloaded(intr)
            print(io, ": overloaded intrinsic")
        else
            print(io, ": \"$(name(intr))\"")
        end
    end
end


text:
########################### Chisq Diagnostic ###########################

function update_hangartner_temp_vars!(u::Matrix{Int64}, X::Matrix{Int64},
                                      t::Int64)
  d = size(X,2)

  for j in 1:d
    u[X[t, j], j] += 1
  end
end

function hangartner_inner(Y::AbstractMatrix, m::Int)
  ## setup temp vars
  n, d = size(Y)

  # Count for each category in each chain
  u = zeros(Int64, m, d)
  v = zeros(Int64, m, d)

  for t in 1:n
    # fill out temp vars
    update_hangartner_temp_vars!(u, Y, t)
  end
  phia, chi_stat, m_tot = weiss_sub(u, v, n)

  return (n * sum(chi_stat), m_tot)
end

function weiss(X::AbstractMatrix{U}) where {U<:Any}
  ## number of iterations, number of chains
  n, d = size(X)

  ## mapping of values to integers
  v_dict = Dict{U, Int}()

  ## max unique categories
  mc = map(c -> length(unique(X[:,c])), 1:d)
  m = length(unique(X))

  ## counter of number of unique values in each chain
  r0 = 0

  ## Count for each category in each chain
  u = zeros(Int, m, d)

  ## Number of times a category did not transition in each chain
  v = zeros(Int, m, d)

  for t in 1:n
    for c in 1:d
      if !(X[t,c] in keys(v_dict))
        r0 += 1
        v_dict[X[t,c]] = r0
      end
      idx1 = v_dict[X[t,c]]
      u[idx1,c] += 1

      if t > 1
        if X[t-1,c] == X[t,c]
          v[idx1, c] += 1
        end
      end
    end
  end
  phia, chi_stat, m_tot = weiss_sub(u, v, n)
  ca = (1 + phia) / (1 - phia)
  stat = (n / ca) * sum(chi_stat)
  pval = NaN
  if ((m_tot - 1) * (d - 1)) >= 1
    pval = 1 - cdf(Chisq((m_tot - 1) * (d - 1)), stat)
  end

  return ( stat, m_tot, pval, ca)
end

function weiss_sub(u::Matrix{Int64}, v::Matrix{Int64}, t::Int)
  m, d = size(u)
  nt = 0.0
  dt = 0.0
  m_tot = 0

  mp = zeros(Float64, m, d)
  ma = zeros(Float64, m)
  phia = 0.0
  ca = 0.0
  df = 0.0

  chi_stat = zeros(Float64, d)

  for j in 1:m
    p1 = 0.0
    p2 = 0.0
    for l in 1:d
      #aggregate
      p1 += v[j, l] / (d * (t - 1))
      p2 += u[j, l] / (d * t)

      #per chain
      mp[j, l] = u[j, l] / t
      ma[j] += u[j, l] / (d * t)
    end
    nt += p1
    dt += p2 ^ 2

    if ma[j] > 0
      m_tot += 1
      for l in 1:d
        chi_stat[l] += (mp[j,l] - ma[j]) ^ 2 / ma[j]
      end
    end
  end
  phia = 1.0 + (1.0 / t) - ((1 - nt) / (1 - dt))
  phia = min(max(phia,0.0),1.0 - eps())
  return (phia, chi_stat, m_tot)
end

function update_billingsley_temp_vars!(f::Array{Int64,3},
                                       X::Matrix{Int64},
                                       t::Int64)
  d = size(X,2)
  for j in 1:d
    if t > 1
      f[X[t-1, j], X[t,j], j] += 1
    end
  end
end

function billingsley_sub(f::Array{Int64,3})
  df = 0.0
  stat = 0.0

  m, d = size(f, 2), size(f, 3)

  # marginal transitions, i.e.
  # number of transitions from each category
  mf = mapslices(sum, f, dims=2)

  # For each category, number of chains for which
  # that category was present
  A = vec(mapslices( (x)-> sum(x.>0), mf, dims=3))

  # For each category, number of categories it
  # transitioned to
  B =  vec(mapslices((x)-> sum(x.>0), mapslices(sum,f,dims=3),dims=2))

  # transition probabilities in each chain
  P = f ./ mf

  # transition probabilities
  mP = reshape((mapslices(sum,f,dims=3) ./ mapslices(sum,mf,dims=3)), Val(2))

  idx = findall(x -> x != 0, A .* B)
  for j in idx
    #df for billingsley
    df += (A[j] - 1) * (B[j] - 1)

    #billingsley
    for k in idx
      if (mP[j,k] > 0.0)
        for l in 1:d
          if mf[j,1,l] > 0 && isfinite(P[j,k,l])
            stat += mf[j,1,l] * (P[j,k,l] - mP[j,k]) ^ 2 / mP[j,k]
          end
        end
      end
    end
  end
  return (stat, df, mP)
end

function bd_inner(Y::AbstractMatrix, m::Int)
  num_iters, num_chains = size(Y)
  # Transition matrix for each chain
  f = zeros(Int64, m, m, num_chains)

  for t in 1:num_iters
    # fill out temp vars
    update_billingsley_temp_vars!(f, Y, t)
  end
  billingsley_sub(f)
end

function simulate_NDARMA(N::Int, p::Int, q::Int, prob::Vector{Float64},
                         phi::Vector{Float64})
  X = zeros(Int64, N)
  X[1:p] = rand(Categorical(prob), p)
  d1 = Multinomial(1, phi)
  d2 = Categorical(prob)
  for t in (p+1):N
    alphabeta = rand(d1)
    eps = rand(d2, q + 1)
    X[t] = sum([X[(t-p):(t-1)]; eps] .* alphabeta)
  end
  return X
end

function simulate_MC(N::Int, P::Matrix{Float64})
  X = zeros(Int64, N)
  n, m = size(P)
  X[1] = sample(1:n)
  for i in 2:N
    X[i] = wsample(1:n, vec(P[X[i-1], :]) )
  end
  return X
end

function diag_all(X::AbstractMatrix{U}, method::Symbol,
                  nsim::Int, start_iter::Int, step_size::Int) where {U<:Any}

  ## number of iterations, number of chains
  n, d = size(X)

  ## mapping of values to integers
  v_dict = Dict{U, Int}()

  ## max unique categories
  mc = map(c -> length(unique(X[:,c])), 1:d)
  m = length(unique(X))

  ## counter of number of unique values in each chain
  r0 = 0

  ## Count for each category in each chain
  u = zeros(Int, m, d)

  ## Number of times a category did not transition in each chain
  v = zeros(Int, m, d)

  ## transition matrix for each chain
  f = zeros(Int, m, m, d)

  result = zeros(Float64, 3, length(start_iter:step_size:n))
  result_iter = 1
  for t in 1:n
    for c in 1:d
      if !(X[t,c] in keys(v_dict))
        r0 += 1
        v_dict[X[t,c]] = r0
      end
      idx1 = v_dict[X[t,c]]
      u[idx1,c] += 1

      if t > 1
        idx2 = v_dict[X[t-1,c]]
        f[idx1, idx2, c] += 1

        if X[t-1,c] == X[t,c]
          v[idx1, c] += 1
        end
      end
    end

    if ((t >= start_iter) && (((t - start_iter) % step_size) == 0))
      phia, chi_stat, m_tot = weiss_sub(u, v, t)
      hot_stat, df, mP = billingsley_sub(f)

      phat = mapslices(sum,u,dims=2)[:,1] / sum(mapslices(sum,u,dims=2))
      ca = (1 + phia) / (1 - phia)
      stat = NaN
      pval = NaN
      df0  = NaN

      if method == :hangartner
        stat = t * sum(chi_stat)
        df0 = (m - 1) * (d - 1)
        if m > 1 && !isnan(stat)
          pval = 1 - cdf(Chisq( (m - 1) * (d - 1)), stat)
        end
      elseif method == :weiss
        stat = (t / ca) * sum(chi_stat)
        df0 = (m - 1) * (d - 1)
        pval = NaN
        if m > 1 && !isnan(stat)
          pval = 1 - cdf(Chisq( (m - 1) * (d - 1)), stat)
        end
      elseif method == :DARBOOT
        stat = t * sum(chi_stat)
        bstats = zeros(Float64, nsim)
        for b in 1:nsim
          Y = hcat([simulate_NDARMA(t, 1, 0, phat, [phia, 1-phia])
                for j in 1:d]...)
          s = hangartner_inner(Y, m)[1]
          bstats[b] = s
        end
        idx = findall(!isnan(bstats))
        df0 = mean(bstats[idx])
        pval = sum(stat .<= bstats[idx])/length(idx)
      elseif method == :MCBOOT
        bstats = zeros(Float64, nsim)
        for b in 1:nsim
          Y = hcat([simulate_MC(t, mP) for j in 1:d]...)
          s = hangartner_inner(Y, m)[1]
          bstats[b] = s
        end
        idx = findall(!isnan(bstats))
        df0 = mean(bstats[idx])
        pval = sum(stat .<= bstats[idx])/length(idx)
      elseif method == :billingsley
        stat = hot_stat
        df0 = df
        if df > 0 && !isnan(hot_stat)
          pval = 1 - cdf(Chisq(df), hot_stat)
        end
      elseif method == :billingsleyBOOT
        stat = hot_stat
        bstats = zeros(Float64, nsim)
        for b in 1:nsim
          Y = hcat([simulate_MC(t, mP) for j in 1:d]...)
          (s,sd) = bd_inner(Y, m)[1:2]
          bstats[b] = s/sd
        end
        idx = findall(!isnan(bstats))
        df0 = mean(bstats[idx])
        pval = sum(stat/df .<= bstats[idx])/length(idx)
      else
        error("Unexpected")
      end
      result[:, result_iter] = [stat, df0, pval]
      result_iter += 1
    end
  end
  return result
end

function discretediag_sub(c::AbstractChains, frac::Real, method::Symbol,
                          nsim::Int, start_iter::Int, step_size::Int)

  num_iters, num_vars, num_chains = size(c.value)

  vals = zeros(Float64, 3 * (num_chains + 1), num_vars)
  plot_vals_stat = zeros(length(start_iter:step_size:num_iters), num_vars)
  plot_vals_pval = zeros(length(start_iter:step_size:num_iters), num_vars)

  ## Between-chain diagnostic
  X = zeros(Int64, num_iters, num_chains)
  for j in 1:length(num_vars)
    X = convert(Array{Int64, 2}, c.value[:,j,:])
    result = diag_all(X, method, nsim, start_iter, step_size)
    plot_vals_stat[:,j] = result[1, :] ./ result[2, :]
    plot_vals_pval[:,j] = result[3, :]
    vals[1:3, j] = result[:, end]
  end

  ## Within-chain diagnostic
  x = zeros(Int64, num_iters)
  Y = zeros(Int64, num_iters, 2)
  for j in 1:num_vars
    for k in 1:num_chains
      x = convert(Array{Int64, 1}, c.value[:,j,k])

      idx1 = 1:round(Int64, frac * num_iters)
      idx2 = round(Int64, num_iters - frac * num_iters + 1):num_iters
      x1 = x[idx1]
      x2 = x[idx2]
      n_min = min(length(x1), length(x2))
      Y = [x1[1:n_min] x2[(end - n_min + 1):end]]

      vals[(3 + 3 * (k - 1) + 1):(3 + 3 * (k - 1) + 3), j] =
        diag_all(Y, method, nsim, n_min, step_size)[:, end]
    end
  end
  return (collect(1:num_vars), vals, plot_vals_stat, plot_vals_pval)

end

function discretediagplot(c::AbstractChains; frac::Real=0.3,
                          method::Symbol=:weiss, nsim::Int=1000,
                          start_iter::Int=100, step_size::Int=10000)

  num_iters, num_vars, num_chains = size(c.value)

  valid_methods = [:hangartner, :weiss, :DARBOOT,
                   :MCBOOT, :billingsley, :billingsleyBOOT]
  if !(method in valid_methods)
    methods_str = join([":$f" for f in valid_methods], ", ")
    throw(ArgumentError("method must be one of ", methods_str))
  end

  if !(0.0 < frac < 1.0)
    throw(ArgumentError("frac must be in (0,1)"))
  end

  if (start_iter > num_iters ) || (step_size > num_iters)
    throw(ArgumentError("start_iter, step_size must be less than $num_iters"))
  end

  V, vals, plot_vals_stat, plot_vals_pval =
    discretediag_sub(c, frac, method, nsim, start_iter, step_size)

  p1 = plot(y=vcat([plot_vals_stat[:,j] for j in 1:length(V)]...),
            x=repeat(collect(c.range[start_iter:step_size:num_iters])/1000,
                     outer=[length(V)]),
            Geom.line,
            Guide.xlabel("Iteration (thousands)", orientation=:horizontal),
            Guide.ylabel("stat/df",orientation=:vertical),
            Scale.color_discrete(), Guide.colorkey(title="Variable"),
            color=repeat(c.names[V],
                         inner=[length(start_iter:step_size:num_iters)]))

  p2 = plot(y=vcat([plot_vals_pval[:,j] for j in 1:length(V)]...),
            x=repeat(collect(c.range[start_iter:step_size:num_iters])/1000,
                     outer=[length(V)]),
            Geom.line,
            Guide.xlabel("Iteration (thousands)", orientation=:horizontal),
            Guide.ylabel("pval",orientation=:vertical),
            Scale.color_discrete(), Guide.colorkey(title="Variable"),
            color=repeat(c.names[V],
                         inner=[length(start_iter:step_size:num_iters)]))

  return [p1, p2]
end

function discretediag(c::AbstractChains; frac::Real=0.3,
                      method::Symbol=:weiss, nsim::Int=1000)

  num_iters, num_vars, num_chains = size(c.value)

  valid_methods = [:hangartner, :weiss, :DARBOOT,
                   :MCBOOT, :billingsley, :billingsleyBOOT]
  if !(method in valid_methods)
    methods_str = join([":$f" for f in valid_methods], ", ")
    throw(ArgumentError("method must be one of ", methods_str))
  end

  if !(0.0 < frac < 1.0)
    throw(ArgumentError("frac must be in (0,1)"))
  end

  V, vals = discretediag_sub(c, frac, method, nsim,
                             size(c.value,1), size(c.value,1))[1:2]

  hdr = header(c) * "\nChisq Diagnostic:\nEnd Fractions = $frac\n" *
  "method = $method\n"
  ChainSummary(copy(round.(vals,digits=3)'), c.names[V],
               convert(Array{AbstractString, 1},
                       vcat([["stat", "df", "p-value"]
                             for k in 1:(num_chains + 1)]...)), hdr)
end


text:
using SpecialPolynomials
using HypergeometricFunctions

export qc, ql, qk
export Qc, Ql, Qk

"MArrays.jl contains the functions and methods for the evolution of initial amplitudes in a one-dimensional masses array with harmonic couplings. Here, we study different coupling laws, their evolution on time, and their energy content (despicted at the EExchanges.jl script)"

# 1. ---------------------------------------------------------------------------

"### Circular arrangement of masses, _m_{j} = 1 ∀j_, for restitution constants _k_{j} = k, ∀j"

"`Ac(N)` gives the normalization term 1/(N + 1)"
Ac(N) = 1/(N + 1)
"`ωc(n, M)` gives an angular frequency who depends on position _m_"
ωc(m, N) = (π * m) * Ac(N)

"`kec` is the kernel function who provide the evolution in time. It is an oscillatory sinusoidal term who is called as `kec(t, k, m, n, l, N)`"
kec(t, k, m, n, l, N) = cos(2 * t * (√k) * sin(ωc(m, N))) * cos(2 * ωc(m, N) * (n - l))

"`qc` is the circular evolution function who integrates the kernel `kec` and their proper normalization terms `Ac` and `ωc`. For a constant _k_, it is called as `qc(k, n, t, N, V::Vector)`"
function qc(k, n, t, N, V::Vector)
    u = zeros(Float64,(N + 1, N + 1))
    for m in 0:N
        for l in 0:N
            u[l + 1, m + 1] = kec(t, k, m, n, l, N) * V[l + 1]
        end
    end
    return Ac(N) * sum(u)
end

"`Qc` is an ad-hoc function to obtain the matrix evolution from the function `qc`. It evaluates an initial _N + 1_ vector _V_ over a time span _t_. The resulting item is an _(N + 1) × `length(t)`_ matrix. It has middle performance. It is called `Qc(k, t, V::Vector)`"
function Qc(k, t, V::Vector)
    N = length(V)-1
    T = length(t)
    v = zeros(Float64,(N + 1,T))
    for l in 0:N
        for m in 1:T
            v[l + 1, m] = qc(k, l, t[m], N, V)
        end
    end
    return v
end

# The initial amplitudes vector V need to be of dimension N + 1 = dim,
# because the loops starts counting at 0.

# 2. ---------------------------------------------------------------------------

"### Lineal arrays with specific conditions in the coupling matrix"

"##### Case m_{j} = 1, k_{j} = k, ∀j"

"We will use the generator method for obtain the second kind Chebyshev polynomial
given by the package SpecialPolynomials.jl"

"`gen(n)` is the generator array for the Chebyshev's. It retuns a vector of _N + 1_ size with zeros everywhere, except at `gen(end) = 1`, who in turn serves as the generator for the n-th Chebyshev U polynomial."
function gen(n)
    gen = zeros(Int64, n + 1)
    gen[end] = 1
    return gen
end

"`chebU(x, n)` gives the n-th Chebyshev polynomial U_{n}(x). It relies on the package SpecialPolynomials.jl"
function chebU(x, n)
    return SpecialPolynomials.ChebyshevU(gen(n))(x)
end

"`ϕ(j, N)` gives the j-th phase of the Chebyshev's U_{n}(x)"
ϕ(j, N) = (π * j) / (N + 2)
"`y(j, N)` gives the j-th zero of the Chebyshev's U_{n}(x)"
y(j, N) = cos(ϕ(j, N))

"`oscl` gives the j-th oscillatory evolution terms at time _t_, it is called as `oscl(k, t, j, N)`"
oscl(k, t, j, N) = cos(2 * t * (√k) * sin(ϕ(j + 1, N)/2))

"`nc(j, N)` gives a kind of _position dependent_ normalization, since it depends on the j-th term."
function nc(j, N)
    nc = zeros(N + 1)
    for l in 0:N
        nc[l + 1] = (chebU(y(j + 1, N), l))^2
    end
    return sum(nc)
end

"`twocheb` provides a compact form to package the multiplication of two mutually excludent Chebyshev's U_{n}(x)×U_{l}(x). It is called as `twocheb(j, l, n, N)`."
twocheb(j, l, n, N) = chebU(y(j + 1, N), l) * chebU(y(j + 1, N), n)

"`kel` is the composition of the `twocheb` and `oscl` functions. It gives the evolution kernel for the lineal array"
kel(t, k, n, l, j, N) = (twocheb(j, l, n, N)/nc(j, N)) * oscl(k, t, j, N)

"`ql` is the main function for the lineal array evolution, when evaluated for some initial _N + 1_ vector V, it gives a _(N + 1) × (N + 1)_ matrix who is a function time _t_ and position index _n_. It is called `ql(k, n, t, N, V::Vector)`"
function ql(k, n, t, N, V::Vector)
    nr = [nc(j, N) for j in 0:N]
    tc(j, l, n) = twocheb(j, l, n, N)
    osc(k, t, j) = oscl(k, t, j, N)
    el = zeros(Float64,(N + 1, N + 1))
    for l in 0:N
        for j in 0:N
            el[l + 1, j + 1] = (tc(j, l, n) / nr[j + 1]) * osc(k, t, j) * V[l + 1]
        end
    end
    return sum(el)
end

"`Ql` is an ad-hoc function to obtain the matrix evolution from the function `ql`. It evaluates an initial _N + 1_ vector _V_ over a time span _t_. The resulting item is an _(N + 1) × `length(t)`_ matrix. It has poor performance. It is called `Ql(k, t, V::Vector)`"
function Ql(k, t, V::Vector)
    dn = length(V)-1
    ell = zeros(Float64, (dn + 1, length(t)))
    for n in 0:dn
        for T in 1:length(t)
            ell[n + 1, T] = ql(k, n, t[T], dn, V)
        end
    end
    return ell
end

# 3. Kravchuk ------------------------------------------------------------------

"### Evolution in Kravchuk coupling. The coupling coefficients belongs to the basis of the discrete and finite harmonic oscillator of _su(2)_"

"`kp` gives the symmetric Kravchuk polynomial, it is evaluated as `kp(i, j, N)` for i,j ∈ [0,N]"
function kp(i, j, N)
    if i == j == N
        return HypergeometricFunctions._₂F₁general2(0, -i, -N, 2) # issue solved when i == j == N
    else
        return HypergeometricFunctions._₂F₁general2(-j, -i, -N, 2)
    end
end

"`ω(j, N)` gives a binomial normalization at the j-th term"
ω(j, N) = 2.0^(-N) * binomial(N, j)
"`h(i, N)` gives a binomial normalization at the i-th term"
h(i, N) = 1 / binomial(N, i)

"`oscc(t, n)` is the evolution kernel at time _t_. It is really simple"
oscc(t, n) = cos(t * sqrt(n))

"`kf` is the symmetric Kravchuk function, who is the composition of `ω` and `h`"
kf(i, j, N) = sqrt(ω(j, N) / h(i, N)) * kp(i, j, N)

"`kfm(N)` evaluates the Kravchuk matrix of `kf`. It returns a matrix of _(N + 1)×(N + 1)_ terms who columns are the functions `kf` at index _n_"
function kfm(N)
    km = zeros(Float64, (N + 1, N + 1))
    for i in 0:N
        for j in 0:N
            km[i + 1, j + 1] = kf(i, j, N)
        end
    end
    return km
end

"`k2` gives the two dimensional Kravchuk function, it is called `k2(m, l, n, N)`"
k2(m, l, n, N) = kf(m, n, N) * kf(l, n, N)

"`qk` is the main function for the Kravchuk array evolution, when evaluated for some initial _N + 1_ vector V, it gives a _(N + 1) × (N + 1)_ matrix who is a function time _t_ and position index _n_. It is called `qk(m, t, N, V::Vector)`"
function qk(m, t, N, V::Vector)
    u = zeros(Float64, (N + 1, N + 1))
    for l in  0:N
        for n in 0:N
            u[l + 1, n + 1] = k2(m, l, n, N) * oscc(t, n) * V[l + 1]
        end
    end
    return sum(u)
end

"`Qk` is an ad-hoc function to obtain the matrix evolution from the function `qk`. It evaluates an initial _N + 1_ vector _V_ over a time span _t_. The resulting item is an _(N + 1) × `length(t)`_ matrix. It has good performance. It is called `Qk(t, V::Vector)`"
function Qk(t, V::Vector)
    dn = length(V)-1
    T = length(t)
    ek = zeros(Float64, (dn + 1, T))
    for m in 0:dn
        for τ in 1:T
            ek[m + 1, τ] = qk(m, t[τ], dn, V)
        end
    end
    return ek
end


text:
## plot.jl : functions for easy visualization of Gaussian random fields

# 1D
"""
plot(grf[, kwargs...])

Plot a sample of the one-dimensional Gaussian random field `grf` and pass optional arguments `kwargs` to the `Plots.plot` command.

# Examples
```jldoctest
julia> cov = CovarianceFunction(1, Exponential(1))
1d exponential covariance function (λ=1.0, σ=1.0, p=2.0)

julia> pts = range(0, stop=1, length=501)
0.0:0.002:1.0

julia> grf = GaussianRandomField(cov, CirculantEmbedding(), pts)
Gaussian random field with 1d exponential covariance function (λ=1.0, σ=1.0, p=2.0) on a 501-point structured grid, using a circulant embedding

julia> plot(grf)
[...]

```
See also: [`plot!`](@ref)
"""
plot

@recipe f(grf::GaussianRandomField{G, C}) where {G, C <: AbstractCovarianceFunction{1}} = (grf.pts[1], sample(grf))

"""
plot!(grf[, kwargs...])

Add a sample of the one-dimensional Gaussian random field `grf` to the existing plot.

See also: [`plot`](@ref)
"""
plot!

# 2D
"""
    surface(grf[, kwargs])

Plot a surface of the two-dimensional Gaussian random field `grf` and pass optional arguments `kwargs` to the `Plots.surface` command.

# Examples
```jldoctest
julia> cov = CovarianceFunction(2, Matern(.3, 1))
2d Matérn covariance function (λ=0.3, ν=1.0, σ=1.0, p=2.0)

julia> pts = range(0, stop=1, length=101)
0.0:0.01:1.0

julia> grf = GaussianRandomField(cov, CirculantEmbedding(), pts, pts, minpadding=157)
Gaussian random field with 2d Matérn covariance function (λ=0.3, ν=1.0, σ=1.0, p=2.0) on a 101×101 structured grid, using a circulant embedding

julia> surface(grf)
[...]

```
See also: [`contour`](@ref), [`contourf`](@ref), [`heatmap`](@ref)
"""
surface

"""
    contour(grf[, kwargs])

Plot a contour map of the two-dimensional Gaussian random field `grf` and pass optional arguments `kwargs` to the `Plots.contour` command.

# Examples
```jldoctest
julia> cov = CovarianceFunction(2, Matern(.3, 1))
2d Matérn covariance function (λ=0.3, ν=1.0, σ=1.0, p=2.0)

julia> pts = range(0, stop=1, length=101)
0.0:0.01:1.0

julia> grf = GaussianRandomField(cov, CirculantEmbedding(), pts, pts, minpadding=157)
Gaussian random field with 2d Matérn covariance function (λ=0.3, ν=1.0, σ=1.0, p=2.0) on a 101×101 structured grid, using a circulant embedding

julia> contour(grf)
[...]

```
See also: [`surface`](@ref), [`contourf`](@ref), [`heatmap`](@ref)
"""
contour

"""
    contourf(grf[, kwargs])

Plot a filled contour map of the two-dimensional Gaussian random field `grf` and pass optional arguments `kwargs` to the `Plots.contourf` command.

# Examples
```jldoctest
julia> cov = CovarianceFunction(2, Matern(.3, 1))
2d Matérn covariance function (λ=0.3, ν=1.0, σ=1.0, p=2.0)

julia> pts = range(0, stop=1, length=101)
0.0:0.01:1.0

julia> grf = GaussianRandomField(cov, CirculantEmbedding(), pts, pts, minpadding=157)
Gaussian random field with 2d Matérn covariance function (λ=0.3, ν=1.0, σ=1.0, p=2.0) on a 101×101 structured grid, using a circulant embedding

julia> contourf(grf)
[...]

```
See also: [`surface`](@ref), [`contour`](@ref), [`heatmap`](@ref)
"""
contourf

"""
    heatmap(grf[, kwargs])

Plot a heatmap of the two-dimensional Gaussian random field `grf` and pass optional arguments `kwargs` to the `Plots.heatmap` command.

# Examples
```jldoctest
julia> cov = CovarianceFunction(2, Matern(.3, 1))
2d Matérn covariance function (λ=0.3, ν=1.0, σ=1.0, p=2.0)

julia> pts = range(0, stop=1, length=101)
0.0:0.01:1.0

julia> grf = GaussianRandomField(cov, CirculantEmbedding(), pts, pts, minpadding=157)
Gaussian random field with 2d Matérn covariance function (λ=0.3, ν=1.0, σ=1.0, p=2.0) on a 101×101 structured grid, using a circulant embedding

julia> heatmap(grf)
[...]

```
See also: [`surface`](@ref), [`contour`](@ref), [`contourf`](@ref)
"""
heatmap

@recipe f(grf::GaussianRandomField{G, C}) where {G, C <: AbstractCovarianceFunction{2}} = begin
    linewidth -> get(plotattributes, :seriestype, :auto) == :contourf ? 0 : :auto
    grf.pts..., sample(grf)
end

# Plot eigenvalues
"""
    plot_eigenvalues(grf)

Log-log plot of the eigenvalues of the Gaussian random field. Only available for Gaussian random field generators of type `Spectral`, `KarhunenLoeve` and `CirculantEmbedding`.

# Examples
```jldoctest
julia> cov = CovarianceFunction(2, Gaussian(.1))
2d Gaussian covariance function (λ=0.1, σ=1.0, p=2.0)

julia> pts = range(0, stop=1, length=51)
0.0:0.02:1.0

julia> grf = GaussianRandomField(cov, KarhunenLoeve(100), pts, pts)
Gaussian random field with 2d Gaussian covariance function (λ=0.1, σ=1.0, p=2.0) on a 51×51 structured grid, using a KL expansion with 100 terms

julia> plot_eigenvalues(grf)
[...]

```
See also: [`plot_eigenfunction`](@ref)
"""
plot_eigenvalues

@userplot Plot_Eigenvalues

recipetype(::Val{:plot_eigenvalues}, args...) = Plot_Eigenvalues(args)

@recipe f(plt::Plot_Eigenvalues) = begin
    grf = plt.args[1]

    if grf isa GaussianRandomField{<:Union{KarhunenLoeve, Spectral},<:CovarianceFunction}
        val = grf.data.eigenval.^2
    elseif grf isa GaussianRandomField{<:CirculantEmbedding}
        val = sort(vec(grf.data[1]), rev=true).^2 * length(grf.data[1])
    elseif grf isa GaussianRandomField{<:KarhunenLoeve,<:SeparableCovarianceFunction}
        order, data = grf.data
        val = Array{Float64}(undef, length(order))
        for (i, o) in enumerate(order)
            val[i] = prod(data[j].eigenval[o[j]] for j in 1:length(data))^2
        end
    end

    @series begin
        seriestype --> :line
        xaxis --> :log
        yaxis --> :log
        xguide --> "eigenvalue number"
        yguide --> "magnitude"
        1:length(val), val
    end
end

# Plot eigenfunctions
"""
    plot_eigenfunction(grf, n)

Contour plot of the `n`th eigenfunction of the Gaussian random field. Only available for Gaussian random field generators of type `Spectral` and `KarhunenLoeve`.

# Examples
```jldoctest
julia> cov = CovarianceFunction(2, Gaussian(.1))
2d Gaussian covariance function (λ=0.1, σ=1.0, p=2.0)

julia> pts = range(0, stop=1, length=51)
0.0:0.02:1.0

julia> grf = GaussianRandomField(cov, KarhunenLoeve(100), pts, pts)
Gaussian random field with 2d Gaussian covariance function (λ=0.1, σ=1.0, p=2.0) on a 51×51 structured grid, using a KL expansion with 100 terms

julia> plot_eigenfunction(grf, 6) # 6th eigenfunction
[...]

```
See also: [`plot_eigenfunction`](@ref)
"""
plot_eigenfunction

@userplot Plot_Eigenfunction

recipetype(::Val{:plot_eigenfunction}, args...) = Plot_Eigenfunction(args)

@recipe f(plt::Plot_Eigenfunction) = begin
    grf = plt.args[1]
    n = plt.args[2]

    if grf isa GaussianRandomField{<:Union{KarhunenLoeve, Spectral}, <:AbstractCovarianceFunction{1}}
        x = grf.pts[1]
        if grf isa GaussianRandomField{G, <:CovarianceFunction{1}} where G
            f = view(grf.data.eigenfunc, :, n)
        elseif grf isa GaussianRandomField{G, <:SeparableCovarianceFunction{1}} where G
            order, data = grf.data
            ordern = order[n]
            f = view(data[1].eigenfunc, :, ordern[1])
        end
        @series begin
            seriestype --> :line
            x, f
        end
    end

    if grf isa GaussianRandomField{<:Union{KarhunenLoeve, Spectral}, <:AbstractCovarianceFunction{2}}
        x, y = grf.pts
        if grf isa GaussianRandomField{G, <:CovarianceFunction{2}} where G
            f = view(grf.data.eigenfunc, :, n)
        elseif grf isa GaussianRandomField{G, <:SeparableCovarianceFunction{2}} where G
            order, data = grf.data
            ordern = order[n]
            f = kron((view(data[j].eigenfunc, :, ordern[j]) for j = 1:length(grf.cov.cov))...)
        end
        @series begin
            seriestype --> :contourf
            linewidth --> 0
            x, y, f
        end
    end
end

# Plot covariance matrix
"""
    plot_covariance_matrix(grf[, pts...])

    Evaluate the covariance function of the Gaussian random field `grf` in the (optional) points `pts` and plot the result.

# Examples
```jldoctest
julia> cov = CovarianceFunction(2, Matern(.3, 1))
2d Matérn covariance function (λ=0.3, ν=1.0, σ=1.0, p=2.0)

julia> pts = range(0, stop=1, length=11)
0.0:0.1:1.0

julia> grf = GaussianRandomField(cov, CirculantEmbedding(), pts, pts, minpadding=7)
Gaussian random field with 2d Matérn covariance function (λ=0.3, ν=1.0, σ=1.0, p=2.0) on a 11×11 structured grid, using a circulant embedding

julia> plot_covariance_matrix(grf)
[...]

julia> pts = range(0, stop=1, length=21)
0.0:0.05:1.0

julia> plot_covariance_matrix(grf, pts, pts)
[...]

```
"""
plot_covariance_matrix

@userplot Plot_Covariance_Matrix

recipetype(::Val{:plot_covariance_matrix}, args...) = Plot_Covariance_Matrix(args)

@recipe function f(plt::Plot_Covariance_Matrix)
    @series begin
        seriestype --> :heatmap
        grf = plt.args[1]
        pts = length(plt.args) > 1 ? plt.args[2] : grf.pts
        Z = apply(grf.cov, pts...)
        x = axes(Z, 1)
        x, x, Z
    end
end


text:
implements(::Mutator, ::Type{<:StaticArrays.StaticArray}) = false
implements(::typeof(setindex!), ::Type{<:StaticArrays.MArray}) = true
copyappendable(src::StaticArrays.StaticVector) = src


text:
using BenchmarkTools
using Random
using Unfold
using CSV
using DataFrames
using CategoricalArrays
const SUITE = BenchmarkGroup()
Random.seed!(3)

# Preparatory code
include("../test/test_utilities.jl"); # to load the simulated data

data,evts = loadtestdata("testCase6",dataPath=(@__DIR__)*"/../test/data") #
transform!(evts, :subject => categorical => :subject)
transform!(evts, :stimulus => categorical => :stimulus)

evts.subjectB = evts.subject;
evts1 = evts[evts.condA.==0,:]
evts2 = evts[evts.condA.==1,:]
data = data.+ 0.5*randn(size(data)) # we have to add minimal noise, else mixed models crashes.

data_r = reshape(data,(1,:)) #make ch x time
data_epochs,times = Unfold.epoch(data=data_r,tbl=evts,τ=(-0.4,0.8),sfreq=100); #epoch
evts,data_epochs = Unfold.dropMissingEpochs(evts,data_epochs) # remove missing

ba1 = firbasis(τ=(0,1),sfreq = 10,name="evts1")
ba2 = firbasis(τ=(0,1),sfreq = 10,name="evts2")

f1_lmm  = @formula 0~1+condB+(1|subject)
f2_lmm  = @formula 0~1+condB+(1|subjectB)

f1  = @formula 0~1+condB
f2  = @formula 0~1+condB

dict_lin = Dict(0 => (f1,ba1),1 => (f2,ba2))
dict_lmm = Dict(0 => (f1_lmm,ba1),1 => (f2_lmm,ba2))

X1_lmm  = designmatrix(UnfoldLinearMixedModel,f1_lmm,evts1,ba1)
X2_lmm  = designmatrix(UnfoldLinearMixedModel,f2_lmm,evts2,ba2)

X1   = designmatrix(UnfoldLinearModelContinuousTime,f1,evts1,ba1)
X2   = designmatrix(UnfoldLinearModelContinuousTime,f2,evts2,ba2)

SUITE = BenchmarkGroup()
SUITE["nodc"] = BenchmarkGroup(
    ["nodc"])
SUITE["dc"] = BenchmarkGroup(["dc"])
# designmatrix generation
SUITE["dc"]["X_gen_lin"] = @benchmarkable designmatrix(UnfoldLinearModelContinuousTime,$f1,$evts1,$ba1)
SUITE["dc"]["X_gen_lmm"] = @benchmarkable designmatrix(UnfoldLinearMixedModelContinuousTime,$f1_lmm,$evts1,$ba1)

# designmat concatination
SUITE["dc"]["X_cat_lin"] = @benchmarkable $X1+$X2
SUITE["dc"]["X_cat_lmm"] = @benchmarkable $X1_lmm+$X2_lmm

# Model Fit
SUITE["nodc"]["fit_lin"] = @benchmarkable fit(UnfoldModel,$f1,$evts1,$data_epochs,$times)
SUITE["nodc"]["fit_lmm"] = @benchmarkable fit(UnfoldModel,$f1_lmm,$evts1,$data_epochs,$times)
SUITE["dc"]["fit_lin"]   = @benchmarkable fit(UnfoldModel,$dict_lin,$evts,$data_r,eventcolumn="condA");
SUITE["dc"]["fit_lmm"] = @benchmarkable   fit(UnfoldModel,$dict_lmm,$evts,$data_r,eventcolumn="condA");

max_stars_repo_name: unfoldtoolbox/jlunfold

max_issues_repo_name: unfoldtoolbox/jlunfold

max_forks_repo_name: unfoldtoolbox/jlunfold


text:
"""
    convert_results(;[specific_run=""])

Convert simulation results for data exchange with other frameworks.

# Arguments
- `specific_run`: Path to a specific .jld file to be converted

See also: [Config](@ref), [cfg_sim](@ref), [cfg_ot](@ref), [cfg_ag](@ref), [cfg_feed](@ref), [cfg_mech](@ref)
"""
function convert_results(
    ;
    specific_run::String=""
)

    if specific_run != ""

        raw_data = load(joinpath("results", specific_run))
        data = raw_data[first(keys(raw_data))]
        filename = specific_run[1:(first(findfirst(".jld2", specific_run)) - 1)]

        if !in(filename, readdir("dataexchange"))
            mkdir(joinpath("dataexchange", filename))
        end

        Base.open(
            joinpath("dataexchange", filename, "config.json"),
            "w"
        ) do io
            Base.write(io, JSON.json(data.config))
        end

        CSV.write(
            joinpath("dataexchange", filename, "agent_log" * ".csv"),
            data.agent_log
        )
        CSV.write(
            joinpath("dataexchange", filename, "post_log" * ".csv"),
            data.post_log
        )

        for i in 1:length(data.graph_list)
            graph_nr = lpad(
                string(i),
                length(string(length(data.graph_list))),
                "0"
            )
            savegraph(
                joinpath("dataexchange", filename, "graph_$graph_nr.gml"),
                data.graph_list[i],
                GraphIO.GML.GMLFormat()
            )
        end

    else

        if !in("dataexchange", readdir())
            mkdir("dataexchange")
        end

        for file in readdir("results")

            if !occursin("jld2", file)
                continue
            end

            raw_data = load(joinpath("results", file))
            data = raw_data[first(keys(raw_data))]
            filename = file[1:(first(findfirst(".jld2", file)) - 1)]

            if !in(filename, readdir("dataexchange"))
                mkdir(joinpath("dataexchange", filename))
            end

            Base.open(
                joinpath("dataexchange", filename, "config.json"),
                "w"
            ) do io
                Base.write(io, JSON.json(data.config))
            end

            CSV.write(
                joinpath("dataexchange", filename, "agent_log" * ".csv"),
                data.agent_log
            )
            CSV.write(
                joinpath("dataexchange", filename, "post_log" * ".csv"),
                data.post_log
            )

            for i in 1:length(data.graph_list)
                graph_nr = lpad(
                    string(i),
                    length(string(length(data.graph_list))),
                    "0"
                )
                savegraph(
                    joinpath("dataexchange", filename, "graph_$graph_nr.gml"),
                    data.graph_list[i],
                    GraphIO.GML.GMLFormat()
                )
            end

        end

    end

end


text:
@testset "Multivariate functions" begin
    # himmilbeau from Daisy benchmarks
    f(x1, x2, x3) = -x1*x2 - 2*x2*x3 - x1 - x3
    dom = Interval(-4.5, -0.3) × Interval(0.4, 0.9) × Interval(3.8, 7.8)
    xref = Interval(-20.786552979420335, -0.540012836551535) # MOSEK deg 6
    for solver in available_solvers
        x = enclose(f, dom, solver)
        r = relative_precision(x, xref)
        @test inf(r) ≤ 5 && sup(r) ≤ 16
    end
end

@testset "Multivariate example from the Quickstart Guide" begin
    g(x, y) = (x + 2y - 7)^2 + (2x + y - 5)^2
    dom = IntervalBox(-10..10, 2)

    x = enclose(g, dom, :IntervalArithmetic)
    xref = Interval(0, 2594)
    r = relative_precision(x, xref)
    @test inf(r) ≤ 1e-5 && sup(r) ≤ 1e-5
end

@testset "Test multivariate polynomial input" begin
    @polyvar x y
    p = (x + 2y - 7)^2 + (2x + y - 5)^2
    dom = IntervalBox(-10..10, 2)

    x = enclose(p, dom)
    xref = Interval(-1446, 2594)
    r = relative_precision(x, xref)
    @test inf(r) ≤ 1e-5 && sup(r) ≤ 1e-5
    # Note: DynamicPolynomials automatically expands p, and evaluation using
    # interval arithmetic gives a worse left bound than the factored expression.
end