text:
module AutoMLPipeline

using AMLPipelineBase
using AMLPipelineBase.AbsTypes
export fit, fit!, transform, transform!,fit_transform, fit_transform!

using AMLPipelineBase
using AMLPipelineBase: AbsTypes, Utils, BaselineModels, Pipelines
using AMLPipelineBase: BaseFilters, FeatureSelectors, DecisionTreeLearners
using AMLPipelineBase: EnsembleMethods, CrossValidators
using AMLPipelineBase: NARemovers

export Machine, Learner, Transformer, Workflow, Computer
export holdout, kfold, score, infer_eltype, nested_dict_to_tuples,
       nested_dict_set!, nested_dict_merge, create_transformer,
       mergedict, getiris, getprofb,
       skipmean,skipmedian,skipstd,
       aggregatorclskipmissing,
       find_catnum_columns,
       train_test_split

export Baseline, Identity
export Imputer,OneHotEncoder,Wrapper
export PrunedTree,RandomForest,Adaboost
export VoteEnsemble, StackEnsemble, BestLearner
export FeatureSelector, CatFeatureSelector, NumFeatureSelector, CatNumDiscriminator
export crossvalidate
export NARemover
export @pipeline, @pipelinex
export +, |>, *, |, >>
export Pipeline, ComboPipeline

import AMLPipelineBase.AbsTypes: fit!, transform!

# --------------------------------------------

include("skpreprocessor.jl")
using .SKPreprocessors
export SKPreprocessor, skpreprocessors

include("sklearners.jl")
using .SKLearners
export SKLearner, sklearners

include("skcrossvalidator.jl")
using .SKCrossValidators
export crossvalidate

export skoperator

function skoperator(name::String; args...)::Machine
   sklr = keys(SKLearners.learner_dict)
   skpr = keys(SKPreprocessors.preprocessor_dict)
   if name ∈ sklr
      obj = SKLearner(name; args...)
   elseif name ∈ skpr
      obj = SKPreprocessor(name; args...)
   else
      skoperator()
      throw(ArgumentError("$name does not exist"))
   end
   return obj
end

function skoperator()
   sklr = keys(SKLearners.learner_dict)
   skpr = keys(SKPreprocessors.preprocessor_dict)
   println("Please choose among these pipeline elements:")
   println([sklr..., skpr...])
end

end # module


text:
N = 10
NG = N#5
NBA = 1
# Set the values of the unconditional means
maxW = 200000
minW = 10000
tmpStrIO = round.([ Vector(linspace(minW,maxW,N));Vector(linspace(minW,maxW ,N))])
N, StrI,StrO = splitVec(tmpStrIO);sum(StrI) == sum(StrO)?():error()
dgpStrIO  =   StaticNets.DegSeq2graphDegSeq(StaticNets.fooErgmDirWCount1,tmpStrIO)
groupsInds = [Utilities.distributeAinVecN(Int.(1:NG),N) , Utilities.distributeAinVecN(Int.(1:1),NG) , Utilities.distributeAinVecN(Int.(1:1),NG)]
StaMod = StaticNets.ErgmDirW1(Int.(dgpStrIO),[zeros(Float64,2N),zeros(Float64,1)],groupsInds[1])
estPar,estIt,estMod = StaticNets.estimate(StaMod)
uBndPar = StaticNets.bndPar2uBndPar(StaMod,estPar)
# Test Dgp
A = ones(1)*0.1
B = ones(1)*0.9
W = uBndPar.*(1-B)
T =300
dgpParArr = [W,B,A];dgpParVec = [W;B;A]
scalingType = "FISHER-DIAG"
mod2Samp = SdErgmDirW1(ones(T,2N), [ zeros(2N) ,0.9  * ones(1), 0.01 * ones(1)], groupsInds,scalingType)
NumberOfGroupsAndABindNodes(mod2Samp,groupsInds)

Y_T, Fitness_TIO = score_driven_filter_or_dgp(mod2Samp,[W;B;A];dgpTN = (T,N))
StrI = sumSq(Y_T,3) ; StrO = sumSq(Y_T,2); StrIO = [StrI  StrO]
Mod2est = SdErgmDirW1(StrIO,dgpParArr, groupsInds,scalingType)
#test filter
Fitness_TIO_Fil, loglike = score_driven_filter_or_dgp(Mod2est,dgpParVec; groupsInds = groupsInds)
Fitness_TIO - Fitness_TIO_Fil
estimateSnapSeq(Mod2est;print_t = true)

hatBndPar,x,xx = StaticNets.estimate(StaticNets.fooErgmDirW1;strIO = StrIO[19,:],groupsInds = groupsInds[1] )
hatUbndPar = StaticNets.bndPar2uBndPar(StaticNets.fooErgmDirW1,hatBndPar)

