struct Control
    iter::Int ## total number of iterations (incl. warmup)
    warmup::Int ## number of warmup iterations
    chains::Int ## number of chains
    save_warmup::Bool ## store and return warmup iterations in Fit
    common_delta::Bool ## common log-HRs across classes
    common_beta::Bool ## common fixed effects across classes
    adapt_delta::Float64 ## target mean acceptance rate for NUTS transitions
end

Control(iter, chains) = Control(iter=iter, chains=chains)

function Control(;iter=2000, warmup=Int(floor(iter/2)), chains=4, save_warmup=false,
                 common_delta=true, common_beta=true, adapt_delta=0.8)
    @argcheck iter > warmup
    Control(iter, Int(floor(iter/2)), chains, save_warmup, common_delta, common_beta, adapt_delta)
end

struct Chain{Tpar<:AbstractParam, Tprior<:AbstractPrior}
    theta::Vector{Tpar}
    accept::Tpar
    loglik_p::Matrix{Float64} ## pointwise log-likelihood
    prior::Tprior
    init::Tpar
end

struct Fit{Tpar<:AbstractParam, Tprior<:AbstractPrior} <: AbstractFit
    control::Control
    chains::Vector{Chain{Tpar, Tprior}}
end

function show(io::IO, x::Data)
    println(io, "Subjects: $(length(x.T))")
    println(io, "Repeated measurements: $(sum(x.n))")
    println(io, "Events: $(sum(x.E))")
end

function show(io::IO, x::Chain{T,S}) where {T,S}
    cm = countmap([par.G for par in x.theta])
    cm0 = countmap([par.G-par.G0 for par in x.theta])

    println(io, "Number of groups: $(collect(zip(keys(cm), values(cm))))")
    println(io, "Number of non-empty groups: $(collect(zip(keys(cm0), values(cm0))))")
end

function show(io::IO, x::Fit{T}) where T <: AbstractParam
    println(io, "Number of chains: $(x.control.chains)")
    println(io, "Iterations: $(x.control.iter)")
    for i in eachindex(x.chains)
        println(io, "Chain $(i):")
        show(io, x.chains[i])
        println(io, "\n")
    end
end

merge(a...) = foldr(merge, a)
merge(a::Chain{T,S}) where {T,S} = a

function merge(a::Chain{T,S}, b::Chain{T,S}) where {T <: AbstractParam, S <: AbstractPrior}
    Chain{T,S}(vcat(a.theta, b.theta), a.accept, vcat(a.loglik_p, b.loglik_p), a.prior, a.init)
end

size(x::Fit) = (length(x.chains), length(first(x.chains)))
slice(x::Chain, t) = x.theta[t]

length(x::Fit) = length(x.chains)
length(x::Chain) = length(x.theta)

function get_samples(x::Fit)
    K, M = size(x)
    res = Matrix{typeof(first(x.chains).init)}(undef, M, K)

    for k in 1:K
        res[:,k] = get_samples(x.chains[k])
    end
    res
end

get_samples(x::Chain) = x.theta

get_loglik_p(x::Chain) = x.loglik_p
get_loglik_p(x::Fit) = get_loglik_p(merge(x.chains...))

get_acceptance_rates(x::Fit) = get_acceptance_rates(merge(x.chains...))
get_acceptance_rates(x::Chain) = x.accept[x.theta[1].G]

