text:
module softbart

using Distributions
using MultivariateStats
using SpecialFunctions
using StatsBase
using LinearAlgebra

mutable struct Node

    is_leaf::Bool
    is_root::Bool
    left::Node
    right::Node
    parent::Node

    ## branch parameters
    var::Int64
    val::Float64

    lower::Float64
    upper::Float64

    ## Leaf parameters
    mu::Float64

    ## Parameters for computing weights efficiently
    current_weight::Float64

    Node() = new()

end

mutable struct Hypers
    alpha::Float64
    beta::Float64
    gamma::Float64
    sigma::Float64
    sigma_mu::Float64
    shape::Float64
    width::Float64
    num_tree::Int64
    s::Vector{Float64}

    ## Double hypers
    sigma_hat::Float64
    sigma_mu_hat::Float64
    alpha_scale::Float64
    alpha_shape_1::Float64
    alpha_shape_2::Float64

    Hypers() = new()
end

mutable struct Opts
    num_burn::Int64
    num_thin::Int64
    num_save::Int64
    num_print::Int64

    update_sigma_mu::Bool
    update_s::Bool
    update_alpha::Bool
    update_beta::Bool
    update_gamma::Bool
end

mutable struct SoftbartOut
    y_hat_train::Array{Float64,2}
    y_hat_test::Array{Float64,2}
    y_hat_train_mean::Vector{Float64}
    y_hat_test_mean::Vector{Float64}
    sigma::Vector{Float64}
    sigma_mu::Vector{Float64}
    s::Array{Float64,2}
    alpha::Vector{Float64}
    beta::Vector{Float64}
    gamma::Vector{Float64}
end

function SoftbartOut(opts::Opts, n_train::Int64, n_test::Int64, hypers::Hypers)

    num_save = opts.num_save

    y_hat_train = zeros(num_save, n_train)
    y_hat_test = zeros(num_save, n_test)
    y_hat_train_mean = zeros(n_train)
    y_hat_test_mean = zeros(n_test)

    sigma = zeros(num_save)

    sigma_mu = opts.update_sigma_mu ? zeros(num_save) : hypers.sigma_mu * ones(1)
    s = opts.update_s ?
        zeros(num_save, length(hypers.s)) :
        reshape(hypers.s, 1, length(hypers.s))
    alpha = opts.update_alpha ? zeros(num_save) : hypers.alpha * ones(1)
    beta = opts.update_beta ? zeros(num_save) : hypers.beta * ones(1)
    gamma = opts.update_gamma ? zeros(num_save) : hypers.gamma * ones(1)

    return SoftbartOut(y_hat_train, y_hat_test, y_hat_train_mean,
                       y_hat_test_mean, sigma, sigma_mu, s, alpha, beta, gamma)
end

function Opts(;num_burn = 1000,
              num_thin = 1,
              num_save = 1000,
              num_print = 100,
              update_sigma_mu = true,
              update_s = true,
              update_alpha = true,
              update_beta = true,
              update_gamma = true)
    return Opts(num_burn, num_thin, num_save, num_print, update_sigma_mu,
                update_s, update_alpha, update_beta, update_gamma)
end

function Hypers(X::Array{Float64, 2},
                Y::Vector{Float64};
                alpha = 0.95,
                beta = 2.0,
                gamma = 0.95,
                k = 2.0,
                width = 0.1,
                shape = 1.0,
                num_tree = 50,
                alpha_scale = 0.0,
                alpha_shape_1 = 0.5,
                alpha_shape_2 = 1.0)

    out = Hypers()
    out.alpha = alpha
    out.beta = beta
    out.gamma = gamma
    out.sigma = GetSigma(X,Y)
    out.sigma_mu = 0.5 / (k * sqrt(num_tree))
    out.shape = shape
    out.width = width
    out.num_tree = num_tree
    out.s = ones(size(X,2)) / size(X,2)

    out.sigma_hat = out.sigma
    out.sigma_mu_hat = out.sigma_mu

    out.alpha_scale = alpha_scale == 0.0 ? size(X,2) : alpha_scale
    out.alpha_shape_1 = alpha_shape_1
    out.alpha_shape_2 = alpha_shape_2

    return out;

end

function Root()
    x = Node()

    x.is_leaf = true
    x.is_root = true
    x.left = x
    x.right = x
    x.parent = x

    x.var = 0
    x.val = 0.0
    x.lower = 0.0
    x.upper = 1.0
    x.mu = 0.0
    x.current_weight = 1.0

    return x
end

function InitRoot(X::Array{Float64,2})
    x = Root()
    return x
end

function GetSigma(X::Array{Float64, 2}, Y::Vector{Float64})
    if size(X,2) < size(X,1) / 5.0
        beta = llsq(X,Y)        # linear least square
        p = size(X,2)
        Yhat = X * beta[1:p] .+ beta[p+1]
        return sqrt(mean((Y - Yhat).^2))
    else
        return sqrt(var(Y))
    end
end

function AddLeaves!(x::Node)
    x.left = Node()
    x.right = Node()
    x.is_leaf = false

    ## Initialize the leafs
    x.left.is_leaf = true
    x.left.parent = x
    x.left.var = 0
    x.left.val = 0.0
    x.left.is_root = false
    x.left.lower = 0.0
    x.left.upper = 1.0
    x.left.mu = 0.0
    x.left.current_weight = 0.0
    x.right.is_leaf = true
    x.right.parent = x
    x.right.var = 0
    x.right.val = 0.0
    x.right.is_root = false
    x.right.lower = 0.0
    x.right.upper = 1.0
    x.right.mu = 0.0
    x.right.current_weight = 0.0

    return nothing
end

function BirthLeaves!(x::Node, hypers::Hypers)
    if !x.is_leaf
        error("Node must be a leaf")
    end
    AddLeaves!(x)
    wv = ProbabilityWeights(hypers.s)
    x.var = sample(wv)
    # x.var = rand(1:P) ## TODO change this to allow non-iid selection
    lower, upper = GetLimits(x)
    x.lower = lower
    x.upper = upper
    x.val = (upper - lower) * rand() + lower
end

function GenTree(hypers::Hypers, P::Int64)
    x = Root()
    GenTree!(x, hypers, P)
    return x
end

function GenTree!(x::Node, hypers::Hypers, P::Int64)
    grow_prob = SplitProb(x, hypers)
    u = rand()
    if u < grow_prob
        BirthLeaves!(x, hypers)
        GenTree!(x.left, hypers, P)
        GenTree!(x.right, hypers, P)
    end
end

function SplitProb(n::Node, hypers::Hypers)
    d = depth(n)
    grow_prob = hypers.gamma / (1 + d)^(hypers.beta)
    return grow_prob
end

function GetLimits(x::Node)
    y = x
    lower = 0.0
    upper = 1.0
    my_bool = y.is_root ? false : true
    while my_bool
        is_left = (y.parent.left == y)
        y = y.parent
        my_bool = y.is_root ? false : true
        if y.var == x.var
            my_bool = false
            if is_left
                upper = y.val
                lower = y.lower
            else
                upper = y.upper
                lower = y.val
            end
        end
    end
    return lower, upper
end

function GetSuffStats(n::Node,
                      y::Vector{Float64},
                      x::Array{Float64,2},
                      hypers::Hypers)
    leafs = leaves(n)
    num_leaves = length(leafs)
    w_i = zeros(num_leaves)
    mu_hat = zeros(num_leaves)
    Lambda = zeros(num_leaves, num_leaves)
    for i in 1:size(x,1)
        GetW!(n, x, i, hypers)
        for j in 1:num_leaves
            w_i[j] = leafs[j].current_weight
        end
        mu_hat += y[i] * w_i
        Lambda += w_i * transpose(w_i)
    end

    Lambda /= hypers.sigma^2
    mu_hat /= hypers.sigma^2

    Omega_inv = Lambda + LinearAlgebra.I  / hypers.sigma_mu^2

    mu_hat = Omega_inv\mu_hat  # TODO: what does this symbol "\" represent???
    # \(x, y) Left division operator: multiplication of y by the inverse of x on the left.
    return mu_hat, Omega_inv

end

function LogLT(n::Node, Y::Vector{Float64}, X::Array{Float64,2}, hypers::Hypers)

    leafs = leaves(n)
    num_leaves = length(leafs)

    ## Get sufficient statistics
    mu_hat, Omega_inv = GetSuffStats(n, Y, X, hypers)

    N = size(Y,1)

    out = -0.5 * N * log(2.0 * pi * hypers.sigma^2)
    out -= 0.5 * num_leaves * log(2.0 * pi * hypers.sigma_mu^2)
    out -= 0.5 * logdet(Omega_inv / (2.0 * pi))
    out -= 0.5 * sum(Y.^2) / hypers.sigma^2
    out += 0.5 * dot(mu_hat, Omega_inv * mu_hat)

    return out

end

function cauchy_jacobian(tau::Float64, sigma_hat::Float64)
    sigma = tau^(-0.5)
    out = logpdf(Cauchy(0,sigma_hat), sigma)

    out = out + log(0.5) - 3.0 / 2.0 * log(tau)

    return out

end

function update_sigma(r::Vector{Float64}, sigma_hat::Float64, sigma_old::Float64)
    SSE = sum(r .* r)
    n = length(r)               # Gamma(shape, scale)
    sigma_prop = rand(Gamma(0.5 * n + 1.0, 2.0 / SSE))^(-0.5)
    tau_prop = sigma_prop^(-2)
    tau_old = sigma_old^(-2)

    loglik_rat = cauchy_jacobian(tau_prop, sigma_hat) -
        cauchy_jacobian(tau_old, sigma_hat)

    return log(rand()) < loglik_rat ? sigma_prop : sigma_old

end

function UpdateSigma!(hypers::Hypers, r::Vector{Float64})
    hypers.sigma = update_sigma(r, hypers.sigma_hat, hypers.sigma)
end

function UpdateSigmaMu!(hypers::Hypers, means::Vector{Float64})
    hypers.sigma_mu = update_sigma(means, hypers.sigma_mu_hat, hypers.sigma_mu)
end

function UpdateMu!(n::Node,
                   Y::Vector{Float64},
                   X::Array{Float64,2},
                   hypers::Hypers)

    leafs = leaves(n)
    num_leaves = length(leafs)

    ## Get mean and covariance
    mu_hat, Omega_inv = GetSuffStats(n,Y,X,hypers)

    ## Update
    mu = mu_hat .+ rand(MvNormalCanon(Omega_inv),1) # sample 1 vector from the distribution, each column is a sample

    # for Multinomial distributin, the canonical parameters:
    # h = Σ^-1 × μ
    # J = Σ^-1
    # MvNormalCanon(J):  Construct a multivariate normal distribution with zero mean (thus zero potential vector) and precision matrix represented by J.
    for i in 1:num_leaves
        leafs[i].mu = mu[i]
    end

end

function predict(forest::Vector{Node},
                 X::Array{Float64,2},
                 hypers::Hypers)

    out = zeros(size(X,1))
    T = length(forest)
    for t in 1:T
        out += predict(forest[t], X, hypers)
    end

    return out

end

function predict(n::Node,
                 X::Array{Float64,2},
                 hypers::Hypers)

    leafs = leaves(n)
    num_leaves = length(leafs)
    N = size(X,1)
    out = zeros(N)

    for i in 1:N
        GetW!(n, X, i, hypers)
        for j in 1:num_leaves
            out[i] += leafs[j].current_weight * leafs[j].mu
        end
    end

    return out

end

function GetW!(n::Node, x::Array{Float64, 2}, i::Int64, hypers::Hypers)

    if !n.is_leaf

        weight = activation(x[i, n.var], n.val, hypers)
        n.left.current_weight = weight * n.current_weight
        n.right.current_weight = (1 - weight) * n.current_weight

        GetW!(n.left, x, i, hypers)
        GetW!(n.right, x, i, hypers)

    end

end

function is_left(n::Node)
    if n == n.parent.left
        return true
    else
        return false
    end
end

function DeleteLeaves!(x::Node)
    x.left = x
    x.right = x
    x.is_leaf = true
    return nothing
end

function depth(x::Node)
    x.is_root ? 0 : 1 + depth(x.parent)
end

function PrintTree(x::Node)
    d = depth(x)
    for i in 0:d
        print("    ")
    end
    print(x.var)
    print("\n")
    if !x.is_leaf
        PrintTree(x.right)
        PrintTree(x.left)
    end
end

function leaves!(x::Node, leaves::Vector{Node})
    if x.is_leaf
        push!(leaves, x);
    else
        leaves!(x.left, leaves);
        leaves!(x.right, leaves);
    end
end

