text:
module Random

using Distributions, Sobol
using LinearAlgebra
using Random: GLOBAL_RNG, MersenneTwister
 #=================================================================================
                        Multinomial *Index* Sampling.

        (Returns the respective *indices* of the generated random nums
        as opp. to the *one-hot* representation as is usual for multinomial.)
 ==================================================================================#

function multinomial_indices(n::Int, p::Vector{Float64})
    #=  Returns the respective *indices* of the generated random nums
        as opp. to the *one-hot* representation as is usual for multinomial
        (adapted from Distributions.jl/src/samplers/multinomial.jl)
        -----------------------------------------------------------
        ===> Uses a binomial for each element of the multinomial
        vector. This is efficient when the number of trials is large,
        but if trials is small, and length is large, use `_linear` version.
    =#

    k = length(p)
    rp = 1.0  # remaining total probability
    i = 0
    km1 = k - 1
    x = zeros(Int32, n)
    op_ix = 1

    while i < km1 && n > 0
        i += 1
        @inbounds pi = p[i]
        if pi < rp
            xi = rand(Binomial(n, pi / rp))
            x[op_ix:(op_ix+xi-1)] .= i
            op_ix += xi
            n -= xi
            rp -= pi
        else
            # In this case, we don't even have to sample
            # from Binomial. Just assign remaining counts
            # to xi.
            x[op_ix:(op_ix+n-1)] .= i
            n = 0
        end
    end

    if i == km1
        x[op_ix:end] .= i+1
    end

    return xx
end

function smp_from_logprob(n_samples::Int, logp::Vector{Float64})
    #= multinomial sampling from log probability vector. Performs
       softmax and then multinomial sampling. Not sure why this
       has its own softmax calc: should be updated to use NNlibs.
    =#
    p = exp.(logp .- maximum(logp))
    p /= sum(p)
    return multinomial_indices(n_samples, p)
end

function multinomial_indices_linear(n::Int, p::AbstractArray)
    #= multinomial sampling for when n is small, and length(p)
       is large. This performs n * (linear) scans of
       the vector `p`, with very little extra overhead.
       More efficiency could be achieved here with larger n.
    =#
    m = length(p)
    x = zeros(Int32, n)

    function linearsearch(p::AbstractArray, m::Int64, rn)
        cs = 0.0
        for ii in 1:m
            @inbounds cs += p[ii]
            if cs > rn
                return ii
            end
        end
        return m
    end

    for i in 1:n
        rn = rand()
        @inbounds x[i] = linearsearch(p, m, rn)
    end
    return x
end

function multinomial_indices_binsearch(n::Int, p::Vector{T}) where T <: AbstractFloat
    #= multinomial sampling that makes sense if n ≈ 10.
       Calculates the cumsum of probabilities and uses binary search on a
       uniform random for each draw to return the index.
    =#
    x = zeros(Int32, n)
    cump = cumsum(p)
    for i in 1:n
        x[i] = searchsortedfirst(cump, rand()*cump[end])
    end
    return x
end

#=================================================================================
                        Quasi Monte Carlo
 ==================================================================================#

randomised_sobol(n, d) = randomised_sobol(GLOBAL_RNG, n, d)
uniform_rand_sobol(n, lims...) = uniform_rand_sobol(GLOBAL_RNG, n, lims...)
sobol_gaussian(n, d) = sobol_gaussian(GLOBAL_RNG, n, d)

function randomised_sobol(rng::MersenneTwister, n, d)
    s = SobolSeq(d)
    p = reduce(hcat, [next!(s) for i = 1:n])
    ϵ = rand(rng, d)
    prand = [(p[j,:] .+ ϵ[j]) .% 1.0 for j in 1:d]
    return hcat(prand...)
end

# randomised Sobol within rectangle
function uniform_rand_sobol(rng::MersenneTwister, n, lims...)
    d = length(lims)
    rsob = randomised_sobol(rng, n, d)
    for (i, interval) in enumerate(lims)
        @assert (length(interval) == 2) format("interval {:d} does not have length 2", i)
        rsob[:,i] .*= diff(interval)
        rsob[:,i] .+= interval[1]
    end
    return rsob
end