## Calculate class-membership probabilities for each subject
function class_prob(xi, Xp)
    cp = exp.(Xp * xi')  ## N x G
    cp ./ sum(cp, dims=2)[:,1] ## normalize each row
end

function build_cov_matrix(S::Vector{T}, L::Matrix{T}) where T
    SD = Diagonal(S)
    Array(Symmetric(SD*L*L'*SD))
end

build_cov_matrix(theta::Tp, g::Int) where Tp <: AbstractParam = buidl_cov_matrix(theta.S[g,:], theta.L[g,:,:])

function post_cind_prob(theta::ParamSpline, data)
    @unpack Y, times, n, Xl, Z, Xp, Xe, T, E, haz_basis, cumhaz_basis = data
    @unpack xi, beta, mu, sigma, w, delta, S, L = theta

    G, q = size(mu)
    N = length(data.T)
    cp = class_prob(xi, Xp)
    res = Matrix{Float64}(undef, N, G)

    for g in 1:G
        Sigma = build_cov_matrix(S[g,:], L[g,:,:])

        for i in 1:N
            nt = 1:n[i]
            lll_vec = loglik_p_long(beta[g,:], mu[g,:], Sigma, sigma, Y[i,1:n[i]], times[i,1:n[i]], n[i], Xl[i,1:n[i],:], Z[i,1:n[i],:])
            lle_vec = loglik_p_surv(T[i], E[i], Xe[i:i,:], haz_basis[i:i,:], cumhaz_basis[i:i,:], w[g,:], delta[g,:])
            res[i,g] = lll_vec + lle_vec + log(cp[g])
        end
    end
    res
end

function post_cind_prob(theta::Param, data)
    @unpack Y, times, n, Xl, Z, Xp, Xe, T, E = data
    @unpack xi, beta, mu, sigma, bl_shape, bl_scale, delta, S, L = theta

    G, q = size(mu)
    N = length(data.T)
    cp = class_prob(xi, Xp)
    res = Matrix{Float64}(undef, N, G)

    for g in 1:G
        Sigma = build_cov_matrix(S[g,:], L[g,:,:])

        for i in 1:N
            nt = 1:n[i]
            lll_vec = loglik_p_long(beta[g,:], mu[g,:], Sigma, sigma, Y[i,1:n[i]], times[i,1:n[i]], n[i], Xl[i,1:n[i],:], Z[i,1:n[i],:])
            lle_vec = loglik_p_surv(T[i], E[i], Xe[i:i,:], bl_shape[g], bl_scale[g], delta[g,:])
            res[i,g] = lll_vec + lle_vec + log(cp[g])
        end
    end
    res
end

## sample from discrete distribution with probabilities given by exp.(alpha)/sum(exp.(alpha)) using Gumbel-max trick
sample_discrete(alpha) = last(findmax(alpha + rand(Gumbel(), length(alpha))))

function sample_cind(theta::Tp, data::Data) where Tp <: AbstractParam
    alpha = post_cind_prob(theta, data)
    [sample_discrete(r) for r in eachrow(alpha)]
end

function sample_sigma(b, beta, mu, sigma, Xl, Z, Y, n, times, cind, prior)
    N = size(Xl)[1]
    q = size(Xl)[3]
    v = 0.0

    for i in 1:N
        g = cind[i]
        u = b[i,:] ##mu[g,:] + b[i,:]

        for j in 1:n[i]
            v += (Y[i,j] - dot(Z[i,j,:], u) - dot(Xl[i,j,:], beta[g,:]))^2
        end
    end
    v /= 2

    log_target(x::Float64) = logpdf(InverseGamma(sum(n)/2 - 1, v), exp(2*x)) + logpdf(prior, exp(x)) + x

    y, ac = metropolis(log_target, log(sigma), rand(Normal(log(sigma), 0.1)))

    exp(y), ac
end

function sample_mu_conj(Sigma, beta, mu, sigma, Xl, Z, Y, n, prior_mu)
    N = size(Y,1)
    q = length(mu)

    V0 = inv(cov(prior_mu))
    ll_mean_mu = zeros(Float64, q)
    ll_inv_cov_mu = zeros(Float64, q, q)

    for i in 1:N
        tt = 1:n[i]
        Zi = Z[i,tt,:]
        Xli = Xl[i,tt,:]
        C = inv(Array(Symmetric(Zi*Sigma*Zi' + sigma^2*diagm(0 => ones(Float64, n[i])))))
        ll_mean_mu += Zi'*C*(Y[i,tt] - Xli*beta)
        ll_inv_cov_mu += Zi'*C*Zi
    end

    post_cov = inv(Symmetric(V0 + ll_inv_cov_mu))
    post_mean = post_cov * (V0*mean(prior_mu) + ll_mean_mu)
    rand(MvNormal(post_mean, post_cov)), 1
end

function sample_beta_conj(theta, data, prior_beta)
    @unpack beta, mu, sigma, S, L, cind = theta
    @unpack Xl, Z, Y, n = data

    N = size(Y,1)
    l = size(beta,2)
    G, q = size(mu)

    V0 = inv(cov(prior_beta))

    ll_mean_beta = zeros(Float64, l)
    ll_inv_cov_beta = zeros(Float64, l, l)

    Sigma = zeros(eltype(beta), G, q, q)
    for g in 1:G
        Sigma[g,:,:] = build_cov_matrix(S[g,:], L[g,:,:])
    end

    for i in 1:N
        tt = 1:n[i]
        Zi = Z[i,tt,:]
        Xli = Xl[i,tt,:]
        g = cind[i]
        C = inv(Array(Symmetric(Zi*Sigma[g,:,:]*Zi' + sigma^2*diagm(0 => ones(Float64, n[i])))))
        ll_mean_beta += Xli'*C*(Y[i,tt] - Zi*mu[g,:])
        ll_inv_cov_beta += Xli'*C*Xli
    end

    post_cov = inv(Symmetric(V0 + ll_inv_cov_beta))
    post_mean = post_cov * (V0*mean(prior_beta) + ll_mean_beta)
    rand(MvNormal(post_mean, post_cov)), 1
end

function sample_mu_beta_conj(Sigma, beta, mu, sigma, Xl, Z, Y, n, prior_mu, prior_beta)
    N = size(Y,1)
    q = length(mu)
    l = length(beta)

    mb0 = vcat(mean(prior_mu), mean(prior_beta))
    V0 = zeros(Float64, q+l, q+l)
    ## assuming independent priors for mu and beta
    V0[1:q, 1:q] = inv(cov(prior_mu))
    V0[(q+1):end, (q+1):end] = inv(cov(prior_beta))

    ll_mean_mu = zeros(Float64, q)
    ll_mean_beta = zeros(Float64, l)
    ll_inv_cov_mu = zeros(Float64, q, q)
    ll_inv_cov_beta = zeros(Float64, l, l)

    for i in 1:N
        tt = 1:n[i]
        Zi = Z[i,tt,:]
        Xli = Xl[i,tt,:]
        C = inv(Array(Symmetric(Zi*Sigma*Zi' + sigma^2*diagm(0 => ones(Float64, n[i])))))
        ll_mean_mu += Zi'*C*(Y[i,tt] - Xli*beta)
        ll_mean_beta += Xli'*C*(Y[i,tt] - Zi*mu)
        ll_inv_cov_mu += Zi'*C*Zi
        ll_inv_cov_beta += Xli'*C*Xli
    end

    ll_inv_cov = zeros(Float64, q+l, q+l)
    ll_inv_cov[1:q, 1:q] = ll_inv_cov_mu
    ll_inv_cov[(q+1):end, (q+1):end] = ll_inv_cov_beta

    post_cov = inv(Symmetric(V0 + ll_inv_cov))
    post_mean = post_cov * (V0*mb0 + vcat(ll_mean_mu, ll_mean_beta))

    mu_beta = rand(MvNormal(post_mean, post_cov))
    mu_beta[1:q], mu_beta[(q+1):end], 1
end

function inv_cov_SL(S, L)
    ##try
        invSL = inv(LowerTriangular(Diagonal(S)*L))
        return invSL'invSL
    ##catch ex
    ##    @info S, L
    ##    error("singular")
    ##end
    ##return diagm(0 => ones(Float64, length(S)))
end

function sample_b_conj(theta, data)
    @unpack G, b, beta, mu, sigma, cind, S, L = theta
    @unpack T, E, Z, Xe, Xl, Z, Y, n, times = data

    N = length(T)
    q = size(b,2)
    z = Matrix{Float64}(undef, N, q)
    invSigma = [inv_cov_SL(S[g,:], L[g,:,:]) for g in 1:G]
    z0 = rand(MvNormal(diagm(0 => ones(eltype(b), q))), N)

    for i in 1:N
        g = cind[i]
        tt = 1:n[i]
        Zi = Z[i,tt,:]
        ll_mean = Zi'*(Y[i,tt] - Xl[i,tt,:]*beta[g,:])
        ll_inv_cov = Zi'*Zi / sigma^2

        V = Symmetric(invSigma[g] + ll_inv_cov)
        V_chol = cholesky!(V; check=false)

        if !issuccess(V_chol)
            z[i,:] = b[i,:]
        else
            VL = inv(V_chol.L)
            post_mean = VL'*VL*(invSigma[g]*mu[g,:] + ll_mean / sigma^2)
            z[i,:] = post_mean + VL'*z0[:,i]
        end
    end
    z, 1
end

function update_lmm_params_common!(theta::T, accept::T, data::Data, prior) where T <: AbstractParam
    G, q = size(theta.mu)

    theta.beta[1,:], ac = sample_beta_conj(theta, data, prior.beta_mv)

    for g in 1:G
        theta.beta[g,:] = theta.beta[1,:]
    end
    accept.beta .+= ac

    for g in 1:G
        sel = theta.cind .== g

        if sum(sel) > 0
            data_g = slice_data(data, sel)

            Sigma = build_cov_matrix(theta.S[g,:], theta.L[g,:,:])

            theta.mu[g,:], ac = sample_mu_conj(Sigma, theta.beta[g,:], theta.mu[g,:], theta.sigma,
                                               data_g.Xl, data_g.Z, data_g.Y, data_g.n, prior.mu_mv)
            accept.mu[g,:] .+= ac
        else
            theta.mu[g,:] = rand(prior.mu_mv)
            accept.mu[g,:] .+= 1
        end
    end

    theta.b, ac = sample_b_conj(theta, data)
    accept.b .+= ac

    theta.sigma, ac = sample_sigma(theta.b, theta.beta, theta.mu, theta.sigma, data.Xl, data.Z, data.Y, data.n, data.times, theta.cind, prior.sigma)
    accept.sigma += ac
end

function update_lmm_params_mix!(theta::T, accept::T, data::Data, prior) where T <: AbstractParam
    G, q = size(theta.mu)

    for g in 1:G
        sel = theta.cind .== g

        if sum(sel) > 0
            data_g = slice_data(data, sel)

            Sigma = build_cov_matrix(theta.S[g,:], theta.L[g,:,:])

            theta.mu[g,:], theta.beta[g,:], ac = sample_mu_beta_conj(Sigma, theta.beta[g,:], theta.mu[g,:], theta.sigma, data_g.Xl, data_g.Z, data_g.Y, data_g.n, prior.mu_mv, prior.beta_mv)
            accept.beta[g,:] .+= ac
            accept.mu[g,:] .+= ac
        else
            theta.mu[g,:] = rand(prior.mu_mv)
            accept.mu[g,:] .+= 1

            theta.beta[g,:] = rand(prior.beta_mv)
            accept.beta[g,:] .+= 1
        end
    end

    theta.b, ac = sample_b_conj(theta, data)
    accept.b .+= ac

    theta.sigma, ac = sample_sigma(theta.b, theta.beta, theta.mu, theta.sigma, data.Xl, data.Z, data.Y, data.n, data.times, theta.cind, prior.sigma)
    accept.sigma += ac
end

function transform_zero_sum(xi)
    G = size(xi,1)
    A = diagm(0 => ones(eltype(xi), G)) .- 1/G
    A*xi
end

function transform_ref_group(xi)
    G, p = size(xi)
    A = diagm(0 => ones(eltype(xi), G-1)) .- 1/G
    inv(A)*xi[1:(G-1),:]
end

function sample_xi_NUTS(sampler, xi, cind, Xp)
    G, p = size(xi)
    ##init = xi[1:(G-1),:]
    init = transform_ref_group(xi)
    res, ac = nuts_mlogit_step!(sampler, cind_to_matrix(cind, G), init)
    ##vcat(res.xi, zeros(eltype(res.xi), 1, p)), ac
    new_xi = transform_zero_sum(vcat(res.xi, zeros(eltype(res.xi), 1, p)))
    new_xi, ac
end

function update_cp_params!(theta::T, accept::T, data::Data, prior, sampler) where T <: AbstractParam
    G, q = size(theta.mu)

    if G == 1
        theta.cind .= 1
        return loglik(theta, data)
    else
        theta.cind = sample_cind(theta, data)
        accept.cind .+= 1

        theta.xi, ac = sample_xi_NUTS(sampler, theta.xi, theta.cind, data.Xp)
        ##theta.xi, ac = sample_xi(theta.xi, theta.cind, data.Xp)
        accept.xi .+= ac
    end
end

function cind_to_matrix(cind::Vector{Int}, G::Int)
    Y = zeros(Int, length(cind), G)
    for i in 1:length(cind)
        Y[i, cind[i]] = 1
    end
    Y
end

function init_NUTS_samplers(warmup, G, theta, data, prior, cp=Dict{Int, NUTSGibbsSampler}(),
                            surv=Dict{Int, NUTSGibbsSampler}(),
                            cov=Dict{Int, NUTSGibbsSampler}(); common_delta=false)
    if !haskey(surv, G)
        if G > 1
            cp[G] = nuts_mlogit_init(cind_to_matrix(theta.cind, G), data.Xp, prior.xi, warmup)
            @debug "Initialized NUTS sampler for class membership regression (G = $(G))"
        end

        if common_delta
            surv[G] = surv_model_init2(theta, data, prior, warmup)
        else
            surv[G] = surv_model_init(theta, data, prior, warmup)
        end
        @debug "Initialized NUTS sampler for survival submodel (G = $(G))"

        ##@info "Initializing NUTS sampler for LMM submodel"
        ##lmm[G] = NUTS_JLCM_lmm_init(theta, data, prior)

        cov[G] = cov_model_init(theta, data, prior, warmup)
        @debug "Initialized NUTS sampler for covariance model (G = $(G))"

        ##@info "Initializing NUTS sampler for JLCM model (G = $(G))"
        ##cov[G] = NUTS_JLCM_init(theta, data, prior)
    end

    cp, surv, cov
end

function sampling(M::Int, data::Data, par_type::Type{Tp}, prior::Tprior=Prior(data), prop::Tprior=prior;
                  chains=1, G=1, warmup=Int(floor(M/2)), common_delta=false, common_beta=false,
                  return_warmup=false) where {Tp <: AbstractParam, Tprior <: AbstractPrior}
    control = Control(M, warmup, chains, return_warmup, common_delta, common_beta, 0.8)
    sampling(control, data, par_type, prior, prop; G=G)
end

function sampling(control::Control, data::Data, par_type::Type{Tp}, prior::Tprior=Prior(data), prop::Tprior=prior; G=1) where {Tp <: AbstractParam, Tprior <: AbstractPrior}

    sampling(control, data, [Tp(data,G) for k in 1:control.chains], prior, prop)
end

function sampling(M, data::Data, x::T; warmup=floor(Int, M/2), from_init=true) where T <: AbstractFit
    c = x.control
    ctrl = Control(M, warmup, c.chains, c.save_warmup, c.common_delta, c.common_beta, c.adapt_delta)

    if from_init
        init = [xx.init for xx in x.chains]
    else
        init = [xx.theta[end] for xx in x.chains]
    end

    sampling(ctrl, data, init, x.chains[1].prior)
end

function sampling(control::Control, data::Data, init::Vector{Tp}, prior::Tprior=Prior(data), prop::Tprior=prior) where {Tp <: AbstractParam, Tprior <: AbstractPrior}

    @assert length(init) == control.chains

    chains = progress_pmap(1:control.chains) do k
        @info "Chain $(k) of $(control.chains)"
        sampling(control, data, init[k]; prior=prior, prop=prop)
    end

    Fit(control, chains)
end

function sampling(control::Control, data::Data, init::Tp=Param(data, 1);
                  prior::Tpr=Prior(), prop::Tpr=prior) where {Tp <: AbstractParam, Tpr <: AbstractPrior}
    G = init.G
    M = control.iter
    warmup = control.warmup

    offset = control.save_warmup ? 0 : warmup

    @info "Taking $(M) samples including $(warmup) warmup iterations in $(G)-class model"

    loglik_ar = Matrix{Float64}(undef, M-offset, length(data.T))
    theta_stor = Vector{Tp}(undef, M-offset)

    theta = deepcopy(init)

    if offset == 0
        theta_stor[1] = deepcopy(theta)
        loglik_ar[1,:] = loglik_p(theta, data)
    end

    accept = Tp(data, G; set_zero=true)

    cp_sampler, surv_sampler, cov_sampler = init_NUTS_samplers(warmup, G, theta, data, prior; common_delta=control.common_delta)
    ac = 0.0

    ##@showprogress 1 "Sampling..."
    for t in 2:M
        if G > 1
            update_cp_params!(theta, accept, data, prior, cp_sampler[G])
        end

        if control.common_beta
            update_lmm_params_common!(theta, accept, data, prior)
        else
            update_lmm_params_mix!(theta, accept, data, prior)
        end

        theta, ac = cov_model_step!(cov_sampler[G], theta)
        accept.S .+= ac
        accept.L .+= ac
        accept.sigma += ac

        if control.common_delta
            theta, ac = surv_model_step2!(surv_sampler[G], theta)
        else
            theta, ac = surv_model_step!(surv_sampler[G], theta)
        end

        accept.delta .+= ac
        accept.bl_shape .+= ac
        accept.bl_scale .+= ac

        theta.G0 = sum([sum(theta.cind .== g) == 0 for g in 1:G])

        if t > offset
            ll = loglik_p(theta, data)
            loglik_ar[t-offset,:] = ll
            theta_stor[t-offset] = deepcopy(theta)
        end
    end

    Chain{Tp,Tpr}(theta_stor, accept, loglik_ar, prior, init)
end


using MPIMeasurements
using Base.Test
using Unitful
using Compat
using HDF5

# define Grid
shape = [5,5,1]
fov = [200.0,200.0,1.0]Unitful.mm
center = [0.0,0.0,0.0]Unitful.mm
positions = MeanderingGridPositions( RegularGridPositions(shape,fov,center) )

# create Scanner
#scanner = MPIScanner("HeadScanner.toml")
scanner = MPIScanner("DummyScanner.toml")
robot = getRobot(scanner)
safety = getSafety(scanner)
gaussmeter = getGaussMeter(scanner)

#TODO mT sollte man hier nicht angeben müssen. Das sollte im Gaussmeter gekapselt sein
mfMeasObj = MagneticFieldMeas(gaussmeter, Unitful.mT, positions)

res = performTour!(robot, safety, positions, mfMeasObj)

filenameField = joinpath(homedir(),"TestBackground.h5")
saveMagneticFieldAsHDF5(mfMeasObj, filenameField)

pos, field = loadMagneticField(filenameField)

MPISimulations.plotMagneticField(field, pos, 3, 1)


using MacroTools: postwalk, @capture

export @dice, compile, rundice

struct DiceCode
    interpret::Function
end

"""
process dice code before running it
currently it makes if-then-else, &&, and || polymorphic
"""
function to_dice(code)

    # manual hygiene
    mgr = gensym(:mgr)
    flip = gensym(:flip)
    ite = gensym(:ite)

    transformed_code = postwalk(esc(code)) do x
        # TODO: replace by custom desugaring that is still lazy for boolean guards
        # this will not work in general
        # for example when control flow has `return`` inside
        if x isa Expr && (x.head == :if || x.head == :elseif)
            @assert length(x.args) == 3 "@dice macro only supports purely functional if-then-else"
            return :($ite($(x.args...)))
        end
        @capture(x, flip(P_)) && return :($flip($P))
        @capture(x, A_ || B_) && return :($A | $B)
        @capture(x, A_ && B_) && return :($A & $B)
        return x
    end

    return quote

        DiceCode( $(esc(mgr)) -> begin

            $(esc(flip))(prob::Number) =
                DistBool($(esc(mgr)), prob)

            $(esc(ite))(args...) =
                ifelse(args...)

            # transformed user code
            $transformed_code
        end)
    end
end

macro dice(code)
    to_dice(code)
end

function compile(code::DiceCode)
    code.interpret(CuddMgr())
end

function to_dice_ir(code::DiceCode)
    interpretation = code.interpret(IrMgr())
    treeify(interpretation.bit)
end

function rundice(code::DiceCode)
    rundice(to_dice_ir(code))
end

function infer(code::DiceCode, algo)
    x = if algo == :ocaml
        to_dice_ir(code)
    elseif algo == :bdd
        compile(code)
    end
    infer(x)
end


using Random
using DABenchmarks

Random.seed!(0)

using Distributed
machine = [("MACHINENAME", 8)]
addprocs(machine, dir = "/home/USER", exename="/path/to/julia", exeflags="--project")
# addprocs(8, exeflags="--project") # For local workers

@everywhere using StochasticPrograms
@everywhere using LShapedSolvers
@everywhere using HydroModels
@everywhere using Gurobi

function prepare_dbenchmark(nscenarios::Integer, nsamples::Integer; timeout::Integer = 1000)
    dtr_1 = LShapedSolver(GurobiSolver(OutputFlag=0), regularization = :tr, log=false, distributed = true, subsolver = ()->GurobiSolver(OutputFlag=0), κ=1.0, bundle = 1)
    dtr_10 = LShapedSolver(GurobiSolver(OutputFlag=0), regularization = :tr, log=false, distributed = true, subsolver = ()->GurobiSolver(OutputFlag=0), κ=1.0, bundle = 10)
    dtr_50 = LShapedSolver(GurobiSolver(OutputFlag=0), regularization = :tr, log=false, distributed = true, subsolver = ()->GurobiSolver(OutputFlag=0), κ = 1.0, bundle = 50)
    dtr_100 = LShapedSolver(GurobiSolver(OutputFlag=0), regularization = :tr, log=false, distributed = true, subsolver = ()->GurobiSolver(OutputFlag=0), κ = 1.0, bundle = 100)
    dtr_125 = LShapedSolver(GurobiSolver(OutputFlag=0), regularization = :tr, log=false, distributed = true, subsolver = ()->GurobiSolver(OutputFlag=0), κ=1.0, bundle = 125)
    solvers = [dtr_1, dtr_10, dtr_50, dtr_100, dtr_125]
    solvernames = ["1", "10", "50", "100", "125"]
    # Create Day-ahead benchmark
    return Scaling(solvernames, solvers, nscenarios, nsamples; timeout = timeout)
end
@info "Running bundled L-shaped benchmarks on 16 workers..."
dab = prepare_dbenchmark(1000, 10); warmup_benchmarks!(dab); run_benchmarks!(dab); save_results!(dab, "tr_bundle_8.json")
exit()


## PARAMETER RANGES

abstract type ParamRange{T} end

Base.isempty(::ParamRange) = false

abstract type Boundedness end

abstract type   Bounded <: Boundedness end
abstract type Unbounded <: Boundedness end

abstract type   LeftUnbounded <: Unbounded end
abstract type  RightUnbounded <: Unbounded end
abstract type DoublyUnbounded <: Unbounded end

struct NumericRange{T,B<:Boundedness,D} <: ParamRange{T}
    field::Union{Symbol,Expr}
    lower::Union{T,Float64}     # Float64 to allow for -Inf
    upper::Union{T,Float64}     # Float64 to allow for Inf
    origin::Float64
    unit::Float64
    scale::D
end

struct NominalRange{T,N} <: ParamRange{T}
    field::Union{Symbol,Expr}
    values::NTuple{N,T}
end

function Base.show(stream::IO,
#                   ::MIME"text/plain",
                   r::NumericRange{T}) where T
    fstr = string(r.field)
    repr = "NumericRange($(r.lower) ≤ $fstr ≤ $(r.upper); "*
        "origin=$(r.origin), unit=$(r.unit))"
    if r.scale isa Symbol
        r.scale !== :linear && (repr *= " on $(r.scale) scale")
    else
        repr = "transformed "*repr
    end
    print(stream, repr)
    return nothing
end

function Base.show(stream::IO,
#                   ::MIME"text/plain",
                   r::NominalRange{T}) where T
    fstr = string(r.field)
    seqstr = sequence_string(collect(r.values))
    repr = "NominalRange($fstr = $seqstr)"
    print(stream, repr)
    return nothing
end

"""
    r = range(model, :hyper; values=nothing)

Define a one-dimensional `NominalRange` object for a field `hyper` of
`model`. Note that `r` is not directly iterable but `iterator(r)` is.

A nested hyperparameter is specified using dot notation. For example,
`:(atom.max_depth)` specifies the `max_depth` hyperparameter of
the submodel `model.atom`.

    r = range(model, :hyper; upper=nothing, lower=nothing,
              scale=nothing, values=nothing)

Assuming `values` is not specified, define a one-dimensional
`NumericRange` object for a `Real` field `hyper` of `model`.  Note
that `r` is not directly iteratable but `iterator(r, n)`is an iterator
of length `n`. To generate random elements from `r`, instead apply
`rand` methods to `sampler(r)`. The supported scales are `:linear`,`
:log`, `:logminus`, `:log10`, `:log2`, or a callable object.

Note that `r` is not directly iterable, but `iterator(r, n)` is, for
given resolution (length) `n`.

By default, the behaviour of the constructed object depends on the
type of the value of the hyperparameter `:hyper` at `model` *at the
time of construction.* To override this behaviour (for instance if
`model` is not available) specify a type in place of `model` so the
behaviour is determined by the value of the specified type.

A nested hyperparameter is specified using dot notation (see above).

If `scale` is unspecified, it is set to `:linear`, `:log`,
`:logminus`, or `:linear`, according to whether the interval `(lower,
upper)` is bounded, right-unbounded, left-unbounded, or doubly
unbounded, respectively.  Note `upper=Inf` and `lower=-Inf` are
allowed.

If `values` is specified, the other keyword arguments are ignored and
a `NominalRange` object is returned (see above).

See also: [`iterator`](@ref), [`sampler`](@ref)

"""
function Base.range(model::Union{Model, Type}, field::Union{Symbol,Expr};
                    values=nothing, lower=nothing, upper=nothing,
                    origin=nothing, unit=nothing, scale::D=nothing) where D
    all(==(nothing), [values, lower, upper, origin, unit]) &&
        throw(ArgumentError("You must specify at least one of these: "*
                            "values=..., lower=..., upper=..., origin=..., "*
                            "unit=..."))

    if model isa Model
        value = recursive_getproperty(model, field)
        T = typeof(value)
    else
        T = model
    end
    if T <: Real && values === nothing
        return numeric_range(T, D, field, lower, upper, origin, unit, scale)
    else
        return nominal_range(T, field, values)
    end
end

function numeric_range(T, D, field, lower, upper, origin, unit, scale)
    lower === Inf &&
        throw(ArgumentError("`lower` must be finite or `-Inf`."))
    upper === -Inf &&
        throw(ArgumentError("`upper` must be finite or `Inf`."))

    lower === nothing && (lower = -Inf)
    upper === nothing && (upper = Inf)

    lower < upper ||
        throw(ArgumentError("`lower` must be strictly less than `upper`."))

    is_unbounded = (lower === -Inf || upper === Inf)

    if origin === nothing
        is_unbounded &&
            throw(DomainError("For an unbounded range you must specify " *
                              "`origin=...` to define a centre.\nTo make " *
                              "the range bounded, specify finite " *
                              "`upper=...` and `lower=...`."))
        origin = (upper + lower)/2
    end
    if unit === nothing
        is_unbounded &&
            throw(DomainError("For an unbounded range you must specify " *
                              "`unit=...` to define a unit of scale.\nTo " *
                              "make the range bounded, specify finite " *
                              "`upper=...` and `lower=...`."))
        unit = (upper - lower)/2
    end
    unit > 0 || throw(DomainError("`unit` must be positive."))
    origin < upper && origin > lower ||
        throw(DomainError("`origin` must lie strictly between `lower` and " *
                          "`upper`."))
    if lower === -Inf
        if upper === Inf
            B = DoublyUnbounded
            scale === nothing && (scale = :linear)
        else
            B = LeftUnbounded
            scale === nothing && (scale = :logminus)
        end
    else
        if upper === Inf
            B = RightUnbounded
            scale === nothing && (scale = :log)
        else
            B = Bounded
            scale === nothing && (scale = :linear)
        end
    end
    lower isa Union{T, Float64} || (lower = convert(T, lower) )
    upper isa Union{T, Float64} || (upper = convert(T, upper) )
    scale isa Symbol && (D = Symbol)
    return NumericRange{T,B,D}(field, lower, upper, origin, unit, scale)
end

nominal_range(T, field, values) = throw(ArgumentError(
   "`$values` must be an instance of type `AbstractVector{<:$T}`."
    * (T <: Model ? "\n Perharps you forgot to instantiate model"
     * "as `$(T)()`" : "") ))

nominal_range(T, field, ::Nothing) = throw(ArgumentError(
"The inferred hyper-parameter type is $T, which is nominal. "*
"If this is true, you must specify values=... "*
"If this is false, specify the correct type as "*
"first argument of `range`, as in "*
"the example, "*
"`range(Int, :dummy, lower=1, upper=10)`. "  ))

function nominal_range(::Type{T}, field, values::AbstractVector{<:T}) where T
    return NominalRange{T,length(values)}(field, Tuple(values))
end

#specific def for T<:AbstractFloat(Allows conversion btw AbstractFloats and Signed types)
function nominal_range(::Type{T}, field,
        values::AbstractVector{<:Union{AbstractFloat,Signed}}) where T<: AbstractFloat
    return NominalRange{T,length(values)}(field, Tuple(values))
end

#specific def for T<:Signed (Allows conversion btw Signed types)
function nominal_range(::Type{T}, field,
               values::AbstractVector{<:Signed}) where T<: Signed
    return NominalRange{T,length(values)}(field, Tuple(values))
end


#
# This file is part of the DiscreteEvents.jl Julia package, MIT license
#
# Paul Bayer, 2020
#
# This is a Julia package for discrete event simulation
#
using DiscreteEvents

println("... basic tests: printing  ...")
resetClock!(𝐶)
println(𝐶)
println(𝐶, stdout, 1, 2, 3)
DiscreteEvents.prettyClock(true)
str = "Clock 1: state=:idle, t=0.0, Δt=0.01, prc:0\n  scheduled ev:0, cev:0, sampl:0\n"
@test repr(𝐶) == str
# DiscreteEvents.prettyClock(false)
# str = "Clock{Array{DiscreteEvents.ClockChannel,1}}(1, DiscreteEvents.ClockChannel[], DiscreteEvents.Idle(), 0.0, , 0.01, DiscreteEvents.Schedule(DataStructures.PriorityQueue{DiscreteEvents.DiscreteEvent,Float64,Base.Order.ForwardOrdering}(), DiscreteEvents.DiscreteCond[], DiscreteEvents.Sample[]), Dict{Any,Prc}(), Channel[], 0.0, 0.0, 0.0, 0, 0)"
# @test repr(𝐶) == str
DiscreteEvents.prettyClock(true)
@test tau() == 0

println("... basic tests: only events  ...")
ex1 = :(1+1)
ex2 = :(1+2)
e() = 123
f(a) = a+3
g(a) = a+4
h(a, b; c = 1, d = 2) = a + b + c + d
i(; a = 1, b = 2) = a + b
j(x) = x == :unknown

@test DiscreteEvents._evaluate(fun(e)) == 123
@test DiscreteEvents._evaluate(fun(f, 1)) == 4
@test DiscreteEvents._evaluate(fun(h, 1, 2, c=3, d=4)) == 10

a = 11; b = 12; c = 13; d = 14;
sf1 = fun(h, a, b, c=c, d=d)
sf2 = fun(h, :a, :b, c=:c, d=:d)
sf3 = fun(h, a, b, c=c, d=d)
sf4 = fun(h, :a, :b, c=:c, d=:d)
sf5 = fun(i, a=a, b=b)
sf6 = fun(event!, fun(f, a), 1, cid=2, spawn=true)
@test sf1() == 50
@test sf2() == 50
@test sf3() == 50
@test sf4() == 50
@test sf5() == 23
sf6()
ev = DiscreteEvents._nextevent(𝐶)
@test ev.ex() == 14
a = 21; b = 22; c = 23; d = 24;
@test sf1() == 50
@test sf2() == 90
@test fun(h, :a, 2, c=:c, d=4)() == 50
@test fun(j, :unknown)()
@test fun(<=, fun(tau), 1)()

@test fun(i, a=10, b=20)() == 30

# one expression
ev = DiscreteEvents.DiscreteEvent(:(1+1), 10.0, nothing, 1)
@test DiscreteEvents._evaluate(ev.ex) == 2
@test ev.t == 10

# two expressions
ev = DiscreteEvents.DiscreteEvent((:(1+1), :(1+2)), 15.0, nothing, 1)
@test DiscreteEvents._evaluate(ev.ex) == (2, 3)
@test ev.t == 15

# one fun
ev = DiscreteEvents.DiscreteEvent(fun(f, 1), 10.0, nothing, 1)
@test DiscreteEvents._evaluate(ev.ex) == 4

# two funs
ev = DiscreteEvents.DiscreteEvent((fun(f, 1), fun(g, 1)), 10.0, nothing, 1)
@test DiscreteEvents._evaluate(ev.ex) == (4, 5)

# expressions and funs mixed in a tuple
ev = DiscreteEvents.DiscreteEvent((:(1+1), fun(g,2), :(1+2), fun(f, 1)), 10.0, nothing, 1)
@test sum([DiscreteEvents._evaluate(ex) for ex in ev.ex if ex isa Function]) == 10
@test sum([eval(ex) for ex in ev.ex if isa(ex, Expr)]) == 5
@test DiscreteEvents._evaluate(ev.ex) == (2, 6, 3, 4)

@test DiscreteEvents._scale(0) == 1
@test DiscreteEvents._scale(pi*1e7) == 1e7

clk = Clock()  # set up clock without sampling
@test_warn "undefined transition" DiscreteEvents.step!(clk, clk.state, DiscreteEvents.Resume())
DiscreteEvents.init!(clk)
@test clk.state == DiscreteEvents.Idle()
@test tau(clk) == 0
clk = Clock(t0=100)
@test tau(clk) == 100
# @test_warn "nothing to _evaluate" incr!(clk)


using Cairo, Compose

function plot_network(graph::AbstractGraph, user_paths, locs_x, locs_y)
    used_edges = []
    used_by = Dict()

    colour_pal = [colorant"lightgrey", colorant"orange", colorant"lightslateblue", colorant"green"]

    # nodecolor = ["blue" for i in 1:nv(graph)]

    this_user = 0
    for paths in user_paths
        this_user += 1
        for path in paths
            for edge in path
                push!(used_edges, edge)
                used_by[(edge.src,edge.dst)] = this_user
            end
        end
    end

    colours = []
    widths = []
    for edge in edges(graph)
        if edge in used_edges
            push!(colours, used_by[(edge.src,edge.dst)] + 1)
            push!(widths, 5)
        else
            push!(colours, 1)
            push!(widths, 1)
        end
    end

    mygplot = gplot(graph, locs_x, locs_y, edgestrokec=colour_pal[colours],
    edgelinewidth=widths, arrowlengthfrac=0.04)#, layout=spring_layout)
    # Save to pdf
    draw(PDF("plots/network_drawing.pdf", 16cm, 16cm), mygplot)
end


struct ByKeyTrait end
const ByKey = ByKeyTrait()

struct ByValueTrait end
const ByValue = ByValueTrait()

"""
    index_by(a::AbstractIndex, i)

Returns `ByKey` if the index of `i` should be identified by searching the keys
and returns `ByValue` if the index of `i` should be identified directly from the
values.
"""
index_by(a::AbstractIndex{name,K}, i::Function) where {name,K} = ByKey

index_by(a::AbstractIndex{name,K}, i::K) where {name,K} = ByKey
index_by(a::AbstractIndex{name,K}, i::AbstractVector{K}) where {name,K} = ByKey

index_by(a::AbstractIndex{name,K}, i::K) where {name,K<:Integer} = ByKey
index_by(a::AbstractIndex{name,K}, i::AbstractVector{K}) where {name,K<:Integer} = ByKey

index_by(a::AbstractIndex{name,K}, i::I) where {name,K,I<:Integer} = ByValue
index_by(a::AbstractIndex{name,K}, i::AbstractVector{I}) where {name,K,I<:Integer} = ByValue

@inline @propagate_inbounds function Base.to_index(a::AbstractIndex, i)
    return _to_index(index_by(a, i), a, i)
end

# _to_index
# TODO find_all should be filter where possible
@propagate_inbounds function _to_index(b::ByKeyTrait, a, i::Function)
    return __to_index(a, i, find_all(i, keys(a)))
end
@propagate_inbounds function _to_index(b::ByKeyTrait, a, i::AbstractVector)
    return __to_index(a, i, find_all(in(i), keys(a)))
end
@propagate_inbounds function _to_index(b::ByKeyTrait, a, i)
    return __to_index(a, i, find_first(==(i), keys(a)))
end
@propagate_inbounds function _to_index(b::ByValueTrait, a, i::Function)
    return __to_index(a, i, find_all(i, values(a)))
end
@propagate_inbounds function _to_index(b::ByValueTrait, a, i)
    @boundscheck if !checkindex(Bool, values(a), i)
        throw(BoundsError(a, i))
    end
    return @inbounds getindex(values(a), i)
end
@propagate_inbounds function _to_index(b::ByValueTrait, a, i::AbstractVector)
    @boundscheck if !checkindex(Bool, values(a), i)
        throw(BoundsError(a, i))
    end
    return @inbounds unsafe_reindex(a, i)
end

# __to_index
@propagate_inbounds function __to_index(a, i, idx::T) where {T<:Union{Integer,Nothing}}
    @boundscheck if T <: Nothing
        throw(BoundsError(a, i))
    end
    return @inbounds getindex(values(a), idx)
end
@propagate_inbounds function __to_index(a, i, idx::AbstractVector{T}) where {T<:Union{Integer,Nothing}}
    @boundscheck if !(T<:Integer)
        throw(BoundsError(a, i))
    end
    return unsafe_reindex(a, idx)
end


struct IIF1MConstantCache{vecuType,rhsType,nl_rhsType} <: StochasticDiffEqConstantCache
  uhold::vecuType
  rhs::rhsType
  nl_rhs::nl_rhsType
end

@cache struct IIF1MCache{uType,vecuType,DiffCacheType,rhsType,nl_rhsType,rateType,rateNoiseType,rateNoiseCollectionType,NoiseTmpType} <: StochasticDiffEqMutableCache
  u::uType
  uprev::uType
  uhold::vecuType
  dual_cache::DiffCacheType
  tmp::uType
  rhs::rhsType
  nl_rhs::nl_rhsType
  rtmp1::rateType
  rtmp2::rateNoiseType
  rtmp3::rateNoiseCollectionType
  noise_tmp::NoiseTmpType
end

function alg_cache(alg::IIF1M,prob,u,ΔW,ΔZ,p,rate_prototype,noise_rate_prototype,jump_rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,f,t,dt,::Type{Val{false}})
  uhold = Vector{typeof(u)}(undef, 1)
  tmp = zero(rate_prototype)
  rhs = RHS_IIF1M_Scalar(f,tmp,t,t,p)
  nl_rhs = alg.nlsolve(Val{:init},rhs,uhold)
  IIF1MConstantCache(uhold,rhs,nl_rhs)
end

function alg_cache(alg::IIF1M,prob,u,ΔW,ΔZ,p,rate_prototype,noise_rate_prototype,jump_rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,f,t,dt,::Type{Val{true}})

  tmp = similar(u,axes(u)); rtmp1 = zero(rate_prototype)
  dual_cache = DiffCache(u,Val{determine_chunksize(u,get_chunksize(alg.nlsolve))})
  uhold = vec(u) # this makes uhold the same values as integrator.u
  rhs = RHS_IIF1(f,tmp,t,t,dual_cache,size(u),p)
  nl_rhs = alg.nlsolve(Val{:init},rhs,uhold)
  noise_tmp = tmp

  rtmp2 = zero(noise_rate_prototype)
  if is_diagonal_noise(prob)
    rtmp3 = rtmp2
  else
    rtmp3 = zero(rate_prototype)
  end
  IIF1MCache(u,uprev,uhold,dual_cache,tmp,rhs,nl_rhs,rtmp1,rtmp2,rtmp3,noise_tmp)
end

struct IIF2MConstantCache{vecuType,rhsType,nl_rhsType} <: StochasticDiffEqConstantCache
  uhold::vecuType
  rhs::rhsType
  nl_rhs::nl_rhsType
end

@cache struct IIF2MCache{uType,vecuType,DiffCacheType,rhsType,nl_rhsType,rateType,rateNoiseType,rateNoiseCollectionType,NoiseTmpType} <: StochasticDiffEqMutableCache
  u::uType
  uprev::uType
  uhold::vecuType
  dual_cache::DiffCacheType
  tmp::uType
  rhs::rhsType
  nl_rhs::nl_rhsType
  rtmp1::rateType
  rtmp2::rateNoiseType
  rtmp3::rateNoiseCollectionType
  noise_tmp::NoiseTmpType
end

function alg_cache(alg::IIF2M,prob,u,ΔW,ΔZ,p,rate_prototype,noise_rate_prototype,jump_rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,f,t,dt,::Type{Val{false}})
  uhold = Vector{typeof(u)}(undef, 1)
  tmp = zero(rate_prototype)
  rhs = RHS_IIF2M_Scalar(f,tmp,t,t,p)
  nl_rhs = alg.nlsolve(Val{:init},rhs,uhold)
  IIF2MConstantCache(uhold,rhs,nl_rhs)
end

function alg_cache(alg::IIF2M,prob,u,ΔW,ΔZ,p,rate_prototype,noise_rate_prototype,jump_rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,f,t,dt,::Type{Val{true}})

  tmp = similar(u,axes(u)); rtmp1 = zero(rate_prototype)
  dual_cache = DiffCache(u,Val{determine_chunksize(u,get_chunksize(alg.nlsolve))})
  uhold = vec(u) # this makes uhold the same values as integrator.u
  rhs = RHS_IIF2(f,tmp,t,t,dual_cache,size(u),p)
  nl_rhs = alg.nlsolve(Val{:init},rhs,uhold)
  noise_tmp = tmp

  rtmp2 = zero(noise_rate_prototype)
  if is_diagonal_noise(prob)
    rtmp3 = rtmp2
  else
    rtmp3 = zero(rate_prototype)
  end
  IIF2MCache(u,uprev,uhold,dual_cache,tmp,rhs,nl_rhs,rtmp1,rtmp2,rtmp3,noise_tmp)
end

struct IIF1MilConstantCache{vecuType,rhsType,nl_rhsType} <: StochasticDiffEqConstantCache
  uhold::vecuType
  rhs::rhsType
  nl_rhs::nl_rhsType
end

@cache struct IIF1MilCache{uType,vecuType,DiffCacheType,rhsType,nl_rhsType,rateType,rateNoiseType,rateNoiseCollectionType,NoiseTmpType} <: StochasticDiffEqMutableCache
  u::uType
  uprev::uType
  uhold::vecuType
  dual_cache::DiffCacheType
  tmp::uType
  rhs::rhsType
  nl_rhs::nl_rhsType
  rtmp1::rateType
  rtmp2::rateNoiseType
  rtmp3::rateNoiseCollectionType
  noise_tmp::NoiseTmpType
  gtmp::rateNoiseType
  gtmp2::rateNoiseType
end

function alg_cache(alg::IIF1Mil,prob,u,ΔW,ΔZ,p,rate_prototype,noise_rate_prototype,jump_rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,f,t,dt,::Type{Val{false}})
  uhold = Vector{typeof(u)}(undef, 1)
  C = zero(rate_prototype)
  rhs = RHS_IIF1M_Scalar(f,C,t,t)
  nl_rhs = alg.nlsolve(Val{:init},rhs,uhold)
  IIF1MilConstantCache(uhold,rhs,nl_rhs)
end

function alg_cache(alg::IIF1Mil,prob,u,ΔW,ΔZ,p,rate_prototype,noise_rate_prototype,jump_rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,f,t,dt,::Type{Val{true}})
  tmp = similar(u,axes(u)); rtmp1 = zero(rate_prototype)
  dual_cache = DiffCache(u,Val{determine_chunksize(u,get_chunksize(alg.nlsolve))})
  uhold = vec(u) # this makes uhold the same values as integrator.u
  rhs = RHS_IIF1(f,tmp,t,t,dual_cache,size(u))
  nl_rhs = alg.nlsolve(Val{:init},rhs,uhold)
  noise_tmp = zero(noise_rate_prototype)
  gtmp = zero(noise_rate_prototype); gtmp2 = zero(noise_rate_prototype)
  rtmp2 = zero(noise_rate_prototype)
  if is_diagonal_noise(prob)
    rtmp3 = rtmp2
  else
    rtmp3 = zero(rate_prototype)
  end
  IIF1MilCache(u,uprev,uhold,dual_cache,tmp,rhs,nl_rhs,rtmp1,rtmp2,rtmp3,noise_tmp,gtmp,gtmp2)
end


#!/usr/bin/env julia

using Drs
using Logging

include(joinpath(Pkg.dir("MathProgBase"),"test","linprog.jl"))
linprogtest(DrsMathProgSolver(; LogLevel=DEBUG))

include(joinpath(Pkg.dir("MathProgBase"),"test","linproginterface.jl"))
linprogsolvertest(DrsMathProgSolver())

include(joinpath(Pkg.dir("MathProgBase"),"test","conicinterface.jl"))
coniclineartest(DrsMathProgSolver())


"""
    is_canonical(::Type{I}) -> StaticBool

Returns `True` if instances of `I` can be used for indexing without any further change
(e.g., `Int`, `StaticInt`, `UnitRange{Int}`)
"""
is_canonical(x) = is_canonical(typeof(x))
is_canonical(::Type{T}) where {T} = static(false)
is_canonical(::Type{<:CanonicalInt}) = static(true)
is_canonical(::Type{StepRange{Int,Int}}) = static(true)
is_canonical(::Type{UnitRange{Int}}) = static(true)
is_canonical(::Type{OneTo{Int}}) = static(true)
is_canonical(::Type{<:OptionallyStaticRange}) = static(true)
is_canonical(::Type{Vector{Int}}) = static(true)
is_canonical(::Type{<:Slice}) = static(true)
@inline is_canonical(::Type{T}) where {N,T<:Tuple{Vararg{Any,N}}} = _is_canon(T, static(N))
_is_canon(::Type{T}, ::StaticInt{0}) where {T} = static(true)
@inline function _is_canon(::Type{T}, n::StaticInt{N}) where {T,N}
    return is_canonical(_get_tuple(T, n)) & _is_canon(T, n - static(1))
end

"""
    canonicalize(x)

Checks if `x` is in a canonical form for indexing. If `x` is already in a canonical form
then it is returned unchanged. If `x` is not in a canonical form then it is passed to
`canonical_convert`.
"""
canonicalize(x) = _canonicalize(is_canonical(x), x)
_canonicalize(::True, x) = x
_canonicalize(::False, x) = canonical_convert(x)

canonical_convert(x::Integer) = convert(Int, x)
function canonical_convert(x::AbstractRange)
    return OptionallyStaticStepRange(static_first(x), static_step(x), static_last(x))
end
function canonical_convert(x::AbstractUnitRange)
    return OptionallyStaticUnitRange(static_first(x), static_last(x))
end

is_linear_indexing(A, args::Tuple{Arg}) where {Arg} = ndims_index(Arg) < 2
is_linear_indexing(A, args::Tuple{Arg,Vararg{Any}}) where {Arg} = false

"""
    to_indices(A, inds::Tuple) -> Tuple

Maps indexing arguments `inds` to the axes of `A`, ensures they are converted to a native
indexing form, and that they are inbounds. Unless all indices in `inds` return `static(true)`
on a call to [`is_canonical`](@ref), then they each are checked at the axis level with
[`to_index`](@ref).
"""
@propagate_inbounds to_indices(A, ::Tuple{}) = to_indices(A, lazy_axes(A), ())
@propagate_inbounds to_indices(A, inds::Tuple) = _to_indices(is_canonical(inds), A, inds)
@propagate_inbounds function to_indices(A, inds::Tuple{LinearIndices})
    to_indices(A, lazy_axes(A), axes(getfield(inds, 1)))
end
@propagate_inbounds function _to_indices(::True, A, inds)
    if isone(sum(ndims_index(inds)))
        @boundscheck if !checkindex(Bool, eachindex(IndexLinear(), A), getfield(inds, 1))
            throw(BoundsError(A, inds))
        end
        return inds
    else
        @boundscheck if !Base.checkbounds_indices(Bool, lazy_axes(A), inds)
            throw(BoundsError(A, inds))
        end
        return inds
    end
end
@propagate_inbounds function _to_indices(::False, A, inds)
    if isone(sum(ndims_index(inds)))
        return (to_index(LazyAxis{:}(A), getfield(inds, 1)),)
    else
        return to_indices(A, lazy_axes(A), inds)
    end
end
@propagate_inbounds function to_indices(A, axs, inds::Tuple{<:AbstractCartesianIndex,Vararg{Any}})
    to_indices(A, axs, (Tuple(getfield(inds, 1))..., tail(inds)...))
end
@propagate_inbounds function to_indices(A, axs, inds::Tuple{I,Vararg{Any}}) where {I}
    _to_indices(ndims_index(I), A, axs, inds)
end

@propagate_inbounds function _to_indices(::StaticInt{1}, A, axs, inds)
    (to_index(_maybe_first(axs), getfield(inds, 1)),
     to_indices(A, _maybe_tail(axs), _maybe_tail(inds))...)
end

@propagate_inbounds function _to_indices(::StaticInt{N}, A, axs, inds) where {N}
    axsfront, axstail = Base.IteratorsMD.split(axs, Val(N))
    if IndexStyle(A) === IndexLinear()
        index = to_index(LinearIndices(axsfront), getfield(inds, 1))
    else
        index = to_index(CartesianIndices(axsfront), getfield(inds, 1))
    end
    return (index, to_indices(A, axstail, _maybe_tail(inds))...)
end
# When used as indices themselves, CartesianIndices can simply become its tuple of ranges
@propagate_inbounds function to_indices(A, axs, inds::Tuple{CartesianIndices, Vararg{Any}})
    to_indices(A, axs, (axes(getfield(inds, 1))..., tail(inds)...))
end
# but preserve CartesianIndices{0} as they consume a dimension.
@propagate_inbounds function to_indices(A, axs, inds::Tuple{CartesianIndices{0},Vararg{Any}})
    (getfield(inds, 1), to_indices(A, _maybe_tail(axs), tail(inds))...)
end
@propagate_inbounds function to_indices(A, axs, ::Tuple{})
    @boundscheck if length(getfield(axs, 1)) != 1
        error("Cannot drop dimension of size $(length(first(axs))).")
    end
    return to_indices(A, _maybe_tail(axs), ())
end
to_indices(A, ::Tuple{}, ::Tuple{}) = ()

# if there aren't anymore axes than we are using trailing dimensions of size one
_maybe_first(::Tuple{}) = static(1):static(1)
_maybe_first(x::Tuple) = getfield(x, 1)
_maybe_tail(::Tuple{}) = ()
_maybe_tail(x::Tuple) = tail(x)

"""
    to_index([::IndexStyle, ]axis, arg) -> index

Convert the argument `arg` that was originally passed to `getindex` for the dimension
corresponding to `axis` into a form for native indexing (`Int`, Vector{Int}, etc.). New
axis types with unique behavior should use an `IndexStyle` trait:

```julia
to_index(axis::MyAxisType, arg) = to_index(IndexStyle(axis), axis, arg)
to_index(::MyIndexStyle, axis, arg) = ...
```
"""
@propagate_inbounds to_index(axis, arg) = to_index(IndexStyle(axis), axis, arg)
function to_index(s, axis, arg)
    throw(ArgumentError("invalid index: IndexStyle $s does not support indices of " *
                        "type $(typeof(arg)) for instances of type $(typeof(axis))."))
end
to_index(::IndexLinear, axis, arg::Colon) = indices(axis) # Colons get converted to slices by `indices`
to_index(::IndexLinear, axis, arg::CartesianIndices{0}) = arg
to_index(::IndexLinear, axis, arg::CartesianIndices{1}) = axes(arg, 1)
@propagate_inbounds function to_index(::IndexLinear, axis, arg::AbstractCartesianIndex{1})
    return to_index(axis, first(Tuple(arg)))
end
@propagate_inbounds function to_index(::IndexLinear, x, arg::Union{Array{Bool}, BitArray})
    @boundscheck checkbounds(x, arg)
    return LogicalIndex{Int}(arg)
end
@propagate_inbounds function to_index(::IndexLinear, x, arg::AbstractArray{<:AbstractCartesianIndex})
    @boundscheck _multi_check_index(axes(x), arg) || throw(BoundsError(x, arg))
    return arg
end
@propagate_inbounds function to_index(::IndexLinear, x, arg::LogicalIndex)
    @boundscheck checkbounds(Bool, x, arg) || throw(BoundsError(x, arg))
    return arg
end
@propagate_inbounds function to_index(::IndexLinear, x, arg::Integer)
    @boundscheck checkindex(Bool, indices(x), arg) || throw(BoundsError(x, arg))
    return canonicalize(arg)
end
@propagate_inbounds function to_index(::IndexLinear, axis, arg::AbstractArray{Bool})
    @boundscheck checkbounds(axis, arg)
    return LogicalIndex(arg)
end
@propagate_inbounds function to_index(::IndexLinear, x, arg::AbstractArray{<:Integer})
    @boundscheck checkindex(Bool, x, arg) || throw(BoundsError(x, arg))
    return arg
end
@propagate_inbounds function to_index(::IndexLinear, x, arg::AbstractRange{Integer})
    @boundscheck checkindex(Bool, indices(axis), arg) || throw(BoundsError(axis, arg))
    return static_first(arg):static_step(arg):static_last(arg)
end

## IndexCartesian ##
to_index(::IndexCartesian, x, arg::Colon) = CartesianIndices(x)
to_index(::IndexCartesian, x, arg::CartesianIndices{0}) = arg
to_index(::IndexCartesian, x, arg::AbstractCartesianIndex) = arg
function to_index(::IndexCartesian, x, arg)
    @boundscheck _multi_check_index(axes(x), arg) || throw(BoundsError(x, arg))
    return arg
end
@propagate_inbounds function to_index(::IndexCartesian, x, arg::AbstractArray{<:AbstractCartesianIndex})
    @boundscheck _multi_check_index(axes(x), arg) || throw(BoundsError(x, arg))
    return arg
end
@propagate_inbounds function to_index(::IndexCartesian, x, arg::AbstractArray{Bool})
    @boundscheck checkbounds(x, arg)
    return LogicalIndex(arg)
end

function _multi_check_index(axs::Tuple, arg::AbstractArray{T}) where {T<:AbstractCartesianIndex}
    b = true
    for i in arg
        b &= Base.checkbounds_indices(Bool, axs, (i,))
    end
    return b
end

@propagate_inbounds function to_index(::IndexCartesian, x, arg::Union{Array{Bool}, BitArray})
    @boundscheck checkbounds(x, arg)
    return LogicalIndex{Int}(arg)
end

# TODO delete this once the layout interface is working
_array_index(::IndexLinear, a, i::CanonicalInt) = i
@inline function _array_index(::IndexStyle, a, i::CanonicalInt)
    CartesianIndices(ntuple(dim -> indices(a, dim), Val(ndims(a))))[i]
end
_array_index(::IndexLinear, a, i::AbstractCartesianIndex{1}) = getfield(Tuple(i), 1)
@inline function _array_index(::IndexLinear, a, i::AbstractCartesianIndex)
    N = ndims(a)
    StrideIndex{N,ntuple(+, Val(N)),nothing}(size_to_strides(size(a), static(1)), offsets(a))[i]
end
_array_index(::IndexStyle, a, i::AbstractCartesianIndex) = i

"""
    unsafe_reconstruct(A, data; kwargs...)

Reconstruct `A` given the values in `data`. New methods using `unsafe_reconstruct`
should only dispatch on `A`.
"""
function unsafe_reconstruct(axis::OneTo, data; kwargs...)
    if axis === data
        return axis
    else
        return OneTo(data)
    end
end
function unsafe_reconstruct(axis::UnitRange, data; kwargs...)
    if axis === data
        return axis
    else
        return UnitRange(first(data), last(data))
    end
end
function unsafe_reconstruct(axis::OptionallyStaticUnitRange, data; kwargs...)
    if axis === data
        return axis
    else
        return OptionallyStaticUnitRange(static_first(data), static_last(data))
    end
end
function unsafe_reconstruct(A::AbstractUnitRange, data; kwargs...)
    return static_first(data):static_last(data)
end

"""
    to_axes(A, inds) -> Tuple

Construct new axes given the corresponding `inds` constructed after
`to_indices(A, args) -> inds`. This method iterates through each pair of axes and
indices calling [`to_axis`](@ref).
"""
@inline function to_axes(A, inds::Tuple)
    if ndims(A) === 1
        return (to_axis(axes(A, 1), first(inds)),)
    elseif isone(sum(ndims_index(inds)))
        return (to_axis(eachindex(IndexLinear(), A), first(inds)),)
    else
        return to_axes(A, axes(A), inds)
    end
end
# drop this dimension
to_axes(A, a::Tuple, i::Tuple{<:Integer,Vararg{Any}}) = to_axes(A, tail(a), tail(i))
to_axes(A, a::Tuple, i::Tuple{I,Vararg{Any}}) where {I} = _to_axes(ndims_index(I), A, a, i)
function _to_axes(::StaticInt{1}, A, axs::Tuple, inds::Tuple)
    return (to_axis(first(axs), first(inds)), to_axes(A, tail(axs), tail(inds))...)
end
@propagate_inbounds function _to_axes(::StaticInt{N}, A, axs::Tuple, inds::Tuple) where {N}
    axes_front, axes_tail = Base.IteratorsMD.split(axs, Val(N))
    if IndexStyle(A) === IndexLinear()
        axis = to_axis(LinearIndices(axes_front), getfield(inds, 1))
    else
        axis = to_axis(CartesianIndices(axes_front), getfield(inds, 1))
    end
    return (axis, to_axes(A, axes_tail, tail(inds))...)
end
to_axes(A, ::Tuple{Ax,Vararg{Any}}, ::Tuple{}) where {Ax} = ()
to_axes(A, ::Tuple{}, ::Tuple{}) = ()

"""
    to_axis(old_axis, index) -> new_axis

Construct an `new_axis` for a newly constructed array that corresponds to the
previously executed `to_index(old_axis, arg) -> index`. `to_axis` assumes that
`index` has already been confirmed to be in bounds. The underlying indices of
`new_axis` begins at one and extends the length of `index` (i.e., one-based indexing).
"""
@inline function to_axis(axis, inds)
    if !can_change_size(axis) &&
       (known_length(inds) !== nothing && known_length(axis) === known_length(inds))
        return axis
    else
        return to_axis(IndexStyle(axis), axis, inds)
    end
end

# don't need to check size b/c slice means it's the entire axis
@inline function to_axis(axis, inds::Slice)
    if can_change_size(axis)
        return copy(axis)
    else
        return axis
    end
end
to_axis(S::IndexLinear, axis, inds) = StaticInt(1):static_length(inds)

################
### getindex ###
################
"""
    ArrayInterface.getindex(A, args...)

Retrieve the value(s) stored at the given key or index within a collection. Creating
another instance of `ArrayInterface.getindex` should only be done by overloading `A`.
Changing indexing based on a given argument from `args` should be done through,
[`to_index`](@ref), or [`to_axis`](@ref).
"""
@propagate_inbounds getindex(A, args...) = unsafe_getindex(A, to_indices(A, args)...)
@propagate_inbounds function getindex(A; kwargs...)
    return unsafe_getindex(A, to_indices(A, order_named_inds(dimnames(A), values(kwargs)))...)
end
@propagate_inbounds getindex(x::Tuple, i::Int) = getfield(x, i)
@propagate_inbounds getindex(x::Tuple, ::StaticInt{i}) where {i} = getfield(x, i)

## unsafe_getindex ##
function unsafe_getindex(a::A) where {A}
    parent_type(A) <: A && throw(MethodError(unsafe_getindex, (A,)))
    return unsafe_getindex(parent(a))
end
function unsafe_getindex(a::A, i::CanonicalInt) where {A}
    idx = _array_index(IndexStyle(A), a, i)
    if idx === i
        parent_type(A) <: A && throw(MethodError(unsafe_getindex, (A, i)))
        return unsafe_getindex(parent(a), i)
    else
        return unsafe_getindex(a, idx)
    end
end
function unsafe_getindex(a::A, i::AbstractCartesianIndex) where {A}
    idx = _array_index(IndexStyle(A), a, i)
    if idx === i
        parent_type(A) <: A && throw(MethodError(unsafe_getindex, (A, i)))
        return unsafe_getindex(parent(a), i)
    else
        return unsafe_getindex(a, idx)
    end
end
function unsafe_getindex(a, i::CanonicalInt, ii::Vararg{CanonicalInt})
    unsafe_getindex(a, NDIndex(i, ii...))
end
unsafe_getindex(a, i::Vararg{Any}) = unsafe_get_collection(a, i)

unsafe_getindex(A::Array) = Base.arrayref(false, A, 1)
unsafe_getindex(A::Array, i::CanonicalInt) = Base.arrayref(false, A, Int(i))

unsafe_getindex(A::LinearIndices, i::CanonicalInt) = Int(i)

unsafe_getindex(A::CartesianIndices, i::AbstractCartesianIndex) = CartesianIndex(i)

unsafe_getindex(A::SubArray, i::CanonicalInt) = @inbounds(A[i])
unsafe_getindex(A::SubArray, i::AbstractCartesianIndex) = @inbounds(A[i])

# This is based on Base._unsafe_getindex from https://github.com/JuliaLang/julia/blob/c5ede45829bf8eb09f2145bfd6f089459d77b2b1/base/multidimensional.jl#L755.
#=
    unsafe_get_collection(A, inds)

Returns a collection of `A` given `inds`. `inds` is assumed to have been bounds-checked.
=#
function unsafe_get_collection(A, inds)
    axs = to_axes(A, inds)
    dest = similar(A, axs)
    if map(Base.unsafe_length, axes(dest)) == map(Base.unsafe_length, axs)
        _unsafe_get_index!(dest, A, inds...) # usually a generated function, don't allow it to impact inference result
    else
        Base.throw_checksize_error(dest, axs)
    end
    return dest
end

function _generate_unsafe_get_index!_body(N::Int)
    quote
        Compat.@inline()
        D = eachindex(dest)
        Dy = iterate(D)
        @inbounds Base.Cartesian.@nloops $N j d -> I[d] begin
            # This condition is never hit, but at the moment
            # the optimizer is not clever enough to split the union without it
            Dy === nothing && return dest
            (idx, state) = Dy
            dest[idx] = unsafe_getindex(src, NDIndex(Base.Cartesian.@ntuple($N, j)))
            Dy = iterate(D, state)
        end
        return dest
    end
end
@generated function _unsafe_get_index!(dest, src, I::Vararg{Any,N}) where {N}
    return _generate_unsafe_get_index!_body(N)
end

_ints2range(x::Integer) = x:x
_ints2range(x::AbstractRange) = x
@inline function unsafe_get_collection(A::CartesianIndices{N}, inds) where {N}
    if (length(inds) === 1 && N > 1) || stride_preserving_index(typeof(inds)) === False()
        return Base._getindex(IndexStyle(A), A, inds...)
    else
        return CartesianIndices(to_axes(A, _ints2range.(inds)))
    end
end
@inline function unsafe_get_collection(A::LinearIndices{N}, inds) where {N}
    if isone(sum(ndims_index(inds)))
        return @inbounds(eachindex(A)[first(inds)])
    elseif stride_preserving_index(typeof(inds)) === True()
        return LinearIndices(to_axes(A, _ints2range.(inds)))
    else
        return Base._getindex(IndexStyle(A), A, inds...)
    end
end

#################
### setindex! ###
#################
"""
    ArrayInterface.setindex!(A, args...)

Store the given values at the given key or index within a collection.
"""
@propagate_inbounds function setindex!(A, val, args...)
    if can_setindex(A)
        return unsafe_setindex!(A, val, to_indices(A, args)...)
    else
        error("Instance of type $(typeof(A)) are not mutable and cannot change elements after construction.")
    end
end
@propagate_inbounds function setindex!(A, val; kwargs...)
    return unsafe_setindex!(A, val, to_indices(A, order_named_inds(dimnames(A), values(kwargs)))...)
end

## unsafe_setindex! ##
function unsafe_setindex!(a::A, v) where {A}
    parent_type(A) <: A && throw(MethodError(unsafe_setindex!, (A, v)))
    return unsafe_setindex!(parent(a), v)
end
function unsafe_setindex!(a::A, v, i::CanonicalInt) where {A}
    idx = _array_index(IndexStyle(A), a, i)
    if idx === i
        parent_type(A) <: A && throw(MethodError(unsafe_setindex!, (A, v, i)))
        return unsafe_setindex!(parent(a), v, i)
    else
        return unsafe_setindex!(a, v, idx)
    end
end
function unsafe_setindex!(a::A, v, i::AbstractCartesianIndex) where {A}
    idx = _array_index(IndexStyle(A), a, i)
    if idx === i
        parent_type(A) <: A && throw(MethodError(unsafe_setindex!, (A, v, i)))
        return unsafe_setindex!(parent(a), v, i)
    else
        return unsafe_setindex!(a, v, idx)
    end
end
function unsafe_setindex!(a, v, i::CanonicalInt, ii::Vararg{CanonicalInt})
    unsafe_setindex!(a, v, NDIndex(i, ii...))
end
function unsafe_setindex!(A::Array{T}, v) where {T}
    Base.arrayset(false, A, convert(T, v)::T, 1)
end
function unsafe_setindex!(A::Array{T}, v, i::CanonicalInt) where {T}
    return Base.arrayset(false, A, convert(T, v)::T, Int(i))
end

unsafe_setindex!(a, v, i::Vararg{Any}) = unsafe_set_collection!(a, v, i)

# This is based on Base._unsafe_setindex!.
#=
    unsafe_set_collection!(A, val, inds)

Sets `inds` of `A` to `val`. `inds` is assumed to have been bounds-checked.
=#
@inline unsafe_set_collection!(A, v, i) = _unsafe_setindex!(A, v, i...)

function _generate_unsafe_setindex!_body(N::Int)
    quote
        x′ = Base.unalias(A, x)
        Base.Cartesian.@nexprs $N d -> (I_d = Base.unalias(A, I[d]))
        idxlens = Base.Cartesian.@ncall $N Base.index_lengths I
        Base.Cartesian.@ncall $N Base.setindex_shape_check x′ (d -> idxlens[d])
        Xy = iterate(x′)
        @inbounds Base.Cartesian.@nloops $N i d->I_d begin
            # This is never reached, but serves as an assumption for
            # the optimizer that it does not need to emit error paths
            Xy === nothing && break
            (val, state) = Xy
            unsafe_setindex!(A, val, NDIndex(Base.Cartesian.@ntuple($N, i)))
            Xy = iterate(x′, state)
        end
        A
    end
end

@generated function _unsafe_setindex!(A, x, I::Vararg{Any,N}) where {N}
    return _generate_unsafe_setindex!_body(N)
end


#
module Spectral

using LinearAlgebra
using FastGaussQuadrature

abstract type AbstractSpectralSpace end

end # module


# ------------------------------------------------------------------
# Licensed under the ISC License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    EstimationSolution

A solution to a spatial estimation problem.
"""
struct EstimationSolution{𝒟,ℳ,𝒱}
  domain::𝒟
  mean::ℳ
  variance::𝒱
end

# -------------
# VARIABLE API
# -------------

Base.getindex(solution::EstimationSolution, var::Symbol) =
  (mean=solution.mean[var], variance=solution.variance[var])

# ------------
# IO methods
# ------------
function Base.show(io::IO, solution::EstimationSolution)
  N = ncoords(solution.domain)
  print(io, "$(N)D EstimationSolution")
end

function Base.show(io::IO, ::MIME"text/plain", solution::EstimationSolution)
  println(io, solution)
  println(io, "  domain: ", solution.domain)
  print(  io, "  variables: ", join(keys(solution.mean), ", ", " and "))
end


################################################# FILE DESCRIPTION #########################################################

# This file contains tests for the Vertex DataFrame.

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

@testset "Edge DataFrame" begin
   g = completegraph(10)

   eit = edges(g)

   float_arr = rand(90)
   int_arr = rand(Int, 90)
   bool_arr = rand(Bool, 90)
   str_arr = map(string, 1:90)
   void_arr = fill(nothing, 90)

   E = g.edata

   # Mutation on emtpy DataFrame
   addedge!(E)
   @test size(E) == (0,0)
   rmedge!(E, 11)
   @test size(E) == (0,0)
   rmedge!(E, 1:90)
   @test size(E) == (0,0)

   ###
   # SETVPROP
   ###
   # Try to create a property only on a subset of properties
   try seteprop!(g, eit[rand(1:90)], 1, :p1); @test false catch @test true end
   try seteprop!(g, eit[rand(1:90, rand(1:90))], 1, :p1); @test false catch @test true end
   try seteprop!(g, :, rand(rand(1:89)), :p1); @test false catch @test true end

   # Create new properties
   seteprop!(g, :, float_arr, :p1)
   seteprop!(g, :, int_arr, :p2)
   seteprop!(g, :, bool_arr, :p3)
   seteprop!(g, :, str_arr, :p4)
   seteprop!(g, :, nothing, :p5)

   # Test property names and types
   @test listeprops(g) == [:p1, :p2, :p3, :p4, :p5]
   @test eltypes(g.edata) == [Float64, Int, Bool, String, Void]

   # Test property values
   @test geteprop(g, :, :p1) == float_arr
   @test geteprop(g, :, :p2) == int_arr
   @test geteprop(g, :, :p3) == bool_arr
   @test geteprop(g, :, :p4) == str_arr
   @test geteprop(g, :, :p5) == void_arr

   # Mutation on DataFrame
   @test addedge!(g, 1=>1) == true
   @test g.indxs[1=>1] == 91
   for eprop in listeprops(g)
      isequal(geteprop(g, 1=>1, eprop), NA)
   end

   rmedge!(g, 1=>1)
   @test size(g.edata) == (90,5)

   # Modify single edge
   e = rand(eit)

   seteprop!(g, e, 0.5, :p1)
   @test geteprop(g, e, :p1) == 0.5

   seteprop!(g, e, "p4", :p4)
   @test geteprop(g, e, :p4) == "p4"

   # Modify a range of edges
   es = rand(eit, 20)
   seteprop!(g, es, 0, :p1)
   @test geteprop(g, es, :p1) == zeros(20)

   es = unique(rand(eit, 50))
   vals = 1 : length(es)
   seteprop!(g, es, vals, :p2)
   @test geteprop(g, es, :p2) == collect(vals)
end


"""
    EncodingDirection

Possible values: `ipos`, `jpos`, `kpos`, `ineg, `jneg`, `kneg` (the axis of the NIfTI data along which
slices were acquired, and the direction in which SliceTiming is defined with
respect to). `ipos`, `jpos`, `kpos` identifiers correspond to the first, second and third axis
of the data in the NIfTI file. `*neg` indicates that the contents of
SliceTiming are defined in reverse order - that is, the first entry corresponds
to the slice with the largest index, and the final entry corresponds to slice
index zero.
"""
@enum EncodingDirection begin
    ipos = 1
    ineg = -1
    jpos = 2
    jneg = -2
    kpos = 3
    kneg = -3
end

EncodingDirection(e::AbstractString) = EncodingDirection(Symbol(e))

function EncodingDirection(e::Symbol)
    if  (e === Symbol("i")) | (e === Symbol("ipos"))
        return ipos
    elseif (e === Symbol("i-")) | (e === Symbol("ineg"))
        return ineg
    elseif  (e === Symbol("j")) | (e === Symbol("jpos"))
        return jpos
    elseif (e === Symbol("j-")) | (e === Symbol("jneg"))
        return jneg
    elseif  (e === Symbol("k")) | (e === Symbol("kpos"))
        return kpos
    elseif (e === Symbol("k-")) | (e === Symbol("kneg"))
        return kneg
    else
        error("$e is not a supported encoding direction.")
    end
end
Base.String(e::EncodingDirection) = String(Symbol(e))

@inline function encoding_name(x::Int)
    if x === 1
        return :left
    elseif x === -1
        return :right
    elseif x === 2
        return :posterior
    elseif x === -2
        return :anterior
    elseif x === 3
        return :inferior
    elseif x === -3
        return :superior
    else
        error("$x does not map to a dimension name.")
    end
end

encoding_name(x) = encoding_name(Int(x))


using BeaData
using Test

@testset "metadata methods" begin
    ds = bea_datasets()
    @test "NIPA" ∈ ds[!, :DatasetName]

    pl = bea_parameterlist("FixedAssets")
    @test "TableName" ∈ pl[!, :ParameterName]

    pv = bea_parametervalues("NIUnderlyingDetail", "TableName")
    @test "U001A" ∈ pv[!, :Value]
end

@testset "data retrieval" begin
    nipa111 = bea_table("NIPA", "T10101", "A", 0 , 0)
    @test typeof(nipa111) <: BeaTable

    ud1a = bea_table("NIUnderlyingDetail", "U001A", "Q", 1996, 1996)
    @test typeof(ud1a) <: BeaTable

    faat101 = bea_table("FixedAssets", "FAAt101", "A", 0, 0)
    @test typeof(faat101) <: BeaTable
end


module PickAxisArrays
export PickAxisArray

struct PickAxisArray{T,N,AT<:AbstractArray,P,PERM} <: AbstractArray{T,N}
    parent::AT
end

struct Window
    i::Int
    pre::Int
    after::Int
end

function PickAxisArray(parent, indmask, perm = nothing)
    f = findall(isequal(true), indmask)
    f2 = findall(isequal(Colon()), indmask)
    f3 = findall(i -> isa(i, Tuple{Int,Int}), indmask)
    o = sort([f; f2; f3])
    o = isempty(f2) ? o : replace(o, map(i -> i => Colon(), f2)...)
    o = isempty(f3) ? o : replace(o, map(i -> i => Window(i, indmask[i]...), f3)...)
    nsub = 0
    for i = 1:length(o)
        if o[i] isa Colon
            nsub += 1
        elseif o[i] isa Window
            o[i] = Window(o[i].i - nsub, o[i].pre, o[i].after)
        else
            o[i] = o[i] - nsub
        end
    end
    if perm !== nothing
        length(perm) != length(f2) + length(f3) && error("Not a valid permutation")
        perm = (perm...,)
    end
    im = (o...,)
    T = anycol(im) ? Array{eltype(parent),ncol(im)} : eltype(parent)
    nd = length(f)
    PickAxisArray{T,nd,typeof(parent),im,perm}(parent)
end
indmask(::PickAxisArray{<:Any,<:Any,<:Any,i}) where {i} = i
getind(i, j) = i[j]
getind(i, j::Colon) = j
function getind(i, j::Window)
    c = i[j.i]
    c-j.pre:c+j.after
end
function Base.size(p::PickAxisArray)
    im = indmask(p)
    ii2 = filter(i->!isa(i,Union{Colon,Window}),im)
    permsize(p,map(i->size(p.parent,i),ii2))
end

permsize(::PickAxisArray{<:Any,<:Any,<:Any,<:Any,P}, x) where {P} = x[P]
permsize(::PickAxisArray{<:Any,<:Any,<:Any,<:Any,nothing}, x) = x
permout(::PickAxisArray{<:Any,<:Any,<:Any,<:Any,P}, x) where {P} = permutedims(x, P)
permout(::PickAxisArray{<:Any,<:Any,<:Any,<:Any,nothing}, x) = x
function Base.view(p::PickAxisArray, i::Integer...)
    inew = map(j -> getind(i, j), indmask(p))
    r = permout(p, view(p.parent, inew...))
    r
end
function Base.getindex(p::PickAxisArray, i::Integer...)
    inew = map(j -> getind(i, j), indmask(p))
    permout(p, getindex(p.parent, inew...))
end
anycol(::Tuple{}) = false
anycol(t::Tuple) = anycol(first(t), Base.tail(t))
anycol(::Union{Colon,Window}, t::Tuple) = true
anycol(i, ::Tuple{}) = false
anycol(::Union{Colon,Window}, ::Tuple{}) = true
anycol(i, t::Tuple) = anycol(first(t), Base.tail(t))

ncol(t::Tuple) = ncol(first(t), Base.tail(t), 0)
ncol(::Union{Colon,Window}, t::Tuple, n) = ncol(first(t), Base.tail(t), n + 1)
ncol(i, ::Tuple{}, n) = n
ncol(::Union{Colon,Window}, ::Tuple{}, n) = n + 1
ncol(i, t::Tuple, n) = ncol(first(t), Base.tail(t), n)

Base.getindex(p::PickAxisArray, i::CartesianIndex) = p[i.I...]
Base.view(p::PickAxisArray, i::CartesianIndex) = view(p,i.I...)
end


## PART 1

println("~~ QUESTION 4 ~~");
println("\nPart 1");

# Read in values from previous simulation (but only after thermalization):
println("\n\tReading in values from MD simulation...");

temps_1069 = readdlm("q4/temps-rho-0.75-temp-1.069.dat")[401:end]
temps_1304 = readdlm("q4/temps-rho-0.75-temp-1.304.dat")[401:end]

pe_1069 = readdlm("q4/pe-rho-0.75-temp-1.069.dat")[401:end]
pe_1304 = readdlm("q4/pe-rho-0.75-temp-1.304.dat")[401:end]

vir_1069 = readdlm("q4/vir-rho-0.75-temp-1.069.dat")[401:end]
vir_1304 = readdlm("q4/vir-rho-0.75-temp-1.304.dat")[401:end]

# Slap it all into a single array so that we can reuuse
#   old code
md_1069 = [temps_1069 pe_1069 vir_1069];
md_1304 = [temps_1304 pe_1304 vir_1304];


module Visualize

export visualize,
       external_visualize

using AbstractPlotting
using ..MDTrajAnalyse

const colours = Dict(
    "C"=>:black,
    " C"=>:black,
    "O"=>:red,
    " O"=>:red,
    "H"=>:white,
    " H"=>:white,
    "Ar"=>:orange,
    "N"=>:blue,
    " N"=>:blue
    )

const sizes = Dict(
    "C"=>0.9,
    " C"=>0.9,
    "O"=>0.9,
    " O"=>0.9,
    "H"=>0.6,
    " H"=>0.6,
    "Ar"=>1.0,
    "N"=>0.9,
    " N"=>0.9
    )

function visualize(traj::AbstractTrajectory, frame)
    acol = map( x -> colours[x], traj.names)
    asizes = map( x -> sizes[x], traj.names)
    sc = meshscatter(traj[frame]; markersize=asizes,color=acol)
    return sc
end

function visualize(traj::AbstractTrajectory; step=100)
    acol = map( x -> colours[x], traj.names)
    asizes = map( x -> sizes[x], traj.names)
    fig = Figure()
    ax = Axis3(fig[1,1])
    i = Slider(fig[2,1]; range=1:step:length(traj), startvalue=1)
    x = lift( j -> traj[j][1,:], i.value)
    y = lift( j -> traj[j][2,:], i.value)
    z = lift( j -> traj[j][3,:], i.value)
    meshscatter!(ax,x,y,z; markersize=asizes, color=acol)
    return fig
end

function external_visualize(traj::AbstractTrajectory;
    cmd="vlc", stdout=devnull, stderr=devnull, command="vmd"
    )
    fname= tempname()*".xyz"
    write_xyz(fname, traj)
    run(pipeline(`$(command) $(fname)`, stdout=stdout, stderr=stderr))
    rm(fname)
end

end # module


using skalman
using Test

@testset "skalman.jl" begin
    # Write your tests here.
    @test legendrepoly(0.0,0) == [1.0]

end


#!/usr/bin/env julia
#=
This is to use a fixed set of ID as reference, and then sample more and more ID
from the rest to mask, and then to impute them back. The purpose is to test what
happes if a lot more ID are masked.
=#

using DelimitedFiles, Plots, Statistics
dat = readdlm("num.txt")
dat = reshape(dat, 2, :)'
n = length(dat[:, 1])

x=50:100:1950
plot(x, dat[:, 1],
     xlabel="Masked set size",
     label="Error rates",
     dpi=300,
     ylabel="Error rates",
     leg=false)

savefig("fix-2-less2more.png")


using ACE1pack, Test, LazyArtifacts

test_train_set = joinpath(artifact"TiAl_tiny_dataset", "TiAl_tiny.xyz")

@info("Test constructing `data_params` and reading data")
params = data_params(fname = test_train_set, energy_key = "energy", force_key = "force", virial_key = "virial")
data = ACE1pack.read_data(params)


using NeuroCore, Test, Unitful

using NeuroCore: F64Sec, F64Tesla, F64Hz, IntDeg, OneF64Sec, OneF64Tesla, OneF64Hz, OneIntDeg
using NeuroCore: UnkownContrast, IODINE, ipos, jpos

include("bids_tests.jl")

#= TODO
    :stream_offset,
    :auxfiles,
    :srcfile,
    :calmax,
    :calmin,
    :freqdim,
    :phasedim,
    :slicedim,
    :slice_start,
=#


export PrecomputedInitialisation, RandomInitialisation, KMeansInitialisation, IncrementalInitialisation

mutable struct PrecomputedInitialisation <: InitialisationType
  Z::Array{Int}
  PrecomputedInitialisation(Z::Array{Int}) = new(Z)
end

mutable struct RandomInitialisation <: InitialisationType
  k::Int
  RandomInitialisation(;k = 2) = new(k)
end

mutable struct KMeansInitialisation <: InitialisationType
  k::Int
  maxiterations::Int
  KMeansInitialisation(;k = 2, maxiterations = 1000) = new(k, maxiterations)
end

mutable struct IncrementalInitialisation <: InitialisationType end


plot_function(plt::Function, grouped) = plt
plot_function(plt::Tuple, grouped) = grouped ? plt[2] : plt[1]

combine_cols(dict, ns) = length(ns) > 1 ? hcat((dict[n] for n in ns)...) : dict[ns[1]]

function dataviewer(t; throttle = 0.1, nbins = 30, nbins_range = 1:100)
    (t isa AbstractObservable) || (t = Observable{Any}(t))

    columns_and_names = @map create_columns_from_iterabletable(getiterator(&t))

    names = @map (&columns_and_names)[2]

    dict = @map Dict((key, convert_missing.(val)) for (val, key)  in zip((&columns_and_names)...))
    x =  Widgets.dropdown(names, placeholder = "First axis", multiple = true)
    y =  Widgets.dropdown(names, placeholder = "Second axis", multiple = true)
    y_toggle = Widgets.togglecontent(y, value = false, label = "Second axis")
    plot_type = Widgets.dropdown(
        OrderedDict(
            "line"         => Plots.plot,
            "scatter"      => Plots.scatter,
            "bar"          => (Plots.bar, StatsPlots.groupedbar),
            "boxplot"      => (StatsPlots.boxplot, StatsPlots.groupedboxplot),
            "corrplot"     => StatsPlots.corrplot,
            "cornerplot"   => StatsPlots.cornerplot,
            "density"      => StatsPlots.density,
            "cdensity"     => StatsPlots.cdensity,
            "histogram"    => StatsPlots.histogram,
            "marginalhist" => StatsPlots.marginalhist,
            "violin"       => (StatsPlots.violin, StatsPlots.groupedviolin),
        ),
        placeholder = "Plot type")

    # Add bins if the plot allows it
    display_nbins = @map (&plot_type) in [corrplot, cornerplot, histogram, marginalhist] ? "block" : "none"
    nbins =  (Widgets.slider(nbins_range, extra_obs = ["display" => display_nbins], value = nbins, label = "number of bins"))
    nbins.scope.dom = Widgets.div(nbins.scope.dom, attributes = Dict("data-bind" => "style: {display: display}"))
    nbins_throttle = Observables.throttle(throttle, nbins)

    by = Widgets.dropdown(names, multiple = true, placeholder="Group by")
    by_toggle = Widgets.togglecontent(by, value = false, label = "Split data")
    plt = Widgets.button("plot")
    output = @map begin
        if (&plt == 0)
            plot()
        else
            args = Any[]
            # add first and maybe second argument
            push!(args, combine_cols(&dict, x[]))
            has_y = y_toggle[] && !isempty(y[])
            has_y && push!(args, combine_cols(&dict, y[]))

            # compute automatic kwargs
            kwargs = Dict()

            # grouping kwarg
            has_by = by_toggle[] && !isempty(by[])
            by_tup = Tuple(getindex(&dict, b) for b in by[])
            has_by && (kwargs[:group] = NamedTuple{Tuple(by[])}(by_tup))

            # label kwarg
            if length(x[]) > 1
                kwargs[:label] = x[]
            elseif y_toggle[] && length(y[]) > 1
                kwargs[:label] = y[]
            end

            # x and y labels
            densityplot1D = plot_type[] in [cdensity, density, histogram]
            (length(x[]) == 1 && (densityplot1D || has_y)) && (kwargs[:xlabel] = x[][1])
            if has_y && length(y[]) == 1
                kwargs[:ylabel] = y[][1]
            elseif !has_y && !densityplot1D && length(x[]) == 1
                kwargs[:ylabel] = x[][1]
            end

            plot_func = plot_function(plot_type[], has_by)
            plot_func(args...; nbins = &nbins_throttle, kwargs...)
        end
    end
    wdg = Widget{:dataviewer}(["x" => x, "y" => y, "y_toggle" => y_toggle, "by" => by, "by_toggle" => by_toggle,
        "plot_type" => plot_type, "plot_button" => plt, "nbins" => nbins], output = output)
    @layout! wdg Widgets.div(
        Widgets.div(
            :x,
            :y_toggle,
            :plot_type,
            :by_toggle,
            :plot_button
        ),
        Widgets.div(style = Dict("width" => "3em")),
        Widgets.div(
            Widgets.observe(_),
            :nbins
        ),
        style = Dict("display" => "flex", "direction" => "row")
    )
end


function link_f(net::FastNet,c,rp)
    @inbounds linkpos=net.kcstart[c]+rp
    @inbounds net.kid[linkpos]
end

function link_f(net::FastNet,rp)
    @inbounds net.kid[rp]
end

function linkstate_f(net::FastNet,kid)
    @inbounds net.kstate[kid]
end

function makelink_f!(net::FastNet,src,dst)
    @inbounds begin
        kid=link_f(net,net.L,1)
        net.ksrc[kid]=src
        net.kdst[kid]=dst
        cls=net.ttable[net.nstate[src],net.nstate[dst]]
        linkstate_f!(net,kid,cls)
        net.knexts[kid]=net.nout[src]
        net.knextd[kid]=net.nin[dst]
        net.nout[src]=kid
        net.nin[dst]=kid
    end
    kid
end

function destroylink_f!(net::FastNet,kid)
    lastsrc=0
    lastdst=0
    @inbounds begin
        srcn=net.ksrc[kid]
        dstn=net.kdst[kid]
        cur=net.nout[srcn]
        while cur!=kid
            lastsrc=cur
            cur=net.knexts[cur]
        end
        cur=net.nin[dstn]
        while cur!=kid
            lastdst=cur
            cur=net.knextd[cur]
        end
        if lastsrc!=0
            net.knexts[lastsrc]=net.knexts[kid]
        else
            net.nout[srcn]=net.knexts[kid]
        end
        if lastdst!=0
            net.knextd[lastdst]=net.knextd[kid]
        else
            net.nin[dstn]=net.knextd[kid]
        end
        linkstate_f!(net,kid,net.L)
    end
    nothing
end

function linkexists_f(net::FastNet,kid)
    @inbounds net.kstate[kid]!=net.L
end

function linksrc_f(net::FastNet,kid)
    @inbounds net.ksrc[kid]
end

function linkdst_f(net::FastNet,kid)
    @inbounds net.kdst[kid]
end

function nextlinkout_f(net::FastNet,kid)
    @inbounds net.knexts[kid]
end

function nextlinkin_f(net::FastNet,kid)
    @inbounds net.knextd[kid]
end

function countlinks_f(net::FastNet)
    @inbounds net.K-net.kclen[net.L]
end

function countlinks_f(net::FastNet,cls::Integer)
    @inbounds net.kclen[cls]
end

function linkcounts_f(net::FastNet)
    n=net.L-2
    @inbounds net.kclen[1:n]
end

function countlinks_f(net::FastNet,cls::Union{Vector,Tuple})
    ret=0
    @inbounds begin
        for c in cls
            ret+=net.kclen[c]
        end
    end
    ret
end

function randomlink_f(net::FastNet)
    link_f(net,rand(net.rng,1:countlinks_f(net)))
end

function randomlink_f(net::FastNet,cls::Integer)
    link_f(net,cls,rand(net.rng,1:countlinks_f(net,cls)))
end

function randomlink_f(net::FastNet,cls::Union{Array,Tuple})
    tot=countlinks_f(net,cls)
    r=rand(net.rng,1:tot)
    old=r
    i=1
    cr=old-countlinks_f(net,cls[i])
    while cr>0
        i+=1
        old=cr
        cr=old-countlinks_f(net,cls[i])
    end
    link_f(net,cls[i],old)
end


"""
    HMCDA(n_iters::Int, n_adapt::Int, delta::Float64, lambda::Float64)

Hamiltonian Monte Carlo sampler wiht Dual Averaging algorithm.

Usage:

```julia
HMCDA(1000, 200, 0.65, 0.3)
```

Example:

```julia
# Define a simple Normal model with unknown mean and variance.
@model gdemo(x) = begin
  s ~ InverseGamma(2,3)
  m ~ Normal(0, sqrt(s))
  x[1] ~ Normal(m, sqrt(s))
  x[2] ~ Normal(m, sqrt(s))
  return s, m
end

sample(gdemo([1.5, 2]), HMCDA(1000, 200, 0.65, 0.3))
```
"""
mutable struct HMCDA{T} <: Hamiltonian
  n_iters   ::  Int       # number of samples
  n_adapt   ::  Int       # number of samples with adaption for epsilon
  delta     ::  Float64   # target accept rate
  lambda    ::  Float64   # target leapfrog length
  space     ::  Set{T}    # sampling space, emtpy means all
  gid       ::  Int       # group ID
end
function HMCDA(n_adapt::Int, delta::Float64, lambda::Float64, space...)
    _space = isa(space, Symbol) ? Set([space]) : Set(space)
    return HMCDA(1, n_adapt, delta, lambda, _space, 0)
end
function HMCDA(n_iters::Int, delta::Float64, lambda::Float64)
    n_adapt_default = Int(round(n_iters / 2))
    n_adapt = n_adapt_default > 1000 ? 1000 : n_adapt_default
    return HMCDA(n_iters, n_adapt, delta, lambda, Set(), 0)
end
function HMCDA(alg::HMCDA, new_gid::Int)
    return HMCDA(alg.n_iters, alg.n_adapt, alg.delta, alg.lambda, alg.space, new_gid)
end
HMCDA{T}(alg::HMCDA, new_gid::Int) where {T} = HMCDA(alg, new_gid)
function HMCDA(n_iters::Int, n_adapt::Int, delta::Float64, lambda::Float64)
    return HMCDA(n_iters, n_adapt, delta, lambda, Set(), 0)
end
function HMCDA(n_iters::Int, n_adapt::Int, delta::Float64, lambda::Float64, space...)
    _space = isa(space, Symbol) ? Set([space]) : Set(space)
    return HMCDA(n_iters, n_adapt, delta, lambda, _space, 0)
end
function HMCDA(n_iters::Int, n_adapt::Int, δ::Float64, λ::Float64, space::Set, gid::Int)
    return HMCDA{eltype(space)}(n_iters, n_adapt, δ, λ, space, gid)
end

function step(model, spl::Sampler{<:HMCDA}, vi::VarInfo, is_first::Bool)
    if is_first
        if ~haskey(spl.info, :wum)
            spl.alg.gid != 0 && link!(vi, spl)

            init_warm_up_params(vi, spl)

            θ = vi[spl]
            ϵ = spl.alg.delta > 0 ?
                find_good_eps(model, vi, spl) :       # heuristically find optimal ϵ
                spl.info[:pre_set_ϵ]
            vi[spl] = θ

            spl.alg.gid != 0 && invlink!(vi, spl)
            push!(spl.info[:wum][:ϵ], ϵ)
            update_da_μ(spl.info[:wum], ϵ)
        end

        push!(spl.info[:accept_his], true)
    else
        # Set parameters
        λ = spl.alg.lambda
        ϵ = spl.info[:wum][:ϵ][end]
        @debug "current ϵ: $ϵ"

        spl.info[:lf_num] = 0   # reset current lf num counter

        @debug "X-> R..."
        if spl.alg.gid != 0
            link!(vi, spl)
            runmodel!(model, vi, spl)
        end

        grad_func = gen_grad_func(vi, spl, model)
        lj_func = gen_lj_func(vi, spl, model)
        rev_func = gen_rev_func(vi, spl)
        log_func = gen_log_func(spl)

        θ, lj, stds = vi[spl], vi.logp, spl.info[:wum][:stds]

        θ_new, lj_new, is_accept, τ_valid, α = _hmc_step(
            θ, lj, lj_func, grad_func, ϵ, λ, stds; rev_func=rev_func, log_func=log_func)

        if PROGRESS[] && spl.alg.gid == 0
            stds_str = string(spl.info[:wum][:stds])
            stds_str = length(stds_str) >= 32 ? stds_str[1:30]*"..." : stds_str
            haskey(spl.info, :progress) && ProgressMeter.update!(
                spl.info[:progress],
                spl.info[:progress].counter;
                showvalues = [(:ϵ, ϵ), (:α, α), (:pre_cond, stds_str)],
            )
        end

        @debug "decide whether to accept..."
        if is_accept
            push!(spl.info[:accept_his], true)
            vi[spl] = θ_new
            setlogp!(vi, lj_new)
        else
            push!(spl.info[:accept_his], false)
            vi[spl] = θ
            setlogp!(vi, lj)
        end

        # QUES: why do we need the 2nd condition here (Kai)
        if spl.alg.delta > 0
            # TODO: figure out whether or not the condition below is needed
            # if spl.alg.delta > 0 && τ_valid > 0    # only do adaption for HMCDA
            adapt!(spl.info[:wum], α, vi[spl], adapt_ϵ = true)
        end

        @debug "R -> X..."
        spl.alg.gid != 0 && invlink!(vi, spl)
    end
    return vi
end

function _hmc_step(
    θ::AbstractVector{<:Real},
    lj::Real,
    lj_func,
    grad_func,
    ϵ::Real,
    λ::Real,
    stds::AbstractVector{<:Real};
    rev_func=nothing,
    log_func=nothing,
)

    θ_dim = length(θ)

    @debug "sampling momentums..."
    p = _sample_momentum(θ_dim, stds)

    @debug "recording old values..."
    H = _find_H(θ, p, lj, stds)

    τ = max(1, round(Int, λ / ϵ))
    @debug "leapfrog for $τ steps with step size $ϵ"
    θ_new, p_new, τ_valid = _leapfrog(θ, p, τ, ϵ, grad_func; rev_func=rev_func, log_func=log_func)

    @debug "computing new H..."
    lj_new = lj_func(θ_new)
    H_new = (τ_valid == 0) ? Inf : _find_H(θ_new, p_new, lj_new, stds)

    @debug "deciding wether to accept and computing accept rate α..."
    is_accept, logα = mh_accept(H, H_new)

    if is_accept
        θ = θ_new
        lj = lj_new
    end

    return θ, lj, is_accept, τ_valid, exp(logα)
end

function _hmc_step(
    θ::AbstractVector{<:Real},
    lj::Real,
    lj_func,
    grad_func,
    τ::Int,
    ϵ::Real,
    stdsstds::AbstractVector{<:Real},
)
    return _hmc_step(θ, lj, lj_func, grad_func, ϵ, τ * ϵ, stds)
end


"""
    format_table_value(width::Int, value; default_precision=2)

Format the table value using the given `width` and value.

For `val::Real`
- return a string representation that fits in the width and return "<<" if val is smaller than 0 but can't be represented or ">>" when it's bigger than 0.

For `val::Integer`
- return simply the string representation and `t.l.` if it doesn't fit

For all others:
- return simply the string representation and `t.l.` if it doesn't fit
"""
function format_table_value(width::Int, val::Real; default_precision=2)
    s_val = fmt("<.10f", val)
    precision = default_precision
    s_val_split = split(s_val, ".")
    if length(s_val_split[1]) == 1 && s_val_split[1] == "0" && length(s_val_split) == 2
        while precision < 10
            if s_val_split[2][precision - 1] != '0'
                precision -= precision > 2 ? 1 : 0
                break
            end
            precision += 1
        end
    end

    prec_fmt = FormatSpec("<.$(precision)f")
    s_val = fmt(prec_fmt, val)
    while length(s_val) > width && precision > 0
        precision -= 1
        prec_fmt = FormatSpec("<.$(precision)f")
        s_val = fmt(prec_fmt, val)
    end
    if length(s_val) > width
        s_val = val > 0 ? ">>" : "<<"
    end
    return s_val
end

function format_table_value(width::Int, value::Integer; default_precision=2)
    s_val = string(value)
    if length(s_val) > width
        s_val = "t.l."
    end
    return s_val
end

function format_table_value(width::Int, value; default_precision=2)
    s_val = string(value)
    if length(s_val) > width
        s_val = "t.l."
    end
    return s_val
end


"""
    `readrun(file)` reading running data.
"""
function readrun(file)
    data = readtable(file, header=false)
    inx = unique([find(x -> isna(x), data[1]);
                  find(x -> isna(x), data[2]);
                  find(x -> isna(x), data[3])])
    deleterows!(data, inx)
    date = DateTime(data[1])
    if length(find(x -> isna(x), data[4])) == 0
        ele = data[4]
    else
        ele = Vector{AbstractFloat}[]
    end
    if length(find(x -> isna(x), data[5])) == 0
        lap = data[5]
    else
        lap = Vector{Integer}[]
    end
    if length(find(x -> isna(x), data[6])) == 0
        hr = data[6]
    else
        hr = Vector{Integer}[]
    end
    run = Run(date, data[2], data[3], ele, lap, hr)
    return run
end

"""
    `readttwatch(file)` reading csv-data generated by ttwatch.

See: https://github.com/ryanbinns/ttwatch
"""
function readttwatch(file)
    data = readtable(file, header=false, skipstart=1)
    inx = unique([find(x -> isna(x), data[12]);
                  find(x -> isna(x), data[8]);
                  find(x -> isna(x), data[7])])
    deleterows!(data, inx)
    date = DateTime(data[12])
    if length(find(x -> isna(x), data[3])) == 0
        lap = data[3]
    else
        lap = Vector{Integer}[]
    end
    if length(find(x -> isna(x), data[9])) == 0
        ele = data[9]
    else
        ele = Vector{AbstractFloat}[]
    end
    if length(find(x -> isna(x), data[10])) == 0
        hr = data[10]
    else
        hr = Vector{Integer}[]
    end
    run = Run(lap, date, data[8], data[7], ele, hr)
    return run
end


# user functions

"""
    identitypose6fg()

Initialize a and exiting or new factor graph with a first pose and prior as specified by default
keywords or user.
"""
function identitypose6fg(;fg=nothing,
      N::Int=100,
      initCov::Array{Float64,2}=0.001*eye(6),
      initpose::SE3=SE3(0) )
  #
  fg = fg != nothing ? fg : initfg()
  N = 100

  println("Adding PriorPose3 and :x0 to graph...")
  # pts = 0.001*randn(6,N)
  initPosePrior = PriorPose3( MvNormal(veeEuler(initpose), initCov)  )
  pts = getSample(initPosePrior,N)[1]
  v0 = addNode!(fg, :x0,  pts,  N=N, labels=["POSE"])
  f1  = addFactor!(fg,[v0], initPosePrior)
  return fg
end

# function solveandvisualize(fg::FactorGraph,
#       vc::DrakeVisualizer.Visualizer;
#       refresh::Number=1.0,
#       densitymeshes::Vector{Symbol}=Symbol[],
#       drawlandms::Bool=true,
#       N::Int=100,
#       drawtype::Symbol=:max,
#       gt::Dict{Symbol, Tuple{Symbol,Vector{Float64}}}=Dict{Symbol, Tuple{Symbol,Vector{Float64}}}() )
#   #
#   deletemeshes!(vc)
#
#   # draw while solving[1]==true
#   solving = [true]
#   @async begin
#     while solving[1]
#       println(".")
#       visualizeallposes!(vc, fg, drawlandms=drawlandms, drawtype=drawtype, gt=gt)
#       i = 1
#       deletemeshes!(vc)
#       for dm in densitymeshes
#         i+=1
#         visualizeDensityMesh!(vc, fg, dm)
#       end
#       sleep(1)
#     end
#   end
#
#   # solve
#   tree = wipeBuildNewTree!(fg)
#   @time inferOverTree!(fg, tree, N=N)
#   solving[1]=false;
#   nothing
# end

function hasval(d::Dict, va)
  for (k,v) in d
    if v == va
      return true
    end
  end
  return false
end

function projectrbe(fgl::FactorGraph, from::Symbol, to::Vector{Float64})
  x1 = getKDEMax(getVertKDE(fgl,from))
  x2 = to

  wTc = SE3(x1[1:3],Euler(x1[4:6]...))

  cX2 = (matrix(wTc) \ [x2[1:3];1.0])[1:3]

  # Dx = x2[1:3]-x1[1:3]
  Dx = cX2
  range = norm(Dx)
  bearing = atan2(Dx[2],Dx[1])
  elev = -atan2(Dx[3],norm(Dx[1:2]))
  return range, bearing, elev
end

function projectrbe(fgl::FactorGraph, from::Symbol, to::Symbol)
  # x1 = getKDEMax(getVertKDE(fg,from))
  x2 = getKDEMax(getVertKDE(fgl,to))
  return projectrbe(fgl, from, x2[1:3])

  # wTc = SE3(x1[1:3],Euler(x1[4:6]...))
  #
  # cX2 = (matrix(wTc) \ [x2;1.0])[1:3]
  #
  # # Dx = x2[1:3]-x1[1:3]
  # Dx = cX2
  # range = norm(Dx)
  # bearing = atan2(Dx[2],Dx[1])
  # elev = -atan2(Dx[3],norm(Dx[1:2]))
  # return range, bearing, elev
end

#