function leaves(x::Node)
    leaves = Node[];
    leaves!(x, leaves);
    return leaves;
end

function get_means(forest::Vector{Node})
    means = Float64[]
    for tree in forest
        get_means!(tree, means)
    end
    return means
end

function get_means!(node::Node, means::Vector{Float64})

    if node.is_leaf
        push!(means, node.mu)
    else
        get_means!(node.left, means)
        get_means!(node.right, means)
    end
end

function test_tree()
    x = Root()
    AddLeaves!(x)
    AddLeaves!(x.left)
    AddLeaves!(x.left.left)
    AddLeaves!(x.left.right)

    x.var = 1
    x.left.var = 2
    x.right.var = 0
    x.left.left.var = 3
    x.left.right.var = 4
    x.left.left.left.var = 0
    x.left.left.right.var = 0
    x.left.right.left.var = 0
    x.left.right.right.var = 0
    return x
end

function test_tree_2(P::Int64, hypers::Hypers)
    x = Root()
    BirthLeaves!(x, hypers)
    BirthLeaves!(x.left, hypers)
    BirthLeaves!(x.left.left, hypers)
    BirthLeaves!(x.left.right, hypers)

    return x
end

function init_forest(X::Array{Float64, 2}, Y::Vector{Float64}, hypers)
    forest = Node[]
    for t in 1:hypers.num_tree
        n = InitRoot(X)
        push!(forest, n)
    end
    return forest
end

function soft_bart(
                        X::Array{Float64, 2},
                        Y::Vector{Float64},
                        X_test::Array{Float64,2},
                        hypers::Hypers,
                        opts::Opts)

    a = minimum(Y)
    b = maximum(Y)
    Y_u = (Y .- a) ./ (b - a) .- 0.5
    hypers_u = deepcopy(hypers)
    hypers_u.sigma /= (b-a)
    hypers_u.sigma_hat /= (b-a)

    out = soft_bart_norm(X,Y_u,X_test,hypers_u,opts)

    ## Unnormalize
    out.y_hat_train = (b - a) * (out.y_hat_train .+ 0.5) .+ a
    out.y_hat_test = (b - a) * (out.y_hat_test .+ 0.5) .+ a
    out.sigma = (b - a) * out.sigma
    out.y_hat_train_mean = (b - a) * (out.y_hat_train_mean .+ 0.5) .+ a
    out.y_hat_test_mean = (b - a) * (out.y_hat_test_mean .+ 0.5) .+ a

    # return (b-a) * (y_hat + 0.5) + a, (b-a) * sigma, sigma_mu, forest, hypers_u

    return out

end

function soft_bart_norm(
                   X::Array{Float64,2},
                   Y::Vector{Float64},
                   X_test::Array{Float64,2},
                   hypers::Hypers,
                   opts::Opts)

    MH_BD = 0.7

    forest = init_forest(X,Y,hypers)

    Y_hat = zeros(size(X,1))

    ## Output
    # Y_hat_test = zeros(opts.num_save, size(X_test, 1))
    # sigma = zeros(opts.num_save)
    # sigma_mu = zeros(opts.num_save)
    out = SoftbartOut(opts, size(X,1), size(X_test,1), hypers)

    for i in 1:opts.num_burn
        for t in 1:hypers.num_tree
            Y_star = Y_hat - predict(forest[t], X, hypers)
            res = Y - Y_star

            if forest[t].is_leaf
                birth_death!(forest[t], X, res, hypers)
            elseif rand() < MH_BD
                birth_death!(forest[t], X, res, hypers)
            else
                change_decision_rule!(forest[t], X, res, hypers)
            end

            UpdateMu!(forest[t],res,X,hypers)
            Y_hat = Y_star .+ predict(forest[t], X, hypers)
        end

        res = Y - Y_hat

        ## Update sigma
        UpdateSigma!(hypers, res)

        ## Update sigma_mu
        if opts.update_sigma_mu
            means = get_means(forest)
            UpdateSigmaMu!(hypers, means)
        end

        ## update s and alpha
        if (i > opts.num_burn / 2.0 && opts.update_s)
            UpdateS!(forest, hypers)
            if(opts.update_alpha)
                UpdateAlpha!(hypers)
            end
        end

        ## Update beta and gamma
        if opts.update_beta
            UpdateBeta!(forest, hypers)
        end
        if opts.update_gamma
            UpdateGamma!(forest, hypers)
        end

        if(i % opts.num_print == 0)
            print("Finishing burn", i, "\n")
        end
    end

    for i in 1:opts.num_save
        for b in 1:opts.num_thin
            for t in 1:hypers.num_tree
                Y_star = Y_hat - predict(forest[t], X, hypers)
                res = Y - Y_star

                if forest[t].is_leaf
                    birth_death!(forest[t], X, res, hypers)
                elseif rand() < MH_BD
                    birth_death!(forest[t], X, res, hypers)
                else
                    change_decision_rule!(forest[t], X, res, hypers)
                end
                UpdateMu!(forest[t],res,X,hypers)
                Y_hat = Y_star .+ predict(forest[t], X, hypers)
            end

            res = Y - Y_hat
            UpdateSigma!(hypers, res)

            ## Update sigma
            UpdateSigma!(hypers, res)

            ## Update sigma_mu
            if opts.update_sigma_mu
                means = get_means(forest)
                UpdateSigmaMu!(hypers, means)
            end

            ## update s and alpha
            if (opts.update_s)
                UpdateS!(forest, hypers)
                if(opts.update_alpha)
                    UpdateAlpha!(hypers)
                end
            end

            ## Update beta and gamma
            if opts.update_beta
                UpdateBeta!(forest, hypers)
            end
            if opts.update_gamma
                UpdateGamma!(forest, hypers)
            end

        end

        ## Record predictions
        out.y_hat_test[i,:] = predict(forest,X_test,hypers)
        out.y_hat_train[i,:] = Y_hat

        ## Record sigma
        out.sigma[i] = hypers.sigma

        ## Record hypers
        if opts.update_sigma_mu
            out.sigma_mu[i] = hypers.sigma_mu
        end
        if opts.update_s
            out.s[i,:] = hypers.s
        end
        if opts.update_alpha
            out.alpha[i] = hypers.alpha
        end
        if opts.update_beta
            out.beta[i] = hypers.beta
        end
        if opts.update_gamma
            out.gamma[i] = hypers.gamma
        end

        out.sigma_mu[i] = hypers.sigma_mu
        out.sigma[i] = hypers.sigma
        if (i % opts.num_print == 0)
            print("Finishing save", i, "\n")
        end
    end

    out.y_hat_train_mean = vec(mean(out.y_hat_train, dims = (1)))
    out.y_hat_test_mean = vec(mean(out.y_hat_test, dims = (1)))

    return out
end

function fried(X::Array{Float64, 2})
    n = size(X,1)
    p = size(X,2)

    Y = zeros(n)
    for i in 1:n
        Y[i] = 10 * sin(2 * pi * X[i,1] .* X[i,2]) .+
            20 * (X[i,3] - 0.5).^2 .+ 10 * X[i,4] .+ 5 * X[i,5]
    end
    return Y
end

function logit(x::Float64)
    return log(x) - log(1 - x)
end

function normalize(y::Vector{Float64})
    a = minimum(y)
    b = maximum(y)
    return (y - a) / (b - a) - 0.5
end

function expit(x::Float64)
    return 1.0 / (1 + exp(-x))
end

function activation(x::Float64, c::Float64, hypers::Hypers)
    return 1- expit((x - c) / hypers.width)
end

function birth_death!(tree::Node,
                     X::Array{Float64, 2},
                     Y::Vector{Float64},
                     hypers::Hypers)

    u = rand()
    p_birth = probability_node_birth(tree)
    if u < p_birth
        node_birth!(tree::Node,
                    X::Array{Float64, 2},
                    Y::Vector{Float64},
                    hypers::Hypers)
    else
        node_death!(tree::Node,
                    X::Array{Float64,2},
                    Y::Vector{Float64},
                    hypers::Hypers)
    end
end

function node_birth!(tree::Node,
                     X::Array{Float64,2},
                     Y::Vector{Float64},
                     hypers::Hypers)

    leaf, leaf_probability = birth_node(tree)

    leaf_depth = depth(leaf)
    leaf_prior = growth_prior(leaf, leaf_depth, hypers)

    ## Get likelihood before
    ll_before = LogLT(tree, Y, X, hypers)
    ll_before += log(1.0 - leaf_prior)

    ## Get transition probability
    p_forward = log(probability_node_birth(tree) * leaf_probability)

    ## Birth new leaves
    BirthLeaves!(leaf, hypers) ## TODO Change selection of feature

    ## Get likelihood after
    ll_after = LogLT(tree, Y, X, hypers)
    ll_after += log(leaf_prior) +
        log(1.0 - growth_prior(leaf.left, leaf_depth + 1, hypers)) +
        log(1.0 - growth_prior(leaf.right, leaf_depth + 1, hypers))

    ## Get probability of reverse transition: TODO not sure if I need to account
    ## for the choice of splitting value; I don't think I need to, but not sure
    num_not_grand_branches = length(not_grand_branches(tree))
    p_not_grand = 1.0/num_not_grand_branches
    p_backward = log((1 - probability_node_birth(tree)) * p_not_grand)

    transition_prob = ll_after + p_backward - ll_before - p_forward

    if log(rand()) > transition_prob
        DeleteLeaves!(leaf)
        leaf.var = 0
    end

end

function node_death!(tree::Node,
                     X::Array{Float64,2},
                     Y::Vector{Float64},
                     hypers::Hypers)

    ## Select branch to kill children
    branch, p_not_grand = death_node(tree)

    ## Compute before likelihood
    leaf_depth = depth(branch.left)
    leaf_prob = growth_prior(branch, leaf_depth - 1, hypers)
    left_prior = growth_prior(branch.left, leaf_depth, hypers)
    right_prior = left_prior
    ll_before = LogLT(tree, Y, X, hypers) +
        log(1.0 - left_prior) + log(1.0 - right_prior) + log(leaf_prob)

    ## Compute forward transition prob
    p_forward = log(p_not_grand * (1.0 - probability_node_birth(tree)))

    ## Save old leafs
    left = branch.left
    right = branch.right

    ## Compute after likelihood
    DeleteLeaves!(branch)

    ll_after = LogLT(tree, Y, X, hypers) + log(1.0 - leaf_prob)

    ## Compute backwards transition
    p_backwards = log(1.0 / length(leaves(tree)) * probability_node_birth(tree))

    ## Do MH
    transition_prob = ll_after + p_backwards - ll_before - p_forward
    if log(rand()) > transition_prob
        branch.left = left
        branch.right = right
        branch.is_leaf = false
    end

end

function change_decision_rule!(tree::Node,
                               X::Array{Float64, 2},
                               Y::Vector{Float64},
                               hypers::Hypers)

    ## Select a node to change
    ngb = not_grand_branches(tree)
    branch = rand(ngb)

    ## Calculate likelihood before proposal
    ll_before = LogLT(tree, Y, X, hypers)
    # ll_before = ll_before

    ## save old split
    old_feature = branch.var
    old_value   = branch.val
    old_lower = branch.lower
    old_upper = branch.upper

    # branch.var = rand(1:size(X,2)) ## TODO change for sparsity
    branch.var = sample(ProbabilityWeights(hypers.s))
    lower, upper = GetLimits(branch)
    branch.lower = lower
    branch.upper = upper
    branch.val = (upper - lower) * rand() + lower

    ll_after = LogLT(tree, Y, X, hypers)

    log_trans_prob = ll_after - ll_before

    if log(rand()) > log_trans_prob
        branch.var = old_feature
        branch.val = old_value
        branch.lower = old_lower
        branch.upper = old_upper
    end

end

function growth_prior(leaf::Node, leaf_depth::Int64, hypers::Hypers)
    return hypers.gamma * (1.0 + leaf_depth)^(-hypers.beta)
end

function birth_node(tree::Node)
    leafs = leaves(tree)
    leaf = rand(leafs)
    num_leafs = length(leafs)
    leaf_node_probability = 1.0 / num_leafs

    return leaf, leaf_node_probability
end

function test_data(n::Int64; sigma = 1.0, p = 10)
    X = rand(n,p)
    Y = fried(X) .+ sigma * randn(n)
    return X, Y
end

function probability_node_birth(tree::Node)
    out = tree.is_leaf ? 1.0 : 0.5
end

function death_node(tree::Node)
    ngb = not_grand_branches(tree)
    branch = rand(ngb)
    p_not_grand = 1.0 / length(ngb)
    return branch, p_not_grand
end