# randomised Sobol Gaussian random variates
function sobol_gaussian(rng::MersenneTwister, n, d)
    (n == 0) && return zeros(d,0)'
    s = SobolSeq(d)
    p = reduce(hcat, [next!(s) for i = 1:n])
    ϵ = rand(rng, d)
    prand = [(p[j,:] .+ ϵ[j]) .% 1.0 for j in 1:d]
    p = reduce(vcat, [quantile.(Normal(), prand[j])' for j in 1:d])'
    return p
end

#=================================================================================
                                Misc stuff
 ==================================================================================#

function psd_matrix(d; ϵ=0.)
    A = randn(d, d)
    U = eigvecs((A + A')/2) # (Don't really have to divide by 2)
    P = U*diagm(0=>abs.(randn(d)) .+ ϵ)*U';
    return (P + P')/2
end

end  # module


text:
"""
    tour = sortperm_msa(msa::AbstractMultipleSequenceAlignment)

Order the sequences in `msa` to minimize the "tour length" visiting each sequence once.
The length between sequences is defined at `100 - percentsimilarity(seq1, seq2)`.

This can be useful for graphical or alignment display by grouping obviously-similar
sequences near one another.
"""
function sortperm_msa(msa::AbstractMultipleSequenceAlignment)
    sim = percentsimilarity(msa)
    D = 100 .- Matrix(sim)
    tour, _ = solve_tsp(D; quality_factor=80)
    pop!(tour)  # returns a closed path, but that duplicates the first/last entry
    return tour
end

"""
    filter_species!(msa::AbstractMultipleSequenceAlignment, speciesname::AbstractString)

Remove all sequences from `msa` except those with [`species(sequencename)`](@ref) equal to `speciesname`.
"""
function filter_species!(msa::AbstractMultipleSequenceAlignment, speciesname::AbstractString)
    mask = map(x -> species(x) == speciesname, sequencenames(msa))
    filtersequences!(msa, mask)
end

"""
    filter_long!(msa::AbstractMultipleSequenceAlignment, minres::Real)

Remove all sequences from `msa` with fewer than `minres` matching residues.
"""
function filter_long!(msa::AbstractMultipleSequenceAlignment, minres::Real)
    # Get rid of short sequences
    nresidues = map(eachrow(msa)) do v
        sum(!=(Residue('-')), v)
    end
    mask = nresidues .> minres
    filtersequences!(msa, mask)
end

# Move this to MIToS?
if !hasmethod(getsequencemapping, Tuple{AnnotatedAlignedSequence})
    function MIToS.MSA.getsequencemapping(seq::AnnotatedAlignedSequence)
        getsequencemapping(seq, sequencenames(seq)[1])
    end
    function MIToS.MSA.getsequencemapping(msa::Union{AnnotatedAlignedSequence,AnnotatedMultipleSequenceAlignment}, seq_id::String)
        MIToS.MSA._str2int_mapping(getannotsequence(msa, seq_id, "SeqMap"))
    end
    function MIToS.MSA.getsequencemapping(msa::AnnotatedMultipleSequenceAlignment, seqid::Regex)
        id = findfirst(str -> occursin(seqid, str), sequencenames(msa))
        getsequencemapping(msa, id)
    end
end

struct SequenceMapping <: AbstractVector{Int}
    seqmap::Vector{Int}
end
SequenceMapping(seq::AnnotatedAlignedSequence) = SequenceMapping(getsequencemapping(seq))

Base.size(sm::SequenceMapping) = size(sm.seqmap)
Base.getindex(sm::SequenceMapping, i::Int) = sm.seqmap[i]
Base.setindex!(sm::SequenceMapping, v, i::Int) = sm.seqmap[i] = v
Base.similar(::SequenceMapping, ::Type{Int}, dims::Dims{1}) = SequenceMapping(Vector{Int}(undef, dims))

# We restrict the following to `Vector` to prevent ambiguities
Base.getindex(fullseqvec::Vector{T}, sm::SequenceMapping) where T = Union{T,Nothing}[i == 0 ? nothing : fullseqvec[i] for i in sm.seqmap]

function Base.setindex!(fullseqvec::Vector, colvalues::AbstractVector, sm::SequenceMapping)
    idxnz = sm.seqmap .> 0
    fullseqvec[sm.seqmap[idxnz]] = colvalues[idxnz]
end

"""
    chimerax_script(filename, uprot_list, msa::AnnotatedMultipleSequenceAlignment, colidxs;
                    dir=pwd(), chain_transparency=80, styles=Dict{Int,String}(), extras=String[])

Create a [chimerax]() visualization script with name `filename`. `uprot_list` is a list of UniProtX names that you
want to visualize.  The AlphaFold PDB files for these proteins should be `dir`. `msa` is a Multiple Sequence alignment
and `colidxs` specifies the column indices in `msa` that you'd like to visualize.

`chain_transparency` sets the transparency on the ribbon diagrams (0 = not transparent).
`styles` can be used to affect the display, e.g., `Dict(k => "@SD sphere")` would cause methionines at column index `k`
to be displayed with the sulfur in sphere mode. `extras` can be used to hand-specify a number of additional commands;
this can be useful if, for example, the `msa` has occasional misalignments.

# Examples

Suppose you have the `msa` for rhodopsin (mouse: P15409), then:

```julia
chimerax_script("myscript.cxc", ["P15409"], msa, [i1, i2, i3])
```

where `i1` through `i3` are column-indices in the `msa` that you'd like to view.
"""
function chimerax_script(filename, uprot_list, msa::AnnotatedMultipleSequenceAlignment, colidxs;
                         dir=pwd(), chain_transparency=80, styles=Dict{Int,String}(), extras=String[])
    open(filename, "w") do io
        for (i, p) in enumerate(uprot_list)
            afname = alphafoldfile(p)
            println(io, "open ", joinpath(dir, afname))
            seqidx = findfirst(str -> startswith(str, p), sequencenames(msa))
            sm = getsequencemapping(msa, seqidx)
            for c in colidxs
                ridx = sm[c]
                if iszero(ridx)
                    @warn "column $c not set in $p"
                    continue
                end
                println(io, "show #", i, " :", ridx)
                style = get(styles, c, nothing)
                if style !== nothing
                    println(io, "style #", i, " :", ridx, " ", style)
                end
            end
        end
        for ex in extras
            println(io, ex)
        end
        println(io, "transparency #1-", length(uprot_list), ' ', chain_transparency, " target c")
        if !any(str -> startswith(str, "align"), extras)
            println(io, "matchmaker #2-", length(uprot_list), " to #1")
        end
    end
end

ext: jl


text:
import Base.println, Base.flush

# Define valid syslog levels and common aliases.
SYSLOG_LEVELS = [:emerg, :alert, :crit, :err, :warning, :notice, :info, :debug]
ALIAS_LEVELS = Dict(:warn=>:warning, :error=>:err, :emergency=>:emerg, :critical=>:crit)

# Define valid syslog facilities (even kern, which isn't accessible from userspace).
FACILITIES = [
    :auth, :authpriv, :cron, :daemon, :ftp, :kern, :lpr, :mail, :news, :syslog, :user,
    :uucp, :local0, :local1, :local2, :local3, :local4, :local5, :local6, :local7, :security
]

type Syslog <: IO
    facility::Symbol
    tag::AbstractString
    pid::Integer

    function Syslog(facility=:local0, tag="julia", tag_pid::Bool=false)
        # Verify that logger is installed.
        if !success(`which logger`)
            error("syslog output only available on systems with logger installed")
        end

        if !(facility in FACILITIES)
            error("invalid logging facility: $(facility)")
        end

        # Appends Julia's process ID to the log tag if show_pid is specified.
        pid = tag_pid ? getpid() : -1

        new(facility, tag, pid)
    end
end

function println(log::Syslog, level::Symbol, msg::AbstractString)
    level = get(ALIAS_LEVELS, level, level)
    if !(level in SYSLOG_LEVELS)
        throw(ErrorException("invalid level: $level"))
    end

    tag = log.tag * (log.pid > -1 ? "[$(log.pid)]" : "")

    # @mendable is required to allow the test cases to modify the run command for testing.
    @mendable run(`logger -t $(tag) -p $(log.facility).$level $msg`)
end

function println(log::Syslog, level::AbstractString, msg::AbstractString)
    println(log, symbol(lowercase(level)), msg)
end

# Defined just in case somebody decides to call flush, which is totally unnecessary.
flush(log::Syslog) = log


text:
using ElectrostaticPIC1D, Random, Test

Random.seed!(0)

@testset "ElectrostaticPIC1D tests" begin
  include("./BasisFunctions.jl")
  include("./Fields.jl")
  include("./Particles.jl")
  #include("./Simulations.jl")
end


text:
"""
    Mlmnet_cv(MLMNets, lambdas, data, rowFolds, colFolds)

Type for storing the results of running cross-validation for `mlmnet`

"""
mutable struct Mlmnet_cv

    # Mlmnet objects
    MLMNets::Array{Mlmnet, 1}
    # Lambda penalties
    lambdas::Array{Float64,1}

    # Response and predictor matrices
    data::RawData
    # Row folds
    rowFolds::Array{Array,1}
    # Column folds
    colFolds::Array{Array,1}

    # These will be generated and should not be supplied as input to the
    # constructor.
    # Test MSE for each of the CV folds for each lambda
    mse::Array{Float64, 2}
    # Proportion of zero interaction coefficients for each of the CV folds
    # for each lambda
    propZero::Array{Float64, 2}

    Mlmnet_cv(MLMNets, lambdas, data, rowFolds, colFolds, dig) =
        new(MLMNets, lambdas, data, rowFolds, colFolds,
            calc_mse(MLMNets, data, lambdas, rowFolds, colFolds),
            calc_prop_zero(MLMNets, lambdas; dig=dig))
end

"""
    mlmnet_cv(fun, data, lambdas, rowFolds, colFolds;
              isXIntercept, isZIntercept, isXReg, isZReg,
              isXInterceptReg, isZInterceptReg, isStandardize, isVerbose,
              stepsize, setStepsize, dig, funArgs...)

Performs cross-validation for `mlmnet` using row and column folds from user
input.

# Arguments

- fun = function that applies an L1-penalty estimate method
- data = RawData object
- lambdas = 1d array of floats consisting of lambda penalties in descending
  order. If they are not in descending order, they will be sorted.
- rowFolds = 1d array of arrays (one array for each fold), each containing
  the indices for a row fold; must be same length as colFolds. Can be
  generated with a call to `make_folds`, which is based on `Kfold` from the
  MLBase package.
- colFolds = 1d array of arrays (one array for each fold), each containing
  the indices for a column fold; must be same length as rowFolds. Can be
  generated with a call to `make_folds`, which is based on `Kfold` from the
  MLBase package

# Keyword arguments

- isXIntercept = boolean flag indicating whether or not to include an `X`
  intercept (row main effects). Defaults to `true`.
- isZIntercept = boolean flag indicating whether or not to include a `Z`
  intercept (column main effects). Defaults to `true`.
- isXReg = 1d array of bit flags indicating whether or not to regularize each
  of the `X` (row) effects. Defaults to 2d array of `true`s with length
  equal to the number of `X` effects (equivalent to `data.p`).
- isZReg = 1d array of bit flags indicating whether or not to regularize each
  of the `Z` (column) effects. Defaults to 2d array of `true`s with length
  equal to the number of `Z` effects (equivalent to `data.q`).
- isXInterceptReg = boolean flag indicating whether or not to regularize the
  `X` intercept Defaults to `false`.
- isZInterceptReg = boolean flag indicating whether or not to regularize the
  `Z` intercept. Defaults to `false`.
- isStandardize = boolean flag indicating if the columns of `X` and `Z`
  should be standardized (to mean 0, standard deviation 1). Defaults to `true`.
- isVerbose = boolean flag indicating whether or not to print messages.
  Defaults to `true`.
- stepsize = float; step size for updates (irrelevant for coordinate
  descent and when `setStepsize` is set to `true` for `ista!` and `fista!`).
  Defaults to `0.01`.
- setStepsize = boolean flag indicating whether the fixed step size should be
  calculated (for `ista!` and `fista!`). Defaults to `true`.
- dig = integer; digits of precision for zero coefficients. Defaults to 12.
- funArgs = variable keyword arguments to be passed into `fun`

# Value

An Mlmnet_cv object.

# Some notes

This is the base `mlmnet_cv` function that all other variants call. Folds
are computed in parallel when possible.

"""
function mlmnet_cv(fun::Function, data::RawData,
                   lambdas::AbstractArray{Float64,1},
                   rowFolds::Array{Array{Int64,1},1},
                   colFolds::Array{Array{Int64,1},1};
                   isXIntercept::Bool=true, isZIntercept::Bool=true,
                   isXReg::BitArray{1}=trues(size(get_X(data),2)),
                   isZReg::BitArray{1}=trues(size(get_Z(data),2)),
                   isXInterceptReg::Bool=false, isZInterceptReg::Bool=false,
                   isStandardize::Bool=true, isVerbose::Bool=true,
                   stepsize::Float64=0.01, setStepsize::Bool=true,
                   dig::Int64=12, funArgs...)

    # Number of folds
    nFolds = length(rowFolds)
    # Ensure that the number of row folds is the same as the number of column
    # folds
    if nFolds != length(colFolds)
        error("Number of row folds must equal number of column folds. ")
    else
        println_verbose(string("Performing ", nFolds,
                               "-fold cross validation."), isVerbose)
    end

    # Generate list of RawData objects that hold out each of the folds
    dataFolds = Array{RawData}(undef, nFolds)
    for i in 1:nFolds
        dataFolds[i] = RawData(Response(get_Y(data)[rowFolds[i],
                                                    colFolds[i]]),
                               Predictors(get_X(data)[rowFolds[i],:],
                                          get_Z(data)[colFolds[i],:],
                                          data.predictors.isXIntercept,
                                          data.predictors.isZIntercept))
    end

    # Run mlmnet on each RawData object, in parallel when possible
    MLMNets = Distributed.pmap(data -> mlmnet(fun, data, lambdas;
                                              isXIntercept=isXIntercept,
                                              isZIntercept=isZIntercept,
                                              isXReg=isXReg, isZReg=isZReg,
                                              isXInterceptReg=isXInterceptReg,
                                              isZInterceptReg=isZInterceptReg,
                                              isStandardize=isStandardize,
                                              isVerbose=isVerbose,
                                              stepsize=stepsize,
                                              setStepsize=setStepsize,
                                              funArgs...), dataFolds)

    return Mlmnet_cv(MLMNets, lambdas, data, rowFolds, colFolds, dig)
end

"""
    mlmnet_cv(fun, data, lambdas, rowFolds, nColFolds;
              isXIntercept, isZIntercept, isXReg, isZReg,
              isXInterceptReg, isZInterceptReg, isStandardize, isVerbose,
              stepsize, setStepsize, dig, funArgs...)

Performs cross-validation for `mlmnet` using row folds from user input and
non-overlapping column folds randomly generated using a call to `make_folds`.
Calls the base `mlmnet_cv` function.

# Arguments

- fun = function that applies an L1-penalty estimate method
- data = RawData object
- lambdas = 1d array of floats consisting of lambda penalties in descending
  order. If they are not in descending order, they will be sorted.
- rowFolds = 1d array of arrays (one array for each fold), each containing
  the indices for a row fold. Can be generated with a call to `make_folds`,
  which is based on `Kfold` from the MLBase package.
- nColFolds = integer corresponding to the number of column folds to randomly
  generate. Can be either `length(rowFolds)` or `1` (to use all columns in
  every fold).

# Keyword arguments

- isXIntercept = boolean flag indicating whether or not to include an `X`
  intercept (row main effects). Defaults to `true`.
- isZIntercept = boolean flag indicating whether or not to include a `Z`
  intercept (column main effects). Defaults to `true`.
- isXReg = 1d array of bit flags indicating whether or not to regularize each
  of the `X` (row) effects. Defaults to 2d array of `true`s with length
  equal to the number of `X` effects (equivalent to `data.p`).
- isZReg = 1d array of bit flags indicating whether or not to regularize each
  of the `Z` (column) effects. Defaults to 2d array of `true`s with length
  equal to the number of `Z` effects (equivalent to `data.q`).
- isXInterceptReg = boolean flag indicating whether or not to regularize the
  `X` intercept Defaults to `false`.
- isZInterceptReg = boolean flag indicating whether or not to regularize the
  `Z` intercept. Defaults to `false`.
- isStandardize = boolean flag indicating if the columns of `X` and `Z`
  should be standardized (to mean 0, standard deviation 1). Defaults to `true`.
- isVerbose = boolean flag indicating whether or not to print messages.
  Defaults to `true`.
- stepsize = float; step size for updates (irrelevant for coordinate
  descent and when `setStepsize` is set to `true` for `ista!` and `fista!`).
  Defaults to `0.01`.
- setStepsize = boolean flag indicating whether the fixed step size should be
  calculated (for `ista!` and `fista!`). Defaults to `true`.
- dig = integer; digits of precision for zero coefficients. Defaults to 12.
- funArgs = variable keyword arguments to be passed into `fun`

# Value

An Mlmnet_cv object.

# Some notes

Folds are computed in parallel when possible.

"""
function mlmnet_cv(fun::Function, data::RawData,
                   lambdas::AbstractArray{Float64,1},
                   rowFolds::Array{Array{Int64,1},1}, nColFolds::Int64;
                   isXIntercept::Bool=true, isZIntercept::Bool=true,
                   isXReg::BitArray{1}=trues(size(get_X(data),2)),
                   isZReg::BitArray{1}=trues(size(get_Z(data),2)),
                   isXInterceptReg::Bool=false, isZInterceptReg::Bool=false,
                   isStandardize::Bool=true, isVerbose::Bool=true,
                   stepsize::Float64=0.01, setStepsize::Bool=true,
                   dig::Int64=12, funArgs...)

    # Generate random column folds
    colFolds = make_folds(data.m, nColFolds, length(rowFolds))

    # Pass in user input row folds and randomly generated column folds to the
    # base mlmnet_cv function
    mlmnet_cv(fun, data, lambdas, rowFolds, colFolds;
              isXIntercept=isXIntercept, isZIntercept=isZIntercept,
              isXReg=isXReg, isZReg=isZReg,
              isXInterceptReg=isXInterceptReg,
              isZInterceptReg=isZInterceptReg,
              isVerbose=isVerbose, isStandardize=isStandardize,
              stepsize=stepsize, setStepsize=setStepsize, dig=dig, funArgs...)
end

"""
    mlmnet_cv(fun, data, lambdas, nRowFolds, colFolds;
              isXIntercept, isZIntercept, isXReg, isZReg,
              isXInterceptReg, isZInterceptReg, isStandardize, isVerbose,
              stepsize, setStepsize, dig, funArgs...)

Performs cross-validation for `mlmnet` using non-overlapping row folds
randomly generated using a call to `make_folds` and column folds from user
input. Calls the base `mlmnet_cv` function.

# Arguments

- fun = function that applies an L1-penalty estimate method
- data = RawData object
- lambdas = 1d array of floats consisting of lambda penalties in descending
  order. If they are not in descending order, they will be sorted.
- nRowFolds = integer corresponding to the number of row folds to randomly
  generate. Can be either `length(colFolds)` or `1`
  (to use all rows in every fold).
- colFolds = 1d array of arrays (one array for each fold), each containing
  the indices for a column fold. Can be generated with a call to `make_folds`,
  which is based on `Kfold` from the MLBase package.

# Keyword arguments

- isXIntercept = boolean flag indicating whether or not to include an `X`
  intercept (row main effects). Defaults to `true`.
- isZIntercept = boolean flag indicating whether or not to include a `Z`
  intercept (column main effects). Defaults to `true`.
- isXReg = 1d array of bit flags indicating whether or not to regularize each
  of the `X` (row) effects. Defaults to 2d array of `true`s with length
  equal to the number of `X` effects (equivalent to `data.p`).
- isZReg = 1d array of bit flags indicating whether or not to regularize each
  of the `Z` (column) effects. Defaults to 2d array of `true`s with length
  equal to the number of `Z` effects (equivalent to `data.q`).
- isXInterceptReg = boolean flag indicating whether or not to regularize the
  `X` intercept Defaults to `false`.
- isZInterceptReg = boolean flag indicating whether or not to regularize the
  `Z` intercept. Defaults to `false`.
- isStandardize = boolean flag indicating if the columns of `X` and `Z`
  should be standardized (to mean 0, standard deviation 1). Defaults to `true`.
- isVerbose = boolean flag indicating whether or not to print messages.
  Defaults to `true`.
- stepsize = float; step size for updates (irrelevant for coordinate
  descent and when `setStepsize` is set to `true` for `ista!` and `fista!`).
  Defaults to `0.01`.
- setStepsize = boolean flag indicating whether the fixed step size should be
  calculated (for `ista!` and `fista!`). Defaults to `true`.
- dig = integer; digits of precision for zero coefficients. Defaults to 12.
- funArgs = variable keyword arguments to be passed into `fun`

# Value

An Mlmnet_cv object.

# Some notes

Folds are computed in parallel when possible.

"""
function mlmnet_cv(fun::Function, data::RawData,
                   lambdas::AbstractArray{Float64,1},
                   nRowFolds::Int64, colFolds::Array{Array{Int64,1},1};
                   isXIntercept::Bool=true, isZIntercept::Bool=true,
                   isXReg::BitArray{1}=trues(size(get_X(data),2)),
                   isZReg::BitArray{1}=trues(size(get_Z(data),2)),
                   isXInterceptReg::Bool=false, isZInterceptReg::Bool=false,
                   isStandardize::Bool=true, isVerbose::Bool=true,
                   stepsize::Float64=0.01, setStepsize::Bool=true,
                   dig::Int64=12, funArgs...)

    # Generate random row folds
    rowFolds = make_folds(data.n, nRowFolds, length(colFolds))

    # Pass in randomly generated row folds and user input column folds to the
    # base mlmnet_cv function
    mlmnet_cv(fun, data, lambdas, rowFolds, colFolds;
              isXIntercept=isXIntercept, isZIntercept=isZIntercept,
              isXReg=isXReg, isZReg=isZReg,
              isXInterceptReg=isXInterceptReg,
              isZInterceptReg=isZInterceptReg,
              isVerbose=isVerbose, isStandardize=isStandardize,
              stepsize=stepsize, setStepsize=setStepsize, dig=dig, funArgs...)
end

"""
    mlmnet_cv(fun, data, lambdas, nRowFolds, nColFolds;
              isXIntercept, isZIntercept, isXReg, isZReg,
              isXInterceptReg, isZInterceptReg, isStandardize, isVerbose,
              stepsize, setStepsize, dig, funArgs...)

Performs cross-validation for `mlmnet` using non-overlapping row and column
folds randomly generated using calls to `make_folds`. Calls the base
`mlmnet_cv` function.

# Arguments

- fun = function that applies an L1-penalty estimate method
- data = RawData object
- lambdas = 1d array of floats consisting of lambda penalties in descending
  order. If they are not in descending order, they will be sorted.
- nRowFolds = integer corresponding to the number of row folds to randomly
  generate. Can be either equal to `nColFolds` or `1`
  (to use all rows in every fold).
- nColFolds = integer corresponding to the number of column folds to randomly
  generate. Can be either equal to `nRowFolds` or `1` (to use all columns in
  every fold).

# Keyword arguments

- isXIntercept = boolean flag indicating whether or not to include an `X`
  intercept (row main effects). Defaults to `true`.
- isZIntercept = boolean flag indicating whether or not to include a `Z`
  intercept (column main effects). Defaults to `true`.
- isXReg = 1d array of bit flags indicating whether or not to regularize each
  of the `X` (row) effects. Defaults to 2d array of `true`s with length
  equal to the number of `X` effects (equivalent to `data.p`).
- isZReg = 1d array of bit flags indicating whether or not to regularize each
  of the `Z` (column) effects. Defaults to 2d array of `true`s with length
  equal to the number of `Z` effects (equivalent to `data.q`).
- isXInterceptReg = boolean flag indicating whether or not to regularize the
  `X` intercept Defaults to `false`.
- isZInterceptReg = boolean flag indicating whether or not to regularize the
  `Z` intercept. Defaults to `false`.
- isStandardize = boolean flag indicating if the columns of `X` and `Z`
  should be standardized (to mean 0, standard deviation 1). Defaults to `true`.
- isVerbose = boolean flag indicating whether or not to print messages.
  Defaults to `true`.
- stepsize = float; step size for updates (irrelevant for coordinate
  descent and when `setStepsize` is set to `true` for `ista!` and `fista!`).
  Defaults to `0.01`.
- setStepsize = boolean flag indicating whether the fixed step size should be
  calculated (for `ista!` and `fista!`). Defaults to `true`.
- dig = integer; digits of precision for zero coefficients. Defaults to 12.
- funArgs = variable keyword arguments to be passed into `fun`

# Value

An Mlmnet_cv object.

# Some notes

Folds are computed in parallel when possible.

The default method for choosing the fixed step size for `fista!` or `ista!`
is to use the reciprocal of the product of the maximum eigenvalues of
`X*transpose(X)` and `Z*transpose(Z)`. This is computed when `fista!` or
`ista!` is passed into the `fun` argument and `setStepsize` is set to `true`.
If `setStepsize` is set to `false`, the value of the `stepsize` argument will
be used as the fixed step size. Note that obtaining the eigenvalues when `X`
and/or `Z` are very large may exceed computational limitations.

Specifying a good starting step size (`stepsize`) and multiplying factor
(`gamma`) when `fista_bt!` is passed into the `fun` argument can be difficult.
Shrinking the step size too gradually can result in slow convergence. Doing so
too quickly can cause the criterion to diverge. We have found that setting
`stepsize` to 0.01 often works well in practice; choice of `gamma` appears to
be less consequential.

"""
function mlmnet_cv(fun::Function, data::RawData, lambdas::Array{Float64,1},
                   nRowFolds::Int64=10, nColFolds::Int64=10;
                   isXIntercept::Bool=true, isZIntercept::Bool=true,
                   isXReg::BitArray{1}=trues(size(get_X(data),2)),
                   isZReg::BitArray{1}=trues(size(get_Z(data),2)),
                   isXInterceptReg::Bool=false, isZInterceptReg::Bool=false,
                   isStandardize::Bool=true, isVerbose::Bool=true,
                   stepsize::Float64=0.01, setStepsize::Bool=true,
                   dig::Int64=12, funArgs...)

    # Generate random row and column folds
    rowFolds = make_folds(data.n, nRowFolds, max(nRowFolds, nColFolds))
    colFolds = make_folds(data.m, nColFolds, max(nRowFolds, nColFolds))

    # Pass in randomly generated row and column folds to the base mlmnet_cv
    # function
    mlmnet_cv(fun, data, lambdas, rowFolds, colFolds;
              isXIntercept=isXIntercept, isZIntercept=isZIntercept,
              isXReg=isXReg, isZReg=isZReg,
              isXInterceptReg=isXInterceptReg,
              isZInterceptReg=isZInterceptReg,
              isVerbose=isVerbose, isStandardize=isStandardize,
              stepsize=stepsize, setStepsize=setStepsize, dig=dig, funArgs...)
end


text:
# Infastructure to support generating overloads from rules.
_package_hook(::Base.PkgId) = refresh_rules()

# Holds all the hook functions that are invokes when a new rule is defined
const RRULE_DEFINITION_HOOKS = Function[]
const FRULE_DEFINITION_HOOKS = Function[]
_hook_list(::typeof(rrule)) = RRULE_DEFINITION_HOOKS
_hook_list(::typeof(frule)) = FRULE_DEFINITION_HOOKS

"""
    on_new_rule(hook, frule | rrule)

Register a `hook` function to run when new rules are defined.
The hook receives a signature type-type as input, and generally will use `eval` to define
an overload of an AD system's overloaded type
For example, using the signature type `Tuple{typeof(+), Real, Real}` to make
`+(::DualNumber, ::DualNumber)` call the `frule` for `+`.
A signature type tuple always has the form:
`Tuple{typeof(operation), typeof{pos_arg1}, typeof{pos_arg2}...}`, where `pos_arg1` is the
first positional argument.

The hooks are automatically run on new rules whenever a package is loaded.
They can be manually triggered by [`refresh_rules`](@ref).
When a hook is first registered with `on_new_rule` it is run on all existing rules.
"""
function on_new_rule(hook_fun, rule_kind)
    # apply the hook to the existing rules
    ret = map(_rule_list(rule_kind)) do method
        sig = _primal_sig(rule_kind, method)
        _safe_hook_fun(hook_fun, sig)
    end

    # register hook for new rules -- so all new rules get this function applied
    push!(_hook_list(rule_kind), hook_fun)
    return ret
end

"""
    clear_new_rule_hooks!(frule|rrule)

Clears all hooks that were registered with corresponding [`on_new_rule`](@ref).
This is useful for while working interactively to define your rule generating hooks.
If you previously wrong an incorrect hook, you can use this to get rid of the old one.

!!! warning
    This absolutely should not be used in a package, as it will break any other AD system
    using the rule hooks that might happen to be loaded.
"""
clear_new_rule_hooks!(rule_kind) = empty!(_hook_list(rule_kind))

###########################################################################################

"""
    _rule_list(frule | rrule)

Returns a list of all the methods of the currently defined rules of the given kind.
Excluding the fallback rule (that return `nothing` for every input) and `@opt_out` opted out
rules, and excluding rules that require a particular `RuleConfig`.
"""
function _rule_list(rule_kind::Union{typeof(frule), typeof(rrule)})
    opted_out = Set(arg_type_tuple(m.sig) for m in _no_rule_list(rule_kind))
    return Iterators.filter(methods(rule_kind)) do m
        return !_requires_config(m) && arg_type_tuple(m.sig) ∉ opted_out
    end
end

"check if this rule requires a particular configuation (`RuleConfig`)"
_requires_config(m::Method) = m.sig <: Tuple{Any, RuleConfig, Vararg}

_no_rule_list(::typeof(rrule)) = methods(ChainRulesCore.no_rrule)
_no_rule_list(::typeof(frule)) = methods(ChainRulesCore.no_frule)

arg_type_tuple(d::DataType) = Tuple{d.parameters[2:end]...}
function arg_type_tuple(d::UnionAll)
    body = Base.unwrap_unionall(d)
    body_tt = arg_type_tuple(body)
    return Base.rewrap_unionall(body_tt, d)
end

######################################################################

const LAST_REFRESH_RRULE = Ref(0)
const LAST_REFRESH_FRULE = Ref(0)
last_refresh(::typeof(frule)) = LAST_REFRESH_FRULE
last_refresh(::typeof(rrule)) = LAST_REFRESH_RRULE

"""
    refresh_rules()
    refresh_rules(frule | rrule)

This triggers all [`on_new_rule`](@ref) hooks to run on any newly defined rules.
It is *automatically* run when ever a package is loaded.
It can also be manually called to run it directly, for example if a rule was defined
in the REPL or within the same file as the AD function.
"""
function refresh_rules()
    refresh_rules(frule);
    refresh_rules(rrule)
end

function refresh_rules(rule_kind)
    isempty(_rule_list(rule_kind)) && return  # if no hooks, exit early, nothing to run
    already_done_world_age = last_refresh(rule_kind)[]
    for method in _rule_list(rule_kind)
        _defined_world(method) < already_done_world_age && continue
        sig = _primal_sig(rule_kind, method)
        _trigger_new_rule_hooks(rule_kind, sig)
    end

    last_refresh(rule_kind)[] = _current_world()
    return nothing
end

@static if VERSION >= v"1.2"
    _current_world() = Base.get_world_counter()
    _defined_world(method) = method.primary_world
else
    _current_world() = ccall(:jl_get_world_counter, UInt, ())
    _defined_world(method) = method.min_world
end

"""
    _primal_sig(frule|rule, rule_method | rule_sig)

Returns the signature as a `Tuple{function_type, arg1_type, arg2_type,...}`.
"""
_primal_sig(rule_kind, method::Method) = _primal_sig(rule_kind, method.sig)
function _primal_sig(::typeof(frule), rule_sig::DataType)
    @assert rule_sig.parameters[1] == typeof(frule)
    # need to skip frule and the deriviative info, so starting from the 3rd
    return Tuple{rule_sig.parameters[3:end]...}
end
function _primal_sig(::typeof(rrule), rule_sig::DataType)
    @assert rule_sig.parameters[1] == typeof(rrule)
    # need to skip rrule so starting from the 2rd
    return Tuple{rule_sig.parameters[2:end]...}
end
function _primal_sig(rule_kind, rule_sig::UnionAll)
    # This looks a lot like Base.unwrap_unionall and Base.rewrap_unionall, but using those
    # seems not to work
    p_sig = _primal_sig(rule_kind, rule_sig.body)
    return UnionAll(rule_sig.var, p_sig)
end

function _trigger_new_rule_hooks(rule_kind, sig)
    for hook_fun in _hook_list(rule_kind)
        _safe_hook_fun(hook_fun, sig)
    end
end

function _safe_hook_fun(hook_fun, sig)
    try
        hook_fun(sig)
    catch err
        @error "Error triggering hook" hook_fun sig exception=(err, catch_backtrace())
    end
end


text:
#----------------------------------------------------------------------------------------------#
#            Package User Interface -- Suitable types must implement the following             #
#----------------------------------------------------------------------------------------------#

import Base: cp

# Empty generic functions to specify interface. (from:
# https://docs.julialang.org/en/v1/manual/methods/#Empty-generic-functions-1).

# Model range functions
# ---------------------

"""
Function to return the minimum allowed temperature for state specification.
"""
function minT end           # Minimum allowed temperature

"""
Function to return the maximum allowed temperature for state specification.
"""
function maxT end           # Maximum allowed temperature

"""
Function to return the minimum allowed pressure for state specification.
"""
function minP end           # Minimum allowed pressure

"""
Function to return the maximum allowed pressure for state specification.
"""
function maxP end           # Maximum allowed pressure

"""
Function to return the minimum allowed specific volume for state specification.
"""
function minv end           # Minimum allowed sp. volume

"""
Function to return the maximum allowed specific volume for state specification.
"""
function maxv end           # Maximum allowed sp. volume

export minT, maxT, minP, maxP, minv, maxv

# EoS-independent properties
# --------------------------

"""
Function to return the molecular "weight" (mass).
"""
function M end              # Substance molecular "weight", i.e., molar mass

"""
Function to return the gas constant.
"""
function R end              # Substance gas constant

"""
Function to return the isobaric specific heat.
"""
function cp end             # Substance/Heat model isobaric specific heat

"""
Function to return the isochoric specific heat.
"""
function cv end             # Substance/Heat model isochoric specific heat

"""
Function to return the isobaric-to-isochoric specific heat ratio.\n
Aliased to `γ`: U+3b3, or `\\gamma<TAB>` on the REPL.
"""
function gamma end          # Substance/Heat model specific heat ratio
γ = gamma                   # γ: U+3b3

"""
Function to return the specific formation enthalpy at the standard state.\n
Aliased to `hf0`.
"""
function hf° end            # Substance formation enthalpy # °: U+b0
hf0 = hf°

"""
Function to return the specific formation internal energy at the standard state.\n
Aliased to `uf0`.
"""
function uf° end            # Substance formation intenrgy # °: U+b0
uf0 = uf°

"""
Function to return the reference state.
"""
function ref end            # Substance model reference state

export M, R, cp, cv, gamma, hf°, uf°, ref
export               γ    , hf0, uf0

# Specific heat temperature integrals
# -----------------------------------

"""
Function to return the process specific internal energy variation.\n
Aliased to `Du`.
"""
function Δu end             # Substance/Heat model cv integral # Δ: U+394
Du = Δu

"""
Function to return the process specific enthalpy variation.\n
Aliased to `Dh`.
"""
function Δh end             # Substance/Heat model cp integral # Δ: U+394
Dh = Δh

"""
Function to return the process specific ideal gas partial entropy variation.\n
Aliased to `Ds0`.
"""
function Δs° end            # Substance/Heat model cp/T integral # Δ: U+394; °: U+b0
Ds0 = Δs°

export Δu, Δh, Δs°
export Du, Dh, Ds0

# Substance properties
# --------------------

"""
Function to return the pressure.
"""
function P end              # Thermodynamic (absolute) pressure

"""
Function to return the temperature.
"""
function T end              # Thermodynamic (absolute) temperature

"""
Function to return the specific volume.
"""
function v end              # Specific volume

"""
Function to return the density.\n
Aliased to `ρ`: `\\rho<TAB>`.
"""
function rho end            # Mass density
ρ = rho                     # ρ: U+3c1

"""
Function to return the specific flow work.
"""
function Pv end             # Specific flow work

"""
Function to return the specific internal energy.
"""
function u end              # Specific internal energy

"""
Function to return the specific enthalpy.
"""
function h end              # Specific enthalpy

"""
Function to return the specific entropy.
"""
function s end              # Specific entropy

"""
Function to return the specific Gibbs energy.
"""
function g end              # Specific Gibbs energy

"""
Function to return the specific Helmholtz energy.
"""
function a end              # Specific Helmholtz energy

"""
Function to return the quality.
"""
function x end              # Quality

"""
Function to return the Massieu function, \$j ≡ -a/T\$.
"""
function j end              # Massieu function j ≡ -a/T

"""
Function to return the Planck function, \$r ≡ -g/T\$.
"""
function r end              # Plank function r ≡ -g/T

export P, T, v, rho, Pv, u, h, s, g, a, x, j, r
export          ρ

# Derived substance properties
# ----------------------------

"""
Function to return the isentropic expansion coefficient, \$k ≡ (-v/P)(∂P/∂v)_s\$.
"""
function k end              # Isentropic expansion exponent,    k ≡ (-v/P)(∂P/∂v)_s

"""
Function to return the coefficient of volume expansion, \$β ≡ (1/v)(∂v/∂T)_P\$.\n
Aliased to `β`: `\\beta<TAB>`.
"""
function beta end           # Coefficient of volume expansion,  β ≡ (1/v)(∂v/∂T)_P
β = beta                    # β: U+3b2

"""
Function to return the isothermal compressibility, \$κ_T ≡ (-1/v)(∂v/∂P)_T\$.\n
Aliased to `κT`: `\\kappa<TAB>T`.
"""
function kappaT end         # Isothermal compressibility,       κT ≡ (-1/v)(∂v/∂P)_T
κT = kappaT                 # κ: U+3ba

"""
Function to return the isentropic compressibility, \$κ_s ≡ (-1/v)(∂v/∂P)_s\$.\n
Aliased to `κs`: `\\kappa<TAB>s`.
"""
function kappas end         # Isentropic compressibility,       κs ≡ (-1/v)(∂v/∂P)_s
κs = kappas                 # κ: U+3ba

"""
Function to return the specific heat difference, \$Δc ≡ cp - cv = vTβ²/κT\$.\n
Aliased to `Dc`.
"""
function Δc end             # Specific heat difference,         Δc ≡ cp - cv = vTβ²/κT
Dc = Δc

"""
Function to return the adiabatic sound velocity, \$cs² ≡ (∂P/∂ρ)_s = -v(∂P/∂v)_s\$.
"""
function cs end             # Adiabatic sound velocity,         cs² ≡ (∂P/∂ρ)_s = -v(∂P/∂v)_s

"""
Function to return the Mach number.
"""
function Ma end             # Mach number

"""
Function to return the Joule-Thomson coefficient, \$μJT ≡ (∂T/∂P)_h = v(Tβ-1)/cp\$.\n
Aliased to `μJT`: `\\mu<TAB>JT`.
"""
function muJT end           # Joule-Thomson coefficient,        μJT ≡ (∂T/∂P)_h = v(Tβ-1)/cp
μJT = muJT                  # μ: U+3bc

"""
Function to return the isentropic expansion coefficient, \$μs ≡ (∂T/∂P)_s = vTβ/cp\$.\n
Aliased to `μs`: `\\mu<TAB>s`.
"""
function mus end            # Isentropic expansion coefficient, μs ≡ (∂T/∂P)_s = vTβ/cp
μs = mus                    # μ: U+3bc

export k, beta, kappaT, kappas, Δc, cs, Ma, muJT, mus
export    β   , κT    , κs    , Dc,         μJT , μs

# System properties
# -----------------

"""
Function to return the system mass.
"""
function m end              # System mass

"""
Function to return the system chemical amount.
"""
function N end              # System chemical amount

export m, N

# Macroscopic properties
# ----------------------

"""
Function to return the specific kinetic energy.
"""
function ek end             # Specific kinetic energy

"""
Function to return the specific potential energy.
"""
function ep end             # Specific potential energy

export ek, ep

# Total properties
# ----------------

"""
Function to return the specific total energy.\n
Aliased to `𝖾`: `\\sanse<TAB>`.
"""
function e end              # Substance specific total energy
𝖾 = e

"""
Function to return the flowing substance specific total energy.\n
Aliased to `θ`: `\\theta<TAB>`.
"""
function theta end          # Flowing substance specific total energy
θ = theta

export e, theta
export 𝖾, θ

# Interactions
# ------------

"""
Function to return the specific heat transfer.
"""
function q end              # Specific heat transfer

"""
Function to return the specific work transfer.
"""
function w end              # Specific work transfer

"""
Function to return the specific energy transfer.\n
Aliased to `De`.
"""
function Δe end             # Specific energy transfer
De = Δe

"""
Function to return the specific entropy transfer.\n
Aliased to `Ds`.
"""
function Δs end             # Specific entropy transfer
Ds = Δs

export q, w, Δe, Δs
export       De, Ds


text:
#!/usr/local/bin/julia
#
const PKGS = ["RDatasets","DataFrames","Plots","PyPlot","PlotlyJS","UnicodePlots","Images",
              "Winston","PGFPlots","Luxor","Gadfly","StatsPlots","GraphPlot","Makie"]

println("Installing required packages ...")

using Pkg
ipks = Pkg.installed()
for p in PKGS
    println("> $p")
    !haskey(ipks,p) && Pkg.add(p)
end
println("Done!")


text:
@testset "TrapezoidalPoints" begin
    @test_throws Exception TrapezoidalPoints{Int}(4)
    @test eltype(@inferred(TrapezoidalPoints(4))) == Float64

    for T in Floats
        @test eltype(@inferred(TrapezoidalPoints{T}(3))) == T
        @test typeof(@inferred(TrapezoidalPoints{T}(3)[1])) == T

        @test TrapezoidalPoints{T}(0) ≈ T[]
        @test TrapezoidalPoints{T}(1) ≈ T[0]  atol = eps(real(T))
        @test TrapezoidalPoints{T}(2) ≈ T[-1,1]
        @test TrapezoidalPoints{T}(3) ≈ T[-1,0,1]   atol = eps(real(T))
        @test TrapezoidalPoints{T}(4) ≈ T[-1,-T(1)/3,T(1)/3,1]
        @test TrapezoidalPoints{T}(5) ≈ T[-1,-T(1)/2,0,T(1)/2,1]   atol = eps(real(T))
    end
end


text:
### A Pluto.jl notebook ###
# v0.14.1

using Markdown
using InteractiveUtils

# ╔═╡ 194aeca0-9c99-11eb-35c7-891390734658
using Reproduce, Plots, RollingFunctions, Statistics, FileIO, PlutoUI, NaNMath

# ╔═╡ f2964d90-9257-4131-924e-1be16e38b162
color_scheme = [
    colorant"#44AA99",
    colorant"#332288",
    colorant"#DDCC77",
    colorant"#999933",
    colorant"#CC6677",
    colorant"#AA4499",
    colorant"#DDDDDD",
	colorant"#117733",
	colorant"#882255",
	colorant"#88CCEE",
]

# ╔═╡ 51cea0d1-bdd3-49e5-8eae-e89c9edbe522
md"""

# Data Functions

"""

# ╔═╡ 2c3a4b8f-0f40-4445-9462-2644b1157c15
function mean_uneven(d::Vector{Array{F, 1}}) where {F}
    ret = zeros(F, maximum(length.(d)))
    n = zeros(Int, maximum(length.(d)))
    for v ∈ d
        ret[1:length(v)] .+= v
        n[1:length(v)] .+= 1
    end
    ret ./ n
end

# ╔═╡ 93549b83-fffc-421c-a23e-376ce5787a9e
function std_uneven(d::Vector{Array{F, 1}}) where {F}

    m = mean_uneven(d)

    ret = zeros(F, maximum(length.(d)))
    n = zeros(Int, maximum(length.(d)))
    for v ∈ d
        ret[1:length(v)] .+= (v .- m[1:length(v)]).^2
        n[1:length(v)] .+= 1
    end
    sqrt.(ret ./ n)
end

# ╔═╡ abb330d0-8144-4395-b22b-7262055cba3f
function Reproduce.search(f, ic::ItemCollection)

    found_items = Vector{Reproduce.Item}()
    for (item_idx, item) in enumerate(ic.items)
        if f(item)
            push!(found_items, item)
        end
    end
    return ItemCollection(found_items)

end

# ╔═╡ f41e07cd-e7e7-4fce-bed2-d310d1998ec9
load_data(item) = FileIO.load(joinpath(item.folder_str, "results.jld2"))
# 	catch
# 		@show "$(f) not loadable"
# 	end
# end

# ╔═╡ 4ed2764d-cf06-4dd8-8458-cac9a8f4d0a5
function get_runs(ic, get_data)
    tmp = get_data(FileIO.load(joinpath(ic[1].folder_str, "results.jld2")))
    d = typeof(tmp)[]
    for (idx, item) ∈ enumerate(ic)
        push!(d, get_data(load_data(item)))
    end
    d
end

# ╔═╡ 5d660aa4-8e6b-464a-bd90-af616a61fc43
function internal_func(
		ic,
		param_keys;
		comp=findmax,
		get_comp_data,
		get_data=get_comp_data)

	# @show length(ic)
    ic_diff = diff(ic)
	# @show ic_diff
    params = if param_keys isa String
        ic_diff[param_keys]
    else
        collect(Iterators.product([ic_diff[k] for k ∈ param_keys]...))
    end

    s = zeros(length(params))
    for (p_idx, p) ∈ enumerate(params)
        sub_ic = if param_keys isa String
            search(ic, Dict(param_keys=>p))
        else
            search(ic, Dict(param_keys[i]=>p[i] for i ∈ 1:length(p)))
        end
        s[p_idx] = mean(get_runs(sub_ic, get_comp_data))
    end

	v, idx = comp(s)

	# get data of best setting
	best_ic = search(
		ic,
		Dict(param_keys[i]=>params[idx][i] for i ∈ 1:length(params[idx])))
	data = get_runs(best_ic, get_data)

    data, v, params[idx]
end

# ╔═╡ f5885dc3-c6e8-4af9-a78e-4f240bf88807
function get_line_data_for(
		ic::ItemCollection,
		line_keys,
		param_keys;
		comp=findmin,
	    get_comp_data,
	    get_data)
	ic_diff = diff(ic)
	params = if line_keys isa String
        ic_diff[line_keys]
    else
        collect(Iterators.product([ic_diff[k] for k ∈ line_keys]...))
    end

	strg = []

	Threads.@threads for p_idx ∈ 1:length(params)
		p = params[p_idx]
		sub_ic = if line_keys isa String
            search(ic, Dict(line_keys=>p))
        else
            search(ic, Dict(line_keys[i]=>p[i] for i ∈ 1:length(p)))
        end
		if length(sub_ic) != 0
			d, c, ps = internal_func(
				sub_ic,
				param_keys;
				get_comp_data=get_comp_data,
				get_data=get_data)
			push!(strg, (params[p_idx], ps, d, c))
		end
	end
	strg
end

# ╔═╡ 883d2bf4-5fa7-452b-8cd2-8ca33886689f
function new_diff(items::Array{Reproduce.Item, 1};
                 exclude_keys::Union{Array{String,1}, Array{Symbol,1}} = Array{String, 1}(),
                 exclude_parse_values::Bool=true)

    kt = keytype(items[1].parsed_args)
    exclude_keys = kt.(exclude_keys)
    if exclude_parse_values == true
        exclude_keys = [exclude_keys; kt.([Reproduce.HASH_KEY, Reproduce.SAVE_NAME_KEY, Reproduce.GIT_INFO_KEY])]
    end
    diff_parsed = Dict{kt, Array}()
    for item in items
        tmp_dict = items[1] - item
        for key in filter((k)->k ∉ exclude_keys, keys(tmp_dict))
            if key ∉ keys(diff_parsed)
				diff_parsed[key] = Array{Any, 1}()
            end

            if tmp_dict[key][1] ∉ diff_parsed[key]
                push!(diff_parsed[key], tmp_dict[key][1])
				new_type = typeof(tmp_dict[key])
            end
            if tmp_dict[key][2] ∉ diff_parsed[key]
                push!(diff_parsed[key], tmp_dict[key][2])
            end
        end
    end
    for key in keys(diff_parsed)
		diff_parsed[key] = collect(promote(diff_parsed[key]...))
        sort!(diff_parsed[key])
    end
    return diff_parsed
end

# ╔═╡ c394afe9-7780-4e07-8e40-4f52265f8afa
new_diff(ic::Reproduce.ItemCollection; kwargs...) = new_diff(ic.items; kwargs...)

# ╔═╡ b749f892-15b8-4512-963a-ebc3e63b3f55
function get_oned_error(d; agg=(x)->mean(x;dims=1)[1, :])
	ret = agg(d["results"][:oned_tmaze_start_error])
	ret[isnan.(ret)] .= Inf
	ret
end

# ╔═╡ d6a4d95b-4953-48b7-a31a-e66a69a1c90c
begin
	get_oned_error_AUC(d) = sum(get_oned_error(d))
	get_oned_error_AUE(d; final_perc = 0.1) = begin
		x = get_oned_error(d)
		len_inc = Int(clamp(round(length(x)*final_perc), 1, Inf))
		sum(x[end-len_inc:end])
	end
	get_oned_error_gvf(d; idx) = get_oned_error(d; agg=(x)->x[idx,:])
end

# ╔═╡ bf5417e3-965b-4bb4-948c-ed954ac18cdb
function get_gvf_data(sub_ic; get_comp_data=get_oned_error_AUC)
	# descent_tb_ic = search(ic, Dict(
	# 		"cumulant_schedule"=>"DrifterDistractor",
	# 		"demon_opt"=>"Auto",
	# 		"demon_update"=>"ESARSA"))
	data, v, ps = internal_func(
		sub_ic,
		["num_tilings", "num_tiles", "demon_eta"];
		comp=findmin,
		get_comp_data = get_comp_data,
		get_data=(x)->[get_oned_error_gvf(x; idx=i) for i ∈ 1:4])
	[(mean(getindex.(data, i)), std(getindex.(data, i))./sqrt(length(data))) for i in 1:4], ps
end

# ╔═╡ 7f05f324-2600-46d1-9468-4ee6b60fef57
md"""
# Data Analysis

Get Item Collections for OneDTMazeQ and OneDTMazeSR and combine.

Sarsa also run w/ Q learners, but failed for some optimizers.

"""

# ╔═╡ ea84781a-cc17-48f7-8cc6-cedb3bc06686
begin
	ic_q = search(ItemCollection("../local_data/OneDTMazeQ/")) do item
		item.parsed_args["demon_update"] ∈ ["ESARSA", "TB"]
	end
	ic_sr = ItemCollection("../local_data/OneDTMazeSR_ideal/")
	ic = ItemCollection([ic_q.items; ic_sr.items])
	diff_dict = new_diff(ic.items)
end

# ╔═╡ eeed36c6-326f-4cc9-93f4-143bb8cae54e
md"""
#### Get data with parameters chosen over AUC for all gvfs.

"""

# ╔═╡ 3c167a75-f726-4579-abb1-7e3308d5b704
begin
	gvf_best_data = get_line_data_for(
		ic,
		["demon_learner", "demon_update", "demon_opt", "num_tiles", "num_tilings", "cumulant_schedule"],
		["demon_eta"];
		comp=findmin,
	    get_comp_data = get_oned_error_AUC,
	    get_data = (x)->[get_oned_error_gvf(x; idx=i) for i ∈ 1:4])
end

# ╔═╡ d92544b9-b571-4bbd-af4b-db83d360dc68
function plot_gvf_data(ic, gvf_idx;	cumulant_schedule = "DrifterDistractor", demon_learner="Q", demon_update="TB")

	sub_ic = search(ic, Dict(
			"cumulant_schedule"=>cumulant_schedule,
			"demon_learner"=>demon_learner,
			"demon_update"=>demon_update))
	plt = nothing
	d = new_diff(sub_ic)
	prms = if "demon_opt" ∈ keys(d)
		d["demon_opt"]
	else
		[sub_ic.items[1].parsed_args["demon_opt"]]
	end
	@show prms
	for (demon_opt) ∈ prms
		sub_ic_2 = search(sub_ic,
			   Dict("demon_opt"=>demon_opt))
		data, params = get_gvf_data(sub_ic_2)
		if plt isa Nothing
			plt = plot(data[gvf_idx][1], ribbon=data[gvf_idx][2], label="$(demon_opt), $(demon_update), $(params)", title="GVF $(gvf_idx)", palette=color_scheme)
		else
			plot!(plt, data[gvf_idx][1], ribbon=data[gvf_idx][2], label="$(demon_opt), $(demon_update), $(params)", palette=color_scheme)
		end
	end
	plt
end

# ╔═╡ 76b499f8-950e-4fb5-9e5e-be7b42cc18e0
# plot([plot_gvf_data(ic, i; demon_learner="Q", demon_update="TB", cumulant_schedule="DrifterDistractor") for i ∈ 1:4]..., size=(800, 800))

# ╔═╡ d5e946b4-b534-46fe-b4f6-e6aaadf18bd4
# plot([plot_gvf_data(ic, i; demon_learner="SR", demon_update="TB") for i ∈ 1:4]..., size=(800, 800))

# ╔═╡ bc2f558d-f4e6-48df-9c1d-06c1f5f4cf69
function get_gvf_plot_data(ic; line_args)
	d = new_diff(ic)
	data_col = []
	for args ∈ Iterators.product([d[la] for la ∈ line_args]...)
		sub_ic_2 = search(ic,
			   Dict(line_args[i]=>args[i] for i ∈ 1:length(line_args)))
		if length(sub_ic_2) != 0
			data, params = get_gvf_data(sub_ic_2)
			push!(data_col, (data, params, args))
		end
	end
	data_col
end

# ╔═╡ b7d30b13-9e5f-4f0e-bb03-c58cf667328a
function plot_all_gvf_data(data_col, gvf_idx; plt_kwargs...)

	plt = nothing
	for d ∈ data_col
			data, params, args = d
			if plt isa Nothing
				plt = plot(data[gvf_idx][1], ribbon=data[gvf_idx][2], label="$(args), $(params)", title="GVF $(gvf_idx)"; plt_kwargs...)
			else
				plot!(plt, data[gvf_idx][1], ribbon=data[gvf_idx][2], label="$(args), $(params)"; plt_kwargs...)
			end
	end
	plt
end

# ╔═╡ 0ee7ee2c-8998-4062-a43d-9e0a5f34f1c6
# gvf_data_dd = get_gvf_plot_data(search(ic, Dict("cumulant_schedule"=>"DrifterDistractor"));
# 	line_args=["demon_opt", "demon_update", "demon_learner"])

# ╔═╡ effda83c-8c65-4e75-b695-22f4644604ea
# begin
# 	for i ∈ 1:4
# 		plot_all_gvf_data(gvf_data_dd, i; palette=color_scheme, size=(600, 600), lw=2, xlims=(0,200), ylims=(0.0,1.5))
# 		savefig("oned_tmaze_dd_round_robin_gvf_$(i).pdf")
# 	end
# 	for i ∈ 1:4
# 		plot_all_gvf_data(gvf_data_dd, i; palette=color_scheme, size=(600, 600), lw=2, xlims=(0,200))
# 		savefig("oned_tmaze_dd_round_robin_gvf_$(i)_nl.pdf")
# 	end
# end

# ╔═╡ 2a50b3e8-d06b-4a1f-b274-4386ef8a4081
# gvf_data_c = get_gvf_plot_data(search(ic, Dict("cumulant_schedule"=>"Constant"));
# 	line_args=["demon_opt", "demon_update", "demon_learner"])

# ╔═╡ b44f8824-72c6-4500-971a-4adf90dafed8
# begin
# 	for i ∈ 1:4
# 		plot_all_gvf_data(gvf_data_c, i; palette=color_scheme, size=(600, 600), lw=2, xlims=(0,200), ylims=(0.0,1.5))
# 		savefig("oned_tmaze_c_round_robin_gvf_$(i).pdf")
# 	end
# 	for i ∈ 1:4
# 		plot_all_gvf_data(gvf_data_c, i; palette=color_scheme, size=(600, 600), lw=2, xlims=(0,200))
# 		savefig("oned_tmaze_c_round_robin_gvf_$(i)_nl.pdf")
# 	end
# end

# ╔═╡ Cell order:
# ╠═194aeca0-9c99-11eb-35c7-891390734658
# ╟─f2964d90-9257-4131-924e-1be16e38b162
# ╟─51cea0d1-bdd3-49e5-8eae-e89c9edbe522
# ╟─2c3a4b8f-0f40-4445-9462-2644b1157c15
# ╟─93549b83-fffc-421c-a23e-376ce5787a9e
# ╟─abb330d0-8144-4395-b22b-7262055cba3f
# ╟─5d660aa4-8e6b-464a-bd90-af616a61fc43
# ╟─f41e07cd-e7e7-4fce-bed2-d310d1998ec9
# ╟─4ed2764d-cf06-4dd8-8458-cac9a8f4d0a5
# ╟─f5885dc3-c6e8-4af9-a78e-4f240bf88807
# ╟─883d2bf4-5fa7-452b-8cd2-8ca33886689f
# ╟─c394afe9-7780-4e07-8e40-4f52265f8afa
# ╟─b749f892-15b8-4512-963a-ebc3e63b3f55
# ╟─d6a4d95b-4953-48b7-a31a-e66a69a1c90c
# ╟─bf5417e3-965b-4bb4-948c-ed954ac18cdb
# ╟─7f05f324-2600-46d1-9468-4ee6b60fef57
# ╟─ea84781a-cc17-48f7-8cc6-cedb3bc06686
# ╟─eeed36c6-326f-4cc9-93f4-143bb8cae54e
# ╠═3c167a75-f726-4579-abb1-7e3308d5b704
# ╟─d92544b9-b571-4bbd-af4b-db83d360dc68
# ╠═76b499f8-950e-4fb5-9e5e-be7b42cc18e0
# ╠═d5e946b4-b534-46fe-b4f6-e6aaadf18bd4
# ╠═bc2f558d-f4e6-48df-9c1d-06c1f5f4cf69
# ╠═b7d30b13-9e5f-4f0e-bb03-c58cf667328a
# ╠═0ee7ee2c-8998-4062-a43d-9e0a5f34f1c6
# ╠═effda83c-8c65-4e75-b695-22f4644604ea
# ╠═2a50b3e8-d06b-4a1f-b274-4386ef8a4081
# ╠═b44f8824-72c6-4500-971a-4adf90dafed8

max_stars_repo_name: MatthewMcLeod/curiosity

max_issues_repo_name: MatthewMcLeod/curiosity

max_forks_repo_name: MatthewMcLeod/curiosity


text:
"""
    responding!(actr, task, key, args...; kwargs...)

Sets the declarative motor module as busy and registers a new event for executing
a key stroke

# Arguments

- `actr`: an ACT-R model object
- `task`: a task that is a subtype of `AbstractTask`
- `key`: a string representing a response key
"""
function responding!(actr, task, key, args...; kwargs...)
    actr.motor.state.busy = true
    description = "Respond"
    type = "model"
    id = get_name(actr)
    tΔ = rnd_time(.060)
    register!(actr, respond!, after, tΔ , actr, task, key;
        id, description, type)
    return tΔ
end

"""
    respond!(actr, task, key, args...; kwargs...)

Executes a motor response with user defined `press_key!` function and sets module state to
busy = false.

# Arguments

- `actr`: an ACT-R model object
- `task`: a task that is a subtype of `AbstractTask`
- `key`: a string representing a response key
"""
function respond!(actr, task, key)
    actr.motor.state.busy = false
    press_key!(task, actr, key)
end


text:
using StatsBase:mean
const PolicyOrSelector = Union{AbstractPolicy, AbstractActionSelector}

"""
    MonteCarloLearner(approximator::Tapp, π::Tp, γ::Float64=1., α::Float64 = 1.0, first_visit::Bool = true) where {Tapp,Tp}

See more details at Section (5.1) on Page 92 of the book *Sutton, Richard S., and Andrew G. Barto. Reinforcement learning: An introduction. MIT press, 2018.*
"""
struct MonteCarloLearner{visit,Tapp <: AbstractApproximator,Tp <:PolicyOrSelector,Tr <: Function} <: AbstractMonteCarloLearner
    approximator::Tapp
    π::Tp
    γ::Float64
    α::Float64
    returns::Tr
    function MonteCarloLearner(approximator::Tapp, π::Tp, γ::Float64=1., α::Float64 = 1.0, first_visit::Bool = true) where {Tapp,Tp}
        returns = cached_sample_avg()
        if first_visit
            new{:FirstVisit,Tapp,Tp,typeof(returns)}(approximator, π, γ, α, returns)
        else
            new{:EveryVisit,Tapp,Tp,typeof(returns)}(approximator, π, γ, α, returns)
        end
    end
end

(learner::MonteCarloLearner{Tvisit, <:AbstractQApproximator, <:AbstractActionSelector})(s) where Tvisit = learner.approximator(s) |> learner.π
(learner::MonteCarloLearner{Tvisit, <:AbstractApproximator, <:AbstractPolicy})(s) where Tvisit = learner.π(s)
(learner::MonteCarloLearner{Tvisit, <:AbstractApproximator, <:AbstractPolicy})(s, ::Val{:dist}) where Tvisit = learner.π(s, Val(:dist))

update!(learner::MonteCarloLearner, buffer::EpisodeSARDBuffer) = isfull(buffer) && update!(learner, @view(buffer.state[1:end - 1]), @view(buffer.action[1:end - 1]), buffer.reward)
update!(learner::MonteCarloLearner, buffer::EpisodeSARDSBuffer) = isfull(buffer) && update!(learner, buffer.state, buffer.action, buffer.reward)
update!(learner::MonteCarloLearner{Tvisit,<:AbstractVApproximator}, states, actions, rewards) where Tvisit = update!(learner, states, rewards)

function update!(learner::MonteCarloLearner{:FirstVisit,<:AbstractVApproximator}, states, rewards)
    α, γ, V, Returns, G = learner.α, learner.γ, learner.approximator, learner.returns, 0.
    for (isfirstvisit, s, r) in Iterators.reverse(zip(IsFirstVisit(states), states, rewards))
        if isfirstvisit
            G = γ * G + r
            update!(V, s, α * (Returns(s, G) - V(s)))
        end
    end
end

function update!(learner::MonteCarloLearner{:EveryVisit,<:AbstractVApproximator}, states, rewards)
    α, γ, V, Returns, G = learner.α, learner.γ, learner.approximator, learner.returns, 0.
    for (s, r) in Iterators.reverse(zip(states, rewards))
        G = γ * G + r
        update!(V, s, α * (Returns(s, G) - V(s)))
    end
end

function update!(learner::MonteCarloLearner{:FirstVisit,<:AbstractQApproximator}, states, actions, rewards)
    α, γ, Q, π, Returns, G = learner.α, learner.γ, learner.approximator, learner.π, learner.returns, 0.
    for (isfirstvisit, s, a, r) in Iterators.reverse(zip(IsFirstVisit(states), states, actions, rewards))
        G = γ * G + r
        if isfirstvisit
            update!(Q, s, a, α * (Returns((s, a), G) - Q(s, a)))
            update!(π, s, Q(s, Val(:argmax)))
        end
    end
end

function update!(learner::MonteCarloLearner{:EveryVisit, <:AbstractQApproximator}, states, actions, rewards)
    α, γ, Q, π, Returns, G = learner.α, learner.γ, learner.approximator, learner.π, learner.returns, 0.
    for (s, a, r) in Iterators.reverse(zip(states, actions, rewards))
        G = γ * G + r
        update!(Q, s, a, α * (Returns((s, a), G) - Q(s, a)))
        update!(π, s, Q(s, Val(:argmax)))
    end
end

"""
    MonteCarloExploringStartLearner(approximator::Tapp, π::Tp, π_start::RandomPolicy, γ::Float64, α::Float64 = 1.0; is_first_visit::Bool = true) where {Tapp,Tp}

See more details at Section (5.3) on Page 99 of the book *Sutton, Richard S., and Andrew G. Barto. Reinforcement learning: An introduction. MIT press, 2018.*
"""
mutable struct MonteCarloExploringStartLearner{visit,Tapp <: AbstractApproximator,Tp <: AbstractPolicy,Tr <: Function} <: AbstractMonteCarloLearner
    approximator::Tapp
    π::Tp
    π_start::RandomPolicy
    γ::Float64
    α::Float64
    returns::Tr
    is_start::Bool
    function MonteCarloExploringStartLearner(approximator::Tapp, π::Tp, π_start::RandomPolicy, γ::Float64, α::Float64 = 1.0; is_first_visit::Bool = true) where {Tapp,Tp}
        returns = cached_sample_avg()
        if is_first_visit
            new{:FirstVisit,Tapp,Tp,typeof(returns)}(approximator, π, π_start, γ, α, returns, true)
        else
            new{:EveryVisit,Tapp,Tp,typeof(returns)}(approximator, π, π_start, γ, α, returns, true)
        end
    end
end

(learner::MonteCarloExploringStartLearner)(s) = learner.is_start ? learner.π_start(s) : learner.π(s)

function update!(learner::MonteCarloExploringStartLearner{Tvisit,<:AbstractQApproximator}, buffer::EpisodeSARDBuffer) where Tvisit
    α, γ, Q, π, Returns = learner.α, learner.γ, learner.approximator, learner.π, learner.returns
    if isfull(buffer)
        learner.is_start = true  # reset start flag
        G, states_actions, rewards = 0, zip(@view(buffer.state[1:end - 1]), @view(buffer.action[1:end - 1])), buffer.reward
        for (isfirstvisit, (s, a), r) in Iterators.reverse(zip(IsFirstVisit(states_actions), states_actions, rewards))
            G = γ * G + r
            if Tvisit == :EveryVisit || (Tvisit == :FirstVisit && isfirstvisit)
                update!(Q, s, a, α * (Returns((s, a), G) - Q(s, a)))
                update!(π, s, Q(s, Val(:argmax)))
            end
        end
    else
        learner.is_start = false  # reset start flag
    end
end

"""
    OffPolicyMonteCarloLearner(approximator::Tapp, π_behavior::Tpb, π_target::Tpt, γ::Float64, α::Float64 = 1.0; isfirstvisit::Bool = true, sampling::Symbol = :OrdinaryImportanceSampling,) where {Tapp,Tpb,Tpt}

See more details at Section (5.7) on Page 111 of the book *Sutton, Richard S., and Andrew G. Barto. Reinforcement learning: An introduction. MIT press, 2018.*
"""
struct OffPolicyMonteCarloLearner{Tvisit,Sampling,Tapp <: AbstractApproximator,Tpb <: AbstractPolicy,Tpt <: AbstractPolicy} <: AbstractMonteCarloLearner
    approximator::Tapp
    π_behavior::Tpb
    π_target::Tpt
    γ::Float64
    α::Float64
    ρ_G::Dict{Int,Vector{Pair{Float64,Float64}}}
    function OffPolicyMonteCarloLearner(approximator::Tapp, π_behavior::Tpb, π_target::Tpt, γ::Float64, α::Float64 = 1.0; isfirstvisit::Bool = true, sampling::Symbol = :OrdinaryImportanceSampling,) where {Tapp,Tpb,Tpt}
        sampling in [:OrdinaryImportanceSampling, :WeightedImportanceSampling] || error("unknown sampling method $method")
        if isfirstvisit
            new{:FirstVisit,sampling,Tapp,Tpb,Tpt}(approximator, π_behavior, π_target, γ, α, Dict{Int,Vector{Pair{Float64,Float64}}}())
        else
            new{:EveryVisit,sampling,Tapp,Tpb,Tpt}(approximator, π_behavior, π_target, γ, α, Dict{Int,Vector{Pair{Float64,Float64}}}())
        end
    end
end

(learner::OffPolicyMonteCarloLearner)(s) = learner.π_behavior(s)

function update!(learner::OffPolicyMonteCarloLearner{Tvisit,Tsampling,<:AbstractVApproximator}, buffer::EpisodeSARDBuffer) where {Tvisit,Tsampling}
    α, γ, V = learner.α, learner.γ, learner.approximator
    if isfull(buffer)
        G, ρ, states, actions, rewards = 0, 1., @view(buffer.state[1:end - 1]), @view(buffer.action[1:end - 1]), buffer.reward
        for (isfirstvisit, s, a, r) in Iterators.reverse(zip(IsFirstVisit(states), states, actions, rewards))
            G = γ * G + r
            ρ *= learner.π_target(s, a) / learner.π_behavior(s, a)
            if Tvisit == :EveryVisit || (Tvisit == :FirstVisit && isfirstvisit)
                if haskey(learner.ρ_G, s)
                    push!(learner.ρ_G[s], ρ => G)
                else
                    learner.ρ_G[s] = [ρ => G]
                end
            end
        end
        if Tsampling == :OrdinaryImportanceSampling
            for s in keys(learner.ρ_G)
                v = mean(ρ * G for (ρ, G) in learner.ρ_G[s])
                update!(V, s, α * (v - V(s)))
            end
        else
            for s in keys(learner.ρ_G)
                numerator = sum(ρ * G for (ρ, G) in learner.ρ_G[s])
                denominator = sum(ρ for (ρ, _) in learner.ρ_G[s])
                v = denominator == 0 ? 0 : numerator / denominator
                update!(V, s, α * (v - V(s)))
            end
        end
    end
end


text:
function jump_sdplib(solver, path; verbose = false, test = false)

    println("running: $(path)")

    n, m, F, c = sdplib_data(path)

    # Build model
    model = Model(with_optimizer(solver))
    @variable(model, X[1:n, 1:n], PSD)

    # Objective function
    @objective(model, Min, sum(F[0][idx...] * X[idx...]
        for idx in zip(findnz(F[0])[1:end-1]...)))

    # Linear equality constraints
    for k = 1:m
        @constraint(model, sum(F[k][idx...] * X[idx...]
            for idx in zip(findnz(F[k])[1:end-1]...)) == c[k])
    end

    teste = @time optimize!(model)

    XX = value.(X)

    verbose && sdplib_eval(F,c,n,m,XX)

    objval = objective_value(model)
    stime = MOI.get(model, MOI.SolveTime())

    # @show tp = typeof(model.moi_backend.optimizer.model.optimizer)
    # @show fieldnames(tp)
    rank = -1
    try
        @show rank = model.moi_backend.optimizer.model.optimizer.sol.final_rank
    catch
    end
    status = 0
    if JuMP.termination_status(model) == MOI.OPTIMAL
        status = 1
    end
    return (objval, stime, rank, status)
    # return (objval, stime)
end


text:
function RRQRDecomposition(arg0::RealMatrix)
    return RRQRDecomposition((RealMatrix,), arg0)
end

function RRQRDecomposition(arg0::RealMatrix, arg1::jdouble)
    return RRQRDecomposition((RealMatrix, jdouble), arg0, arg1)
end

function get_p(obj::RRQRDecomposition)
    return jcall(obj, "getP", RealMatrix, ())
end

function get_rank(obj::RRQRDecomposition, arg0::jdouble)
    return jcall(obj, "getRank", jint, (jdouble,), arg0)
end

function get_solver(obj::RRQRDecomposition)
    return jcall(obj, "getSolver", DecompositionSolver, ())
end


text:
using Base.Threads

function getInput()
    println("Enter start, end, n")
    collect([parse(Int64, x) for x in split(readline())])
end

function trap(
    left_endpt::Float64,
    right_endpt::Float64,
    trap_count::Int64,
    base_len::Float64)

    f(x) = x^2

    estimate = (f(left_endpt) + f(right_endpt)) / 2.0
    for i=1:trap_count - 1
        estimate += f(left_endpt + i * base_len) # type instability
    end
    return estimate * base_len
end

function main(a=0, b=131072, n=32768)
    # a, b, n = getInput()
    h = (b - a) / n

    n_threads = nthreads()
    if n_threads == 1
        println("
        Number of threads detected to be =1\n
        environment variables may not be properly set up!"
        )
    end
    total = Atomic{Float64}(0)
    local_n::Int64 = div(n, n_threads)
    @threads for rank in 0:n_threads - 1
        local_a = a + rank * local_n * h
        local_b = local_a + local_n * h
        local_ans = trap(local_a, local_b, local_n, h)
        println("local_a = $local_a, local_b = $local_b, ans = $local_ans")
        atomic_add!(total, local_ans)
    end

    println("the answer for a = $a, b = $b, n = $n: $(total[])")
end

for i=1:10000
    main()
end


text:
using VibronicToolkit: a, Basis, id, mkid, mkop, n, operators, q

using LinearAlgebra: I

let sys = DenseSystem(test_params_trivial...),
    size = 3,
    basis = Basis(sys, size)

    @test id(basis) == I
    @test isapprox(a(basis)' * a(basis), n(basis))

    q_ = q(basis)
    ip_ = -(a(basis)' - a(basis)) / sqrt(2.0)
    ho = (-ip_^2 + q_^2)/2
    @test isapprox(ho[1:(end-1),1:(end-1)], n(basis)[1:(end-1),1:(end-1)]+0.5I)

    @test mkid(basis) == I
    @test mkop(basis, a, 1) == kron(a(basis), id(basis))
    @test mkop(basis, q, 2) == kron(id(basis), q(basis))

    h0, V = operators(basis, sys)
    h0_exp = zeros(basis.dim, basis.dim)
    for s in 1:test_S
        idx = ((s-1)*basis.dim1+1):(s*basis.dim1)
        h0_exp[idx, idx] .+= I + kron(n(basis)+0.5I, id(basis)) + kron(id(basis), n(basis)+0.5I)
    end
    @test h0 == h0_exp
    V_exp = zeros(basis.dim, basis.dim)
    for s1 in 1:test_S
        for s2 in 1:test_S
            s1 == s2 && continue
            idx1 = ((s1-1)*basis.dim1+1):(s1*basis.dim1)
            idx2 = ((s2-1)*basis.dim1+1):(s2*basis.dim1)
            V_exp[idx2, idx1] .+= Matrix(I, basis.dim1, basis.dim1)
        end
    end
    @test V == V_exp
end


text:
module BenchDataFrames

using BangBang: push!!
using BenchmarkTools: @benchmarkable, BenchmarkGroup
using CategoricalArrays: CategoricalVector
using DataFrames: DataFrame
using DataStructures: LittleDict

function repeat_push(df, row, n = 10)
    row = LittleDict(keys(row), values(row))
    for _ = 1:n
        foreach(empty!, eachcol(df))
        for i = 1:2^10
            push!!(df, row)
        end
    end
end

mkrow(x) = (; tuple.(Symbol.('a':'a'+(8-1)), x)...)

suite = BenchmarkGroup()

suite["Int-Int"] = @benchmarkable repeat_push($(DataFrame([mkrow(0)])), $(mkrow(0)))
suite["Missing-Int"] = @benchmarkable repeat_push(
    $(DataFrame([mkrow(missing)])),
    $(mkrow(0)),
)
suite["String-String"] = @benchmarkable repeat_push(
    $(DataFrame([mkrow("x")])),
    $(mkrow("x")),
)
suite["Missing-String"] = @benchmarkable repeat_push(
    $(DataFrame([mkrow(missing)])),
    $(mkrow("x")),
)
suite["Categorical"] = @benchmarkable repeat_push(
    $(DataFrame(mkrow(Ref(CategoricalVector(["x"]))))),
    $(mkrow("x")),
)

end  # module
BenchDataFrames.suite


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

JLLWrappers.@generate_main_file_header("Libglvnd")
JLLWrappers.@generate_main_file("Libglvnd", UUID("7e76a0d4-f3c7-5321-8279-8d96eeed0f29"))
end  # module Libglvnd_jll


text:
#Include KG_Module
include("../src/KG_Module.jl")

#name of file to analyse
filepath = "./data/binary_run_1/binary_bh_rotating_radius_1.5_alpha_10.0_Rorbit_5.0.h5"

#Get field at certain positions
time_vector , pos_vector_1 = get_field_point( filepath , 0.0 , 0.0)
_           , pos_vector_2  = get_field_point(filepath, 0.0 , 10.0)
_           , pos_vector_3  = get_field_point(filepath, 0.0 , 24.0)

#Write to file
open("dislpacement_at_00_00.txt", "w") do io
    writedlm(io , [time_vector pos_vector_1])
end

open("dislpacement_at_00_10.txt", "w") do io
    writedlm(io , [time_vector pos_vector_2])
end

open("dislpacement_at_00_24.txt", "w") do io
    writedlm(io , [time_vector pos_vector_3])
end


text:
# This file is a part of Julia. License is MIT: https://julialang.org/license

module AccessorsTest

using Dates
using Test

@testset "yearmonthday/yearmonth/monthday" begin
    # yearmonthday is the opposite of totaldays
    # taking Rata Die Day # and returning proleptic Gregorian date
    @test DatesPlus.yearmonthday(-306) == (0, 2, 29)
    @test DatesPlus.yearmonth(-306) == (0, 2)
    @test DatesPlus.monthday(-306) == (2, 29)
    @test DatesPlus.yearmonthday(-305) == (0, 3, 1)
    @test DatesPlus.yearmonth(-305) == (0, 3)
    @test DatesPlus.monthday(-305) == (3, 1)
    @test DatesPlus.yearmonthday(-2) == (0, 12, 29)
    @test DatesPlus.yearmonth(-2) == (0, 12)
    @test DatesPlus.monthday(-2) == (12, 29)
    @test DatesPlus.yearmonthday(-1) == (0, 12, 30)
    @test DatesPlus.yearmonth(-1) == (0, 12)
    @test DatesPlus.monthday(-1) == (12, 30)
    @test DatesPlus.yearmonthday(0) == (0, 12, 31)
    @test DatesPlus.yearmonth(-0) == (0, 12)
    @test DatesPlus.monthday(-0) == (12, 31)
    @test DatesPlus.yearmonthday(1) == (1, 1, 1)
    @test DatesPlus.yearmonth(1) == (1, 1)
    @test DatesPlus.monthday(1) == (1, 1)
    @test DatesPlus.yearmonthday(730120) == (2000, 1, 1)
end
@testset "year/month/day" begin
    # year, month, and day return the indivial components
    # of yearmonthday, avoiding additional calculations when possible
    @test DatesPlus.year(-1) == 0
    @test DatesPlus.month(-1) == 12
    @test DatesPlus.day(-1) == 30
    @test DatesPlus.year(0) == 0
    @test DatesPlus.month(0) == 12
    @test DatesPlus.day(0) == 31
    @test DatesPlus.year(1) == 1
    @test DatesPlus.month(1) == 1
    @test DatesPlus.day(1) == 1
    @test DatesPlus.year(730120) == 2000
    @test DatesPlus.month(730120) == 1
    @test DatesPlus.day(730120) == 1
end
@testset "totaldays/yearmonthday over many years" begin
    # Test totaldays and yearmonthday from January 1st of "from" to December 31st of "to"
    # test_dates(-10000, 10000) takes about 15 seconds
    # test_dates(year(typemin(Date)), year(typemax(Date))) is full range
    # and would take.......a really long time
    let from=0, to=2100, y=0, m=0, d=0
        test_day = DatesPlus.totaldays(from, 1, 1)
        for y in from:to
            for m = 1:12
                for d = 1:DatesPlus.daysinmonth(y, m)
                    days = DatesPlus.totaldays(y, m, d)
                    @test days == test_day
                    @test (y, m, d) == DatesPlus.yearmonthday(days)
                    test_day += 1
                end
            end
        end
    end
end

@testset "year, month, day, hour, minute" begin
    let y=0, m=0, d=0, h=0, mi=0
        for m = 1:12
            for d = 1:DatesPlus.daysinmonth(y, m)
                for h = 0:23
                    for mi = 0:59
                        dt = DatesPlus.DateTime(y, m, d, h, mi)
                        @test y == DatesPlus.year(dt)
                        @test m == DatesPlus.month(dt)
                        @test d == DatesPlus.day(dt)
                        @test d == DatesPlus.dayofmonth(dt)
                        @test h == DatesPlus.hour(dt)
                        @test mi == DatesPlus.minute(dt)
                        @test (m, d) == DatesPlus.monthday(dt)
                    end
                end
            end
        end
    end
end
@testset "second, millisecond" begin
    let y=0, m=0, d=0, h=0, mi=0, s=0, ms=0
        for y in [-2013, -1, 0, 1, 2013]
            for m in [1, 6, 12]
                for d in [1, 15, DatesPlus.daysinmonth(y, m)]
                    for h in [0, 12, 23]
                        for s = 0:59
                            for ms in [0, 1, 500, 999]
                                dt = DatesPlus.DateTime(y, m, d, h, mi, s, ms)
                                @test y == DatesPlus.year(dt)
                                @test m == DatesPlus.month(dt)
                                @test d == DatesPlus.day(dt)
                                @test h == DatesPlus.hour(dt)
                                @test s == DatesPlus.second(dt)
                                @test ms == DatesPlus.millisecond(dt)
                            end
                        end
                    end
                end
            end
        end
    end
end
@testset "year, month, day, hour, minute, second over many years" begin
    let from=0, to=2100, y=0, m=0, d=0
        for y in from:to
            for m = 1:12
                for d = 1:DatesPlus.daysinmonth(y, m)
                    dt = DatesPlus.Date(y, m, d)
                    @test y == DatesPlus.year(dt)
                    @test m == DatesPlus.month(dt)
                    @test d == DatesPlus.day(dt)
                end
            end
        end
    end

    # test hour, minute, second
    let h=0, mi=0, s=0, ms=0, us=0, ns=0
        for h = (0, 23), mi = (0, 59), s = (0, 59),
            ms in (0, 1, 500, 999), us in (0, 1, 500, 999), ns in (0, 1, 500, 999)
            t = DatesPlus.Time(h, mi, s, ms, us, ns)
            @test h == DatesPlus.hour(t)
            @test mi == DatesPlus.minute(t)
            @test s == DatesPlus.second(t)
            @test ms == DatesPlus.millisecond(t)
            @test us == DatesPlus.microsecond(t)
            @test ns == DatesPlus.nanosecond(t)
        end
    end
end
@testset "week" begin
    # Tests from https://en.wikipedia.org/wiki/ISO_week_date
    @test DatesPlus.week(DatesPlus.Date(2005, 1, 1)) == 53
    @test DatesPlus.week(DatesPlus.Date(2005, 1, 2)) == 53
    @test DatesPlus.week(DatesPlus.Date(2005, 12, 31)) == 52
    @test DatesPlus.week(DatesPlus.Date(2007, 1, 1)) == 1
    @test DatesPlus.week(DatesPlus.Date(2007, 12, 30)) == 52
    @test DatesPlus.week(DatesPlus.Date(2007, 12, 31)) == 1
    @test DatesPlus.week(DatesPlus.Date(2008, 1, 1)) == 1
    @test DatesPlus.week(DatesPlus.Date(2008, 12, 28)) == 52
    @test DatesPlus.week(DatesPlus.Date(2008, 12, 29)) == 1
    @test DatesPlus.week(DatesPlus.Date(2008, 12, 30)) == 1
    @test DatesPlus.week(DatesPlus.Date(2008, 12, 31)) == 1
    @test DatesPlus.week(DatesPlus.Date(2009, 1, 1)) == 1
    @test DatesPlus.week(DatesPlus.Date(2009, 12, 31)) == 53
    @test DatesPlus.week(DatesPlus.Date(2010, 1, 1)) == 53
    @test DatesPlus.week(DatesPlus.Date(2010, 1, 2)) == 53
    @test DatesPlus.week(DatesPlus.Date(2010, 1, 2)) == 53
    # Tests from http://www.epochconverter.com/date-and-time/weeknumbers-by-year.php?year=1999
    dt = DatesPlus.DateTime(1999, 12, 27)
    dt1 = DatesPlus.Date(1999, 12, 27)
    check = (52, 52, 52, 52, 52, 52, 52, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2)
    for i = 1:21
        @test DatesPlus.week(dt) == check[i]
        @test DatesPlus.week(dt1) == check[i]
        dt = dt + DatesPlus.Day(1)
        dt1 = dt1 + DatesPlus.Day(1)
    end
    # Tests from http://www.epochconverter.com/date-and-time/weeknumbers-by-year.php?year=2000
    dt = DatesPlus.DateTime(2000, 12, 25)
    dt1 = DatesPlus.Date(2000, 12, 25)
    for i = 1:21
        @test DatesPlus.week(dt) == check[i]
        @test DatesPlus.week(dt1) == check[i]
        dt = dt + DatesPlus.Day(1)
        dt1 = dt1 + DatesPlus.Day(1)
    end
    # Test from http://www.epochconverter.com/date-and-time/weeknumbers-by-year.php?year=2030
    dt = DatesPlus.DateTime(2030, 12, 23)
    dt1 = DatesPlus.Date(2030, 12, 23)
    for i = 1:21
        @test DatesPlus.week(dt) == check[i]
        @test DatesPlus.week(dt1) == check[i]
        dt = dt + DatesPlus.Day(1)
        dt1 = dt1 + DatesPlus.Day(1)
    end
    # Tests from http://www.epochconverter.com/date-and-time/weeknumbers-by-year.php?year=2004
    dt = DatesPlus.DateTime(2004, 12, 20)
    dt1 = DatesPlus.Date(2004, 12, 20)
    check = (52, 52, 52, 52, 52, 52, 52, 53, 53, 53, 53, 53, 53, 53, 1, 1, 1, 1, 1, 1, 1)
    for i = 1:21
        @test DatesPlus.week(dt) == check[i]
        @test DatesPlus.week(dt1) == check[i]
        dt = dt + DatesPlus.Day(1)
        dt1 = dt1 + DatesPlus.Day(1)
    end
end
@testset "Vectorized accessors" begin
    a = DatesPlus.Date(2014, 1, 1)
    dr = [a, a, a, a, a, a, a, a, a, a]
    @test DatesPlus.year.(dr) == repeat([2014], 10)
    @test DatesPlus.month.(dr) == repeat([1], 10)
    @test DatesPlus.day.(dr) == repeat([1], 10)

    a = DatesPlus.DateTime(2014, 1, 1)
    dr = [a, a, a, a, a, a, a, a, a, a]
    @test DatesPlus.year.(dr) == repeat([2014], 10)
    @test DatesPlus.month.(dr) == repeat([1], 10)
    @test DatesPlus.day.(dr) == repeat([1], 10)
    @test DatesPlus.hour.(dr) == repeat([0], 10)
    @test DatesPlus.minute.(dr) == repeat([0], 10)
    @test DatesPlus.second.(dr) == repeat([0], 10)
    @test DatesPlus.millisecond.(dr) == repeat([0], 10)

    b = DatesPlus.Time(1, 2, 3, 4, 5, 6)
    tr = [b, b, b, b, b, b, b, b, b, b]
    @test DatesPlus.hour.(tr) == repeat([1], 10)
    @test DatesPlus.minute.(tr) == repeat([2], 10)
    @test DatesPlus.second.(tr) == repeat([3], 10)
    @test DatesPlus.millisecond.(tr) == repeat([4], 10)
    @test DatesPlus.microsecond.(tr) == repeat([5], 10)
    @test DatesPlus.nanosecond.(tr) == repeat([6], 10)
end

end


text:
# generate correlation dummy for u1-dm symmetry
function get_χ_u1_dm_empty(
    r :: Reduced_lattice,
    m :: Mesh
    ) :: Vector{Matrix{Float64}}

    χxx = zeros(Float64, length(r.sites), m.num_χ)
    χzz = zeros(Float64, length(r.sites), m.num_χ)
    χxy = zeros(Float64, length(r.sites), m.num_χ)
    χ   = Matrix{Float64}[χxx, χzz, χxy]

    return χ
end

# xx kernel for double integral
function compute_χ_kernel_xx!(
    Λ    :: Float64,
    site :: Int64,
    w    :: Float64,
    vv   :: Matrix{Float64},
    buff :: Vector{Float64},
    r    :: Reduced_lattice,
    m    :: Mesh,
    a    :: Action_u1_dm
    )    :: Nothing

    for i in eachindex(buff)
        # map integration arguments and modify increments
        v   = (2.0 * vv[1, i] - 1.0) / ((1.0 - vv[1, i]) * vv[1, i])
        vp  = (2.0 * vv[2, i] - 1.0) / ((1.0 - vv[2, i]) * vv[2, i])
        dv  = (2.0 * vv[1, i] * vv[1, i] - 2.0 * vv[1, i] + 1) / ((1.0 - vv[1, i]) * (1.0 - vv[1, i]) * vv[1, i] * vv[1, i])
        dvp = (2.0 * vv[2, i] * vv[2, i] - 2.0 * vv[2, i] + 1) / ((1.0 - vv[2, i]) * (1.0 - vv[2, i]) * vv[2, i] * vv[2, i])

        # get buffers for non-local term
        bs1 = get_buffer_s(v + vp, 0.5 * (v - w - vp), 0.5 * (-v - w + vp), m)
        bt1 = get_buffer_t(w, v, vp, m)
        bu1 = get_buffer_u(v - vp, 0.5 * (v - w + vp), 0.5 * ( v + w + vp), m)

        # get buffers for local term
        bs2 = get_buffer_s( v + vp, 0.5 * (v - w - vp), 0.5 * (v + w - vp), m)
        bt2 = get_buffer_t(-v + vp, 0.5 * (v - w + vp), 0.5 * (v + w + vp), m)
        bu2 = get_buffer_u(-w, v, vp, m)

        # compute value
        buff[i] = -get_Γ_comp(1, site, bs1, bt1, bu1, r, a, apply_flags_u1_dm) / (2.0 * pi)^2

        if site == 1
            vzz      = get_Γ_comp(2, site, bs2, bt2, bu2, r, a, apply_flags_u1_dm)
            vdd      = get_Γ_comp(4, site, bs2, bt2, bu2, r, a, apply_flags_u1_dm)
            buff[i] -= (vzz - vdd) / (2.0 * (2.0 * pi)^2)
        end

        buff[i] *= get_G(Λ,  v - 0.5 * w, m, a) * get_G(Λ,  v + 0.5 * w, m, a) * dv
        buff[i] *= get_G(Λ, vp - 0.5 * w, m, a) * get_G(Λ, vp + 0.5 * w, m, a) * dvp
    end

    return nothing
end

# zz kernel for double integral
function compute_χ_kernel_zz!(
    Λ    :: Float64,
    site :: Int64,
    w    :: Float64,
    vv   :: Matrix{Float64},
    buff :: Vector{Float64},
    r    :: Reduced_lattice,
    m    :: Mesh,
    a    :: Action_u1_dm
    )    :: Nothing

    for i in eachindex(buff)
        # map integration arguments and modify increments
        v   = (2.0 * vv[1, i] - 1.0) / ((1.0 - vv[1, i]) * vv[1, i])
        vp  = (2.0 * vv[2, i] - 1.0) / ((1.0 - vv[2, i]) * vv[2, i])
        dv  = (2.0 * vv[1, i] * vv[1, i] - 2.0 * vv[1, i] + 1) / ((1.0 - vv[1, i]) * (1.0 - vv[1, i]) * vv[1, i] * vv[1, i])
        dvp = (2.0 * vv[2, i] * vv[2, i] - 2.0 * vv[2, i] + 1) / ((1.0 - vv[2, i]) * (1.0 - vv[2, i]) * vv[2, i] * vv[2, i])

        # get buffers for non-local term
        bs1 = get_buffer_s(v + vp, 0.5 * (v - w - vp), 0.5 * (-v - w + vp), m)
        bt1 = get_buffer_t(w, v, vp, m)
        bu1 = get_buffer_u(v - vp, 0.5 * (v - w + vp), 0.5 * ( v + w + vp), m)

        # get buffers for local term
        bs2 = get_buffer_s( v + vp, 0.5 * (v - w - vp), 0.5 * (v + w - vp), m)
        bt2 = get_buffer_t(-v + vp, 0.5 * (v - w + vp), 0.5 * (v + w + vp), m)
        bu2 = get_buffer_u(-w, v, vp, m)

        # compute value
        buff[i] = -get_Γ_comp(2, site, bs1, bt1, bu1, r, a, apply_flags_u1_dm) / (2.0 * pi)^2

        if site == 1
            vxx      = get_Γ_comp(1, site, bs2, bt2, bu2, r, a, apply_flags_u1_dm)
            vzz      = get_Γ_comp(2, site, bs2, bt2, bu2, r, a, apply_flags_u1_dm)
            vdd      = get_Γ_comp(4, site, bs2, bt2, bu2, r, a, apply_flags_u1_dm)
            buff[i] -= (2.0 * vxx - vzz - vdd) / (2.0 * (2.0 * pi)^2)
        end

        buff[i] *= get_G(Λ,  v - 0.5 * w, m, a) * get_G(Λ,  v + 0.5 * w, m, a) * dv
        buff[i] *= get_G(Λ, vp - 0.5 * w, m, a) * get_G(Λ, vp + 0.5 * w, m, a) * dvp
    end

    return nothing
end

# xy kernel for double integral
function compute_χ_kernel_xy!(
    Λ    :: Float64,
    site :: Int64,
    w    :: Float64,
    vv   :: Matrix{Float64},
    buff :: Vector{Float64},
    r    :: Reduced_lattice,
    m    :: Mesh,
    a    :: Action_u1_dm
    )    :: Nothing

    for i in eachindex(buff)
        # map integration arguments and modify increments
        v   = (2.0 * vv[1, i] - 1.0) / ((1.0 - vv[1, i]) * vv[1, i])
        vp  = (2.0 * vv[2, i] - 1.0) / ((1.0 - vv[2, i]) * vv[2, i])
        dv  = (2.0 * vv[1, i] * vv[1, i] - 2.0 * vv[1, i] + 1) / ((1.0 - vv[1, i]) * (1.0 - vv[1, i]) * vv[1, i] * vv[1, i])
        dvp = (2.0 * vv[2, i] * vv[2, i] - 2.0 * vv[2, i] + 1) / ((1.0 - vv[2, i]) * (1.0 - vv[2, i]) * vv[2, i] * vv[2, i])

        # get buffers for non-local term
        bs1 = get_buffer_s(v + vp, 0.5 * (v - w - vp), 0.5 * (-v - w + vp), m)
        bt1 = get_buffer_t(w, v, vp, m)
        bu1 = get_buffer_u(v - vp, 0.5 * (v - w + vp), 0.5 * ( v + w + vp), m)

        # get buffers for local term
        bs2 = get_buffer_s( v + vp, 0.5 * (v - w - vp), 0.5 * (v + w - vp), m)
        bt2 = get_buffer_t(-v + vp, 0.5 * (v - w + vp), 0.5 * (v + w + vp), m)
        bu2 = get_buffer_u(-w, v, vp, m)

        # compute value
        buff[i] = -get_Γ_comp(3, site, bs1, bt1, bu1, r, a, apply_flags_u1_dm) / (2.0 * pi)^2

        if site == 1
            vzd      = get_Γ_comp(5, site, bs2, bt2, bu2, r, a, apply_flags_u1_dm)
            vdz      = get_Γ_comp(6, site, bs2, bt2, bu2, r, a, apply_flags_u1_dm)
            buff[i] -= (vzd - vdz) / (2.0 * (2.0 * pi)^2)
        end

        buff[i] *= get_G(Λ,  v - 0.5 * w, m, a) * get_G(Λ,  v + 0.5 * w, m, a) * dv
        buff[i] *= get_G(Λ, vp - 0.5 * w, m, a) * get_G(Λ, vp + 0.5 * w, m, a) * dvp
    end

    return nothing
end

# compute correlations in real and Matsubara space
function compute_χ!(
    Λ     :: Float64,
    r     :: Reduced_lattice,
    m     :: Mesh,
    a     :: Action_u1_dm,
    χ     :: Vector{Matrix{Float64}},
    χ_tol :: NTuple{2, Float64}
    )     :: Nothing

    @sync for w in 1 : m.num_χ
        for i in eachindex(r.sites)
            Threads.@spawn begin
                # determine reference scale
                ref = Λ + 0.5 * m.χ[w]

                # compute xx vertex contribution
                integrand  = (vv, buff) -> compute_χ_kernel_xx!(Λ, i, m.χ[w], vv, buff, r, m, a)
                χ[1][i, w] = integrate_χ_boxes(integrand, 4.0 * ref, χ_tol)

                # compute zz vertex contribution
                integrand  = (vv, buff) -> compute_χ_kernel_zz!(Λ, i, m.χ[w], vv, buff, r, m, a)
                χ[2][i, w] = integrate_χ_boxes(integrand, 4.0 * ref, χ_tol)

                # compute xy vertex contribution
                integrand  = (vv, buff) -> compute_χ_kernel_xy!(Λ, i, m.χ[w], vv, buff, r, m, a)
                χ[3][i, w] = integrate_χ_boxes(integrand, 4.0 * ref, χ_tol)

                # compute propagator contribution
                if i == 1
                    integrand   = v -> get_G(Λ, v - 0.5 * m.χ[w], m, a) * get_G(Λ, v + 0.5 * m.χ[w], m, a) / (4.0 * pi)
                    res         = quadgk(integrand, -Inf, -2.0 * ref, 0.0, 2.0 * ref, Inf, atol = χ_tol[1], rtol = χ_tol[2])[1]
                    χ[1][i, w] += res
                    χ[2][i, w] += res
                end
            end
        end
    end

    return nothing
end


text:
@inline function eig_promote(f, X::Hermitian)
    evals, U = eigen(X)
    return Hermitian(U * Diagonal(f(evals)) * U')
end

"""
    majmin(ρ::Hermitian, ϵ)

Returns the minimum in majorization order over the trace distance ball (of density matrices) of radius `ϵ` around a density matrix `ρ`.
"""
majmin(ρ::Hermitian, ϵ) = eig_promote(p -> majmin(p, ϵ), ρ)

"""
    majmax(ρ::Hermitian, ϵ)

Returns the maximum in majorization order over the trace distance ball (of density matrices) of radius `ϵ` around a density matrix `ρ`.
"""
majmax(ρ::Hermitian, ϵ) = eig_promote(p -> majmax(p, ϵ), ρ)

"""
    ≺(ρ::Hermitian, σ::Hermitian; tol=1e-8)

Returns true if `σ` majorizes `ρ` and false otherwise. The keyword argument `tol` specifies a tolerance for the comparisons. Can be used in infix form, i.e. `ρ ≺ σ`.
"""
≺(ρ::Hermitian, σ::Hermitian; tol = 1e-8) = ≺(eigvals(ρ), eigvals(σ); tol = tol)

"""
    tracedist(ρ::AbstractMatrix, σ::AbstractMatrix)

Computes the trace distance between two matrices `ρ` and `σ`.
"""
tracedist(ρ::AbstractMatrix, σ::AbstractMatrix) = 0.5 * sum(abs, eigvals(ρ - σ))

tracedist(ρ::Hermitian, σ::Hermitian) = 0.5 * sum(svdvals(ρ - σ))

# error checking and wrapping for non `Hermitian`-type matrices:

function ≺(ρ::AbstractMatrix, σ::AbstractMatrix; tol = 1e-8)
    ishermitian(ρ) && ishermitian(σ) ||
    throw(ArgumentError("Inputs must be Hermitian matricies"))
    ≺(Hermitian(ρ), Hermitian(σ); tol = tol)
end

function majmax(ρ::AbstractMatrix, ϵ)
    ishermitian(ρ) || throw(ArgumentError("The input must be a Hermitian matrix"))
    return majmax(Hermitian(ρ), ϵ)
end

function majmin(ρ::AbstractMatrix, ϵ)
    ishermitian(ρ) || throw(ArgumentError("The input must be a Hermitian matrix"))
    return majmin(Hermitian(ρ), ϵ)
end


text:
###
### Random Sequence Generator
###
###
### Random LongSequence generator.
###
### This file is a part of BioJulia.
### License is MIT: https://github.com/BioJulia/BioSequences.jl/blob/master/LICENSE.md

import Random: Sampler, rand!
const DEFAULT_RNG = StableRNG(rand(Int) & typemax(Int))

"""
    SamplerUniform{T}

Uniform sampler of type T. Instantiate with a collection of eltype T containing
the elements to sample.

# Examples
```
julia> sp = SamplerUniform(rna"ACGU");
```
"""
struct SamplerUniform{T} <: Sampler{T}
    elems::Vector{T}

    function SamplerUniform{T}(elems) where {T}
        elemsvector = convert(Vector{T}, collect(elems))
        if length(elemsvector) < 1
            throw(ArgumentError("elements collection must be non-empty"))
        end
        return new(elemsvector)
    end
end

SamplerUniform(elems) = SamplerUniform{eltype(elems)}(elems)
Base.eltype(::Type{SamplerUniform{T}}) where {T} = T
Base.rand(rng::AbstractRNG, sp::SamplerUniform) = rand(rng, sp.elems)
Base.rand(sp::SamplerUniform) = rand(DEFAULT_RNG, sp.elems)
const DefaultAASampler = SamplerUniform(aa"ACDEFGHIKLMNPQRSTVWY")

"""
    SamplerWeighted{T}

Weighted sampler of type T. Instantiate with a collection of eltype T containing
the elements to sample, and an orderen collection of probabilities to sample
each element except the last. The last probability is the remaining probability
up to 1.

# Examples
```
julia> sp = SamplerWeighted(rna"ACGUN", fill(0.2475, 4));
```
"""
struct SamplerWeighted{T} <: Sampler{T}
    elems::Vector{T}
    probs::Vector{Float64}

    function SamplerWeighted{T}(elems, probs) where {T}
        elemsvector = convert(Vector{T}, collect(elems))
        probsvector = convert(Vector{Float64}, collect(probs))
        if !isempty(probsvector)
            probsum = sum(probsvector)
            if probsum > 1.0
                throw(ArgumentError("sum of probabilties cannot exceed 1.0"))
            elseif minimum(probsvector) < 0.0
                throw(ArgumentError("probabilities must be non-negative"))
            end
        else
            probsum = 0.0
        end
        if length(elemsvector) != length(probsvector) + 1
            throw(ArgumentError("length of elems must be length of probs + 1"))
        end
        # Even with float weirdness,  we can guarantee x + (1.0 - x) == 1.0,
        # when 0 ≤ x ≤ 1, as there's no exponent, and it works like int addition
        return new(elemsvector, push!(probsvector, 1.0 - probsum))
    end
end

SamplerWeighted(elems, probs) = SamplerWeighted{eltype(elems)}(elems, probs)
Base.eltype(::Type{SamplerWeighted{T}}) where {T} = T

function Base.rand(rng::AbstractRNG, sp::SamplerWeighted)
    r = rand(rng)
    j = 1
    @inbounds cumulative_prob = sp.probs[j]
    while cumulative_prob < r
        j += 1
        @inbounds cumulative_prob += sp.probs[j]
    end
    return @inbounds sp.elems[j]
end
Base.rand(sp::SamplerWeighted) = rand(DEFAULT_RNG, sp)

"""
    randseq([rng::AbstractRNG], A::Alphabet, sp::Sampler, len::Integer)

Generate a LongSequence{A} of length `len` with elements drawn from
the given sampler.

# Example:
```
# Generate 1000-length RNA with 4% chance of N, 24% for A, C, G, or U
julia> sp = SamplerWeighted(rna"ACGUN", fill(0.24, 4))
julia> seq = randseq(RNAAlphabet{4}(), sp, 50)
50nt RNA Sequence:
CUNGGGCCCGGGNAAACGUGGUACACCCUGUUAAUAUCAACNNGCGCUNU
```
"""
function randseq(rng::AbstractRNG, A::Alphabet, sp::Sampler, len::Integer)
    seq = LongSequence{typeof(A)}(len)
    @inbounds for i in 1:len
        letter = rand(rng, sp)
        unsafe_setindex!(seq, letter, i)
    end
    return seq
end
randseq(A::Alphabet, sp::Sampler, len::Integer) = randseq(DEFAULT_RNG, A, sp, len)

"""
    randseq([rng::AbstractRNG], A::Alphabet, len::Integer)

Generate a LongSequence{A} of length `len` from the specified alphabet, drawn
from the default distribution. User-defined alphabets should implement this
method to implement random LongSequence generation.

For RNA and DNA alphabets, the default distribution is uniform across A, C, G,
and T/U.
For AminoAcidAlphabet, it is uniform across the 20 standard amino acids.
For a user-defined alphabet A, default is uniform across all elements of
`symbols(A)`.

# Example:
```
julia> seq = randseq(AminoAcidAlphabet(), 50)
50aa Amino Acid Sequence:
VFMHSIRMIRLMVHRSWKMHSARHVNFIRCQDKKWKSADGIYTDICKYSM
```
"""
function randseq(rng::AbstractRNG, A::NucleicAcidAlphabet{4}, len::Integer)
    seq = LongSequence{typeof(A)}(len)
    data = seq.data
    rand!(rng, data)
    @inbounds for i in eachindex(data)
        nuc = 0x1111111111111111
        mask = data[i]
        nuc = ((nuc & mask) << 1) | (nuc & ~mask)
        mask >>>= 1
        nuc = ((nuc & mask) << 2) | (nuc & ~mask)
        data[i] = nuc
    end
    return seq
end
randseq(A::NucleicAcidAlphabet{4}, len::Integer) = randseq(DEFAULT_RNG, A, len)

function randseq(rng::AbstractRNG, ::AminoAcidAlphabet, len::Integer)
    return randseq(rng, AminoAcidAlphabet(), DefaultAASampler, len)
end
randseq(A::AminoAcidAlphabet, len::Integer) = randseq(DEFAULT_RNG, A, len)

# Generic fallback method.
function randseq(rng::AbstractRNG, A::Alphabet, len::Integer)
    if length(A) == 1 << bits_per_symbol(A)
        return randseq_allbits(rng, A, len)
    else
        letters = symbols(A)
        sampler = SamplerUniform{eltype(A)}(letters)
        return randseq(rng, A, sampler, len)
    end
end
randseq(A::Alphabet, len::Integer) = randseq(DEFAULT_RNG, A, len)

function randseq_allbits(rng::AbstractRNG, A::Alphabet, len::Integer)
    vector = rand(rng, UInt64, seq_data_len(typeof(A), len))
    return LongSequence{typeof(A)}(vector, 1:len, false)
end

"""
    randdnaseq([rng::AbstractRNG], len::Integer)

Generate a random LongSequence{DNAAlphabet{4}} sequence of length `len`,
with bases sampled uniformly from [A, C, G, T]
"""
randdnaseq(rng::AbstractRNG, len::Integer) = randseq(rng, DNAAlphabet{4}(), len)
randdnaseq(len::Integer) = randseq(DEFAULT_RNG, DNAAlphabet{4}(), len)

"""
    randrnaseq([rng::AbstractRNG], len::Integer)

Generate a random LongSequence{RNAAlphabet{4}} sequence of length `len`,
with bases sampled uniformly from [A, C, G, U]
"""
randrnaseq(rng::AbstractRNG, len::Integer) = randseq(rng, RNAAlphabet{4}(), len)
randrnaseq(len::Integer) = randseq(DEFAULT_RNG, RNAAlphabet{4}(), len)

"""
    randaaseq([rng::AbstractRNG], len::Integer)

Generate a random LongSequence{AminoAcidAlphabet} sequence of length `len`,
with amino acids sampled uniformly from the 20 standard amino acids.
"""
randaaseq(rng::AbstractRNG, len::Integer) = randseq(rng, AminoAcidAlphabet(), len)
randaaseq(len::Integer) = randseq(DEFAULT_RNG, AminoAcidAlphabet(), len)