(hatBndPar[1:N] .+ hatBndPar[N+1:2N]' .>=0)
minInd = indmin(hatBndPar[N+1:2N])
hatBndPar[1+N:2N][minInd] .+ hatBndPar[1:N]
## Test Estimate

estimate(Mod2est)
##plots
1;
    N,Fitness_TI,Fitness_TO = splitMat(Fitness_TIO)
    Fitness_TO[1:5,:]
    using Plots
    plotly()
    pparI = plot(Fitness_TI[:,2:end])
    pparO = plot(Fitness_TO);
    plot(pparI,pparO,layout = (2,1),legend=:none,size = (1200,600))
    pstrI = plot(StrI,title = "  Unc Means = $(round(dgpStrIO[1:N]))");
    pstrO = plot(StrO,title = "  Unc Means = $(round(dgpStrIO[N+1:2N]))");
    plot(pstrI,pstrO,layout = (2,1),legend=:none,size = (1200,600))
    dgpStrIOmat =[dgpStrIO[1:N] dgpStrIO[N+1:2N]]
    meanStrIO =  [meanSq(StrI,1)  meanSq(StrO,1)]
    [meanStrIO  dgpStrIOmat ]
    (meanStrIO .- dgpStrIOmat )./dgpStrIOmat

##
A_prob = 0.00123004
B_prob=  0.840832
W_prob_tmp=    [-Inf         -324.547
                5.33316     -0.274072
               -0.412405    -0.378851
               -0.396107    -0.464246
               -0.357238    -0.401917
               -0.41319     -0.428446
               -0.428308    -0.997147
               -0.496042    -0.361475
               -0.433907    -0.36438
               -0.499994    -0.370849]
W_prob = [W_prob_tmp[1:N];W_prob_tmp[N+1:2N]]
Umean_prob =StaticNets.uBndPar2bndPar(StaMod, W_prob./(1-B_prob))
display(StaticNets.expMatrix(StaMod,Umean_prob))
Y_T, Fitness_TIO =score_driven_filter_or_dgp(Mod2est,[W_prob;B_prob;A_prob]; groupsInds = groupsInds)

##
MeanStrT =  meanSq(Mod2est.obsT,1)
[MeanStrT[1:N] MeanStrT[N+1:2N]]
staMod =  StaticNets.ErgmDirW1((MeanStrT),[zeros(Float64,2N),zeros(Float64,1)],groupsInds[1])
StaticNets.estimate(staMod)
MeanStrT[N+1:end] = -MeanStrT[N+1:end]
sum(MeanStrT)
##

##

##

#

ext: jl


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

__precompile__(true)

"""
Utilities for reading and writing delimited files, for example ".csv".
See [`readdlm`](@ref) and [`writedlm`](@ref).
"""
module DelimitedFiles

using Mmap

import Base: _default_delims, tryparse_internal, show

export readdlm, writedlm

Base.@deprecate readcsv(io; opts...) readdlm(io, ','; opts...)
Base.@deprecate readcsv(io, T::Type; opts...) readdlm(io, ',', T; opts...)
Base.@deprecate writecsv(io, a; opts...) writedlm(io, a, ','; opts...)

invalid_dlm(::Type{Char})   = reinterpret(Char, 0xfffffffe)
invalid_dlm(::Type{UInt8})  = 0xfe
invalid_dlm(::Type{UInt16}) = 0xfffe
invalid_dlm(::Type{UInt32}) = 0xfffffffe

const offs_chunk_size = 5000

"""
    readdlm(source, T::Type; options...)

The columns are assumed to be separated by one or more whitespaces. The end of line
delimiter is taken as `\\n`.

# Examples
```jldoctest
julia> using DelimitedFiles

julia> x = [1; 2; 3; 4];

julia> y = [5; 6; 7; 8];

julia> open("delim_file.txt", "w") do io
           writedlm(io, [x y])
       end;

julia> readdlm("delim_file.txt", Int64)
4×2 Array{Int64,2}:
 1  5
 2  6
 3  7
 4  8

julia> readdlm("delim_file.txt", Float64)
4×2 Array{Float64,2}:
 1.0  5.0
 2.0  6.0
 3.0  7.0
 4.0  8.0

julia> rm("delim_file.txt")
```
"""
readdlm(input, T::Type; opts...) = readdlm(input, invalid_dlm(Char), T, '\n'; opts...)

"""
    readdlm(source, delim::Char, T::Type; options...)

The end of line delimiter is taken as `\\n`.

# Examples
```jldoctest
julia> using DelimitedFiles

julia> x = [1; 2; 3; 4];

julia> y = [1.1; 2.2; 3.3; 4.4];

julia> open("delim_file.txt", "w") do io
           writedlm(io, [x y], ',')
       end;

julia> readdlm("delim_file.txt", ',', Float64)
4×2 Array{Float64,2}:
 1.0  1.1
 2.0  2.2
 3.0  3.3
 4.0  4.4

julia> rm("delim_file.txt")
```
"""
readdlm(input, dlm::Char, T::Type; opts...) = readdlm(input, dlm, T, '\n'; opts...)

"""
    readdlm(source; options...)

The columns are assumed to be separated by one or more whitespaces. The end of line
delimiter is taken as `\\n`. If all data is numeric, the result will be a numeric array. If
some elements cannot be parsed as numbers, a heterogeneous array of numbers and strings
is returned.

# Examples
```jldoctest
julia> using DelimitedFiles

julia> x = [1; 2; 3; 4];

julia> y = ["a"; "b"; "c"; "d"];

julia> open("delim_file.txt", "w") do io
           writedlm(io, [x y])
       end;

julia> readdlm("delim_file.txt")
4×2 Array{Any,2}:
 1  "a"
 2  "b"
 3  "c"
 4  "d"

julia> rm("delim_file.txt")
```
"""
readdlm(input; opts...) = readdlm(input, invalid_dlm(Char), '\n'; opts...)

"""
    readdlm(source, delim::Char; options...)

The end of line delimiter is taken as `\\n`. If all data is numeric, the result will be a
numeric array. If some elements cannot be parsed as numbers, a heterogeneous array of
numbers and strings is returned.

# Examples
```jldoctest
julia> using DelimitedFiles

julia> x = [1; 2; 3; 4];

julia> y = [1.1; 2.2; 3.3; 4.4];

julia> open("delim_file.txt", "w") do io
           writedlm(io, [x y], ',')
       end;

julia> readdlm("delim_file.txt", ',')
4×2 Array{Float64,2}:
 1.0  1.1
 2.0  2.2
 3.0  3.3
 4.0  4.4

julia> rm("delim_file.txt")

julia> z = ["a"; "b"; "c"; "d"];

julia> open("delim_file.txt", "w") do io
           writedlm(io, [x z], ',')
       end;

julia> readdlm("delim_file.txt", ',')
4×2 Array{Any,2}:
 1  "a"
 2  "b"
 3  "c"
 4  "d"

julia> rm("delim_file.txt")
```
"""
readdlm(input, dlm::Char; opts...) = readdlm(input, dlm, '\n'; opts...)

"""
    readdlm(source, delim::Char, eol::Char; options...)

If all data is numeric, the result will be a numeric array. If some elements cannot be
parsed as numbers, a heterogeneous array of numbers and strings is returned.
"""
readdlm(input, dlm::Char, eol::Char; opts...) =
    readdlm_auto(input, dlm, Float64, eol, true; opts...)

"""
    readdlm(source, delim::Char, T::Type, eol::Char; header=false, skipstart=0, skipblanks=true, use_mmap, quotes=true, dims, comments=true, comment_char='#')

Read a matrix from the source where each line (separated by `eol`) gives one row, with
elements separated by the given delimiter. The source can be a text file, stream or byte
array. Memory mapped files can be used by passing the byte array representation of the
mapped segment as source.

If `T` is a numeric type, the result is an array of that type, with any non-numeric elements
as `NaN` for floating-point types, or zero. Other useful values of `T` include
`String`, `AbstractString`, and `Any`.

If `header` is `true`, the first row of data will be read as header and the tuple
`(data_cells, header_cells)` is returned instead of only `data_cells`.

Specifying `skipstart` will ignore the corresponding number of initial lines from the input.

If `skipblanks` is `true`, blank lines in the input will be ignored.

If `use_mmap` is `true`, the file specified by `source` is memory mapped for potential
speedups. Default is `true` except on Windows. On Windows, you may want to specify `true` if
the file is large, and is only read once and not written to.

If `quotes` is `true`, columns enclosed within double-quote (\") characters are allowed to
contain new lines and column delimiters. Double-quote characters within a quoted field must
be escaped with another double-quote.  Specifying `dims` as a tuple of the expected rows and
columns (including header, if any) may speed up reading of large files.  If `comments` is
`true`, lines beginning with `comment_char` and text following `comment_char` in any line
are ignored.

# Examples
```jldoctest
julia> using DelimitedFiles

julia> x = [1; 2; 3; 4];

julia> y = [5; 6; 7; 8];

julia> open("delim_file.txt", "w") do io
           writedlm(io, [x y])
       end

julia> readdlm("delim_file.txt", '\\t', Int, '\\n')
4×2 Array{Int64,2}:
 1  5
 2  6
 3  7
 4  8
```
"""
readdlm(input, dlm::Char, T::Type, eol::Char; opts...) =
    readdlm_auto(input, dlm, T, eol, false; opts...)

readdlm_auto(input::Vector{UInt8}, dlm::Char, T::Type, eol::Char, auto::Bool; opts...) =
    readdlm_string(String(input), dlm, T, eol, auto, val_opts(opts))
readdlm_auto(input::IO, dlm::Char, T::Type, eol::Char, auto::Bool; opts...) =
    readdlm_string(read(input, String), dlm, T, eol, auto, val_opts(opts))
function readdlm_auto(input::AbstractString, dlm::Char, T::Type, eol::Char, auto::Bool; opts...)
    isfile(input) || throw(ArgumentError("Cannot open \'$input\': not a file"))
    optsd = val_opts(opts)
    use_mmap = get(optsd, :use_mmap, Sys.iswindows() ? false : true)
    fsz = filesize(input)
    if use_mmap && fsz > 0 && fsz < typemax(Int)
        a = open(input, "r") do f
            Mmap.mmap(f, Vector{UInt8}, (Int(fsz),))
        end
        # TODO: It would be nicer to use String(a) without making a copy,
        # but because the mmap'ed array is not NUL-terminated this causes
        # jl_try_substrtod to segfault below.
        return readdlm_string(GC.@preserve(a, unsafe_string(pointer(a),length(a))), dlm, T, eol, auto, optsd)
    else
        return readdlm_string(read(input, String), dlm, T, eol, auto, optsd)
    end
end

#
# Handlers act on events generated by the parser.
# Parser calls store_cell on the handler to pass events.
#
# DLMOffsets: Keep offsets (when result dimensions are not known)
# DLMStore: Store values directly into a result store (when result dimensions are known)
abstract type DLMHandler end

mutable struct DLMOffsets <: DLMHandler
    oarr::Vector{Vector{Int}}
    offidx::Int
    thresh::Int
    bufflen::Int

    function DLMOffsets(sbuff::String)
        offsets = Vector{Vector{Int}}(uninitialized, 1)
        offsets[1] = Vector{Int}(uninitialized, offs_chunk_size)
        thresh = ceil(min(typemax(UInt), Base.Sys.total_memory()) / sizeof(Int) / 5)
        new(offsets, 1, thresh, sizeof(sbuff))
    end
end

function store_cell(dlmoffsets::DLMOffsets, row::Int, col::Int,
        quoted::Bool, startpos::Int, endpos::Int)
    offidx = dlmoffsets.offidx
    (offidx == 0) && return     # offset collection stopped to avoid choking on memory

    oarr = dlmoffsets.oarr
    offsets = oarr[end]
    if length(offsets) < offidx
        offlen = offs_chunk_size * length(oarr)
        if (offlen + offs_chunk_size) > dlmoffsets.thresh
            est_tot = round(Int, offlen * dlmoffsets.bufflen / endpos)
            if (est_tot - offlen) > offs_chunk_size    # allow another chunk
                # abandon offset collection
                dlmoffsets.oarr = Vector{Int}[]
                dlmoffsets.offidx = 0
                return
            end
        end
        offsets = Vector{Int}(uninitialized, offs_chunk_size)
        push!(oarr, offsets)
        offidx = 1
    end
    offsets[offidx] = row
    offsets[offidx+1] = col
    offsets[offidx+2] = Int(quoted)
    offsets[offidx+3] = startpos
    offsets[offidx+4] = endpos
    dlmoffsets.offidx = offidx + 5
    nothing
end

function result(dlmoffsets::DLMOffsets)
    trimsz = (dlmoffsets.offidx-1) % offs_chunk_size
    ((trimsz > 0) || (dlmoffsets.offidx == 1)) && resize!(dlmoffsets.oarr[end], trimsz)
    dlmoffsets.oarr
end

mutable struct DLMStore{T} <: DLMHandler
    hdr::Array{AbstractString, 2}
    data::Array{T, 2}

    nrows::Int
    ncols::Int
    lastrow::Int
    lastcol::Int
    hdr_offset::Int
    sbuff::String
    auto::Bool
    eol::Char
end

function DLMStore(::Type{T}, dims::NTuple{2,Integer},
                  has_header::Bool, sbuff::String, auto::Bool, eol::Char) where T
    (nrows,ncols) = dims
    nrows <= 0 && throw(ArgumentError("number of rows in dims must be > 0, got $nrows"))
    ncols <= 0 && throw(ArgumentError("number of columns in dims must be > 0, got $ncols"))
    hdr_offset = has_header ? 1 : 0
    DLMStore{T}(fill(SubString(sbuff,1,0), 1, ncols), Matrix{T}(uninitialized, nrows-hdr_offset, ncols),
        nrows, ncols, 0, 0, hdr_offset, sbuff, auto, eol)
end

_chrinstr(sbuff::String, chr::UInt8, startpos::Int, endpos::Int) =
    GC.@preserve sbuff (endpos >= startpos) && (C_NULL != ccall(:memchr, Ptr{UInt8},
    (Ptr{UInt8}, Int32, Csize_t), pointer(sbuff)+startpos-1, chr, endpos-startpos+1))

function store_cell(dlmstore::DLMStore{T}, row::Int, col::Int,
                    quoted::Bool, startpos::Int, endpos::Int) where T
    drow = row - dlmstore.hdr_offset

    ncols = dlmstore.ncols
    lastcol = dlmstore.lastcol
    lastrow = dlmstore.lastrow
    cells::Matrix{T} = dlmstore.data
    sbuff = dlmstore.sbuff

    endpos = prevind(sbuff, nextind(sbuff,endpos))
    if (endpos > 0) && ('\n' == dlmstore.eol) && ('\r' == Char(sbuff[endpos]))
        endpos = prevind(sbuff, endpos)
    end
    if quoted
        startpos += 1
        endpos = prevind(sbuff, endpos)
    end

    if drow > 0
        # fill missing elements
        while ((drow - lastrow) > 1) || ((drow > lastrow > 0) && (lastcol < ncols))
            if (lastcol == ncols) || (lastrow == 0)
                lastcol = 0
                lastrow += 1
            end
            for cidx in (lastcol+1):ncols
                if (T <: AbstractString) || (T == Any)
                    cells[lastrow, cidx] = SubString(sbuff, 1, 0)
                elseif ((T <: Number) || (T <: Char)) && dlmstore.auto
                    throw(TypeError(:store_cell, "", Any, T))
                else
                    error("missing value at row $lastrow column $cidx")
                end
            end
            lastcol = ncols
        end

        # fill data
        if quoted && _chrinstr(sbuff, UInt8('"'), startpos, endpos)
            unescaped = replace(SubString(sbuff, startpos, endpos), r"\"\"" => "\"")
            fail = colval(unescaped, 1, lastindex(unescaped), cells, drow, col)
        else
            fail = colval(sbuff, startpos, endpos, cells, drow, col)
        end
        if fail
            sval = SubString(sbuff, startpos, endpos)
            if (T <: Number) && dlmstore.auto
                throw(TypeError(:store_cell, "", Any, T))
            else
                error("file entry \"$(sval)\" cannot be converted to $T")
            end
        end

        dlmstore.lastrow = drow
        dlmstore.lastcol = col
    else
        # fill header
        if quoted && _chrinstr(sbuff, UInt8('"'), startpos, endpos)
            unescaped = replace(SubString(sbuff, startpos, endpos), r"\"\"" => "\"")
            colval(unescaped, 1, lastindex(unescaped), dlmstore.hdr, 1, col)
        else
            colval(sbuff, startpos, endpos, dlmstore.hdr, 1, col)
        end
    end

    nothing
end

function result(dlmstore::DLMStore{T}) where T
    nrows = dlmstore.nrows - dlmstore.hdr_offset
    ncols = dlmstore.ncols
    lastcol = dlmstore.lastcol
    lastrow = dlmstore.lastrow
    cells = dlmstore.data
    sbuff = dlmstore.sbuff

    if (nrows > 0) && ((lastcol < ncols) || (lastrow < nrows))
        while lastrow <= nrows
            (lastcol == ncols) && (lastcol = 0; lastrow += 1)
            for cidx in (lastcol+1):ncols
                if (T <: AbstractString) || (T == Any)
                    cells[lastrow, cidx] = SubString(sbuff, 1, 0)
                elseif ((T <: Number) || (T <: Char)) && dlmstore.auto
                    throw(TypeError(:store_cell, "", Any, T))
                else
                    error("missing value at row $lastrow column $cidx")
                end
            end
            lastcol = ncols
            (lastrow == nrows) && break
        end
        dlmstore.lastrow = lastrow
        dlmstore.lastcol = ncols
    end
    (dlmstore.hdr_offset > 0) ? (dlmstore.data, dlmstore.hdr) : dlmstore.data
end

function readdlm_string(sbuff::String, dlm::Char, T::Type, eol::Char, auto::Bool, optsd::Dict)
    ign_empty = (dlm == invalid_dlm(Char))
    quotes = get(optsd, :quotes, true)
    comments = get(optsd, :comments, true)
    comment_char = get(optsd, :comment_char, '#')
    dims = get(optsd, :dims, nothing)

    has_header = get(optsd, :header, get(optsd, :has_header, false))
    haskey(optsd, :has_header) && (optsd[:has_header] != has_header) && throw(ArgumentError("conflicting values for header and has_header"))

    skipstart = get(optsd, :skipstart, 0)
    (skipstart >= 0) || throw(ArgumentError("skipstart must be ≥ 0, got $skipstart"))

    skipblanks = get(optsd, :skipblanks, true)

    offset_handler = (dims === nothing) ? DLMOffsets(sbuff) : DLMStore(T, dims, has_header, sbuff, auto, eol)

    for retry in 1:2
        try
            dims = dlm_parse(sbuff, eol, dlm, '"', comment_char, ign_empty, quotes, comments, skipstart, skipblanks, offset_handler)
            break
        catch ex
            if isa(ex, TypeError) && (ex.func == :store_cell)
                T = ex.expected
            else
                rethrow(ex)
            end
            offset_handler = (dims === nothing) ? DLMOffsets(sbuff) : DLMStore(T, dims, has_header, sbuff, auto, eol)
        end
    end

    isa(offset_handler, DLMStore) && (return result(offset_handler))

    offsets = result(offset_handler)
    !isempty(offsets) && (return dlm_fill(T, offsets, dims, has_header, sbuff, auto, eol))

    optsd[:dims] = dims
    return readdlm_string(sbuff, dlm, T, eol, auto, optsd)
end

const valid_opts = [:header, :has_header, :use_mmap, :quotes, :comments, :dims, :comment_char, :skipstart, :skipblanks]
const valid_opt_types = [Bool, Bool, Bool, Bool, Bool, NTuple{2,Integer}, Char, Integer, Bool]

function val_opts(opts)
    d = Dict{Symbol, Union{Bool, NTuple{2, Integer}, Char, Integer}}()
    for (opt_name, opt_val) in opts
        in(opt_name, valid_opts) ||
            throw(ArgumentError("unknown option $opt_name"))
        opt_typ = valid_opt_types[findfirst(equalto(opt_name), valid_opts)::Int]
        isa(opt_val, opt_typ) ||
            throw(ArgumentError("$opt_name should be of type $opt_typ, got $(typeof(opt_val))"))
        d[opt_name] = opt_val
    end
    return d
end

function dlm_fill(T::DataType, offarr::Vector{Vector{Int}}, dims::NTuple{2,Integer}, has_header::Bool, sbuff::String, auto::Bool, eol::Char)
    idx = 1
    offidx = 1
    offsets = offarr[1]
    row = 0
    col = 0
    try
        dh = DLMStore(T, dims, has_header, sbuff, auto, eol)
        while idx <= length(offsets)
            row = offsets[idx]
            col = offsets[idx+1]
            quoted = offsets[idx+2] != 0
            startpos = offsets[idx+3]
            endpos = offsets[idx+4]

            ((idx += 5) > offs_chunk_size) && (offidx < length(offarr)) && (idx = 1; offsets = offarr[offidx += 1])

            store_cell(dh, row, col, quoted, startpos, endpos)
        end
        return result(dh)
    catch ex
        isa(ex, TypeError) && (ex.func == :store_cell) && (return dlm_fill(ex.expected, offarr, dims, has_header, sbuff, auto, eol))
        error("at row $row, column $col : $ex")
    end
end

function colval(sbuff::String, startpos::Int, endpos::Int, cells::Array{Bool,2}, row::Int, col::Int)
    n = tryparse_internal(Bool, sbuff, startpos, endpos, 0, false)
    n === nothing || (cells[row, col] = n)
    n === nothing
end
function colval(sbuff::String, startpos::Int, endpos::Int, cells::Array{T,2}, row::Int, col::Int) where T<:Integer
    n = tryparse_internal(T, sbuff, startpos, endpos, 0, false)
    n === nothing || (cells[row, col] = n)
    n === nothing
end
function colval(sbuff::String, startpos::Int, endpos::Int, cells::Array{T,2}, row::Int, col::Int) where T<:Union{Real,Complex}
    n = tryparse_internal(T, sbuff, startpos, endpos, false)
    n === nothing || (cells[row, col] = n)
    n === nothing
end
function colval(sbuff::String, startpos::Int, endpos::Int, cells::Array{<:AbstractString,2}, row::Int, col::Int)
    cells[row, col] = SubString(sbuff, startpos, endpos)
    return false
end
function colval(sbuff::String, startpos::Int, endpos::Int, cells::Array{Any,2}, row::Int, col::Int)
    # if array is of Any type, attempt parsing only the most common types: Int, Bool, Float64 and fallback to SubString
    len = endpos-startpos+1
    if len > 0
        # check Inteter
        ni64 = tryparse_internal(Int, sbuff, startpos, endpos, 0, false)
        ni64 === nothing || (cells[row, col] = ni64; return false)

        # check Bool
        nb = tryparse_internal(Bool, sbuff, startpos, endpos, 0, false)
        nb === nothing || (cells[row, col] = nb; return false)

        # check float64
        hasvalue, valf64 = ccall(:jl_try_substrtod, Tuple{Bool, Float64},
                                 (Ptr{UInt8}, Csize_t, Csize_t), sbuff, startpos-1, endpos-startpos+1)
        hasvalue && (cells[row, col] = valf64; return false)
    end
    cells[row, col] = SubString(sbuff, startpos, endpos)
    false
end
function colval(sbuff::String, startpos::Int, endpos::Int, cells::Array{<:Char,2}, row::Int, col::Int)
    if startpos == endpos
        cells[row, col] = next(sbuff, startpos)[1]
        return false
    else
        return true
    end
end
colval(sbuff::String, startpos::Int, endpos::Int, cells::Array, row::Int, col::Int) = true

function dlm_parse(dbuff::String, eol::D, dlm::D, qchar::D, cchar::D,
                   ign_adj_dlm::Bool, allow_quote::Bool, allow_comments::Bool,
                   skipstart::Int, skipblanks::Bool, dh::DLMHandler) where D
    ncols = nrows = col = 0
    is_default_dlm = (dlm == invalid_dlm(D))
    error_str = ""
    # 0: begin field, 1: quoted field, 2: unquoted field,
    # 3: second quote (could either be end of field or escape character),
    # 4: comment, 5: skipstart
    state = (skipstart > 0) ? 5 : 0
    is_eol = is_dlm = is_cr = is_quote = is_comment = expct_col = false
    idx = 1
    try
        slen = sizeof(dbuff)
        col_start_idx = 1
        was_cr = false
        while idx <= slen
            val,idx = next(dbuff, idx)
            if (is_eol = (Char(val) == Char(eol)))
                is_dlm = is_comment = is_cr = is_quote = false
            elseif (is_dlm = (is_default_dlm ? in(Char(val), _default_delims) : (Char(val) == Char(dlm))))
                is_comment = is_cr = is_quote = false
            elseif (is_quote = (Char(val) == Char(qchar)))
                is_comment = is_cr = false
            elseif (is_comment = (Char(val) == Char(cchar)))
                is_cr = false
            else
                is_cr = (Char(eol) == '\n') && (Char(val) == '\r')
            end

            if 2 == state   # unquoted field
                if is_dlm
                    state = 0
                    col += 1
                    store_cell(dh, nrows+1, col, false, col_start_idx, idx-2)
                    col_start_idx = idx
                    !ign_adj_dlm && (expct_col = true)
                elseif is_eol
                    nrows += 1
                    col += 1
                    store_cell(dh, nrows, col, false, col_start_idx, idx - (was_cr ? 3 : 2))
                    col_start_idx = idx
                    ncols = max(ncols, col)
                    col = 0
                    state = 0
                elseif (is_comment && allow_comments)
                    nrows += 1
                    col += 1
                    store_cell(dh, nrows, col, false, col_start_idx, idx - 2)
                    ncols = max(ncols, col)
                    col = 0
                    state = 4
                end
            elseif 1 == state   # quoted field
                is_quote && (state = 3)
            elseif 4 == state   # comment line
                if is_eol
                    col_start_idx = idx
                    state = 0
                end
            elseif 0 == state   # begin field
                if is_quote
                    state = (allow_quote && !was_cr) ? 1 : 2
                    expct_col = false
                elseif is_dlm
                    if !ign_adj_dlm
                        expct_col = true
                        col += 1
                        store_cell(dh, nrows+1, col, false, col_start_idx, idx-2)
                    end
                    col_start_idx = idx
                elseif is_eol
                    if (col > 0) || !skipblanks
                        nrows += 1
                        if expct_col
                            col += 1
                            store_cell(dh, nrows, col, false, col_start_idx, idx - (was_cr ? 3 : 2))
                        end
                        ncols = max(ncols, col)
                        col = 0
                    end
                    col_start_idx = idx
                    expct_col = false
                elseif is_comment && allow_comments
                    if col > 0
                        nrows += 1
                        if expct_col
                            col += 1
                            store_cell(dh, nrows, col, false, col_start_idx, idx - 2)
                        end
                        ncols = max(ncols, col)
                        col = 0
                    end
                    expct_col = false
                    state = 4
                elseif !is_cr
                    state = 2
                    expct_col = false
                end
            elseif 3 == state   # second quote
                if is_quote && !was_cr
                    state = 1
                elseif is_dlm && !was_cr
                    state = 0
                    col += 1
                    store_cell(dh, nrows+1, col, true, col_start_idx, idx-2)
                    col_start_idx = idx
                    !ign_adj_dlm && (expct_col = true)
                elseif is_eol
                    nrows += 1
                    col += 1
                    store_cell(dh, nrows, col, true, col_start_idx, idx - (was_cr ? 3 : 2))
                    col_start_idx = idx
                    ncols = max(ncols, col)
                    col = 0
                    state = 0
                elseif is_comment && allow_comments && !was_cr
                    nrows += 1
                    col += 1
                    store_cell(dh, nrows, col, true, col_start_idx, idx - 2)
                    ncols = max(ncols, col)
                    col = 0
                    state = 4
                elseif (is_cr && was_cr) || !is_cr
                    error_str = escape_string("unexpected character '$(Char(val))' after quoted field at row $(nrows+1) column $(col+1)")
                    break
                end
            elseif 5 == state # skip start
                if is_eol
                    col_start_idx = idx
                    skipstart -= 1
                    (0 == skipstart) && (state = 0)
                end
            end
            was_cr = is_cr
        end

        if isempty(error_str)
            if 1 == state       # quoted field
                error_str = "truncated column at row $(nrows+1) column $(col+1)"
            elseif (2 == state) || (3 == state) || ((0 == state) && is_dlm)   # unquoted field, second quote, or begin field with last character as delimiter
                col += 1
                nrows += 1
                store_cell(dh, nrows, col, (3 == state), col_start_idx, idx-1)
                ncols = max(ncols, col)
            end
        end
    catch ex
        if isa(ex, TypeError) && (ex.func == :store_cell)
            rethrow(ex)
        else
            error("at row $(nrows+1), column $col : $ex)")
        end
    end
    !isempty(error_str) && error(error_str)

    return (nrows, ncols)
end

# todo: keyword argument for # of digits to print
writedlm_cell(io::IO, elt::AbstractFloat, dlm, quotes) = print(io, elt)
function writedlm_cell(io::IO, elt::AbstractString, dlm::T, quotes::Bool) where T
    if quotes && !isempty(elt) && (('"' in elt) || ('\n' in elt) || ((T <: Char) ? (dlm in elt) : contains(elt, dlm)))
        print(io, '"', replace(elt, r"\"" => "\"\""), '"')
    else
        print(io, elt)
    end
end
writedlm_cell(io::IO, elt, dlm, quotes) = print(io, elt)
function writedlm(io::IO, a::AbstractMatrix, dlm; opts...)
    optsd = val_opts(opts)
    quotes = get(optsd, :quotes, true)
    pb = PipeBuffer()
    lastc = last(axes(a, 2))
    for i = axes(a, 1)
        for j = axes(a, 2)
            writedlm_cell(pb, a[i, j], dlm, quotes)
            j == lastc ? write(pb,'\n') : print(pb,dlm)
        end
        (bytesavailable(pb) > (16*1024)) && write(io, take!(pb))
    end
    write(io, take!(pb))
    nothing
end

writedlm(io::IO, a::AbstractArray{<:Any,0}, dlm; opts...) = writedlm(io, reshape(a,1), dlm; opts...)

# write an iterable row as dlm-separated items
function writedlm_row(io::IO, row, dlm, quotes)
    state = start(row)
    while !done(row, state)
        (x, state) = next(row, state)
        writedlm_cell(io, x, dlm, quotes)
        done(row, state) ? write(io,'\n') : print(io,dlm)
    end
end

# If the row is a single string, write it as a string rather than
# iterating over characters. Also, include the common case of
# a Number (handled correctly by the generic writedlm_row above)
# purely as an optimization.
function writedlm_row(io::IO, row::Union{Number,AbstractString}, dlm, quotes)
    writedlm_cell(io, row, dlm, quotes)
    write(io, '\n')
end

# write an iterable collection of iterable rows
function writedlm(io::IO, itr, dlm; opts...)
    optsd = val_opts(opts)
    quotes = get(optsd, :quotes, true)
    pb = PipeBuffer()
    for row in itr
        writedlm_row(pb, row, dlm, quotes)
        (bytesavailable(pb) > (16*1024)) && write(io, take!(pb))
    end
    write(io, take!(pb))
    nothing
end

function writedlm(fname::AbstractString, a, dlm; opts...)
    open(fname, "w") do io
        writedlm(io, a, dlm; opts...)
    end
end

"""
    writedlm(f, A, delim='\\t'; opts)

Write `A` (a vector, matrix, or an iterable collection of iterable rows) as text to `f`
(either a filename string or an `IO` stream) using the given delimiter
`delim` (which defaults to tab, but can be any printable Julia object, typically a `Char` or
`AbstractString`).

For example, two vectors `x` and `y` of the same length can be written as two columns of
tab-delimited text to `f` by either `writedlm(f, [x y])` or by `writedlm(f, zip(x, y))`.

# Examples
```jldoctest
julia> using DelimitedFiles

julia> x = [1; 2; 3; 4];

julia> y = [5; 6; 7; 8];

julia> open("delim_file.txt", "w") do io
           writedlm(io, [x y])
       end

julia> readdlm("delim_file.txt", '\\t', Int, '\\n')
4×2 Array{Int64,2}:
 1  5
 2  6
 3  7
 4  8
```
"""
writedlm(io, a; opts...) = writedlm(io, a, '\t'; opts...)

show(io::IO, ::MIME"text/csv", a) = writedlm(io, a, ',')
show(io::IO, ::MIME"text/tab-separated-values", a) = writedlm(io, a, '\t')

end # module DelimitedFiles

max_stars_repo_name: BenChung/Julia-BTI

max_issues_repo_name: BenChung/Julia-BTI

max_forks_repo_name: BenChung/Julia-BTI


text:
"""
    get_line(table::Table)

Get the next line of the table by using `table.current_values`.
Call [`format_table_value`](@ref) to format each value and use the alignments to create the line such that it fits to [`get_header`](@ref).
"""
function get_line(table::Table)
    setup = table.setup
    ln = ""
    for c in 1:length(setup.ids)
        width = setup.widths[c]
        values = table.current_values
        default_precision = setup.precisions[c]
        if isassigned(values, c)
            val = values[c]
            s_val = format_table_value(width-2, get_value(val); default_precision)
        else
            s_val = "-"
        end
        padding = width - 2 - length(s_val)
        alignment = setup.alignments[c]
        if alignment == :center
            left_padding = repeat(" ", fld(padding, 2) + 1)
            right_padding = repeat(" ", cld(padding, 2) + 1)
            ln = "$(ln)$(left_padding)$(s_val)$(right_padding)"
        elseif alignment == :left
            right_padding = repeat(" ", padding + 1)
            ln = "$(ln) $(s_val)$(right_padding)"
        elseif alignment == :right
            left_padding = repeat(" ", padding + 1)
            ln = "$(ln)$(left_padding)$(s_val) "
        else
            @error "Only the alignments :left, :right and :center are defined. $alignment isn't defined."
        end
    end
    return ln
end

"""
    fill_from_prev!(table::Table)

If a value isn't given by a new called [`set_value!`](@ref) since the last call to [`print_line`](@ref) the previous value will be used.
This function overwrites `table.current_values` to set unassigned values to `table.prev_values`.
"""
function fill_from_prev!(table::Table)
    for i in 1:length(table.current_values)
        if !isassigned(table.current_values, i) || isnothing(table.current_values[i])
            table.current_values[i] = table.prev_values[i]
        end
    end
end

"""
    shall_print_line(table::Table; force=false)

Return whether the new line shall be printed. If `force = true` return true immediately.
Otherwise check if at least one value differs enough from the previous value by calling [`differs_enough`](@ref).
"""
function shall_print_line(table::Table; force=false)
    force && return true

    # check if a current value differs enough from the previous value
    shall_update = false
    for i in 1:length(table.current_values)
        !isassigned(table.current_values, i) && continue
        value = table.current_values[i]
        if !isassigned(table.prev_values, i) || differs_enough(value, table.prev_values[i])
            shall_update = true
            break
        end
    end
    return shall_update
end

"""
    print_line(table::Table; force=false)

Print the new line of the table if it differs enough from the previous line or if `force = true`.
If the new line gets printed set the `prev_values` to `current_values` and the `current_values` to an `nothing`.
"""
function print_line(table::Table; force=false)
    fill_from_prev!(table)

    shall_print_line(table; force) || return
    println(get_line(table))
    update_for_new_row(table)
    return
end


text:
"""
    multinomial_hyper(α::AbstractArray{Float32,1})

[Dirichlet Distribution](https://en.wikipedia.org/wiki/Dirichlet_distribution)
"""
struct multinomial_hyper <: distribution_hyper_params
    α::AbstractArray{Float32,1}
end

mutable struct multinomial_sufficient_statistics <: sufficient_statistics
    N::Float32
    points_sum::AbstractArray{Float32,1}
end

function calc_posterior(prior:: multinomial_hyper, suff_statistics::multinomial_sufficient_statistics)
    if suff_statistics.N == 0
        return prior
    end
    return multinomial_hyper(prior.α + suff_statistics.points_sum)
end

function sample_distribution(hyperparams::multinomial_hyper)
    return multinomial_dist(log.(rand(Dirichlet(Float64.(hyperparams.α)))))
end

function create_sufficient_statistics(hyper::multinomial_hyper,posterior::multinomial_hyper,points::AbstractArray{Float32,2}, pts_to_group = 0)
    pts = copy(points)
    points_sum = sum(pts, dims = 2)[:]
    S = pts * pts'
    return multinomial_sufficient_statistics(size(points,2),points_sum)
end

function log_marginal_likelihood(hyper::multinomial_hyper, posterior_hyper::multinomial_hyper, suff_stats::multinomial_sufficient_statistics)
    D = length(suff_stats.points_sum)
    logpi = log(pi)
    val = logabsgamma(sum(hyper.α))[1] -logabsgamma(sum(posterior_hyper.α))[1] + sum((x-> logabsgamma(x)[1]).(posterior_hyper.α) - (x-> logabsgamma(x)[1]).(hyper.α))
    return val
end

function aggregate_suff_stats(suff_l::multinomial_sufficient_statistics, suff_r::multinomial_sufficient_statistics)
    return multinomial_sufficient_statistics(suff_l.N+suff_r.N, suff_l.points_sum + suff_r.points_sum)
end


text:
# The thing that most internal types should probably subtype.
abstract type SthenoAbstractGP <: AbstractGP end

# TYPE-PIRACY
function var(f::GP, x::AbstractVector, x′::AbstractVector)
    return kernelmatrix_diag(f.kernel, x, x′)
end

# Ensure that this package gets to handle the covariance between its own GPs.
# AbstractGPs doesn't support this in general because it's unclear how it ought to be
# implemented, but we have a clear way to implement it here.
function cov(fx::FiniteGP{<:SthenoAbstractGP}, gx::FiniteGP{<:SthenoAbstractGP})
    return cov(fx.f, gx.f, fx.x, gx.x)
end

# A collection of GPs (GPC == "GP Collection"). Used to keep track of GPs.
mutable struct GPC
    n::Int
    GPC() = new(0)
end

ChainRulesCore.@non_differentiable GPC()

next_index(gpc::GPC) = gpc.n + 1

# TYPE PIRACY!
LinearAlgebra.cholesky(D::Diagonal{<:Real, <:Fill}) = AbstractGPs._cholesky(D)


text:
@testset "Objective" begin

    # Test GameObjective
    v = [1,2,3]
    inds = [1,3,5]
    n = 5
    V = [1,0,2,0,3]
    @test V == Algames.expand_vector(v,inds,n)

    # Test Game Objective
    T = Float64
    N = 10
    n = 12
    m = 6
    p = 3
    model = UnicycleGame(p=p)
    Q = [Diagonal(rand(SVector{model.ni[i],T})) for i=1:p]
    R = [Diagonal(rand(SVector{model.mi[i],T})) for i=1:p]
    xf = [i*ones(SVector{model.ni[i],T}) for i=1:p]
    uf = [2i*ones(SVector{model.mi[i],T}) for i=1:p]

    game_obj = GameObjective(Q,R,xf,uf,N,model)

    X = [1.0*SVector{model.n,T}([1,2,3,1,2,3,1,2,3,1,2,3]) for k=1:N]
    U = [1.0*SVector{model.m,T}([2,4,6,2,4,6]) for k=1:N-1]
    dt = 0.1
    Dt = dt*ones(N-1)
    traj = Traj(X, U,Dt)

    @test abs(cost(game_obj.obj[1][1], traj)) <= 1e-10
    @test abs(cost(game_obj.obj[2][1], traj)) <= 1e-10
    @test abs(cost(game_obj.obj[3][1], traj)) <= 1e-10

    # Test cost_gradient! and cost_hessian!
    probsize = ProblemSize(N, model)
    pdtraj = PrimalDualTraj(probsize, dt, f=rand, amplitude=1e1)
    for i = 1:p
        for k = 1:N
            game_obj.E[i][1].cost[k].Q += rand(n,n)
            game_obj.E[i][1].cost[k].R += rand(m,m)
            game_obj.E[i][1].cost[k].H += rand(m,n)
            game_obj.E[i][1].cost[k].q += rand(n)
            game_obj.E[i][1].cost[k].r += rand(m)
            game_obj.E[i][1].cost[k].c += rand()
        end
    end
    cost_gradient!(game_obj, pdtraj)
    cost_hessian!(game_obj, pdtraj)

    # Gradient
    # Stage cost
    @test norm(game_obj.E[1][1].cost[1].q - (game_obj.obj[1][1].cost[1].Q*state(pdtraj.pr[1]) + game_obj.obj[1][1].cost[1].q)*dt, 1) < 1e-10
    @test norm(game_obj.E[1][1].cost[1].r - (game_obj.obj[1][1].cost[1].R*control(pdtraj.pr[1]) + game_obj.obj[1][1].cost[1].r)*dt, 1) < 1e-10
    # Terminal cost
    @test norm(game_obj.E[1][1].cost[end].q - (game_obj.obj[1][1].cost[end].Q*state(pdtraj.pr[end]) + game_obj.obj[1][1].cost[end].q), 1) < 1e-10
    @test norm(game_obj.E[1][1].cost[end].r, 1) < 1e-10

    # Hessian
    # Stage cost
    @test norm(game_obj.E[1][1].cost[1].Q - game_obj.obj[1][1].cost[1].Q*dt, 1) < 1e-10
    @test norm(game_obj.E[1][1].cost[1].R - game_obj.obj[1][1].cost[1].R*dt, 1) < 1e-10
    # Terminal cost
    @test norm(game_obj.E[1][1].cost[end].Q - game_obj.obj[1][1].cost[end].Q, 1) < 1e-10
    @test norm(game_obj.E[1][1].cost[end].R, 1) < 1e-10

    # Test add_collision_cost!
    T = Float64
    dt = 0.1
    N = 10
    p = 3
    model = UnicycleGame(p=p)
    probsize = ProblemSize(N,model)
    n = model.n
    m = model.m
    Q = [Diagonal(rand(SVector{model.ni[i],T})) for i=1:p]
    R = [Diagonal(rand(SVector{model.mi[i],T})) for i=1:p]
    xf = [i*ones(SVector{model.ni[i],T}) for i=1:p]
    uf = [2i*ones(SVector{model.mi[i],T}) for i=1:p]

    game_obj = GameObjective(Q,R,xf,uf,N,model)

    radius = [1.0, 2.0, 3.0]
    μ = [10.0, 20.0, 30.0]
    i = 1
    j = 2
    c = CollisionCost{n,m,T,length(model.px[i])}(μ[i], radius[i], model.px[i], model.px[j])
    add_collision_cost!(game_obj, radius, μ)

    length(game_obj.obj[1]) == 1 + p-1
    length(game_obj.obj[2]) == 1 + p-1
    length(game_obj.obj[3]) == 1 + p-1
    add_collision_cost!(game_obj, radius, μ)
    length(game_obj.obj[1]) == 1 + 2*(p-1)
    length(game_obj.obj[2]) == 1 + 2*(p-1)
    length(game_obj.obj[3]) == 1 + 2*(p-1)

    typeof(game_obj.obj[1][2].cost) <: Vector{<:CollisionCost}
    game_obj.obj[1][2].cost[1].μ == 10.0
    game_obj.obj[2][2].cost[1].μ == 20.0
    game_obj.obj[3][2].cost[1].μ == 30.0

    game_obj.obj[1][2].cost[1].r == 1.0
    game_obj.obj[2][2].cost[1].r == 2.0
    game_obj.obj[3][2].cost[1].r == 3.0

    # Test CollisionCost
    T = Float64
    dt = 0.1
    model = DoubleIntegratorGame(p=2)
    μ = 10.0
    r = 0.2
    pxi = SVector{length(model.px[1]),Int}(model.px[1])
    pxj = SVector{length(model.px[2]),Int}(model.px[2])
    cs = CollisionCost{model.n,model.m,T,length(model.px[1])}(μ,r,pxi,pxj)

    cs0 = copy(cs)
    @test cs0.μ == cs.μ
    @test cs0.r == cs.r
    @test cs0.terminal == cs.terminal
    @test all(cs0.pxi .== cs.pxi)
    @test all(cs0.pxj .== cs.pxj)
    @test state_dim(cs) == model.n
    @test control_dim(cs) == model.m

    @test cs.μ == μ
    @test cs.r == r
    @test cs.pxi == pxi
    @test cs.pxj == pxj
    @test cs.terminal == false
    @test typeof(cs) <: CollisionCost{model.n,model.m,T,length(model.px[1])}
    @test typeof(cs) <: TrajectoryOptimization.CostFunction

    x = @SVector [1.0, 1.1, 2.0, 2.0, 0.0, 0.0, 0.0, 0.0]
    u = rand(SVector{model.m,T})
    z = KnotPoint(x,u,dt)
    @test typeof(z) <: TrajectoryOptimization.AbstractKnotPoint

    # Test Stage Cost
    @test abs(TrajectoryOptimization.stage_cost(cs, x) - 0.05) < 1e-10
    @test abs(TrajectoryOptimization.stage_cost(cs, state(z)) - 0.05) < 1e-10
    @test abs(TrajectoryOptimization.stage_cost(cs, x, u) - 0.05) < 1e-10
    @test abs(TrajectoryOptimization.stage_cost(cs, state(z), control(z)) - 0.05) < 1e-10
    @test abs(TrajectoryOptimization.stage_cost(cs, x) -  0.05) < 1e-10
    @test abs(TrajectoryOptimization.stage_cost(cs, x, u) - 0.05) < 1e-10
    @test abs(TrajectoryOptimization.stage_cost(cs, z) - 0.005) < 1e-10

    @test (@ballocated TrajectoryOptimization.stage_cost($cs, $x)) == 0
    @test (@ballocated TrajectoryOptimization.stage_cost($cs, $x, $u)) == 0
    @test (@ballocated TrajectoryOptimization.stage_cost($cs, $z)) == 0

    # Test Gradient
    function easy_gradient(cs, z)
        function local_eval(x)
            z_ = KnotPoint(x, control(z), z.dt)
            return TrajectoryOptimization.stage_cost(cs,state(z_))
        end
        return ForwardDiff.gradient(local_eval, state(z))
    end

    zr = KnotPoint(rand(SVector{model.n,T}), rand(SVector{model.m,T}), dt)
    cs_active = CollisionCost{model.n,model.m,T,length(model.px[1])}(μ,1e3,pxi,pxj)
    cs_notactive = CollisionCost{model.n,model.m,T,length(model.px[1])}(μ,1e-3,pxi,pxj)
    E = TrajectoryOptimization.QuadraticCost{T}(model.n,model.m)

    TrajectoryOptimization.gradient!(E,cs_active,state(zr))
    grad_x = easy_gradient(cs_active, zr)
    @test norm(E.q - grad_x, 1)/norm(E.q) < 1e-7
    @test E.r == zeros(model.m)

    TrajectoryOptimization.gradient!(E,cs_notactive,state(zr))
    grad_x = easy_gradient(cs_notactive, zr)
    @test norm(E.q - grad_x, 1) < 1e-7
    @test E.r == zeros(model.m)

    @test (@ballocated TrajectoryOptimization.gradient!($E, $cs, $x)) == 0
    @test (@ballocated TrajectoryOptimization.gradient!($E, $cs, $x, $u)) == 0

    # Test Hessian
    function easy_hessian(cs, z)
        function local_eval(x)
            z_ = KnotPoint(x, control(z), z.dt)
            return TrajectoryOptimization.stage_cost(cs,state(z_))
        end
        return ForwardDiff.hessian(local_eval, state(z))
    end

    TrajectoryOptimization.hessian!(E,cs_active,state(zr))
    hess_x = easy_hessian(cs_active, zr)
    E.Q
    hess_x
    @test norm(E.Q - hess_x, 1) < 1e-2
    @test E.r == zeros(model.m)

    TrajectoryOptimization.hessian!(E,cs_notactive,state(zr))
    hess_x = easy_hessian(cs_notactive, zr)
    E.Q
    @test norm(E.Q - hess_x, 1) < 1e-2
    @test E.r == zeros(model.m)

    @test (@ballocated TrajectoryOptimization.hessian!($E, $cs, $x)) == 0
    @test (@ballocated TrajectoryOptimization.hessian!($E, $cs, $x, $u)) == 0

end


text:
module Titan

export @kernel, @gpu

using CUDArt
using MacroTools

# package code goes here
function _kernel(fn)
    ir = codegen(fn)
    ptx = compile_ir(ir)
    mod = create_module(ptx)
    @capture fn function name_(args__)
        body_
    end
    MODULES[name] = mod
    fn′ = wrap_kernel(fn)
    fn′
end

macro kernel(fn)
    _kernel(fn)
end

function codegen(fn)
    ir = readall("/Users/malmaud/Dropbox/Titan/src/kernel.ll")
    ir
end

function compile_ir(ir)
    d = mktempdir()
    kernel_path = joinpath(d, "kernel.ll")
    ptx_path = joinpath(d, "kernel.ptx")
    open(kernel_path, "w") do file
        write(file, ir)
    end
    cmd = `$(LLVM_PATH)/llc $kernel_path -o $ptx_path`
    run(cmd)
    local ptx
    open(ptx_path) do file
        ptx = readall(file)
    end
    ptx
end

function create_module(ptx)
    d = mktempdir()
    path = joinpath(d, "kernel.ptx")
    open(path, "w") do file
        write(file, ptx)
    end

    mod = CuModule(path, false)
end

function wrap_kernel(fn)
    @capture fn function name_(args__)
        body_
    end
    q_name = QuoteNode(name)
    quote
        function $(esc(name))($(args...); griddim=1, blockdim=1)
            mod = MODULES[$q_name]
            fn = CuFunction(mod, string($q_name))
            CUDArt.launch(fn, griddim, blockdim, $(Expr(:tuple, args...)))
        end
    end
end

macro gpu(expr)
    @capture expr name_<griddim_,blockdim_>(args__)
    quote
        $name(($args...); griddim=$griddim, blockdim=$blockdim)
    end
end

function __init__()
    const global LLVM_PATH = get(ENV, "LLVM_PATH", "/Users/malmaud/llvm/install/bin")
    const global MODULES = Dict()
    device(0)
    device_reset(0)
    CUDArt.free(CUDArt.malloc(UInt8, 1))
end

end # module


text:
RecipesBase.@recipe function plot(t::ColoredRootedTree)
  # Compute x and y coordinates recursively
  width = 2.0
  height = 1.0
  color_idx = unique(t.color_sequence)
  if eltype(t.color_sequence) == Bool
    color_val = [:white, :black]
  else
    color_val = Plots.distinguishable_colors(length(color_idx))
  end
  colormap = Dict{eltype(color_idx), eltype(color_val)}()
  for (idx, val) in zip(color_idx, color_val)
    colormap[idx] = val
  end
  x, y, colors = _plot_coordinates(t, 0.0, 0.0, width, height, colormap)

  # Series properties
  linecolor --> :black
  markercolor --> reshape(colors, 1, :)
  background_color --> "LightGray"
  markershape --> :circle
  markersize  --> 6
  linewidth   --> 2

  # Geometric properties
  grid --> false
  ticks --> false
  foreground_color_border --> :white
  legend --> :bottomright

  # We need to `flatten` nested arrays
  if !isempty(t)
    x_min, x_max = extrema(Iterators.flatten(x))
  else
    x_min = x_max = zero(eltype(eltype(x)))
  end
  if x_min ≈ x_max
    xlims --> (x_min - width/2, x_max + width/2)
  end

  if !isempty(t)
    y_min, y_max = extrema(Iterators.flatten(y))
  else
    y_min = y_max = zero(eltype(eltype(x)))
  end
  if y_min ≈ y_max
    ylims --> (y_min - height/2, y_max + height/2)
  end

  # Annotations
  labels = ones(String, 1, length(colors))
  if !isempty(t)
    labels[1] = "tree " * string((t.level_sequence, t.color_sequence))
  end
  label --> labels

  x, y
end

function _plot_coordinates(t::ColoredRootedTree,
                           x_root::T, y_root::T,
                           width::T,  height::T,
                           colormap) where {T}
  # Initialize vectors of return values
  x = Vector{Vector{T}}()
  y = Vector{Vector{T}}()
  colors = Vector{Vector{valtype(colormap)}}()

  if isempty(t)
    return x, y, colors
  end

  push!(x, [x_root])
  push!(y, [y_root])
  color_root = colormap[first(t.color_sequence)]
  push!(colors, [color_root])

  # We cannot indicate a new line series by `NaN` since that doesn't work with
  # colors
  # ┌ Warning: Indices Base.OneTo(9) of attribute `seriescolor` does not match data indices 3:9.
  # └ @ Plots ~/.julia/packages/Plots/AJMX6/src/utils.jl:132
  # ┌ Info: Data contains NaNs or missing values, and indices of `seriescolor` vector do not match data indices.
  # │ If you intend elements of `seriescolor` to apply to individual NaN-separated segements in the data,
  # │ pass each segment in a separate vector instead, and use a row vector for `seriescolor`. Legend entries
  # │ may be suppressed by passing an empty label.
  # │ For example,
  # └     plot([1:2,1:3], [[4,5],[3,4,5]], label=["y" ""], seriescolor=[1 2])
  # ┌ Warning: Indices Base.OneTo(9) of attribute `linecolor` does not match data indices 3:9.
  # └ @ Plots ~/.julia/packages/Plots/AJMX6/src/utils.jl:132
  # ┌ Info: Data contains NaNs or missing values, and indices of `linecolor` vector do not match data indices.
  # │ If you intend elements of `linecolor` to apply to individual NaN-separated segements in the data,
  # │ pass each segment in a separate vector instead, and use a row vector for `linecolor`. Legend entries
  # │ may be suppressed by passing an empty label.
  # │ For example,
  # └     plot([1:2,1:3], [[4,5],[3,4,5]], label=["y" ""], linecolor=[1 2])
  # ┌ Warning: Indices Base.OneTo(9) of attribute `fillcolor` does not match data indices 3:9.
  # └ @ Plots ~/.julia/packages/Plots/AJMX6/src/utils.jl:132
  # ┌ Info: Data contains NaNs or missing values, and indices of `fillcolor` vector do not match data indices.
  # │ If you intend elements of `fillcolor` to apply to individual NaN-separated segements in the data,
  # │ pass each segment in a separate vector instead, and use a row vector for `fillcolor`. Legend entries
  # │ may be suppressed by passing an empty label.
  # │ For example,
  # └     plot([1:2,1:3], [[4,5],[3,4,5]], label=["y" ""], fillcolor=[1 2])
  # ┌ Warning: Indices Base.OneTo(9) of attribute `markercolor` does not match data indices 3:9.
  # └ @ Plots ~/.julia/packages/Plots/AJMX6/src/utils.jl:132
  # ┌ Info: Data contains NaNs or missing values, and indices of `markercolor` vector do not match data indices.
  # │ If you intend elements of `markercolor` to apply to individual NaN-separated segements in the data,
  # │ pass each segment in a separate vector instead, and use a row vector for `markercolor`. Legend entries
  # │ may be suppressed by passing an empty label.
  # │ For example,
  # └     plot([1:2,1:3], [[4,5],[3,4,5]], label=["y" ""], markercolor=[1 2])

  # Compute plot coordinates recursively
  subtr = subtrees(t)

  # Distribute children uniformly in x and at equal y coordinates
  y_child = y_root + height
  num_children = length(subtr)
  distance = width * (num_children - 1) / 2

  x_children = range(x_root - distance, x_root + distance, length=num_children)
  for idx in eachindex(subtr)
    x_child = x_children[idx]
    push!(x, [x_root, x_child])
    push!(y, [y_root, y_child])
    push!(colors, [color_root, colormap[first(subtr[idx].color_sequence)]])
    x_recursive, y_recursive, colors_recursive = _plot_coordinates(subtr[idx],
      x_child, y_child, width / 3, height, colormap)
    append!(x, x_recursive)
    append!(y, y_recursive)
    append!(colors, colors_recursive)
  end

  return x, y, colors
end


text:
export ft_getlibraryversion

function ft_getlibraryversion()
  dllversion = Ref{Cuint}()
  ft_status = ccall((:FT_GetLibraryVersion, d2xx),Cuint,(Ref{Cuint},),dllversion)
  checkstatus(ft_status)
  build = (dllversion[] & 0x000000ff)
  minor = (dllversion[] & 0x0000ff00)>>8
  major = (dllversion[] & 0x00ff0000)>>16
  return VersionNumber(major,minor,build)
end


text:
# Test of multinetwork dimensions and multiple planning years

## Import packages and load common code
import PowerModels; const _PM = PowerModels
import PowerModelsACDC; const _PMACDC = PowerModelsACDC
import FlexPlan; const _FP = FlexPlan
include("../io/create_profile.jl")
include("../io/multiple_years.jl")
# Solvers are imported later, select name here
solver = "Cbc"

## Import and set solver
if solver == "Cbc"
    import Cbc
    optimizer = _FP.optimizer_with_attributes(Cbc.Optimizer,
        "logLevel" => 0, # ∈ {0,1}, default: 0
    ) # Solver options: <https://github.com/jump-dev/Cbc.jl#using-with-jump>
elseif solver == "Gurobi"
    import Gurobi
    optimizer = _FP.optimizer_with_attributes(Gurobi.Optimizer) # Solver options: <https://www.gurobi.com/documentation/current/refman/parameters.html>
end

## Input parameters

test_case = "case6" # Available test cases: "case6", "case67"
number_of_hours = 4 # Number of hourly optimization periods
number_of_scenarios = 2 # Number of scenarios (different generation/load profiles)
number_of_years = 3 # Number of years (different investments)
cost_scale_factor = 1.0 # Scale factor for all costs
model_type = _PM.DCPPowerModel

## Generate test case

data = create_multi_year_network_data(test_case, number_of_hours, number_of_scenarios, number_of_years; cost_scale_factor)

## Solve problem

s = Dict("output" => Dict("branch_flows" => true), "conv_losses_mp" => false, "add_co2_cost" => false)

# One-step alternative, does not expose pm
#result = _FP.stoch_flex_tnep(data, model_type, optimizer, multinetwork=_PM._IM.ismultinetwork(data); setting=s)

# Two-step alternative, exposes pm
pm = _PM.instantiate_model(data, model_type, _FP.post_stoch_flex_tnep; ref_extensions=[_FP.ref_add_gen!, _PMACDC.add_ref_dcgrid!, _PMACDC.add_candidate_dcgrid!, _FP.ref_add_storage!, _FP.ref_add_ne_storage!, _FP.ref_add_flex_load!, _PM.ref_add_on_off_va_bounds!, _PM.ref_add_ne_branch!], setting=s)
result = _PM.optimize_model!(pm; optimizer)

@assert result["termination_status"] ∈ (_PM.OPTIMAL, _PM.LOCALLY_SOLVED) "$(result["optimizer"]) termination status: $(result["termination_status"])"

println("Test completed")


text:
mutable struct ThreadLocalStorage{T,Factory}
    @const factory::Factory
    @atomic storages::Vector{T}
    @const cond::Threads.Condition
end

ThreadLocalStorage{T}(factory) where {T} =
    ThreadLocalStorage{T,_typeof(factory)}(factory, T[], Threads.Condition())

function ThreadLocalStorage(factory)
    value = factory()
    return ThreadLocalStorage{typeof(value),_typeof(factory)}(
        factory,
        [value],
        Threads.Condition(),
    )
end

function ConcurrentUtils.unsafe_takestorages!(tls::ThreadLocalStorage{T}) where {T}
    storages = @atomic :monotonic tls.storages
    @atomic :monotonic tls.storages = T[]
    return storages
end

function Base.empty!(tls::ThreadLocalStorage{T}) where {T}
    @atomic :release tls.storages = T[]
    return tls
end

function getstorages!(tls::ThreadLocalStorage{T}) where {T}
    storages = @atomic :acquire tls.storages
    if length(storages) >= Threads.nthreads()
        return storages
    end
    lock(tls.cond) do
        local storages = @atomic :acquire tls.storages
        n = length(storages)
        if n >= Threads.nthreads()
            return storages
        end
        newstorages = similar(storages, Threads.nthreads())
        copyto!(newstorages, storages)
        local factory = tls.factory
        for i in n+1:Threads.nthreads()
            local value = factory()::T
            newstorages[i] = value
        end
        @atomic :release tls.storages = newstorages
        return newstorages
    end
end

Base.getindex(tls::ThreadLocalStorage) = getstorages!(tls)[Threads.threadid()]

const THREAD_LOCAL_RNG = ThreadLocalStorage(Xoshiro)


text:
# Given a corpus, return a term-document matrix of tf-idf values.
function tfidf(c :: Corpus)
  # D = The number of documents.
  D = documents(c)

  # tf(w, d) = The number of occurrences of term w in document d.
  tf = c.c

  # df(w) = number of documents in which term w occurs
  df = sum(tf .≠ 0, 2)[:]

  # idf(w) = inverse document frequency of term w
  idf = log(D) .- log(df)

  # tfidf(w, d) = tf(w, d) * idf(w)
  tfidf = tf .* idf

  return tfidf
end


text:
function normalize_select(idx, @nospecialize(cols...))
    selected_cols = [Int[], Int[]]
    for i in 1:length(cols)
        normalize_select!(selected_cols, idx, cols[i])
    end
    unique!(selected_cols)
end
function normalize_select!(selected_cols, idx, cols::ColumnIndex)
    push!(selected_cols[1], idx[cols])
end

function normalize_select!(selected_cols, idx, cols::MultiColumnIndex)
    colsidx = idx[cols]
    for i in 1:length(colsidx)
        push!(selected_cols[1], colsidx[i])
    end
    if cols isa Not
        colsidxnot = idx[cols.skip]
        for i in 1:length(colsidxnot)
            push!(selected_cols[2], colsidxnot[i])
        end
    end
end

# sell and sell! will replace select and select!
# Dataset shouldn't support copycols since it causes modifying a data set without telling other alias data sets
"""
    select(ds, args...)

Select columns of `ds` based on `args...`. `args` can be any column selector: column index, column Name, `:`, `Between`, `Not`, Vector of column indices or column names, a regular expression.

It makes a copy of `ds`. See [`select!`](@ref) if an in-place modification is desired.

# Examples
```jldoctest
julia> ds = Dataset(x = [1,2,3,4,5], y = [1.5,2.3,-1,0,2.0], z = Bool[1,0,1,0,1])
5×3 Dataset
 Row │ x         y         z
     │ identity  identity  identity
     │ Int64?    Float64?  Bool?
─────┼──────────────────────────────
   1 │        1       1.5      true
   2 │        2       2.3     false
   3 │        3      -1.0      true
   4 │        4       0.0     false
   5 │        5       2.0      true

julia> select(ds, 2, :)
5×3 Dataset
 Row │ y         x         z
     │ identity  identity  identity
     │ Float64?  Int64?    Bool?
─────┼──────────────────────────────
   1 │      1.5         1      true
   2 │      2.3         2     false
   3 │     -1.0         3      true
   4 │      0.0         4     false
   5 │      2.0         5      true

julia> select(ds, Not(:y))
5×2 Dataset
 Row │ x         z
     │ identity  identity
     │ Int64?    Bool?
─────┼────────────────────
   1 │        1      true
   2 │        2     false
   3 │        3      true
   4 │        4     false
   5 │        5      true
```
"""
function select(ds::Dataset, @nospecialize(args...))
    selected_cols_all = normalize_select(index(ds), args...)
    selected_cols = setdiff!(selected_cols_all[1], selected_cols_all[2])
    res = AbstractVector[]

    for j in 1:length(selected_cols)
        push!(res, copy(_columns(ds)[selected_cols[j]]))
    end
    newds = Dataset(res, _names(ds)[selected_cols], copycols = false)
    # else
    #     for j in 1:length(selected_cols)
    #         push!(res, _columns(ds)[selected_cols[j]])
    #     end
    #     newds = Dataset(res, _names(ds)[selected_cols], copycols = false)
    # end
    for j in 1:length(selected_cols)
        setformat!(newds, j => getformat(ds, selected_cols[j]))
    end
    if all(index(ds).sortedcols .∈ Ref(selected_cols))
        scols = Int[]
        revs = copy(index(ds).rev)
        scols_ds = index(ds).sortedcols
        for k in 1:length(scols_ds)
            push!(scols, findfirst(isequal(scols_ds[k]), selected_cols))
        end
        _copy_grouping_info!(newds, ds)
        # need to reorder things
        empty!(index(newds).sortedcols)
        empty!(index(newds).rev)
        append!(index(newds).sortedcols, scols)
        append!(index(newds).rev, revs)
    # else
    #     _reset_grouping_info!(ds)
    end
    setinfo!(newds, _attributes(ds).meta.info[])
    return newds
end
"""
    select!(ds, args...)

Select columns of `ds` based on `args...`. `args` can be any column selector: column index, column Name, `:`, `Between`, `Not`, Vector of column indices or column names, a regular expression.

It modifies the data set in-place. See [`select`](@ref) if a copy of selected columns is desired.
"""
function select!(ds, @nospecialize(args...))
    selected_cols_all = normalize_select(index(ds), args...)
    selected_cols = setdiff!(selected_cols_all[1], selected_cols_all[2])
    unwanted_cols = setdiff(1:ncol(ds), selected_cols)
    sort!(unwanted_cols, rev = true)

    newnames = _names(ds)[selected_cols]
    newlookup = newnames .=> 1:length(selected_cols)
    res = AbstractVector[]

    for j in 1:length(selected_cols)
        push!(res, _columns(ds)[selected_cols[j]])
    end
    newformat = Dict{Int, Function}()
    for j in 1:length(selected_cols)
        push!(newformat, j => getformat(ds, selected_cols[j]))
    end
    if all(index(ds).sortedcols .∈ Ref(selected_cols))
        scols = Int[]
        # revs = copy(index(ds).rev)
        scols_ds = index(ds).sortedcols
        for k in 1:length(scols_ds)
            push!(scols, findfirst(isequal(scols_ds[k]), selected_cols))
        end
        # _copy_grouping_info!(ds, ds)
        # need to reorder things
        empty!(index(ds).sortedcols)
        # empty!(index(ds).rev)
        append!(index(ds).sortedcols, scols)
        # append!(index(ds).rev, revs)
    # else
    #     _reset_grouping_info!(ds)
    end
    for i in unwanted_cols
      popat!(_columns(ds), i)
      delete!(index(ds), i)
    end
    for j in 1:length(selected_cols)
      _columns(ds)[j] = res[j]
    end
    empty!(index(ds).lookup)
    empty!(index(ds).names)
    empty!(index(ds).format)
    for j in newlookup
      push!(index(ds).lookup, j)
    end
    for j in newnames
      push!(index(ds).names, j)
    end
    for j in newformat
      push!(index(ds).format, j)
    end
    _modified(_attributes(ds))
    ds
end


text:
using VarianceComponentsHDFE
using Random, BenchmarkTools, Test
using CSV
using ArtifactUtils, Pkg.Artifacts
# Make sure LDLFactorizations is version 0.5.0 and that the `multiple-rhs` branch is checked out

pkg_dir = pkgdir(VarianceComponentsHDFE)
artifacts_path = joinpath(pkgdir(VarianceComponentsHDFE), "benchmark\\Artifacts.toml")

#Set Random seed
Random.seed!(1234)

# 1) Medium sized network benchmarks
data = CSV.read(joinpath(pkg_dir, "test.csv");header=false)

y = [data[:,4]...]
id = [data[:,1]...]
firmid = [data[:,2]...]
obs_id = collect(1:length(y))

settings_JLA = VCHDFESettings(leverage_algorithm = JLAAlgorithm(num_simulations=3155), first_id_effects=true, cov_effects=true)
@btime compute_whole(y,id,firmid,nothing,settings_JLA) #55.669s Julia vs 152.97s Matlab (32GB Ram and 4 Threads)

# 2) Large sized network benchmarks
data_hash = artifact_hash("large_nocontrols_pruned",artifacts_path)
data = CSV.read(joinpath(artifact_path(data_hash), "large_nocontrols_pruned.csv"))

y = [data[:,4]...]
id = [data[:,1]...]
firmid = [data[:,2]...]
obs_id = collect(1:length(y))

settings_JLA = VCHDFESettings(leverage_algorithm = JLAAlgorithm(num_simulations=100), first_id_effects=true, cov_effects=true)

@btime compute_whole($y,$id,$firmid,nothing,settings_JLA) #978.056s Julia vs 606.07s Matlab (32GB Ram and 4 Threads, 100 simulations)

settings_JLA = VCHDFESettings(leverage_algorithm = JLAAlgorithm(num_simulations=200), first_id_effects=true, cov_effects=true)

@btime compute_whole($y,$id,$firmid,nothing,settings_JLA) #1702.217 ss Julia vs 1115.04s Matlab (32GB Ram and 4 Threads, 200 simulations)

settings_JLA = VCHDFESettings(leverage_algorithm = JLAAlgorithm(num_simulations=500), first_id_effects=true, cov_effects=true)

@btime compute_whole(y,id,firmid,nothing,settings_JLA) # 4280.219173 s in Julia vs 3328.10s in Matlab (32GB Ram and 4 Threads, 500 simulations)

#Simulations on cluster for 250 rademachers: 949.138618 seconds Julia vs 379.719083  (256GB and 32cores )
#Simulations on cluster for 500 rademachers: 1463.095847 seconds Julia vs 628.135439 seconds Matlab (256GB and 32cores)
#Simulations on cluster for heuristic - 1441 rademachers : 3392.993448 seconds Julia vs 1651.819884 seconds Matlab  (256GB and 32cores)

# 3) Huge sized network benchmarks
data_hash = artifact_hash("huge_nocontrols_pruned",artifacts_path)
data = CSV.read(joinpath(artifact_path(data_hash), "huge_pruned_main.csv"))

y = [data[:,4]...]
id = [data[:,1]...]
firmid = [data[:,2]...]
obs_id = collect(1:length(y))

settings_JLA = VCHDFESettings(leverage_algorithm = JLAAlgorithm(num_simulations=100), first_id_effects=true, cov_effects=true)


text:
using Serialization
import Serialization: serialize, deserialize
import Serialization: serialize_type

const _pickle = PyNULL()

pickle() = ispynull(_pickle) ? copy!(_pickle, pyimport(PyCall.pyversion.major ≥ 3 ? "pickle" : "cPickle")) : _pickle

function serialize(s::AbstractSerializer, pyo::PyObject)
    serialize_type(s, PyObject)
    if ispynull(pyo)
        serialize(s, PyPtr_NULL)
    else
        b = PyBuffer(pycall(pickle()."dumps", PyObject, pyo))
        serialize(s, unsafe_wrap(Array, Ptr{UInt8}(pointer(b)), sizeof(b)))
    end
end

"""
    pybytes(b::Union{String,DenseVector{UInt8}})

Convert `b` to a Python `bytes` object.   This differs from the default
`PyObject(b)` conversion of `String` to a Python string (which may fail if `b`
does not contain valid Unicode), or from the default conversion of a
`Vector{UInt8}` to a `bytearray` object (which is mutable, unlike `bytes`).
"""
function pybytes(b::Union{String,DenseVector{UInt8}})
    b isa String || stride(b,1) == 1 || throw(ArgumentError("pybytes requires stride-1 byte arrays"))
    PyObject(@pycheckn ccall(@pysym(PyString_FromStringAndSize),
                             PyPtr, (Ptr{UInt8}, Int),
                             b, sizeof(b)))
end

function deserialize(s::AbstractSerializer, t::Type{PyObject})
    b = deserialize(s)
    if isa(b, PyPtr)
        @assert b == C_NULL
        return PyNULL()
    else
        return pycall(pickle()."loads", PyObject, pybytes(b))
    end
end


text:
using RR
using RR: current_session
using ASTInterpreter

"""
Attempt to reverse step until the entry to the current function.
"""
function ASTInterpreter.execute_command(state, stack::Union{Gallium.NativeStack,Gallium.CStackFrame}, ::Val{:rf}, command)
    timeline = state.top_interp.session
    RR.silence!(timeline)

    # Algorithm:
    # 1. Determine the location of the function entry point.
    # 2. Determine how to compute the CFA, both here and at the function entry
    #    point.
    # 3. Continue unless the CFA matches

    # Determine module
    stack = isa(stack, Gallium.NativeStack) ? stack.stack[end] : stack
    mod, base, ip = Gallium.modbaseip_for_stack(state, stack)
    modrel = UInt(ip - base)
    loc, fde = Gallium.Unwinder.find_fde(mod, modrel)

    # Compute CFI
    ciecache = nothing
    isa(mod, Module) && (ciecache = mod.ciecache)
    cie::DWARF.CallFrameInfo.CIE, ccoff = realize_cieoff(fde, ciecache)
    target_delta = modrel - loc

    entry_rs = CallFrameInfo.evaluate_program(fde, UInt(0), cie = cie, ciecache = ciecache, ccoff=ccoff)
    here_rs =  CallFrameInfo.evaluate_program(fde, target_delta, cie = cie, ciecache = ciecache, ccoff=ccoff)

    # Compute the CFA here
    regs = icxx"$(current_task(current_session(timeline)))->regs();"
    here_cfa = Gallium.Unwinder.compute_cfa_addr(current_task(current_session(timeline)), regs, here_rs)
    here_rsp = Gallium.get_dwarf(regs, :rsp)

    # Set a breakpoint at function entry
    bp = Gallium.breakpoint(timeline, base + loc)

    # Reverse continue until the breakpoint is hit at a matching CFA, or until
    # we're at a breakpoint higher up the stack (which would imply that we missed it)
    target_tuid = RR.tuid(current_task(timeline))
    while true
        RR.reverse_continue!(timeline)
        if !(RR.tuid(current_task(timeline)) == target_tuid)
            continue
        end
        # TODO: Check that we're at the right breakpoint
        new_regs = icxx"$(current_task(current_session(timeline)))->regs();"
        new_cfa = Gallium.Unwinder.compute_cfa_addr(current_task(current_session(timeline)), new_regs, entry_rs)
        if here_cfa == new_cfa
            break
        elseif Gallium.get_dwarf(new_regs, :rsp) > here_rsp
            println("WARNING: May have missed function call.")
            break
        end
    end
    Gallium.disable(bp)

    # Step once more to get out of the function
    RR.reverse_single_step!(current_session(timeline),current_task(current_session(timeline)),timeline)

    update_stack!(state)
    return true
end

"""
Reverse continue until a breakpoint is hit
"""
function ASTInterpreter.execute_command(state, stack::Union{Gallium.NativeStack,Gallium.CStackFrame}, ::Val{:rc}, command)
    timeline = state.top_interp.session
    RR.silence!(timeline)
    target_tgid = icxx"$(current_task(timeline))->tgid();"
    while true
        RR.reverse_continue!(timeline)
        if icxx"$(current_task(timeline))->tgid();" != target_tgid
            continue
        end
        break
    end
    update_stack!(state)
    return true
end

function task_step_until_bkpt!(timeline::Union{RR.ReplayTimeline, RR.ReplaySession, RR.ReplayTask})
    target_tgid = icxx"$(current_task(timeline))->tgid();"
    while true
        Gallium.step_until_bkpt!(timeline)
        if icxx"$(current_task(timeline))->tgid();" != target_tgid
            continue
        end
        break
    end
end

function ASTInterpreter.execute_command(state, stack, ::Val{:when}, command)
    timeline = state.top_interp.session
    print(STDOUT, "Ticks: ")
    show(STDOUT, UInt64(icxx"$(current_task(current_session(timeline)))->tick_count();"))
    println(STDOUT);
    print(STDOUT, "Time: ", global_time(timeline))
    println(STDOUT)
    return false
end

using Colors
#=
using TerminalExtensions; using Gadfly; using Colors
const repl_theme = Gadfly.Theme(
    panel_fill=colorant"black", default_color=colorant"orange", major_label_color=colorant"white",
    minor_label_color=colorant"white", key_label_color=colorant"white", key_title_color=colorant"white",
    line_width=1mm
   )
eval(Gadfly,quote
   function writemime(io::IO, ::MIME"image/png", p::Plot)
       draw(PNG(io, Compose.default_graphic_width,
                Compose.default_graphic_height), p)
   end
end)
=#

function collect_mark_ticks()
    map(unsafe_load,icxx"""
    std::vector<uint64_t> ticks;
    for(auto it : $timeline->reverse_exec_checkpoints) {
        ticks.push_back(it.first.ptr->key.ticks);
    }
    ticks;
    """)
end

function collect_mark_ticks(marks)
    ticks = UInt[]
    for mark in marks
        push!(ticks, icxx"$(mark.mark).ptr->proto.key.ticks;")
    end
    ticks
end

immutable TimelineEvent
    t::Int
    label::Symbol
end

start_time(timeline) = TimelineEvent(0, :start)
end_time(timeline) = TimelineEvent(RR.count_total_ticks(timeline), :end)
const timeline_default_colors = [colorant"orange",colorant"green",colorant"purple",colorant"blue"]
function timeline_layers(events; timeline_id = 0, pointfilterset=[])
    y = [timeline_id for _ in 1:length(events)]
    filterids = find(x->!(x.label in pointfilterset),events)
    [layer(x=map(x->x.t, events[filterids]),
        y=y[filterids],
        color = map(x->x.label, events[filterids]), Geom.point),
        layer(x=map(x->x.t, events),
            y=y, Geom.line)]
end

function plot_timeline(events, colors = timeline_default_colors; timeline_id = 0)
    display(
    plot(timeline_layers(events, timeline_id = timeline_id)...,
        repl_theme,Guide.xlabel("Time"),
        Guide.ylabel("Timeline"),
        Scale.color_discrete_manual(colors...))
    )
end

function ASTInterpreter.execute_command(state, stack, ::Val{:timeline}, command)
    me = TimelineEvent(UInt64(icxx"$(current_task(current_session(timeline)))->tick_count();"),
        :me)
    ticks = []
    contains(command, "internal") && (ticks = collect_mark_ticks())
    ticks = [TimelineEvent(t, :internal_tick) for t in ticks]
    explicit_ticks = [TimelineEvent(t, :explicit_tick) for t in collect_mark_ticks(mark_stack)]
    colors = [colorant"orange",colorant"green",colorant"purple"]
    (length(ticks) != 0) && unshift!(colors, colorant"blue")
    display(plot_timeline([ticks; start_time(timeline);
        me; end_time(timeline); explicit_ticks], colors))
    println("The timeline is intact.")
    return false
end

immutable AnnotatedMark
    mark
    annotation::String
end

const mark_stack = AnnotatedMark[]
function ASTInterpreter.execute_command(state, stack, ::Val{:mark}, command)
    timeline = state.top_interp.session
    push!(mark_stack,AnnotatedMark(icxx"$timeline->mark();",command[5:end]))
    return false
end

using JLD
"""
Manipulate marks.
"""
function ASTInterpreter.execute_command(state, stack, ::Val{:marks}, command)
    subcmds = split(command," ")[2:end]
    if isempty(subcmds) || subcmds[1] == "list"
        for (i,mark) in enumerate(mark_stack)
            println("[$i] Mark (",mark.annotation,")")
        end
    elseif subcmds[1] == "save"
        annotations = map(x->x.annotation,mark_stack)
        marks = map(x->reinterpret(UInt8,[icxx"$(x.mark).ptr->proto;".data]),mark_stack)
        @save "marks.jld" annotations marks
    elseif subcmds[1] == "load"
        timeline = state.top_interp.session
        @load "marks.jld" annotations marks
        pms = map(x->cxxt"rr::ReplayTimeline::ProtoMark"{312}(reinterpret(NTuple{312,UInt8},x)[]),marks)
        println("Recreating marks. One moment please...")
        for (annotation, pm) in zip(annotations, pms)
            RR.seek(timeline, pm)
            push!(mark_stack,AnnotatedMark(icxx"$timeline->mark();",annotation))
        end
        update_stack!(state)
        return true
    else
        print_with_color(:red, "Unknown subcommand\n")
    end
    return false
end

using ProgressMeter
import RR: when
when(timeline::RR.ReplayTimeline) = when(current_session(timeline))
global_time(timeline) = icxx"$(current_task(timeline))->trace_time();"
function ASTInterpreter.execute_command(state, stack, ::Val{:timejump}, command)
    timeline = state.top_interp.session
    subcmd = split(command)[2:end]
    local target_is_event = false
    if startswith(subcmd[1],"@")
        n = parse(Int, subcmd[1][2:end])
        icxx"$timeline->seek_to_mark($(mark_stack[n].mark));"
        icxx"$timeline->apply_breakpoints_and_watchpoints();"
        #global stack_remap = compute_remap()
        println("We have arrived.")
        update_stack!(state)
        return true
    elseif startswith(subcmd[1], "e")
        target = parse(Int, subcmd[1][2:end])
        me = global_time(timeline)
        n = target - me
        target_is_event = true
    else
        me = when(timeline)
        target = me + n
        n = parse(Int, subcmd[1])
    end
    p = Progress(n, 1, "Time travel in progress (forwards)...", 50)
    function check_for_breakpoint(res)
        if icxx"$res.break_status.breakpoint_hit == true;"
            regs = icxx"$(current_task(current_session(timeline)))->regs();"
            if RR.process_lowlevel_conditionals(Location(timeline, Gallium.ip(regs)), regs)
                println("Interrupted by breakpoint.")
                update_stack!(state)
                return true
            end
        end
        false
    end
    while (target_is_event ? global_time(timeline) : when(timeline)) < target
        # Step past any breakpoints
        if RR.at_breakpoint(timeline)
            RR.emulate_single_step!(timeline, current_vm()) || RR.single_step!(timeline)
        end
        res = RR.step!(current_session(timeline), target; target_is_event = target_is_event)
        if !target_is_event && icxx"$res.break_status.approaching_ticks_target;"
            break
        end
        check_for_breakpoint(res) && return true
        icxx"$timeline->maybe_add_reverse_exec_checkpoint(rr::ReplayTimeline::LOW_OVERHEAD);"
        now = (target_is_event ? global_time(timeline) : when(timeline))
        now != me && ProgressMeter.update!(p, Int64(now) - Int(me))
    end
    if !target_is_event
        while when() < target
            res = RR.single_step!(timeline)
            check_for_breakpoint(res) && return true
            now = when()
            now != me && ProgressMeter.update!(p, Int64(now) - Int(me))
        end
    end
    icxx"$timeline->apply_breakpoints_and_watchpoints();"
    println("We have arrived.")
    update_stack!(state)
    return true
end

function ASTInterpreter.execute_command(state, stack, ::Val{:maps}, command)
    timeline = state.top_interp.session
    parts = split(command, ' ')
    if length(parts) > 1 && parts[2] == "rr"
      icxx"$(current_task(timeline))->vm()->dump();"
    else
      pid = icxx"$(current_task(timeline))->real_tgid();"
      print(readstring("/proc/$pid/maps"))
    end
    return false
end

function ASTInterpreter.execute_command(state, stack, ::Val{:map}, command)
    timeline = state.top_interp.session
    parts = split(command, ' ')
    if length(parts) > 1
        ip = parse(Int, parts[2], 16)
    else
        ip = (isa(stack, Gallium.NativeStack) ? stack.stack[end] : stack).ip
    end
    pid = icxx"$(current_task(timeline))->real_tgid();"
    println(first(filter(map->UInt64(ip) ∈ map[1], Gallium.Ptrace.MapsIterator("/proc/$pid/maps")))[3])
    return false
end

function ASTInterpreter.execute_command(state, stack, ::Val{:gdb}, command)
    icxx"rr::GdbServer::emergency_debug($(current_task(state.top_interp.session)));"
    false
end

function ASTInterpreter.execute_command(state, stack, ::Val{:execjump}, command)
    target_pid = parse(Int, split(command,' ')[2])
    timeline = state.top_interp.session
    println("Time travel in progress (forwards)... Please wait.")
    while icxx"$(current_task(timeline))->tgid();" != target_pid ||
            !icxx"$(current_task(timeline))->execed();"
        # Step past any breakpoints
        if RR.at_breakpoint(timeline)
            RR.emulate_single_step!(timeline, current_vm()) || RR.single_step!(timeline)
        end
        res = RR.step!(current_session(timeline))
        icxx"$timeline->maybe_add_reverse_exec_checkpoint(rr::ReplayTimeline::LOW_OVERHEAD);"
    end
    println("We have arrived.")
    update_stack!(state)
    return true
end

function prepare_remote_execution(timeline::RR.ReplayTimeline)
    session = icxx"$(current_session(timeline))->clone_diversion();"
end

function prepare_remote_execution(session::RR.ReplaySession)
    session = icxx"$session->clone_diversion();"
end

function prepare_remote_execution(task::RR.ReplayTask)
    session = prepare_remote_execution(icxx"&$task->session();")
end

function compute_stack(modules, session::RR.ReplayTimeline; cficache=nothing, ip_only=false)
    task = current_task(current_session(session))
    did_fixup, regs = RR.fixup_RC(task, icxx"$task->regs();")
    stack, RCs = Gallium.stackwalk(regs, task, modules, ip_only=ip_only, rich_c = true, collectRCs = true, cficache=cficache)
    if length(stack) != 0
        stack[end].stacktop = !did_fixup
    end
    Gallium.NativeStack(stack,RCs,modules,session)
end

function ASTInterpreter.execute_command(state, stack, ::Val{:task}, command)
    timeline = state.top_interp.session
    modules = state.top_interp.modules
    subcommand = split(command, " ")[2:end]
    if subcommand[1] == "list"
        icxx"""
            for (auto &task : $(current_session(timeline))->tasks())
                $:(println(IOContext(STDOUT,:modules=>modules),
                    icxx"return task.second;"); nothing);
        """
        println(STDOUT)
    elseif subcommand[1] == "select"
        n = parse(Int, subcommand[2])
        (n < 1 || n > icxx"$(current_session(timeline))->tasks().size();") &&
            (print_with_color(:red, STDERR, "Not a valid task"); return false)
        it = icxx"$(current_session(timeline))->tasks().begin();"
        while n > 1
            icxx"++$it;";
            n -= 1
        end
        update_stack!(state, icxx"$it->second;")
        return true
    end
    return false
end

function Gallium.retrieve_obj_data(session::Union{RR.ReplayTimeline, RR.ReplaySession, RR.ReplayTask}, modules, ip)
    run_function(session, modules, :jl_get_dobj_data, ip) do task
        regs = icxx"$task->regs();"
        @assert UInt(Gallium.ip(regs)) == 0
        array_ptr = Gallium.get_dwarf(regs, :rax)
        data_ptr = Gallium.load(task, RemotePtr{RemotePtr{UInt8,UInt64}}(array_ptr))
        data_size = Gallium.load(task, RemotePtr{Csize_t}(array_ptr+8))
        Gallium.load(task, data_ptr, data_size)
    end
end

function Gallium.retrieve_section_start(session::Union{RR.ReplayTimeline, RR.ReplaySession, RR.ReplayTask}, modules, ip)
    isempty(Gallium.lookup_syms(session, modules, :jl_get_section_start)) &&
        return RemotePtr{Void}(0)
    run_function(session, modules, :jl_get_section_start, ip) do task
        regs = icxx"$task->regs();"
        (UInt(Gallium.ip(regs)) == 0) || return RemotePtr{Void}(0)
        addr = Gallium.get_dwarf(regs, :rax)
        RemotePtr{Void}(addr)
    end
end

Gallium.retrieve_obj_data(task::RR.ReplayTask, ip) =
    Gallium.retrieve_obj_data(icxx"&$task->session();", ip)

Gallium.retrieve_section_start(task::RR.ReplayTask, ip) =
    Gallium.retrieve_section_start(icxx"&$task->session();", ip)

## `rr ps` reimplementation

function find_exit_code(pid, events, current_event, tid_to_pid)
    for (i,e) in enumerate(events[current_event:end])
        # Get the view before this event
        current_tid_to_pid = versioned_view(tid_to_pid, current_event+i-2)
        if icxx"$e.type() == rr::TraceTaskEvent::EXIT;" &&
                current_tid_to_pid[icxx"$e.tid();"] == pid &&
                count(tid->tid==pid, values(current_tid_to_pid)) == 1
            status = icxx"$e.exit_status();"
            icxx"$status.type() == rr::WaitStatus::EXIT;" &&
                return icxx"$status.exit_code();"
            assert(icxx"$status.type() == rr::WaitStatus::FATAL_SIGNAL;")
            return icxx"-$status.fatal_sig();";
        end
    end
    return icxx"-SIGKILL;"
end

# TODO: Replace by proper implementation from DataStructures
type VersionedDict{Key,Value,Version}
    versions::Vector{Tuple{Version,Dict{Key,Value}}}
    current_version::Version
    VersionedDict() = new(Vector{Tuple{Version,Dict{Key,Value}}}(),0)
end
function advance!(dict::VersionedDict, new_version)
    if new_version < dict.current_version
        error("Versions may only be advanced forward")
    end
    dict.current_version = new_version
end
function Base.setindex!{K,V,VER}(dict::VersionedDict{K,V,VER}, value, key)
    if isempty(dict.versions) || dict.versions[end][1] != dict.current_version
        new_dict = isempty(dict.versions) ? Dict{K,V}() : copy(dict.versions[end][2])
        push!(dict.versions, (dict.current_version, new_dict))
    end
    setindex!(dict.versions[end][2], value, key)
    dict
end
Base.getindex(dict::VersionedDict, key) = dict.versions[end][2][key]
function Base.delete!(dict::VersionedDict, key)
    if isempty(dict.versions) || dict.versions[end][1] != dict.current_version
        new_dict = isempty(dict.versions) ? Dict{K,V}() : copy(dict.versions[end][2])
        push!(dict.versions, (dict.current_version, new_dict))
    end
    delete!(dict.versions[end][2], key)
end

function versioned_view(dict::VersionedDict, key)
    dict.versions[findlast(x->x[1]<=key, dict.versions)][2]
end

function update_tid_to_pid_map!(tid_to_pid, e)
    if icxx"$e.type() == rr::TraceTaskEvent::CLONE;"
        if icxx"($e.clone_flags() & CLONE_THREAD) != 0;"
            # thread clone. Record thread's pid.
            tid_to_pid[icxx"$e.tid();"] = tid_to_pid[icxx"$e.parent_tid();"];
        else
            # Some kind of fork. This task is its own pid.
            tid_to_pid[icxx"$e.tid();"] = icxx"$e.tid();"
        end
    elseif icxx"$e.type() == rr::TraceTaskEvent::EXIT;"
        delete!(tid_to_pid, icxx"$e.tid();")
    end
end

function find_cmd_line(pid, events, current_event, tid_to_pid)
    for (i,e) in enumerate(events[current_event:end])
        # Get the view before this event
        current_tid_to_pid = versioned_view(tid_to_pid, current_event+i-2)
        if icxx"$e.type() == rr::TraceTaskEvent::EXEC;" &&
                current_tid_to_pid[icxx"$e.tid();"] == pid
            return current_event + i -1
        elseif icxx"$e.type() == rr::TraceTaskEvent::EXIT;" &&
                current_tid_to_pid[icxx"$e.tid();"] == pid
            return -1;
        end
    end
    -1
end

"""
A reimplementation of `rr ps`
"""
function ASTInterpreter.execute_command(state, stack, ::Val{:ps}, command)
    events = cxxt"rr::TraceTaskEvent"[]
    timeline = state.top_interp.session
    session = current_session(timeline)
    trace = icxx"rr::TraceReader{$session->trace_reader()};"
    icxx"$trace.rewind();"
    while icxx"$trace.good();"
        push!(events, icxx"$trace.read_task_event();")
    end
    isempty(events) || icxx"$(first(events)).type() != rr::TraceTaskEvent::EXEC;" &&
        error("Invalid trace")

    # tid_to_pid is a versi
    tid_to_pid = VersionedDict{cxxt"pid_t",cxxt"pid_t",Int}()
    initial_tid = icxx"$(first(events)).tid();"
    tid_to_pid[initial_tid] = initial_tid;
    for (i, e) in enumerate(events)
        advance!(tid_to_pid, i)
        update_tid_to_pid_map!(tid_to_pid, e)
    end

    print("$initial_tid\t--\t")
    println(join(map(unsafe_string,collect(icxx"$(first(events)).cmd_line();")),' '))

    for (i, e) in enumerate(events)
        if icxx"$e.type() == rr::TraceTaskEvent::CLONE;" &&
                icxx"!($e.clone_flags() & CLONE_THREAD);"
            tid = icxx"$e.tid();"
            pid = versioned_view(tid_to_pid,i)[tid]
            parent_pid = versioned_view(tid_to_pid,i)[icxx"$e.parent_tid();"]
            print("$tid\t$parent_pid\t",find_exit_code(pid, events, i, tid_to_pid)," ")
            idx = find_cmd_line(pid, events, i, tid_to_pid)
            println(idx == -1 ? "(forked but did not exec)" :
                join(map(unsafe_string,collect(icxx"$(events[idx]).cmd_line();")),' '))
        end
    end
    return false
end

function ASTInterpreter.execute_command(state, stack::Union{Gallium.NativeStack,Gallium.CStackFrame}, ::Val{:rsi}, command)
    timeline = state.top_interp.session
    task = isa(stack, Gallium.NativeStack) ? stack.session : current_task(current_session(timeline))
    isa(task, RR.ReplayTimeline) && (task = current_task(current_session(task)))
    RR.reverse_single_step!(current_session(timeline),task,timeline)
    update_stack!(state)
    return true
end


text:
include("TypesTest.jl")
include("BasisTest.jl")
include("NuclearRepulsionTest.jl")
include("OverlapTest.jl")
include("KineticEnergyTest.jl")
include("NuclearAttractionTest.jl")
include("ElecElecRepulsionTest.jl")

run_typesparticles_test()
run_typesbasis_test()
run_readbasis_test()
run_retreadbasis_test()
run_nuclearrepulsion_test()
run_buildoverlap_test()
run_buildkineticenergy_test()
run_buildnuclrattract_test()
run_buildelecelecrepul_test()


text:
@doc raw"""
    MultinomialSymmetric{n} <: AbstractMultinomialDoublyStochastic{N}

The multinomial symmetric matrices manifold consists of all symmetric $n×n$ matrices with
positive entries such that each column sums to one, i.e.

````math
\begin{aligned}
\mathcal{SP}(n) \coloneqq \bigl\{p ∈ ℝ^{n×n}\ \big|\ &p_{i,j} > 0 \text{ for all } i=1,…,n, j=1,…,m,\\
& p^\mathrm{T} = p,\\
& p\mathbf{1}_n = \mathbf{1}_n
\bigr\},
\end{aligned}
````

where $\mathbf{1}_n$ is the vector of length $n$ containing ones.

It is modeled as an [`DefaultIsometricEmbeddingType`](@ref), [`AbstractEmbeddedManifold`](@ref)
via the [`AbstractMultinomialDoublyStochastic`](@ref) type, since it shares a few functions
also with [`AbstractMultinomialDoublyStochastic`](@ref), most and foremost projection of
a point from the embedding onto the manifold.

The tangent space can be written as

````math
T_p\mathcal{SP}(n) \coloneqq \bigl\{
X ∈ ℝ^{n×n}\ \big|\ X = X^{\mathrm{T}} \text{ and }
X\mathbf{1}_n = \mathbf{0}_n
\bigr\},
````

where $\mathbf{0}_n$ is the vector of length $n$ containing zeros.

More details can be found in Section IV[^DouikHassibi2019].

# Constructor

    MultinomialSymmetric(n)

Generate the manifold of matrices $\mathbb R^{n×n}$ that are doubly stochastic and symmetric.

[^DouikHassibi2019]:
    > A. Douik, B. Hassibi:
    > AbstractManifold Optimization Over the Set of Doubly Stochastic Matrices: A Second-Order Geometry,
    > IEEE Transactions on Signal Processing 67(22), pp. 5761–5774, 2019.
    > doi: [10.1109/tsp.2019.2946024](http://doi.org/10.1109/tsp.2019.2946024),
    > arXiv: [1802.02628](https://arxiv.org/abs/1802.02628).
"""
struct MultinomialSymmetric{N} <: AbstractMultinomialDoublyStochastic{N} end

function MultinomialSymmetric(n::Int)
    return MultinomialSymmetric{n}()
end

@doc raw"""
    check_point(M::MultinomialSymmetric, p)

Checks whether `p` is a valid point on the [`MultinomialSymmetric`](@ref)`(m,n)` `M`,
i.e. is a symmetric matrix with positive entries whose rows sum to one.
"""
function check_point(M::MultinomialSymmetric{n}, p; kwargs...) where {n}
    mpv = invoke(check_point, Tuple{supertype(typeof(M)),typeof(p)}, M, p; kwargs...)
    mpv === nothing || return mpv
    # the embedding checks for positivity and unit sum columns, by symmetry we would get
    # the same for the rows, so checking symmetry is the only thing left, we can just use
    # the corresponding manifold for that
    return check_point(SymmetricMatrices(n, ℝ), p)
end
@doc raw"""
    check_vector(M::MultinomialSymmetric p, X; kwargs...)

Checks whether `X` is a valid tangent vector to `p` on the [`MultinomialSymmetric`](@ref) `M`.
This means, that `p` is valid, that `X` is of correct dimension, symmetric, and sums to zero
along any row.
"""
function check_vector(M::MultinomialSymmetric{n}, p, X; kwargs...) where {n}
    mpv = invoke(
        check_vector,
        Tuple{supertype(typeof(M)),typeof(p),typeof(X)},
        M,
        p,
        X;
        kwargs...,
    )
    mpv === nothing || return mpv
    # from the embedding we know that columns sum to zero, only symmety is left, i.e.
    return check_vector(SymmetricMatrices(n, ℝ), p, X; kwargs...)
end

function decorated_manifold(::MultinomialSymmetric{N}) where {N}
    return MultinomialMatrices(N, N)
end

embed!(::MultinomialSymmetric, q, p) = copyto!(q, p)
embed!(::MultinomialSymmetric, Y, ::Any, X) = copyto!(Y, X)

@doc raw"""
    manifold_dimension(M::MultinomialSymmetric{n}) where {n}

returns the dimension of the [`MultinomialSymmetric`](@ref) manifold
namely
````math
\operatorname{dim}_{\mathcal{SP}(n)} = \frac{n(n-1)}{2}.
````
"""
@generated function manifold_dimension(::MultinomialSymmetric{n}) where {n}
    return div(n * (n - 1), 2)
end

@doc raw"""
    project(M::MultinomialSymmetric{n}, p, Y) where {n}

Project `Y` onto the tangent space at `p` on the [`MultinomialSymmetric`](@ref) `M`, return the result in `X`.
The formula reads
````math
    \operatorname{proj}_p(Y) = Y - (α\mathbf{1}_n^{\mathrm{T}} + \mathbf{1}_n α^{\mathrm{T}}) ⊙ p,
````
where $⊙$ denotes the Hadamard or elementwise product and $\mathbb{1}_n$ is the vector of length $n$ containing ones.
The two vector $α ∈ ℝ^{n×n}$ is given by solving
````math
    (I_n+p)α =  Y\mathbf{1},
````
where $I_n$ is teh $n×n$ unit matrix and $\mathbf{1}_n$ is the vector of length $n$ containing ones.

"""
project(::MultinomialSymmetric, ::Any, ::Any)

function project!(::MultinomialSymmetric{n}, X, p, Y) where {n}
    α = (I + p) \ sum(Y, dims=2) # Formula (49) from 1802.02628
    return X .= Y .- (repeat(α, 1, 3) .+ repeat(α', 3, 1)) .* p
end

@generated function representation_size(::MultinomialSymmetric{n}) where {n}
    return (n, n)
end

@doc raw"""
    retract(M::MultinomialSymmetric, p, X, ::ProjectionRetraction)

compute a projection based retraction by projecting $p\odot\exp(X⨸p)$ back onto the manifold,
where $⊙,⨸$ are elementwise multiplication and division, respectively. Similarly, $\exp$
refers to the elementwise exponentiation.
"""
retract(::MultinomialSymmetric, ::Any, ::Any, ::ProjectionRetraction)

function retract!(M::MultinomialSymmetric, q, p, X, ::ProjectionRetraction)
    return project!(M, q, p .* exp.(X ./ p))
end

"""
    vector_transport_to(M::MultinomialSymmetric, p, X, q)

transport the tangent vector `X` at `p` to `q` by projecting it onto the tangent space
at `q`.
"""
vector_transport_to(::MultinomialSymmetric, ::Any, ::Any, ::Any, ::ProjectionTransport)

function vector_transport_to!(M::MultinomialSymmetric, Y, p, X, q, ::ProjectionTransport)
    project!(M, Y, q, X)
    return Y
end

function Base.show(io::IO, ::MultinomialSymmetric{n}) where {n}
    return print(io, "MultinomialSymmetric($(n))")
end