function not_grand_branches(tree::Node)
    ngb = Node[]
    not_grand_branches!(ngb, tree)
    return ngb
end

function not_grand_branches!(ngb::Vector{Node}, node::Node)

    if !node.is_leaf
        left_is_leaf = node.left.is_leaf
        right_is_leaf = node.right.is_leaf
        if left_is_leaf && right_is_leaf
            push!(ngb, node)
        else
            not_grand_branches!(ngb, node.left)
            not_grand_branches!(ngb, node.right)
        end
    end

end

function get_var_counts(forest::Vector{Node}, hypers::Hypers)
    counts = zeros(Int64, size(hypers.s,1))
    num_tree = length(forest)
    for t in 1:num_tree
        get_var_counts!(forest[t], counts)
    end
    return counts
end

function get_var_counts!(node::Node, counts::Vector{Int64})
    if !node.is_leaf
        counts[node.var] += 1
        get_var_counts!(node.left, counts)
        get_var_counts!(node.right, counts)
    end
end

function UpdateS!(forest::Vector{Node}, hypers::Hypers)
    shape_up = get_var_counts(forest, hypers) .+ hypers.alpha / length(hypers.s)
    hypers.s = rand(Dirichlet(shape_up))
end

function alpha_to_rho(alpha::Float64, scale::Float64)
    return alpha / (alpha + scale)
end

function rho_to_alpha(rho::Float64, scale::Float64)
    return scale * rho / (1 - rho)
end

function UpdateAlpha!(hypers::Hypers)
    alpha_current = hypers.alpha

    rho_current = alpha_to_rho(alpha_current, hypers.alpha_scale)
    psi = mean([log(hypers.s[i]) for i in 1:length(hypers.s)])    # psi = mean(log.(s))
    p = 1.0 * length(hypers.s)

    # SpecialFunctions.lgamma(x) = |log(gamma(x::Real))|
    # log_lik_current = log(π(ρ | -)) in S.1
    loglik_current = alpha_current * psi + lgamma(alpha_current) -
        p * lgamma(alpha_current / p) +
        logpdf(Beta(hypers.alpha_shape_1, hypers.alpha_shape_2), rho_current)

    for i in 1:50
        rho_propose = rand(Beta(hypers.alpha_shape_1, hypers.alpha_shape_2))
        alpha_propose = rho_to_alpha(rho_propose, hypers.alpha_scale)

        loglik_propose = alpha_propose * psi + lgamma(alpha_propose) -
            p * lgamma(alpha_propose / p) +
            logpdf(Beta(hypers.alpha_shape_1, hypers.alpha_shape_2), rho_propose)

        if log(rand()) < loglik_propose - loglik_current
            alpha_current = alpha_propose
            rho_current = rho_propose
            loglik_current = loglik_propose
        end
    end

    hypers.alpha = alpha_current

end

function growth_prior(node_depth::Int64, gamma::Float64, beta::Float64)
    return gamma * (1.0 + node_depth)^(-beta)
end

function forest_loglik(forest::Vector{Node}, gamma::Float64, beta::Float64)
    out = 0.0
    for t in 1:length(forest)
        out += tree_loglik(forest[t], 0, gamma, beta)
    end
    return out
end

function tree_loglik(node::Node, node_depth::Int64, gamma::Float64, beta::Float64)
    out = 0.0
    if node.is_leaf
        out += log(1.0 - growth_prior(node_depth, gamma, beta))
    else
        out += log(growth_prior(node_depth, gamma, beta))
        out += tree_loglik(node.left, node_depth + 1, gamma, beta)
        out += tree_loglik(node.right, node_depth + 1, gamma, beta)
    end
    return out
end

function UpdateGamma!(forest::Vector{Node}, hypers::Hypers)
    gamma_current = hypers.gamma
    loglik_current = forest_loglik(forest, gamma_current, hypers.beta)

    for i in 1:10
        gamma_prop = 0.5 * rand() + 0.5
        loglik_prop = forest_loglik(forest, gamma_prop, hypers.beta)
        if log(rand()) < loglik_prop - loglik_current
            gamma_current = gamma_prop
            loglik_current = loglik_prop
        end
    end

    hypers.gamma = gamma_current

end

function UpdateBeta!(forest::Vector{Node}, hypers::Hypers)
    beta_current = hypers.beta
    loglik_current = forest_loglik(forest, hypers.gamma, beta_current)

    for i in 1:10
        beta_prop = 2 * abs(randn())
        loglik_prop = forest_loglik(forest, hypers.gamma, beta_prop)
        if log(rand()) < loglik_prop - loglik_current
            beta_current = beta_prop
            loglik_current = loglik_prop
        end
    end

    hypers.beta = beta_current

end

end
import Base.print
print(io::IO, x::softbart.Node) = print(io, softbart.PrintTree(x));
import Base.show
show(io::IO, x::softbart.Node) = print(io, x);

####### Yinpu added #####
import Base.show
function show(io::IO, fit :: SoftbartOut)
    println(io, "\nFit:")
    println(io, "Updated α: $(fit.alpha[end])")
    println(io, "Updated β: $(fit.beta[end])")
    println(io, "Updated γ: $(fit.gamma[end])")
    println(io, "Updated σ: $(fit,.igma[end])")
    println(io, "Updated σᵤ: $(fit.sigma_mu[end])")
    println(io, "Updated s: $(fit.s[end,:])\n")
end

##### end of Yinpu added ######

max_stars_repo_name: YinpuLi/SparsityBoosting

max_issues_repo_name: YinpuLi/SparsityBoosting

max_forks_repo_name: YinpuLi/SparsityBoosting


text:
using Documenter, CuthillMcKee

makedocs(;
    modules=[CuthillMcKee],
    format=Documenter.HTML(),
    pages=[
        "Home" => "index.md",
    ],
    repo="https://github.com/rleegates/CuthillMcKee.jl/blob/{commit}{path}#L{line}",
    sitename="CuthillMcKee.jl",
    authors="Robert Lee Gates (robert.gates@gmail.com)",
    assets=[],
)

deploydocs(;
    repo="github.com/rleegates/CuthillMcKee.jl",
)

ext: jl


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

# This file tests the new Julia 1.6+ extension interface of TerminalMenus
# To trigger the new interface, at least one configuration keyword argument must be supplied.

# Check to make sure types are imported properly
@test MultiSelectMenu{TerminalMenus.MultiSelectConfig} <: TerminalMenus.ConfiguredMenu  # TODO Julia 2.0: delete parameter

# Constructor
@test MultiSelectMenu(["one", "two", "three"], charset=:ascii).pagesize == 3
@test MultiSelectMenu(string.(1:30), pagesize=-1, charset=:ascii).pagesize == 30
@test MultiSelectMenu(string.(1:4), pagesize=10, charset=:ascii).pagesize == 4
@test MultiSelectMenu(string.(1:100), charset=:ascii).pagesize == 10

multi_menu = MultiSelectMenu(string.(1:20), charset=:ascii)
@test TerminalMenus.options(multi_menu) == string.(1:20)
@test TerminalMenus.header(multi_menu) == "[press: Enter=toggle, a=all, n=none, d=done, q=abort]"

# Output
for kws in ((charset=:ascii,),
            (charset=:unicode,),
            (cursor='@', checked="c", unchecked="u",))
    local multi_menu
    multi_menu = MultiSelectMenu(string.(1:10); kws...)
    cur = isdefined(kws, :cursor) ? kws.cursor : (kws.charset === :ascii ? '>' : '→')
    chk = isdefined(kws, :checked) ? kws.checked : (kws.charset === :ascii ? "[X]" : "✓")
    uck = isdefined(kws, :unchecked) ? kws.unchecked : (kws.charset === :ascii ? "[ ]" : "⬚")

    buf = IOBuffer()
    TerminalMenus.writeline(buf, multi_menu, 1, true)
    @test String(take!(buf)) == "$uck 1"
    TerminalMenus.printmenu(buf, multi_menu, 1; init=true)
    @test startswith(String(take!(buf)), string("\e[2K[press: Enter=toggle, a=all, n=none, d=done, q=abort]\r\n\e[2K $cur $uck 1"))
    push!(multi_menu.selected, 1)
    TerminalMenus.printmenu(buf, multi_menu, 2; init=true)
    @test startswith(String(take!(buf)), string("\e[2K[press: Enter=toggle, a=all, n=none, d=done, q=abort]\r\n\e[2K   $chk 1\r\n\e[2K $cur $uck 2"))
end

# Preselection
sel = [2,5]
multi_menu = MultiSelectMenu(string.(1:20), pagesize=6, selected=sel, charset=:ascii)
buf = IOBuffer()
TerminalMenus.printmenu(buf, multi_menu, 1; init=true)
str = String(take!(buf))
for i = 1:multi_menu.pagesize
    if i ∈ sel
        @test occursin("[X] $i", str)
    else
        @test occursin("[ ] $i", str)
    end
end

# Test SDTIN
multi_menu = MultiSelectMenu(string.(1:10), charset=:ascii)
@test simulate_input(Set([1,2]), multi_menu, :enter, :down, :enter, 'd')
multi_menu = MultiSelectMenu(["single option"], charset=:ascii)
@test simulate_input(Set([1]), multi_menu, :up, :up, :down, :enter, 'd')

max_stars_repo_name: collinwarner/julia

max_issues_repo_name: collinwarner/julia

max_forks_repo_name: collinwarner/julia


text:
## quadrature.jl : nystrom integration methods

abstract type QuadratureRule end

struct GaussLegendre <: QuadratureRule end
struct EOLE <: QuadratureRule end
struct Simpson <: QuadratureRule end
struct Midpoint <: QuadratureRule end
struct Trapezoidal <: QuadratureRule end

# get nodes and weights of Gauss-Legendre quadrature on [a,b]
function get_nodes_and_weights(n::Integer, a, b, q::GaussLegendre)
    nodes, weights = gausslegendre(n)
    weights = (b-a)/2*weights
    nodes = (b-a)/2*nodes .+ (a+b)/2
    return nodes, weights
end

# get nodes and weights of structured grid on [a,b]
function get_nodes_and_weights(n::Integer, a, b, q::EOLE)
    nodes = n == 1 ? [(a+b)/2] : range(a; stop = b,length = n)
    weights = fill((b-a) / n, n)
    return nodes, weights
end

# get nodes and weights of Simpson's rule on [a,b]
function get_nodes_and_weights(n::Integer, a, b, q::Simpson)
    iseven(n) || begin
        @warn "to use Simpson's rule, n must be even (received $(n)). I will continue with n = $(n+1)"
        n += 1
    end
    Δx = (b-a)/n
    nodes = a:Δx:b
    weights = repeat(2:2:4, outer=Int(n/2))
    weights[1] = 1
    push!(weights,1)
    weights *= Δx/3
    return nodes, weights
end

# get nodes and weights of Midpoint rule on [a,b]
function get_nodes_and_weights(n::Integer, a, b, q::Midpoint)
    Δx = (b-a)/n
    nodes = (a+Δx/2):Δx:(b-Δx/2)
    weights = fill(Δx, size(nodes))
    return nodes, weights
end

# get nodes and weights of Trapezoidal rule on [a,b]
function get_nodes_and_weights(n::Integer, a, b, q::Trapezoidal)
    Δx = (b-a)/n
    nodes = a:Δx:b
    weights = fill(Δx, length(nodes))
    weights[1] /= 2
    weights[end] /= 2
    return nodes, weights
end

## EigenSolvers
abstract type AbstractEigenSolver end

struct EigenSolver <: AbstractEigenSolver end

compute(A, n, ::EigenSolver) = eigen(A, sortby=λ->-abs(real(λ)))

struct EigsSolver <: AbstractEigenSolver end

compute(A, n, ::EigsSolver) = eigs(A, nev=n, ritzvec=true, which=:LM, v0=randn(size(A,1)))


text:
# # Debugging

using SwiftShader_jll # hide
using Vulkan
@set_driver :SwiftShader # hide

const debug_callback_c = @cfunction(default_debug_callback, UInt32, (DebugUtilsMessageSeverityFlagEXT, DebugUtilsMessageTypeFlagEXT, Ptr{vk.VkDebugUtilsMessengerCallbackDataEXT}, Ptr{Cvoid}))

create_info = DebugUtilsMessengerCreateInfoEXT(
    |(Vulkan.message_severities...),
    |(
        DEBUG_UTILS_MESSAGE_TYPE_GENERAL_BIT_EXT,
        DEBUG_UTILS_MESSAGE_TYPE_VALIDATION_BIT_EXT,
        DEBUG_UTILS_MESSAGE_TYPE_PERFORMANCE_BIT_EXT
    ),
    debug_callback_c,
)

instance = Instance(["VK_LAYER_KHRONOS_validation"], ["VK_EXT_debug_utils"]; next=create_info)

messenger = unwrap(create_debug_utils_messenger_ext(instance, create_info))


text:
# perf or animation ?
const INFO = true
const record_animation = false

using Printf

using CUDAdrv
using CUDAnative
using CuArrays
CuArrays.allowscalar(false)

using Statistics
#using Makie
using DelimitedFiles

include("main1d-gpu-kernels.jl")

# the type for the domain
struct Domain
    x_start::Float64
    x_end::Float64
end

# Initial solution
function InitSolution!(V::Array{Float64, 2}, domain::Domain)
    nx = size(V, 2)

    hx = (domain.x_end - domain.x_start) / nx
    xc = 0.5  * (domain.x_start + domain.x_end) + nx/8 * hx

    for i=1:nx
        x = domain.x_start + (i-0.5) * hx
        if abs(x - xc) < 0.2
            V[1, i] = 1.
        else
            V[1, i] = 0.
        end

        V[2, i] = 0.
    end
end

function compute_max(mx::Array{Float64}, mx_d::CuArray{Float64}, x_d::CuArray{Float64}, nthreads, nblocks, smem)
    @cuda threads=nthreads blocks=nblocks shmem=smem reduce_max(x_d, mx_d, smem)
    copyto!(mx, 1, mx_d, 1, nblocks)
    return maximum(mx)
end

# initialization
const domain = Domain(0., 1.)
const Nx = 32768

# GPU parameters
const num_gpu_threads = 256
const num_gpu_blocks = ifelse(Nx % num_gpu_threads == 0, div(Nx, num_gpu_threads), div(Nx, num_gpu_threads) + 1)

const nblocks_max = div(num_gpu_blocks, 2) + rem(num_gpu_blocks, 2)
const smem_max = ifelse(num_gpu_threads <= 32, 64 * sizeof(Float64), num_gpu_threads * sizeof(Float64))
const smem_scheme = 3 * (num_gpu_threads + 2) * sizeof(Float64)

println("number of GPU blocks = ", num_gpu_blocks)
println("number of GPU threads per block = ", num_gpu_threads)

# loop in time
const T = 1e-3
const dx = (domain.x_end - domain.x_start) / Nx

# tolerance for dry region
const tol = eps(Float64)

let t=0
    # only V is store on cpu
    V = Array{Float64, 2}(undef, 2, Nx)
    InitSolution!(V, domain)

    # no need for Vold and lambdas to be on cpu
    V_d = CuArray{Float64, 1}(undef, 2*Nx)
    Vold_d = CuArray{Float64, 1}(undef, 2*Nx)

    lambdas_d = CuArray{Float64, 1}(undef, Nx)
    lmax_d = CuArray{Float64, 1}(undef, nblocks_max)
    lmax = Array{Float64, 1}(undef, nblocks_max)

    copyto!(V_d, 1, V, 1, length(V))

    # create kernels (we don't use @cuda macro because there is no need to convert the type
    # of the arguments before each call. We always use the same arguments)
    ker_args_eig = cudaconvert.([lambdas_d, V_d, tol])
    ker_tt_eig = Tuple{Core.Typeof.(ker_args_eig)...}
    ker_eig = cufunction(kernel_UpdateEigenvalues!, ker_tt_eig)

    ker_args_scheme = cudaconvert.([V_d, Vold_d, lambdas_d, 0., tol])
    ker_tt_scheme = Tuple{Core.Typeof.(ker_args_scheme)...}
    ker_scheme = cufunction(kernel_LaxFriedrich!, ker_tt_scheme)

    function update_to_time(tframe::Float64)
        while(t < tframe)
            # update the eigenvalues (first time iteration = initialization)
            #CuArrays.@sync begin
            @cuda threads=num_gpu_threads blocks=num_gpu_blocks kernel_UpdateEigenvalues!(lambdas_d, V_d, tol)
            #ker_eig(ker_args_eig...; threads=num_gpu_threads, blocks=num_gpu_blocks)
            #CUDAdrv.cudacall(ker_eig.fun, ker_tt_eig, ker_args_eig...; threads=num_gpu_threads, blocks=num_gpu_blocks)
            #end

            # compute time step (reduction operation on gpu because we are using CuArrays)
            #Cmax = compute_max(lmax_d, lambdas_d)
            Cmax = compute_max(lmax, lmax_d, lambdas_d, num_gpu_threads, nblocks_max, smem_max)
            dt = min(0.5 * dx / Cmax, tframe - t)

            # copy of V to Vold on gpu (first time iteration = initialization)
            copyto!(Vold_d, 1, V_d, 1, 2*Nx)

            # performs one time step
            #CuArrays.@sync begin
            #@cuda threads=num_gpu_threads blocks=num_gpu_blocks shmem=smem_scheme kernel_LaxFriedrich!(V_d, Vold_d, lambdas_d, dt / dx, tol)
            ker_args_scheme[4] = dt / dx
            ker_scheme(ker_args_scheme...; threads=num_gpu_threads, blocks=num_gpu_blocks, shmem=smem_scheme)
            #CUDAdrv.cudacall(ker_scheme.fun, ker_tt_scheme, ker_args_scheme..., threads=num_gpu_threads, blocks=num_gpu_blocks, shmem=smem_scheme)
            #end

            t += dt
            INFO && @printf("\t%.5f / %.5f\n", t, tframe)
        end
    end

    print("Warming up... ")
    update_to_time(eps(Float64))
    println("Done")

    t = 0
    InitSolution!(V, domain)
    copyto!(V_d, 1, V, 1, length(V))

    println("Initial time t = ", t)
    CUDAdrv.@profile update_to_time(T)
    #CuArrays.@time update_to_time(T)
    println("End of simulation, t = ", t)

    copyto!(V, 1, V_d, 1, length(V))
    println("mean(h) = ", mean(V[1, :]))
    writedlm("sol-gpu", V[1, :], "\n")

    #scene = lines(range(0, stop=1, length=Nx), V[1, :])
    #Makie.save("plot_h_gpu.png", scene)
end

max_stars_repo_name: danielcort/benchmarks-python-julia-c

max_issues_repo_name: danielcort/benchmarks-python-julia-c

max_forks_repo_name: danielcort/benchmarks-python-julia-c


text:
"Inference Algorithms"
module Inference

using Spec
using Random
using Lens
using ..Omega: RandVar, applytrackerr, indomainₛ, logerr,
               UTuple, Ω, applynotrackerr, SimpleΩ, LinearΩ,
               update, cond, randtuple, nelem,
               gradient, linearize, unlinearize, err, indomainₛ, indomain, apl
import ..Omega
using Omega.Space: tagrng

using ..Gradient: value, gradient

using ProgressMeter
import Flux
import ForwardDiff
using Callbacks
using Lens
import UnicodePlots
using DocStringExtensions: SIGNATURES

"Lens called at every iteration of MCMC algorithm"
abstract type Loop end

"Optimization Algorithm"
abstract type OptimAlgorithm end

"Sampling Algorithm"
abstract type SamplingAlgorithm end

"Is the inference algorithm approximate?"
function isapproximate end

"Default probability space type to use"
function defΩ end

include("transforms.jl")# Transformations from [0, 1] to R, etc
include("callbacks.jl") # Common Inference Functions

# Sampling
include("rand.jl")      # Sampling
include("rs.jl")        # Rejection Sampling
include("ssmh.jl")      # Single Site Metropolis Hastings
# include("hmc.jl")       # Hamiltonian Monte Carlo
include("hmcfast.jl")   # Faster Hamiltonian Monte Carlo
# include("hmcfastg.jl")   # Faster Hamiltonian Monte Carlo
include("replica.jl")   # Replica Exchange
# include("dynamichmc.jl")# Dynamic Hamiltonion Monte Carlo
include("fail.jl")# Dynamic Hamiltonion Monte Carlo

# include("sghmc.jl")     # Stochastic Gradient Hamiltonian Monte Carlo
# include("relandscape.jl")  # Variantional Sampling through relandscape

# Optimization
include("argmax.jl")     # NLopt based optimization
include("nlopt.jl")     # NLopt based optimization

export  isapproximate,

        RejectionSample,
        SSMH,
        HMC,
        # SGHMC,
        HMCFAST,
        Replica,

        RejectionSampleAlg,
        SSMHAlg,
        HMCAlg,
        # SGHMCAlg,
        HMCFASTAlg,
        RelandscapeAlg,
        Relandscape,
        NUTS,
        NUTSAlg,
        FailUnsat,

        defalg,
        defcb,
        defΩ,
        defΩProj,

        plotrv,
        plotscalar,
        default_cbs,
        default_cbs_tpl,
        default_cbs,

        # Lenses
        Loop,
        SSMHLoop,
        HMCFASTLoop,

        autorand

end


text:
using DimensionalData, Test, Unitful
using DimensionalData.LookupArrays, DimensionalData.Dimensions

@dim TestDim "Testname"

@testset "dims creation macro" begin
    @test parent(TestDim(1:10)) == 1:10
    @test val(TestDim(1:10)) == 1:10
    @test name(TestDim) == :Testname
    @test label(TestDim) == "Testname"
    @test val(TestDim(:testval)) == :testval
    @test metadata(TestDim(Sampled(1:1; metadata=Metadata(a=1)))) == Metadata(a=1)
    @test units(TestDim) == nothing
    @test label(TestDim) == "Testname"
    @test eltype(TestDim(1)) == Int
    @test eltype(TestDim([1, 2, 3])) <: Int
    @test length(TestDim(1)) == 1
    @test length(TestDim([1, 2, 3])) == 3
    @test eachindex(TestDim(10:10:30)) == Base.OneTo(3)
    @test eachindex(TestDim(1)) == Base.OneTo(1)
    @test step(TestDim(1:2:3)) == 2
    @test step(TestDim(Base.OneTo(10))) == 1
    @test step(TestDim(1:2:3)) == 2
    @test step(TestDim(Sampled(1:2:3; span=Regular(2)))) == 2
    @test first(TestDim(5)) == 5
    @test last(TestDim(5)) == 5
    @test first(TestDim(10:20)) == 10
    @test last(TestDim(10:20)) == 20
    @test firstindex(TestDim(5)) == 1
    @test lastindex(TestDim(5)) == 1
    @test firstindex(TestDim(10:20)) == 1
    @test lastindex(TestDim(10:20)) == 11
    @test eachindex(TestDim(10:20)) == Base.OneTo(11)
    @test size(TestDim(10:20)) == (11,)
    @test ndims(TestDim(10:20)) == 1
    @test ndims(TestDim(1)) == 0
    @test Array(TestDim(10:15)) == [10, 11, 12, 13, 14, 15]
    @test iterate(TestDim(10:20)) == iterate(10:20)
    @test TestDim(5.0:7.0)[2] == 6.0
    @test TestDim(10)[] == 10
    @test TestDim(Sampled(5.0:7.0, ForwardOrdered(), Regular(1.0), Points(), NoMetadata()))[At(6.0)] == 6.0
end

@testset "format" begin
    A = [1 2 3; 4 5 6]
    @test format((X, Y), A) == (X(NoLookup(Base.OneTo(2))), Y(NoLookup(Base.OneTo(3))))
    @test format(Ti, zeros(3)) == (Ti(NoLookup(Base.OneTo(3))),)
    @test format((:a, :b), A) == (Dim{:a}(NoLookup(Base.OneTo(2))),
                                  Dim{:b}(NoLookup(Base.OneTo(3))))
    @test format(:c, 51:100) == (Dim{:c}(NoLookup(Base.OneTo(50))),)
    @test format((a=[:A, :B], b=(10.0:10.0:30.0)), A) ==
        (Dim{:a}(Categorical([:A, :B], ForwardOrdered(), NoMetadata())),
         Dim{:b}(Sampled(10.0:10:30.0, ForwardOrdered(), Regular(10.0), Points(), NoMetadata())))
    @test format((X([:A, :B]; metadata=Metadata(a=5)),
           Y(Categorical(10.0:10.0:30.0, ForwardOrdered(), Metadata("metadata"=>1)))), A) ==
          (X(Categorical([:A, :B], ForwardOrdered(), Metadata(a=5))),
           Y(Categorical(10.0:10.0:30.0, ForwardOrdered(), Metadata("metadata"=>1))))
     @test format((X(Sampled(Base.OneTo(2); order=ForwardOrdered(), span=Regular(), sampling=Points())), Y), A) ==
        (X(Sampled(Base.OneTo(2), ForwardOrdered(), Regular(1), Points(), NoMetadata())),
         Y(NoLookup(Base.OneTo(3))))
end

@testset "Val" begin
    @test dims(Val{X}()) == Val{X}()
    @test lookup(Val{X}()) == NoLookup()
end

@testset "AnonDim" begin
    @test val(AnonDim()) == Colon()
    @test lookup(AnonDim()) == NoLookup()
    @test metadata(AnonDim()) == NoMetadata()
    @test name(AnonDim()) == :Anon
end

@testset "Basic dim and array initialisation and methods" begin
    a = ones(5, 4)

    @test_throws DimensionMismatch DimArray(a, X)
    @test_throws DimensionMismatch DimArray(a, (X, Y, Z))

    da = DimArray(a, (X(LinRange(140, 148, 5)), Y(LinRange(2, 11, 4))))
    dimz = dims(da)

    for args in ((dimz,), (dimz, (X(), Y())), (dimz, X(), Y()),
                 (dimz, (X, Y)), (dimz, X, Y),
                 (dimz, (1, 2)), (dimz, 1, 2))
        @test index(args...) == (LinRange(140, 148, 5), LinRange(2, 11, 4))
        @test name(args...) == (:X, :Y)
        @test units(args...) == (nothing, nothing)
        @test label(args...) == ("X", "Y")
        @test sampling(args...) == (Points(), Points())
        @test span(args...) == (Regular(2.0), Regular(3.0))
        @test locus(args...) == (Center(), Center())
        @test order(args...) == (ForwardOrdered(), ForwardOrdered())
        @test bounds(args...) == ((140, 148), (2, 11))
        @test lookup(args...) == (Sampled(LinRange(140, 148, 5), ForwardOrdered(), Regular(2.0), Points(), NoMetadata()),
                                Sampled(LinRange(2, 11, 4), ForwardOrdered(), Regular(3.0), Points(), NoMetadata()))
    end

    @test val(dimz, ()) == index(dimz, ()) == ()
    @test val(dimz, 1) == val(dimz, X) == val(dimz, X()) == val(dimz[1])

    @test dims(dimz, Y) === dimz[2]
    @test slicedims(dimz, (2:4, 3)) ==
        ((X(Sampled(LinRange(142, 146, 3), ForwardOrdered(), Regular(2.0), Points(), NoMetadata())),),
         (Y(Sampled(LinRange(8.0, 8.0, 1), ForwardOrdered(), Regular(3.0), Points(), NoMetadata())),))

    @test lookup(X) == NoLookup()
end

@testset "repeating dims of the same type is allowed" begin
    dimz = X((143, 145)), X((-38, -35))
    @test dims(dimz, X) === dimz[1]
    @test dims(dimz, (X, X)) === dimz
    @test dimnum(dimz, (X, X)) === (1, 2)
    @test hasdim(dimz, (X, X)) === (true, true)
    @test sortdims(dimz, (X, X)) === dimz
end

@testset "applying function on a dimension" begin
    d = X(0:0.01:2π)
    a = DimArray(cos, d)
    @test length(dims(a)) == 1
    @test typeof(dims(a)[1]) <: X
    @test a.data == cos.(d.val)
end


text:
module C
const NAMES = [
    "kscyce"
    "kdcyce"
    "kdcycee"
    "kdcycea"
    "kasse"
    "kdise"
    "kscyca"
    "kdcyca"
    "kdcycac1"
    "kassa"
    "kdisa"
    "ks27"
    "kd27"
    "kd27e"
    "kd27a"
    "ksemi1"
    "kdemi1"
    "Cdh1T"
    "kacdh1"
    "kicdh1e"
    "kicdh1a"
    "kasec"
    "kdiec"
    "ksskp2"
    "kdskp2"
    "kdskp2c1"
    "Inhibitor"
    "CycElevel"
    "Emi1T"
]

#name2idx
for (idx,name) in enumerate(NAMES)
    eval(Meta.parse("const $name = $idx"))
end

const NUM = length(NAMES)

end  # module

max_stars_repo_name: himoto/Bifurcation

max_issues_repo_name: himoto/Bifurcation

max_forks_repo_name: himoto/Bifurcation


text:
# using Grammatical Evolution to solve the Ant problem
using GrammaticalEvolution
import GrammaticalEvolution.evaluate!
import GrammaticalEvolution.isless
import Base.copy

const food = convert(Int64, '1')
const gap = convert(Int64, '0')
const empty = 46 #convert(Int64, '.')
const startpos = 88 #convert(Int64, 'X')

world1 =
"""
..#..
..#..
.##..
.#...
.X...
"""

function read_map(s)
#     data = [strip(line) for line in split(s)]
    data = [strip(line) for line in s]
    height = length(data)
    width = length(data[1])
    start = ones(Int64, 2)

    map = zeros(Int64, width, height)
    for i=1:height
      for j=1:width
        if data[i][j] != startpos
          map[i, j] = convert(Int64, data[i][j])
        else
          start = Any[i, j]
        end
      end
    end

    return (map, start)
end

# println(read_map(world1))

type World
  width::Int64
  height::Int64
  value::Array{Int64, 2}
  start::Array{Int64, 1}

  function World(width::Int64, height::Int64)
    value = zeros(width, height)
    return new(width, height, value)
  end

  function World(map::Array{Int64, 2}, start::Array)
    return new(size(map, 1), size(map, 2), map, start)
  end
end

copy(world::World) = World(copy(world.value), world.start)

immutable Direction
  value::Int64
end

import Base: +, -, isless
+{T <: Number}(d::Direction, v::T) = Direction(d.value + v)
-{T <: Number}(d::Direction, v::T) = Direction(d.value - v)
isless(d1::Direction, d2::Direction) = isless(d1.value, d2.value)
isless(d1::Direction, v::Int64) = isless(d1.value, value)

const north = Direction(1)
const east = Direction(2)
const south = Direction(3)
const west = Direction(4)

type AntIndividual <: Individual
  genome::Array{Int64, 1}
  fitness::Union{Float64,Void}
  code

  function AntIndividual(size::Int64, max_value::Int64)
    genome = rand(1:max_value, size)
    return new(genome, nothing, nothing)
  end

  AntIndividual(genome::Array{Int64, 1}) = new(genome, nothing, nothing)
end

type AntPopulation <: Population
  individuals::Array{AntIndividual, 1}

  function AntPopulation(population_size::Int64, genome_size::Int64)
    individuals = Array(AntIndividual, 0)
    for i=1:population_size
      push!(individuals, AntIndividual(genome_size, 1000))
    end

    return new(individuals)
  end
end

type Ant
  x::Int64
  y::Int64
  direction::Direction
  eaten::Int64

  function Ant(x::Int64, y::Int64)
    return new(x, y, north, 0)
  end
end

function calc_forward(world::World, ant::Ant)
  nx = ant.x
  ny = ant.y

  if ant.direction === north
    ny -= 1
    if ny < 1 ny = world.height end
  elseif ant.direction === east
    nx += 1
    if nx > world.width nx = 1 end
  elseif ant.direction === south
    ny += 1
    if ny > world.height ny = 1 end
  else
    nx -= 1
    if nx < 1 nx = world.width end
  end

  return (nx, ny)
end

function food_ahead(world::World, ant::Ant)
  (x, y) = calc_forward(world, ant)
  return world.value[x, y] == food
end

function forward(world::World, ant::Ant)
  (ant.x, ant.y) = calc_forward(world, ant)
  if world.value[ant.x, ant.y] == food
    world.value[ant.x, ant.y] = empty
    ant.eaten += 1
  elseif world.value[ant.x, ant.y] == empty
    ant.eaten -= 1
  end
end

function turn_right(ant::Ant)
  #ant.direction = ant.direction + 1 > south ? north : ant.direction+1
  ndir = ant.direction + 1
  ant.direction = ndir.value > south.value ? north : ndir;
end

function turn_left(ant::Ant)
  #ant.direction = ant.direction - 1 < 1 ? south : ant.direction-1
  ndir = ant.direction - 1
  ant.direction = ndir.value < 1 ? south : ndir
end

make_block(lst::Array) = Expr(:block, lst...)
make_loop(values::Array) = Expr(:for, :(=), :i, Expr(:(:), 1, values[1]), values[2])
make_if(values::Array) = Expr(:if, values[1], values[2], values[3])

macro make_if(condition, true_block, false_block)
  Expr(:if, condition, true_block, false_block)
end

macro make_call(fn, args...)
  Expr(:call, fn, args...)
end

macro make_for(start, stop, block)
  Expr(:for, Expr(:(=), :i, Expr(:(:), start, stop)), block)
end

@grammar ant_grammar begin
  start = block
  command = turnleft | turnright | forward | if_statement | loop
  turnleft = Expr(:call, :turn_left, :ant)
  turnright = Expr(:call, :turn_right, :ant)
  forward = Expr(:call, :forward, :world, :ant)
  if_statement = Expr(:if, condition, block, block)
  condition = food_ahead
  food_ahead = Expr(:call, :food_ahead, :world, :ant)
  loop = Expr(:for, Expr(:(=), :i, Expr(:(:), 1, digit)), block)
  block[make_block] = (command)^(1:5)
  digit = 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9
end

function evaluate!(grammar::Grammar, ind::ExampleIndividual, pop::ExamplePopulation, world::World)
  # copy world
  world = copy(world)

  # create ant
  ant = Ant(world.start[1], world.start[2])

  # generate code
  try
    ind.code = transform(grammar, ind)
    @eval fn(world::World, ant::Ant) = $(ind.code)
    pop.totalevals += 1
    if ind.fitness < pop.best_fitness
      pop.best_fitness = ind.fitness
      pop.best_ind = ind
      pop.best_at_eval = pop.totalevals
    end
  catch e
    if typeof(e) !== MaxWrapException
      Base.error_show(STDERR, e, catch_backtrace())
    end
    ind.fitness = realmax(Float64)
    return
  end

  # run code
  fn(world, ant)

  # fitness is the amount of food eaten
  ind.fitness = -convert(Float64, ant.eaten)
end

# need to redefine 'isless' (eating more is better)
#Base.isless(ind1::ExampleIndividual, ind2::ExampleIndividual) = ind1.fitness > ind2.fitness

function main()
  data = open("../examples/santefe.trail")
  (map, start) = read_map(readlines(data))
  world = World(map, start)

  # create population
  pop = ExamplePopulation(500, 500)

  fitness = 0
  generation = 1

  evaluate!(ant_grammar, pop, world)
  while generation < 1000
    tic()
    # generate a new population (based off of fitness)
    pop = generate(ant_grammar, pop, 0.1, 0.2, 0.2, world)

    # population is sorted, so first entry it the best
    fitness = pop[1].fitness
    println("generation: $generation, $(length(pop)), max fitness=$fitness, time=$(signif(toq(),5))\n$(pop[1].code)")
    generation += 1
  end
end

main()

max_stars_repo_name: rcnlee/GrammaticalEvolution

max_issues_repo_name: rcnlee/GrammaticalEvolution

max_forks_repo_name: rcnlee/GrammaticalEvolution


text:
using MATT
using Test

@testset "MATT.jl" begin
    # Write your own tests here.
end


text:
struct RMP3{T} <: AbstractMPWavefunction# where T <: AbstractFloat
    CorrelationEnergy::T
    T2::_T2 where _T2 <: Fermi.AbstractTensor
end

include("DF-RMP3.jl")

function RMP3()
    refwfn = Fermi.HartreeFock.RHF()
    RMP3(refwfn)
end
function RMP3(ref::Fermi.HartreeFock.RHF)
    precision_selector = Dict{Any,Any}("single" => Float32,
                                       "double" => Float64)
    prec = precision_selector[Fermi.CurrentOptions["precision"]]
    RMP3{prec}(ref)
end

function RMP3{T}(ref::Fermi.HartreeFock.RHF) where T <: AbstractFloat
    # convert some options -> singletons
    alg_selector = Dict{Any,Any}("direct" => Direct(),
                                 "df" => DF(),
                                 "conv" => Conventional())
    alg = alg_selector[Fermi.CurrentOptions["mp2_type"]]
    RMP3{T}(ref,alg)
end


text:
## Root finding for Chebyshev expansions
#
#  Contains code that is based in part on Chebfun v5's chebfun/@chebteck/roots.m,
# which is distributed with the following license:

# Copyright (c) 2015, The Chancellor, Masters and Scholars of the University
# of Oxford, and the Chebfun Developers. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#     * Redistributions of source code must retain the above copyright
#       notice, this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#     * Neither the name of the University of Oxford nor the names of its
#       contributors may be used to endorse or promote products derived from
#       this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
# ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

function complexroots(f::Fun{C}) where C<:Chebyshev
    if ncoefficients(f)==0 || (ncoefficients(f)==1 && isapprox(f.coefficients[1],0))
        throw(ArgumentError("Tried to take roots of a zero function."))
    elseif ncoefficients(f)==1
        ComplexF64[]
    elseif ncoefficients(f)==2
        ComplexF64[-f.coefficients[1]/f.coefficients[2]]
    else
        fromcanonical.(f,colleague_eigvals(f.coefficients))
    end
end

#function roots(f::Fun)
#    irts=map(real,filter!(x->abs(x)<=1.+10eps(),filter!#(isreal,complexroots(f.coefficients))))
#
#    map!(x->x>1. ? 1. : x,irts)
#    map!(x->x<-1. ? -1. : x,irts)
#
#    if length(irts)==0
#        Float64[]
#    else
#        fromcanonical(f,irts)
#    end
#end

function roots(f::Fun)
    f2=Fun(f,domain(f)) # default is to convert to Chebyshev/Fourier
    if space(f2)==space(f)
        error("roots not implemented for "*string(typeof(f)))
    else
        roots(f2)
    end
end

roots(f::Fun{<:Chebyshev}) = fromcanonical.(f,roots(setcanonicaldomain(f)))

for (BF,FF) in ((BigFloat,Float64),(Complex{BigFloat},ComplexF64))
    @eval function roots( f::Fun{C,$BF} ) where C<:Chebyshev
    # FIND THE ROOTS OF AN IFUN.

        d = domain(f)
        c = f.coefficients
        vscale = maximum(abs,values(f))
        if vscale == 0
            throw(ArgumentError("Tried to take roots of a zero function."))
        end

        hscale = maximum( [abs(leftendpoint(d)), abs(rightendpoint(d))] )
        htol = eps(2000.)*max(hscale, 1)  # TODO: choose tolerance better

        # calculate Flaot64 roots
        r = Array{$BF}(rootsunit_coeffs(convert(Vector{$FF},c./vscale), Float64(htol)))
        # Map roots from [-1,1] to domain of f:
        rts = fromcanonical.(Ref(d),r)
        fp = differentiate(f)

        # do Newton 3 time
        for _ = 1:3
            rts .-=extrapolate.(Ref(f),rts)./extrapolate.(Ref(fp),rts)
        end

        return rts
    end
end

function roots( f::Fun{C,TT} ) where {C<:Chebyshev,TT<:Union{Float64,ComplexF64}}
# FIND THE ROOTS OF AN IFUN.
    if iszero(f)
        throw(ArgumentError("Tried to take roots of a zero function."))
    end

    d = domain(f)
    c = f.coefficients
    vscale = maximum(abs,values(f))

    hscale = max(abs(leftendpoint(d)), abs(rightendpoint(d)) )
    htol = eps(2000.)*max(hscale, 1)  # TODO: choose tolerance better

    cvscale=c./vscale
    r = rootsunit_coeffs(cvscale, Float64(htol))

    # Check endpoints, as these are prone to inaccuracy
    # which can be deadly.
    if (isempty(r) || !isapprox(last(r),1.)) && abs(sum(cvscale)) < htol
        push!(r,1.)
    end
    if (isempty(r) || !isapprox(first(r),-1.)) && abs(alternatingsum(cvscale)) < htol
        insert!(r,1,-1.)
    end

    # Map roots from [-1,1] to domain of f:
    return fromcanonical.(Ref(d),r)
end

function colleague_matrix( c::Vector{T} ) where T<:Number
#TODO: This is command isn't typed correctly
# COMPUTE THE ROOTS OF A LOW DEGREE POLYNOMIAL BY USING THE COLLEAGUE MATRIX:
    c = chop(c,0) # prune exact zeros
    n = length(c) - 1
    A=zeros(T,n,n)

    for k=1:n-1
        A[k+1,k]=0.5
        A[k,k+1]=0.5
    end
    for k=1:n
        A[1,end-k+1]-=0.5*c[k]/c[end]
    end
    A[n,n-1]=one(T)

    # TODO: can we speed things up because A is lower-Hessenberg
    # Standard colleague matrix (See [Good, 1961]):
    A
end

function colleague_balance!(M)
    n=size(M,1)
    if n<2
        return M
    end

    d=1+abs(M[1,2])
    M[1,2]/=d;M[2,1]*=d

    for k=3:n
        M[k-1,k]*=d;M[k,k-1]/=d
        d+=abs(M[1,k])
        M[1,k]/=d;M[k-1,k]/=d;M[k,k-1]*=d
    end
    M
end

# colleague_eigvals( c::Vector )=hesseneigvals(colleague_balance!(colleague_matrix(c)))

colleague_eigvals( c::Vector )=eigvals(colleague_matrix(c))

function PruneOptions( r, htol::Float64 )
# ONLY KEEP ROOTS IN THE INTERVAL

    # Remove dangling imaginary parts:
    r = real( r[ abs.(imag.(r)) .< htol ] )
    # Keep roots inside [-1 1]:
    r = sort( r[ abs.(r) .<= 1+htol ] )
    # Put roots near ends onto the domain:
    r = min.( max.( r, -1 ), 1)

    # Return the pruned roots:
    return r
end

rootsunit_coeffs(c::Vector{T}, htol::Float64) where {T<:Number} =
    rootsunit_coeffs(c,htol,ClenshawPlan(T,Chebyshev(),length(c),length(c)))
function rootsunit_coeffs(c::Vector{T}, htol::Float64,clplan::ClenshawPlan{S,T}) where {S,T<:Number}
# Computes the roots of the polynomial given by the coefficients c on the unit interval.

    # If recursive subdivision is used, then subdivide [-1,1] into [-1,splitPoint] and [splitPoint,1].
    splitPoint = -0.004849834917525

    # Simplify the coefficients by chopping off the tail:
    nrmc=norm(c, 1)
    @assert nrmc > 0  # There's an error if we make it this far
    c = chop(c,eps()*nrmc)
    n = length(c)

    if n == 0

        # EMPTY FUNCTION
        r = Float64[]

    elseif n == 1

        # CONSTANT FUNCTION
        r = ( c[1] == 0.0 ) ? Float64[0.0] : Float64[]

    elseif n == 2

        # LINEAR POLYNOMIAL
        r1 = -c[1]/c[2];
        r = ( (abs(imag(r1))>htol) | (abs(real(r1))>(1+htol)) ) ? Float64[] : Float64[max(min(real(r1),1),-1)]

    elseif n <= 70

        # COLLEAGUE MATRIX
        # The recursion subdividing below will keep going until we have a piecewise polynomial
        # of degree at most 50 and we likely end up here for each piece.

        # Adjust the coefficients for the colleague matrix
        # Prune roots depending on preferences:
        r = PruneOptions( colleague_eigvals(c), htol )::Vector{Float64}
    else

        #  RECURSIVE SUBDIVISION:
        # If n > 50, then split the interval [-1,1] into [-1,splitPoint], [splitPoint,1].
        # Find the roots of the polynomial on each piece and then concatenate. Recurse if necessary.

        # Evaluate the polynomial at Chebyshev grids on both intervals:
        #(clenshaw! overwrites points, which only makes sence if c is real)

        v1 = isa(c,Vector{Float64}) ? clenshaw!( c, points(-1..splitPoint,n),clplan) : clenshaw( c, points(-1..splitPoint,n),clplan)
        v2 = isa(c,Vector{Float64}) ? clenshaw!( c, points(splitPoint..1 ,n),clplan) : clenshaw( c, points(splitPoint..1 ,n),clplan)

        # Recurse (and map roots back to original interval):
        p = plan_chebyshevtransform( v1 )
        r = [ (splitPoint - 1)/2 .+ (splitPoint + 1)/2 .* rootsunit_coeffs( p*v1, 2*htol,clplan) ;
                 (splitPoint + 1)/2 .+ (1 - splitPoint)/2 .* rootsunit_coeffs( p*v2, 2*htol,clplan) ]

    end

    # Return the roots:
    r
end

extremal_args(f::Fun{S}) where {S<:ContinuousSpace} = cat(1,[extremal_args(fp) for fp in components(f)]..., dims=1)

for op in (:(maximum),:(minimum))
    @eval begin

        $op(f::Fun{ContinuousSpace{T,R},T}) where {T<:Real,R<:Real} =
            $op(map($op,components(f)))
        $op(::typeof(abs), f::Fun{ContinuousSpace{T,R},T}) where {T<:Real,R<:Real} =
            $op(abs, map(g -> $op(abs, g),components(f)))
    end
end

extrema(f::Fun{ContinuousSpace{T,R},T}) where {T<:Real,R<:Real} =
    mapreduce(extrema,(x,y)->extrema([x...;y...]),components(f))

function roots(f::Fun{P}) where P<:ContinuousSpace
    rts=mapreduce(roots,vcat,components(f))
    k=1
    while k < length(rts)
        if isapprox(rts[k],rts[k+1])
            rts=rts[[1:k;k+2:end]]
        else
            k+=1
        end
    end

    rts
end


text:
@doc """
    init_jsi(commands, modeldirs, noises)

!!! note "Advanced"
        init_jsi(commands, modeldirs, noises; <keyword arguments>)

Initialize the package JustSayIt.

# Arguments
- `commands::Dict{String, <:Any}`: the commands to be recognized with their mapping to a function or to a keyboard key or shortcut.
- `modeldirs::Dict{String, String}`: the directories where the unziped speech recognition models to be used are located. Models are downloadable from here: https://alphacephei.com/vosk/models
- `noises::Dict{String, <:AbstractArray{String}}`: for each model, an array of strings with noises (tokens that are to be ignored in the speech as, e.g., "huh").
!!! note "Advanced keyword arguments"
    - `vosk_log_level::Integer=-1`: the vosk log level (see in the Vosk documentation for details).

See also: [`finalize_jsi`](@ref)
"""
init_jsi

let
    global init_jsi, command, command_names, model, noises, noises_names, recognizer, controller, set_controller
    _commands::Dict{String, Union{Function, PyKey, NTuple{N,PyKey} where N}}                            = Dict{String, Union{Function, PyKey, NTuple{N,PyKey} where N}}()
    _models::Dict{String, PyObject}                                                                     = Dict{String, PyObject}()
    _noises::Dict{String, <:AbstractArray{String}}                                                      = Dict{String, Array{String}}()
    _recognizers::Dict{String, PyObject}                                                                = Dict{String, PyObject}()
    _controllers::Dict{String, PyObject}                                                                = Dict{String, PyObject}()
    command(name::AbstractString)                                                                       = _commands[name]
    command_names()                                                                                = keys(_commands)
    model(name::AbstractString=DEFAULT_MODEL_NAME)::PyObject                                            = _models[name]
    noises(modelname::AbstractString)                                                                   = _noises[modelname]  # NOTE: no return value declaration as would be <:AbstractArray{String} which is not possible.
    noises_names()                                                                                = keys(_noises)
    recognizer(id::AbstractString)::PyObject                                                            = _recognizers[id]
    controller(name::AbstractString)::PyObject                                                          = if (name in keys(_controllers)) return _controllers[name] else @APIUsageError("The controller for $name is not available as it has not been set up in init_jsi.") end
    set_controller(name::AbstractString, c::PyObject)                                                   = (_controllers[name] = c; return)

    function init_jsi(commands::Dict{String, <:Any}, modeldirs::Dict{String, String}, noises::Dict{String, <:AbstractArray{String}}; vosk_log_level::Integer=-1)
        Vosk.SetLogLevel(vosk_log_level)

        # Validate and store the commands, adding the help command to it.
        if haskey(commands, COMMAND_NAME_SLEEP) @ArgumentError("the command name $COMMAND_NAME_SLEEP is reserved for putting JustSayIt to sleep. Please choose another command name for your command.") end
        if haskey(commands, COMMAND_NAME_AWAKE) @ArgumentError("the command name $COMMAND_NAME_AWAKE is reserved for awaking JustSayIt. Please choose another command name for your command.") end
        for cmd_name in keys(commands)
            if !(typeof(commands[cmd_name]) <: eltype(values(_commands))) @ArgumentError("the command belonging to commmand name $command_name is of an invalid. Valid are functions (e.g., Keyboard.type), keys (e.g., Key.ctrl or 'f') and tuples of keys (e.g., (Key.ctrl, 'c') )") end
        end
        _commands = commands

        # Verify that DEFAULT_MODEL_NAME is listed in modeldirs and noises. Set the values for the  other surely required models (i.e. which are used in the Commands submodule) to the same as the default if not available.
        if !haskey(modeldirs, DEFAULT_MODEL_NAME) @ArgumentError("a directory for the 'default' model is mandatory.") end
        if !haskey(modeldirs, TYPE_MODEL_NAME) @ArgumentError("a directory for the 'type' model is mandatory.") end
        if haskey(modeldirs, "") @ArgumentError("an empty string is not valid as model identifier.") end
        if !haskey(noises, DEFAULT_MODEL_NAME) @ArgumentError("a noises list for 'default' model is mandatory (e.g. `NOISES_ENGLISH`).") end
        if !haskey(noises, TYPE_MODEL_NAME)
            noises[TYPE_MODEL_NAME] = noises[DEFAULT_MODEL_NAME]
            @warn("no noises given for model \"type\" - falling back to noises configuration of model 'default' when typing.")
        end
        _noises = noises

        # If the modeldirs point to the default path, download a model if none is present (asumed present if the folder is present)
        if modeldirs[DEFAULT_MODEL_NAME] == DEFAULT_MODELDIRS[DEFAULT_MODEL_NAME]
            if !isdir(modeldirs[DEFAULT_MODEL_NAME])
                @info "No 'default' model found in its default location: downloading small english model ($DEFAULT_ENGLISH_MODEL_ARCHIVE) from '$DEFAULT_MODEL_REPO' (~40 MB)..."
                modeldepot = joinpath(modeldirs[DEFAULT_MODEL_NAME], "..")
                download_and_unzip(modeldepot, DEFAULT_ENGLISH_MODEL_ARCHIVE, DEFAULT_MODEL_REPO)
            end
        end
        if modeldirs[TYPE_MODEL_NAME] == DEFAULT_MODELDIRS[TYPE_MODEL_NAME]
            if !isdir(modeldirs[TYPE_MODEL_NAME])
                @info "No 'type' model found in its default location: download (optional) accurate large english model ($DEFAULT_ENGLISH_TYPE_MODEL_ARCHIVE) from '$DEFAULT_MODEL_REPO' (~1 GB)?"
                answer = ""
                while !(answer in ["yes", "no"]) println("Type \"yes\" or \"no\":")
                    answer = readline()
                end
                if answer == "yes"
                    modeldepot = joinpath(modeldirs[TYPE_MODEL_NAME], "..")
                    download_and_unzip(modeldepot, DEFAULT_ENGLISH_TYPE_MODEL_ARCHIVE, DEFAULT_MODEL_REPO)
                else
                    modeldirs[TYPE_MODEL_NAME] = modeldirs[DEFAULT_MODEL_NAME]
                    @warn("Not downloading `type` model: falling back to model `default` for typing.")
                end
            end
        end

        # Set up a default recognizer for each model as well as the command name recognizer.
        for modelname in keys(modeldirs)
            if !isdir(modeldirs[modelname])
                tilde_errmsg = ""
                if startswith(modeldirs[modelname], "~") tilde_errmsg = " Tildes ('~') in strings are not expanded for portability reasons (consider using 'homedir()')." end
                @ArgumentError("directory $(modeldirs[modelname]) does not exist.$tilde_errmsg")
            end
            _models[modelname] = Vosk.Model(modeldirs[modelname])
            _recognizers[modelname] = Vosk.KaldiRecognizer(model(modelname), SAMPLERATE)
            if modelname == DEFAULT_MODEL_NAME
                grammar = json([keys(commands)..., COMMAND_NAME_SLEEP, COMMAND_NAME_AWAKE, noises[modelname]..., UNKNOWN_TOKEN])
                _recognizers[COMMAND_RECOGNIZER_ID] = Vosk.KaldiRecognizer(model(modelname), SAMPLERATE, grammar)
            end
        end

        # Set up the special recognizers defined by voiceargs.
        for f_name in voicearg_f_names()
            for voicearg in keys(voiceargs(f_name))
                kwargs = voiceargs(f_name)[voicearg]
                if haskey(kwargs, :model) && !haskey(modeldirs, kwargs[:model]) @ArgumentError("no directory was given for the model $(kwargs[:model]) required for voicearg $voicearg in function $f_name.") end # NOTE: this error should only ever occur for user defined @voicearg functions; all funtions in submodule Commands must use
                if haskey(kwargs, :valid_input)
                    modelname = haskey(kwargs, :model) ? kwargs[:model] : DEFAULT_MODEL_NAME
                    grammar = json([kwargs[:valid_input]..., noises[modelname]..., UNKNOWN_TOKEN])
                    set_recognizer(f_name, voicearg, Vosk.KaldiRecognizer(model(modelname), SAMPLERATE, grammar))
                end
            end
            @debug "Voiceargs of function `$f_name`:" voiceargs(f_name)
        end

        # Call the garbage collector to avoid having it running right after starting to use JustSayIt.
        GC.gc()
    end

end

_noises(args...) = noises(args...)

function download_and_unzip(destination, filename, repository)
    progress = nothing
    previous = 0
    function show_progress(total::Integer, now::Integer)
        if (total > 0) && (now != previous)
            if isnothing(progress) progress = ProgressMeter.Progress(total; dt=0.1, desc="Download ($(Base.format_bytes(total))): ", color=:magenta, output=stderr) end
            ProgressMeter.update!(progress, now)
        end
        previous = now
    end
    mkpath(destination)
    filepath = joinpath(destination, filename)
    Downloads.download(repository * "/" * filename, filepath; progress=show_progress)
    @pywith Zipfile.ZipFile(filepath, "r") as archive begin
        archive.extractall(destination)
    end
end


text:
@testset "PolyPairBasis" begin

@info("-------- Test PolyPairBasis Implementation ---------")

##

using PolyPairPots, JuLIP, LinearAlgebra, Test
using JuLIP.Testing

##

at = bulk(:W, cubic=true) * 3
rattle!(at, 0.03)
r0 = rnn(:W)
X = copy(positions(at))

trans = PolyTransform(2, r0)
fcut = PolyCutoff1s(2, 0.5*r0, 2.5*r0)
pB = PolyPairBasis(:W, 10, trans, fcut)

##

@info("test (de-)dictionisation of PairBasis")
println(@test decode_dict(Dict(pB)) == pB)

E = energy(pB, at)
DE = - forces(pB, at)

@info("Finite-difference test on PolyPairBasis forces")
for ntest = 1:20
   U = rand(JVecF, length(at)) .- 0.5
   DExU = dot.(DE, Ref(U))
   errs = Float64[]
   for p = 2:10
      h = 0.1^p
      Eh = energy(pB, set_positions!(at, X+h*U))
      DEhxU = (Eh-E) / h
      push!(errs, norm(DExU - DEhxU, Inf))
   end
   success = (/(extrema(errs)...) < 1e-3) || (minimum(errs) < 1e-10)
   print_tf(@test success)
end
println()
##

end

@testset "PolyPairPot" begin

@info("--------------- PolyPairPot Implementation ---------------")

##

using PolyPairPots, JuLIP, LinearAlgebra, Test
using JuLIP.Testing, LinearAlgebra
using JuLIP.MLIPs

##

at = bulk(:W, cubic=true) * 3
rattle!(at, 0.03)
r0 = rnn(:W)

trans = PolyTransform(2, r0)
fcut = PolyCutoff2s(2, 0.5*r0, 1.95*r0)
B = PolyPairBasis(:W, 10, trans, fcut)

##

@info("Testing correctness of `PolyPairPot` against `PolyPairBasis`")

@info("    test `combine`")
coeffs = randcoeffs(B)
pot = combine(B, coeffs)
println(@test pot == PolyPairPot(B, coeffs))

##
@info("   test (de-)dictionisation")
println(@test decode_dict(Dict(pot)) == pot)

@info("      check that PolyPairBasis ≈ PolyPairPot")
for ntest = 1:30
   rattle!(at, 0.01)

   E_pot = energy(pot, at)
   E_b = dot(energy(B, at), coeffs)
   print_tf(@test E_pot ≈ E_b)

   F_pot = forces(pot, at)
   F_b = sum(coeffs .* forces(B, at))
   print_tf(@test F_pot ≈ F_b)

   V_pot = virial(pot, at)
   V_b = sum(coeffs .* virial(B, at))
   print_tf(@test V_pot ≈ V_b)
end
println()

@info("      Standard JuLIP Consistency Test")
variablecell!(at)
rattle!(at, 0.03)
JuLIP.Testing.fdtest(pot, at)

##
end


text:
using Indicators
using Temporal
using Test
using Random
using Statistics

const N = 1_000
const X0 = 50.0
const seed = 1

Random.seed!(seed)
x = cumsum(randn(N), dims=1) .+ X0
ohlc = cumsum(randn(N,4), dims=1) .+ X0
hlc = ohlc[:,2:4]
hl = ohlc[:,2:3]

count_nans(x) = sum(isnan.(x))

@testset "Utilities" begin
    y = x + randn(N)
    cxo = crossover(x, y)
    cxu = crossunder(x, y)
    @test any(cxo)
    @test any(cxu)
    @test !any(cxo .* cxu)  # ensure crossovers and crossunders never coincide
end

# trendy
@testset "Trendlines" begin
    tmp = resistance(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = support(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = minima(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = maxima(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
end

# moving regressions
@testset "Regressions" begin
    tmp = mlr_beta(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 2
    @test count_nans(tmp) != N
    tmp = mlr_slope(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = mlr_intercept(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = mlr(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = mlr_se(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = mlr_ub(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = mlr_lb(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = mlr_bands(tmp)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 3
    @test count_nans(tmp) != N
    tmp = mlr_rsq(x, adjusted=true)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = mlr_rsq(x, adjusted=false)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
end

@testset "Running Calculations" begin
    tmp = diffn(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = diffn(ohlc)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == size(ohlc,2)
    @test count_nans(tmp) != N
    tmp = runmean(x, cumulative=true)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = runmean(x, cumulative=false)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = runsum(x, cumulative=true)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = wilder_sum(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = runmad(x, cumulative=true)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = runmad(x, cumulative=false)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = runvar(x, cumulative=true)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = runvar(x, cumulative=false)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = runsd(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = runcov(x, x.*rand(N), cumulative=true)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = runcov(x, x.*rand(N), cumulative=false)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = runcor(x, x.*rand(N), cumulative=true)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = runcor(x, x.*rand(N), cumulative=false)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = runmin(x, cumulative=true, inclusive=true)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = runmin(x, cumulative=true, inclusive=false)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = runmin(x, cumulative=false, inclusive=true)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = runmin(x, cumulative=false, inclusive=false)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = runmax(x, cumulative=true, inclusive=true)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = runmax(x, cumulative=true, inclusive=false)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = runmax(x, cumulative=false, inclusive=true)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = runmax(x, cumulative=false, inclusive=false)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = mode(map(xi->round(xi), x))
    @test size(tmp, 1) == 1
    @test size(tmp, 2) == 1
    @test !isnan(tmp)
    tmp = runquantile(x, cumulative=true)
    @test !isnan(tmp[2]) && isnan(tmp[1])
    @test tmp[10] == quantile(x[1:10], 0.05)
    tmp = runquantile(x, cumulative=false)
    @test tmp[10] == quantile(x[1:10], 0.05)
    n = 20
    tmp = runacf(x, n=n, maxlag=15, cumulative=true)
    @test all(tmp[n:end,1] .== 1.0)
    tmp = runacf(x, n=n, maxlag=15, cumulative=false)
    @test all(tmp[n:end,1] .== 1.0)
end

# moving average functions
@testset "Moving Averages" begin
    tmp = sma(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = mama(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 2
    @test count_nans(tmp) != N
    tmp = ema(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = wma(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = hma(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = trima(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = mma(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = tema(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = dema(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = swma(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = kama(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = alma(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = zlema(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
end

# momentum function
@testset "Momentum" begin
    tmp = aroon(hl)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 3
    @test count_nans(tmp) != N
    tmp = donch(hl)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 3
    @test count_nans(tmp) != N
    tmp = momentum(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = roc(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = macd(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 3
    @test count_nans(tmp) != N
    tmp = rsi(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = adx(hlc)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 3
    @test count_nans(tmp) != N
    tmp = adx(hlc, wilder=true)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 3
    @test count_nans(tmp) != N
    tmp = psar(hl)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = kst(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = wpr(hlc)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = cci(hlc)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = stoch(hlc, kind=:fast)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 2
    @test count_nans(tmp) != N
    tmp = stoch(hlc, kind=:slow)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 2
    @test count_nans(tmp) != N
    tmp = smi(hlc)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 2
    @test count_nans(tmp) != N
end

# volatility functions
@testset "Volatility" begin
    tmp = bbands(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 3
    @test count_nans(tmp) != N
    tmp = tr(hlc)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = atr(hlc)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = keltner(hlc)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 3
    @test count_nans(tmp) != N
end

# chart patterns functions
@testset "Charting" begin
    tmp = renko(hlc, use_atr=true)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
    tmp = renko(hlc, use_atr=false)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    @test count_nans(tmp) != N
end

# ==== TEMPORAL INTERACTIONS ====
@testset "Temporal" begin
    x = TS(cumsum(randn(N), dims=1) .+ X0)
    x.fields = [:Close]
    ohlc = TS(cumsum(randn(N,4), dims=1))
    ohlc.fields = [:Open, :High, :Low, :Close]
    hlc = ohlc[:,2:4]
    # moving regressions
    tmp = mlr_beta(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 2
    tmp = mlr_slope(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = mlr_intercept(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = mlr(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = mlr_se(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = mlr_ub(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = mlr_lb(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = mlr_bands(tmp)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 3
    tmp = mlr_rsq(x, adjusted=true)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = mlr_rsq(x, adjusted=false)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    # running calculations
    tmp = runmean(x, cumulative=true)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = runmean(x, cumulative=false)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = runsum(x, cumulative=true)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = wilder_sum(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = runmad(x, cumulative=true)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = runmad(x, cumulative=false)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = runvar(x, cumulative=true)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = runvar(x, cumulative=false)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = runsd(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = runcov(ohlc[:,1], ohlc[:,4], cumulative=true)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = runcov(ohlc[:,1], ohlc[:,4], cumulative=false)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = runcor(ohlc[:,1], ohlc[:,4], cumulative=true)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = runcor(ohlc[:,1], ohlc[:,4], cumulative=false)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = runmin(x, cumulative=true, inclusive=true)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = runmin(x, cumulative=true, inclusive=false)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = runmin(x, cumulative=false, inclusive=true)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = runmin(x, cumulative=false, inclusive=false)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = runmax(x, cumulative=true, inclusive=true)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = runmax(x, cumulative=true, inclusive=false)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = runmax(x, cumulative=false, inclusive=true)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = runmax(x, cumulative=false, inclusive=false)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = runquantile(x, cumulative=true)
    @test !isnan(tmp.values[2,1]) && isnan(tmp.values[1,1])
    @test tmp.values[10,1] == quantile(x.values[1:10,1], 0.05)
    tmp = runquantile(x, cumulative=false)
    @test tmp.values[10,1] == quantile(x.values[1:10,1], 0.05)
    n = 20
    tmp = runacf(x, n=n, maxlag=15, cumulative=true)
    @test all(tmp.values[n:end,1] .== 1.0)
    tmp = runacf(x, n=n, maxlag=15, cumulative=false)
    @test all(tmp.values[n:end,1] .== 1.0)
    # moving average functions
    tmp = sma(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = mama(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 2
    tmp = ema(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = wma(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = hma(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = trima(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = mma(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = tema(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = dema(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = swma(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = kama(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = alma(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = zlema(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    # momentum function
    tmp = aroon(hl)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 3
    tmp = donch(hl)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 3
    tmp = momentum(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = roc(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = macd(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 3
    tmp = rsi(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = adx(hlc)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 3
    tmp = adx(hlc, wilder=true)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 3
    tmp = psar(hl)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = kst(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = wpr(hlc)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = cci(hlc)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = stoch(hlc, kind=:fast)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 2
    tmp = stoch(hlc, kind=:slow)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 2
    tmp = smi(hlc)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 2
    # volatility functions
    tmp = bbands(x)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 3
    tmp = tr(hlc)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = atr(hlc)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 1
    tmp = keltner(hlc)
    @test size(tmp, 1) == N
    @test size(tmp, 2) == 3
end


text:
using LinearAlgebra: transpose!
using AstroBase: NAIFId, Epoch, TDBEpoch, from_naifid, julian_twopart, value, SECONDS_PER_DAY

import AstroBase.Ephemerides:
    AbstractEphemeris,
    position,
    position!,
    velocity,
    velocity!,
    state,
    state!

export SPK,
    segments,
    print_segments,
    position,
    position!,
    velocity,
    velocity!,
    state,
    state!

const SIZE_FLOAT64 = sizeof(Float64)

struct OutOfRangeError <: Exception
    date::Float64
    startdate::Float64
    finaldate::Float64
end

Base.showerror(io::IO, err::OutOfRangeError) = print(io,
   "The requested date $(err.date) is outside the interval ($(err.startdate), $(err.finaldate)).")

mutable struct Segment
    name::String
    firstsec::Float64
    lastsec::Float64
    firstdate::Float64
    lastdate::Float64
    target::Int
    center::Int
    frame::Int
    spktype::Int
    firstaddr::Int
    lastaddr::Int
    firstword::Int
    lastword::Int
    initialsecond::Float64
    intlen::Float64
    rsize::Int
    n_records::Int
    order::Int
    cached_record::Int
    cache::Matrix{Float64}
    x::Vector{Float64}
    t::Vector{Float64}
end

jd(sec) = 2451545.0 + sec / SECONDS_PER_DAY
seconds(jd) = (jd - 2451545.0) * SECONDS_PER_DAY

function Segment(daf, name, record)
    firstsec, lastsec = reinterpret_getindex(Float64, record, (1, 9), daf.little)
    target, center, frame, spktype, firstaddr, lastaddr =
        reinterpret_getindex(Int32, record, (17, 21, 25, 29, 33, 37), daf.little)
    if spktype != 2
        throw(ArgumentError("Type $spktype SPK file detected. Only Type 2 SPK files are supported."))
    end
    i0 = lastaddr * SIZE_FLOAT64 - 4 * SIZE_FLOAT64 + 1
    init, intlen, rsize, n_records =
        reinterpret_getindex(Float64, daf.array, (i0, i0 + 8, i0 + 16, i0 + 24), daf.little)
    n_records = round(Int32, n_records)
    order = Int((rsize - 2) ÷ 3)
    Segment(
        name,
        firstsec,
        lastsec,
        jd(firstsec),
        jd(lastsec),
        target,
        center,
        frame,
        spktype,
        firstaddr,
        lastaddr,
        firstaddr*SIZE_FLOAT64 - SIZE_FLOAT64 + 1,
        lastaddr*SIZE_FLOAT64 - SIZE_FLOAT64*4,
        init,
        intlen,
        round(Int32, rsize),
        n_records,
        order,
        -1,
        zeros(3, order),
        zeros(order),
        zeros(order),
    )
end

struct SPK <: AbstractEphemeris
    daf::DAF
    segments::Dict{Int,Dict{Int,Segment}}
end

Base.show(io::IO, spk::SPK) = print(io, "SPK($(spk.segments[0][1].name))")

function SPK(filename)
    daf = DAF(filename)
    segments = Dict{Int, Dict{Int, Segment}}()
    for (name, summary) in getsummaries(daf)
        seg = Segment(daf, name, summary)
        if haskey(segments, seg.center)
            merge!(segments[seg.center], Dict(seg.target=>seg))
        else
            merge!(segments, Dict(seg.center=>Dict(seg.target=>seg)))
        end
    end
    SPK(daf, segments)
end

segments(spk::SPK) = spk.segments

function list_segments(spk::SPK)
    s = String[]
    for (k,v) in spk.segments
        for l in keys(v)
            push!(s, "$(from_naifid(k)) ($k) => $(from_naifid(l)) ($l)")
        end
    end
    sort!(s, lt=segstrlt)
end

function print_segments(spk::SPK)
    s = list_segments(spk)
    println(join(s, "\n"))
end

function segstrlt(a::String, b::String)
   rex = r"\([0-9]*\)$"
   ma = match(rex, a)
   mb = match(rex, b)
   ia = parse(Int, a[ma.offset+1:end-1])
   ib = parse(Int, b[mb.offset+1:end-1])
   ia < ib
end

@inline function checkdate(seg::Segment, tdb::Float64)
    if !(seg.firstdate <= tdb <= seg.lastdate)
        throw(OutOfRangeError(tdb, seg.firstdate, seg.lastdate))
    end
end

@inline function getrecordnum(seg, tdb, tdb2)
    checkdate(seg, tdb+tdb2)
    secs = (seconds(tdb) - seg.initialsecond) + tdb2 * SECONDS_PER_DAY
    recordnum, frac = divrem(secs, seg.intlen)
    recordnum = round(Int, recordnum)
    if recordnum == seg.n_records
        recordnum -= 1
        frac = seg.intlen
    end
    recordnum, frac
end

@inline function update_cache!(spk::SPK, seg::Segment, recordnum)
    components = 3
    seg.cached_record = recordnum
    # Drop the MID and RADIUS values
    first = seg.firstword + SIZE_FLOAT64 * seg.rsize * recordnum + SIZE_FLOAT64 * 2
    ptr = Ptr{Float64}(pointer(spk.daf.array, first))

    cache = unsafe_wrap(Array, Ptr{Float64}(ptr), (seg.order, components), own=false)
    if !spk.daf.little
        transpose!(seg.cache, ntoh.(copy(cache)))
    else
        transpose!(seg.cache, cache)
    end
end

@inline function chebyshev!(seg, frac)
    seg.x[1] = 1.0
    seg.x[2] = 2.0 * frac / seg.intlen - 1.0
    @inbounds for i = 3:seg.order
        seg.x[i] = 2.0 * seg.x[2] * seg.x[i-1] - seg.x[i-2]
    end
end

@inline function chebyshev_deriv!(seg)
    seg.t[2] = 1.0
    if seg.order > 2
        seg.t[3] = 4.0 * seg.x[2]
        @inbounds for i = 4:seg.order
            seg.t[i] = 2.0 * seg.x[2] * seg.t[i-1] - seg.t[i-2] +
                seg.x[i-1] + seg.x[i-1]
        end
    end
    seg.t .*= 2.0
    seg.t ./= seg.intlen
end

@inline function position!(r, seg::Segment, sign::Float64)
    @inbounds @simd for i = 1:3
        for j = 1:seg.order
            r[i] += sign * seg.cache[i, j] * seg.x[j]
        end
    end
    r
end

@inline function position!(r, spk::SPK, seg::Segment, sign::Float64, tdb::Float64, tdb2::Float64=0.0)
    recordnum, frac = getrecordnum(seg, tdb, tdb2)
    if recordnum != seg.cached_record
        update_cache!(spk, seg, recordnum)
    end
    chebyshev!(seg, frac)
    position!(r, seg, sign)
end

@inline function velocity!(v, seg::Segment, sign::Float64)
    chebyshev_deriv!(seg)
    @inbounds @simd for i = 1:3
        for j = 1:seg.order
            v[i] += sign * seg.cache[i, j] * seg.t[j]
        end
    end
    v
end

@inline function velocity!(v, spk::SPK, seg::Segment, sign::Float64, tdb::Float64, tdb2::Float64=0.0)
    recordnum, frac = getrecordnum(seg, tdb, tdb2)
    if recordnum != seg.cached_record
        update_cache!(spk, seg, recordnum)
    end
    chebyshev!(seg, frac)
    velocity!(v, seg, sign)
end

@inline function state!(pos,
                                    vel,
                                    spk::SPK,
                                    seg::Segment,
                                    sign::Float64,
                                    tdb::Float64,
                                    tdb2::Float64=0.0)
    recordnum, frac = getrecordnum(seg, tdb, tdb2)
    if recordnum != seg.cached_record
        update_cache!(spk, seg, recordnum)
    end
    chebyshev!(seg, frac)
    position!(pos, seg, sign)
    velocity!(vel, seg, sign)
    pos, vel
end

@inline function findsegment(segments, origin, target)
    if !(origin in keys(segments) || target in keys(segments))
        throw(ArgumentError("No segment '$origin'->'$target' available."))
    end
    sign = 1.0
    if target < origin
        origin, target = target, origin
        sign = -1.0
    end
    segments[origin][target], sign
end

function position!(pos, spk::SPK, ep::Epoch, from::NAIFId, to::NAIFId)
    seg, sign = findsegment(spk.segments, from, to)
    jd1, jd2 = value.(julian_twopart(TDBEpoch(ep)))
    position!(pos, spk, seg, sign, jd1, jd2)
end

function velocity!(vel, spk::SPK, ep::Epoch, from::NAIFId, to::NAIFId)
    seg, sign = findsegment(spk.segments, from, to)
    jd1, jd2 = value.(julian_twopart(TDBEpoch(ep)))
    velocity!(vel, spk, seg, sign, jd1, jd2)
end

function state!(pos, vel, spk::SPK, ep::Epoch, from::NAIFId, to::NAIFId)
    seg, sign = findsegment(spk.segments, from, to)
    jd1, jd2 = value.(julian_twopart(TDBEpoch(ep)))
    state!(pos, vel, spk, seg, sign, jd1, jd2)
end