text:
# ----------------------------------------- #
# functions to handle IHT output

# an object that houses results returned from an IHT run
immutable IHTResults{T <: Float, V <: DenseVector}
    time :: T
    loss :: T
    iter :: Int
    beta :: V

    #IHTResults{T,V}(time::T, loss::T, iter::Int, beta::V) where {T <: Float, V <: DenseVector{T}} = new{T,V}(time, loss, iter, beta)
    IHTResults{T,V}(time, loss, iter, beta) where {T <: Float, V <: DenseVector{T}} = new{T,V}(time, loss, iter, beta)
end

# strongly typed external constructor for IHTResults
IHTResults(time::T, loss::T, iter::Int, beta::V) where {T <: Float, V <: DenseVector{T}} = IHTResults{T, V}(time, loss, iter, beta)

# function to display IHTResults object
function Base.show(io::IO, x::IHTResults)
    println(io, "IHT results:")
    println(io, "\nCompute time (sec):   ", x.time)
    println(io, "Final loss:           ", x.loss)
    println(io, "Iterations:           ", x.iter)
    println(io, "IHT estimated ", countnz(x.beta), " nonzero coefficients.")
    print(io, DataFrame(Predictor=find(x.beta), Estimated_β=x.beta[find(x.beta)]))
    return nothing
end

# ----------------------------------------- #
# functions for handling temporary arrays

# an object to contain intermediate variables and temporary arrays
type IHTVariables{T <: Float, V <: DenseVector}
    b    :: V
    b0   :: Vector{T}
    xb   :: V
    xb0  :: Vector{T}
    xk   :: Matrix{T}
    gk   :: Vector{T}
    xgk  :: Vector{T}
    idx  :: BitArray{1}
    idx0 :: BitArray{1}
    r    :: V
    df   :: V

    IHTVariables{T,V}(b::V, b0::Vector{T}, xb::V, xb0::Vector{T}, xk::Matrix{T}, gk::Vector{T}, xgk::Vector{T}, idx::BitArray{1}, idx0::BitArray{1}, r::V, df::V) where {T <: Float, V <: DenseVector{T}} = new{T,V}(b, b0, xb, xb0, xk, gk, xgk, idx, idx0, r, df)
end

## strong type construction of IHTVariables
#IHTVariables{T <: Float}(
#    b    :: DenseVector{T},
#    b0   :: Vector{T},
#    xb   :: DenseVector{T},
#    xb0  :: Vector{T},
#    xk   :: Matrix{T},
#    gk   :: Vector{T},
#    xgk  :: Vector{T},
#    idx  :: BitArray{1},
#    idx0 :: BitArray{1},
#    r    :: DenseVector{T},
#    df   :: DenseVector{T}
#) where {T <: Float, V <: DenseVector{T}} = IHTVariables{T, V}(b, b0, xb, xb0, xk, gk, xgk, idx, idx0, r, df)

# minor type instability with SharedArray constructors
# a small optimization for the future, but will not dramatically change performance
# try defining b0, xb0 first and then calling "SharedVector(b0)", "SharedVector(xb0)"
function IHTVariables{T <: Float}(
    x :: SharedMatrix{T},
    y :: SharedVector{T},
    k :: Int
)
    pids = procs(x)
    V    = typeof(y)
    n, p = size(x)
    b    = SharedArray{T}((p,), pids=pids) :: V
    df   = SharedArray{T}((p,), pids=pids) :: V
    xb   = SharedArray{T}((n,), pids=pids) :: V
    r    = SharedArray{T}((n,), pids=pids) :: V
    b0   = zeros(T, p)
    xb0  = zeros(T, n)
    xk   = zeros(T, n, k)
    xgk  = zeros(T, n)
    gk   = zeros(T, k)
    idx  = falses(p)
    idx0 = falses(p)
    return IHTVariables{T, V}(b, b0, xb, xb0, xk, gk, xgk, idx, idx0, r, df)
end

# strongly typed constructor for IHTVariables using regular floating point arrays
function IHTVariables{T <: Float}(
    x :: Matrix{T},
    y :: Vector{T},
    k :: Int
)
    n, p = size(x)
    b    = zeros(T, p)
    b0   = zeros(T, p)
    df   = zeros(T, p)
    xk   = zeros(T, n, k)
    xb   = zeros(T, n)
    xb0  = zeros(T, n)
    r    = zeros(T, n)
    xgk  = zeros(T, n)
    gk   = zeros(T, k)
    idx  = falses(p)
    idx0 = falses(p)
    return IHTVariables{T, typeof(y)}(b, b0, xb, xb0, xk, gk, xgk, idx, idx0, r, df)
end

# strongly typed constructor for IHTVariables using BEDFile + SharedVector
function IHTVariables{T <: Float}(
    x :: BEDFile{T},
    y :: SharedVector{T},
    k :: Int
)
    n, p = size(x)
    pids = procs(x)
    V    = typeof(y)
    b    = SharedArray{T}((p,), pids=pids) :: V
    df   = SharedArray{T}((p,), pids=pids) :: V
    xb   = SharedArray{T}((n,), pids=pids) :: V
    r    = SharedArray{T}((n,), pids=pids) :: V
    b0   = zeros(T, p)
    xb0  = zeros(T, n)
    xk   = zeros(T, n, k)
    xgk  = zeros(T, n)
    gk   = zeros(T, k)
    idx  = falses(p)
    idx0 = falses(p)
    return IHTVariables{T, V}(b, b0, xb, xb0, xk, gk, xgk, idx, idx0, r, df)
end

# function to modify fields of IHTVariables that depend on k
function update_variables!{T <: Float}(
    v :: IHTVariables{T},
    x :: DenseMatrix{T},
    k :: Int
)
    n    = size(x,1)
    v.xk = zeros(T, n, k)
    v.gk = zeros(T, k)
    return nothing
end

function update_variables!{T <: Float}(
    v :: IHTVariables{T},
    x :: BEDFile{T},
    k :: Int
)
    n    = size(x,1)
    v.xk = zeros(T, n, k)
    v.gk = zeros(T, k)
    return nothing
end

# ----------------------------------------- #
# crossvalidation routines

# subroutine to compute a default number of folds
@inline cv_get_num_folds(nmin::Int, nmax::Int) = max(nmin, min(Sys.CPU_CORES::Int, nmax))

# subroutine to refit preditors after crossvalidation
function refit_iht{T <: Float}(
    x        :: DenseMatrix{T},
    y        :: DenseVector{T},
    k        :: Int;
    tol      :: T    = convert(T, 1e-6),
    max_iter :: Int  = 100,
    max_step :: Int  = 50,
    quiet    :: Bool = true,
)
    # initialize β vector and temporary arrays
    v = IHTVariables(sdata(x), sdata(y), k)

    # first use exchange algorithm to extract model
    output = L0_reg(sdata(x), sdata(y), k, v=v, max_iter=max_iter, max_step=max_step, quiet=quiet, tol=tol)

    # which components of β are nonzero?
    # cannot use binary indices here since we need to return Int indices
    bidx = find(output.beta) :: Vector{Int}
    k2 = length(bidx)

    # allocate the submatrix of x corresponding to the inferred model
    # cannot use SubArray with BLAS since result is not StridedArray
    # issue is that bidx is Vector{Int} and not a Range object
    # use of SubArray is more memory efficient, so use it with Julia general multiplications,
    # e.g. At_mul_B and others
    x_inferred = view(sdata(x), :, bidx)

    # now estimate β with the ordinary least squares estimator β = inv(x'x)x'y
    # return it with the vector of MSEs
    xty = Base.At_mul_B(x_inferred, sdata(y)) :: Vector{T}
    xtx = Base.At_mul_B(x_inferred, x_inferred) :: Matrix{T}
    b   = zeros(T, k2)
    try
        b = (xtx \ xty) :: Vector{T}
    catch e
        warn("caught error: ", e, "\nSetting returned values of b to -Inf")
        fill!(b, -Inf)
    end

    return b, bidx
end

# refitting routine for GWAS data with x', mean, prec files
function refit_iht(
    T        :: Type,
    xfile    :: String,
    xtfile   :: String,
    x2file   :: String,
    yfile    :: String,
    meanfile :: String,
    precfile :: String,
    k        :: Int;
    pids     :: Vector{Int} = procs(),
    tol      :: Float = convert(T, 1e-6),
    max_iter :: Int   = 100,
    max_step :: Int   = 50,
    quiet    :: Bool  = true,
    header   :: Bool  = false
)

    # initialize all variables
    x = BEDFile(T, xfile, xtfile, x2file, meanfile, precfile, pids=pids, header=header) :: BEDFile{T}
    y = SharedArray{T}(abspath(yfile), (x.geno.n,), pids=pids) :: SharedVector{T}

    # extract model with IHT
    output = L0_reg(x, y, k, max_iter=max_iter, max_step=max_step, quiet=quiet, tol=tol)

    # which components of β are nonzero?
    inferred_model = output.beta .!= zero(T)
    bidx = find(inferred_model)

    # allocate the submatrix of x corresponding to the inferred model
    x_inferred = zeros(T, x.geno.n, sum(inferred_model))
    decompress_genotypes!(x_inferred, x, inferred_model)

    # now estimate b with the ordinary least squares estimator b = inv(x'x)x'y
    # return it with the vector of MSEs
    xty = Base.At_mul_B(x_inferred, sdata(y))   :: Vector{T}
    xtx = Base.At_mul_B(x_inferred, x_inferred) :: Matrix{T}
    b   = zeros(T, length(bidx))
    try
        b = (xtx \ xty) :: Vector{T}
    catch e
        warn("caught error: ", e, "\nSetting returned values of b to -Inf")
        fill!(b, -Inf)
    end

    # get predictor names before returning
    bids = prednames(x)[bidx]

    return b, bidx, bids
end

# refitting routine for GWAS data with just genotypes, covariates, y
function refit_iht(
    T        :: Type,
    xfile    :: String,
    x2file   :: String,
    yfile    :: String,
    k        :: Int;
    pids     :: Vector{Int} = procs(),
    tol      :: Float = convert(T, 1e-6),
    max_iter :: Int   = 100,
    quiet    :: Bool  = true,
    header   :: Bool  = false
)

    # initialize all variables
    x = BEDFile(T, xfile, x2file, pids=pids, header=header) :: BEDFile{T}
    y = SharedArray{T}(abspath(yfile), (x.geno.n,), pids=pids) :: SharedVector{T}

    # first use exchange algorithm to extract model
    output = L0_reg(x, y, k, max_iter=max_iter, quiet=quiet, tol=tol, window=k)

    # which components of β are nonzero?
    inferred_model = output.beta .!= zero(T)
    bidx = find(inferred_model)

    # allocate the submatrix of x corresponding to the inferred model
    x_inferred = zeros(T, x.geno.n, sum(inferred_model))
    decompress_genotypes!(x_inferred, x, inferred_model)

    # now estimate β with the ordinary least squares estimator b = inv(x'x)x'y
    # return it with the vector of MSEs
    xty = Base.At_mul_B(x_inferred, sdata(y))   :: Vector{T}
    xtx = Base.At_mul_B(x_inferred, x_inferred) :: Matrix{T}
    b   = zeros(T, length(bidx))
    try
        b = (xtx \ xty) :: Vector{T}
    catch e
        warn("caught error: ", e, "\nSetting returned values of b to -Inf")
        fill!(b, -Inf)
    end

    # get predictor names before returning
    bids = prednames(x)[bidx]

    return b, bidx, bids
end

# return type for crossvalidation
immutable IHTCrossvalidationResults{T <: Float}
    mses :: Vector{T}
    path :: Vector{Int}
    b    :: Vector{T}
    bidx :: Vector{Int}
    k    :: Int
    bids :: Vector{String}

    #IHTCrossvalidationResults(mses::Vector{T}, path::Vector{Int}, b::Vector{T}, bidx::Vector{Int}, k::Int, bids::Vector{String}) where {T <: Float} = new(mses, path, b, bidx, k, bids)
    IHTCrossvalidationResults(mses::Vector{T}, path::Vector{Int}, b::Vector{T}, bidx::Vector{Int}, k::Int, bids::Vector{String}) where {T <: Float} = new{T}(mses, path, b, bidx, k, bids)
end

# constructor for when bids are not available
# simply makes vector of "V$i" where $i are drawn from bidx
function IHTCrossvalidationResults{T <: Float}(
    mses :: Vector{T},
    path :: Vector{Int},
    b    :: Vector{T},
    bidx :: Vector{Int},
    k    :: Int
)
    bids = convert(Vector{String}, ["V" * "$i" for i in bidx]) :: Vector{String}
    IHTCrossvalidationResults(mses, path, b, bidx, k, bids) :: IHTCrossvalidationResults{T}
end

# function to view an IHTCrossvalidationResults object
function Base.show(io::IO, x::IHTCrossvalidationResults)
    println(io, "Crossvalidation results:")
    println(io, "Minimum MSE ", minimum(x.mses), " occurs at k = $(x.k).")
    println(io, "Best model β has the following nonzero coefficients:")
    println(io, DataFrame(Predictor=x.bidx, Name=x.bids, Estimated_β=x.b))
    return nothing
end

function Gadfly.plot(x::IHTCrossvalidationResults)
    df = DataFrame(ModelSize=x.path, MSE=x.mses)
    plot(df, x="ModelSize", y="MSE", xintercept=[x.k], Geom.line, Geom.vline(color=colorant"red"), Guide.xlabel("Model size"), Guide.ylabel("MSE"), Guide.title("MSE versus model size"))
end

# ----------------------------------------- #
# subroutines for L0_reg

#function update_xb!{T <: Float}(
#   v :: IHTVariables{T},
#   x :: DenseMatrix{T},
#   k :: Int
#)
#    sum(v.idx) <= k || throw(ArgumentError("Argument indices with $(sum(indices)) trues should have at most $k of them"))
#    fill!(v.xb, zero(T))
#    numtrue = 0
#    @inbounds for j in eachindex(v.idx)
#        if v.idx[j]
#            numtrue += 1
#            @inbounds for i in eachindex(v.xb)
#                v.xb[i] += v.b[j]*x[i,j]
#            end
#        end
#        numtrue >= k && break
#    end
#    return nothing
#end

# subroutine to update residuals and gradient from data
function update_r_grad!{T}(
    v :: IHTVariables{T},
    x :: DenseMatrix{T},
    y :: DenseVector{T}
)
    v.r .= y .- v.xb
    Base.At_mul_B!(v.df, x, v.r) # v.df = x' * v.r
    return nothing
end

function initialize_xb_r_grad!{T <: Float}(
    v :: IHTVariables{T},
    x :: DenseMatrix{T},
    y :: DenseVector{T},
    k :: Int
)
    # update x*beta
    if sum(v.idx) == 0
        fill!(v.xb, zero(T))
    else
        #update_indices!(v.idx, v.b)
        v.idx .= v.b .!= 0
        update_xb!(v.xb, x, v.b, v.idx, k)
        #A_mul_B!(v.xb, view(x :, v.idx), view(v.b, v.idx) )
    end

    # update r and gradient
    update_r_grad!(v, x, y)
end

#function update_r_grad!{T}(
#    v    :: IHTVariables{T},
#    x    :: BEDFile{T},
#    y    :: DenseVector{T};
#    pids :: DenseVector{Int} = procs(x)
#)
#    #difference!(v.r, y, v.xb)
#    broadcast!(-, v.r, y, v.xb)
#    PLINK.At_mul_B!(v.df, x, v.r, pids=pids)
#    return nothing
#end
#
#function initialize_xb_r_grad!{T <: Float}(
#    v    :: IHTVariables{T},
#    x    :: BEDFile{T},
#    y    :: DenseVector{T},
#    k    :: Int;
#    pids :: DenseVector{Int} = procs(x)
#)
#    if sum(v.idx) == 0
#        fill!(v.xb, zero(T))
#        copy!(v.r, y)
#        At_mul_B!(v.df, x, v.r, pids=pids)
#    else
#        update_indices!(v.idx, v.b)
#        v.idx .= v.b .!= 0
#        A_mul_B!(v.xb, x, v.b, v.idx, k, pids=pids)
#        update_r_grad!(v, x, y, pids=pids)
#    end
#    return nothing
#end

# ----------------------------------------- #
# common subroutines for IHT stepping

# this function updates the BitArray indices for b
function _iht_indices{T <: Float}(
    v :: IHTVariables{T},
    k :: Int;
)
    # which components of beta are nonzero?
    #update_indices!(v.idx, v.b)
    v.idx .= v.b .!= 0

    # if current vector is 0,
    # then take largest elements of d as nonzero components for b
    if sum(v.idx) == 0
        a = select(v.df, k, by=abs, rev=true) :: T
#        threshold!(IDX, g, abs(a), n=p)
        v.idx[abs.(v.df) .>= abs(a)-2*eps()] = true
        v.gk = zeros(T, sum(v.idx))
    end

    return nothing
end

# this function computes the step size for one update with iht()
function _iht_stepsize{T <: Float}(
    v :: IHTVariables{T},
    k :: Int
)
    # store relevant components of gradient
    v.gk = v.df[v.idx]

    # compute xk' * gk
    A_mul_B!(v.xgk, v.xk, v.gk)

    # warn if xgk only contains zeros
#    all(v.xgk .== zero(T)) && warn("Entire active set has values equal to 0")

    # return step size
    a = sumabs2(v.gk)  :: T
    b = sumabs2(v.xgk) :: T
    return a / b :: T
end

# this function computes one gradient step in iht()
function _iht_gradstep{T <: Float}(
    v  :: IHTVariables{T},
    mu :: T,
    k  :: Int;
)
    # take gradient step
    BLAS.axpy!(mu, sdata(v.df), sdata(v.b))

    # preserve top k components of b
    project_k!(v.b, k)

    # which indices of new beta are nonzero?
    #update_indices!(v.idx, sdata(v.b))
    v.idx .= v.b .!= 0

    # must correct for equal entries at kth pivot of b
    # **note**: this uses Base.unsafe_setindex! to circumvent type stability issues
    # this is a VERY DANGEROUS DARK SIDE HACK!
    # hack randomly permutes indices of duplicates and retains one
    if sum(v.idx) > k
        a = select(v.b, k, by=abs, rev=true) :: T    # compute kth pivot
        dupes = (abs(v.b) .== abs(a)) :: BitArray{1} # find duplicates
        l = sum(dupes) :: Int                        # how many duplicates?
        l <= 1 && return nothing                     # if no duplicates, then simply return
        d = find(dupes) :: Vector{Int}               # find duplicates
        shuffle!(d)                                  # permute duplicates
        deleteat!(d, 1)                              # save first duplicate
        Base.unsafe_setindex!(v.b, zero(T), d)       # zero out other duplicates
        Base.unsafe_setindex!(v.idx, false, d)       # set corresponding indices to false
    end

    return nothing
end

# this function calculates the omega (here a / b) used for determining backtracking
function _iht_omega{T <: Float}(
    v :: IHTVariables{T}
)
    a = sqeuclidean(v.b, v.b0::Vector{T}) :: T
    b = sqeuclidean(v.xb, v.xb0::Vector{T}) :: T
    return a, b
end

# a function for determining whether or not to backtrack
function _iht_backtrack{T <: Float}(
    v       :: IHTVariables{T},
    ot      :: T,
    ob      :: T,
    mu      :: T,
    mu_step :: Int,
    nstep   :: Int
)
    mu*ob > 0.99*ot              &&
    sum(v.idx) != 0              &&
    sum(xor.(v.idx,v.idx0)) != 0 &&
    mu_step < nstep
end

# ----------------------------------------- #
# printing routines

# this prints the start of the algo
function print_header(io::IO)
     println(io, "\nBegin IHT algorithm\n")
     println(io, "Iter\tHalves\tMu\t\tNorm\t\tObjective")
     println(io, "0\t0\tInf\t\tInf\t\tInf")
end

# default IO for print_header is STDOUT
print_header() = print_header(STDOUT)

# alert when a descent error is found
function print_descent_error{T <: Float}(io::IO, iter::Int, loss::T, next_loss::T)
    print_with_color(:red, io, "\nIHT algorithm fails to descend!\n")
    print_with_color(:red, io, "Iteration: $(iter)\n")
    print_with_color(:red, io, "Current Objective: $(loss)\n")
    print_with_color(:red, io, "Next Objective: $(next_loss)\n")
    print_with_color(:red, io, "Difference in objectives: $(abs(next_loss - loss))\n")
end

# default IO for print_descent_error is STDOUT
print_descent_error{T <: Float}(iter::Int, loss::T, next_loss::T) = print_descent_error(iter, loss, next_loss)

# announce algo convergence
function print_convergence{T <: Float}(io::IO, iter::Int, loss::T, ctime::T)
    println(io, "\nIHT algorithm has converged successfully.")
    println(io, "Results:\nIterations: $(iter)")
    println(io, "Final Loss: $(loss)")
    println(io, "Total Compute Time: $(ctime)")
end

# default IO for print_convergence is STDOUT
print_convergence{T <: Float}(iter::Int, loss::T, ctime::T) = print_convergence(STDOUT, iter, loss, ctime)

# check the finiteness of an objective function value
# throw an error if value is not finite
function check_finiteness{T <: Float}(x::T)
    isnan(x) && throw(error("Objective function is NaN, aborting..."))
    isinf(x) && throw(error("Objective function is Inf, aborting..."))
end

# alert if iteration limit is reached
function print_maxiter{T <: Float}(io::IO, max_iter::Int, loss::T)
    print_with_color(:red, io, "IHT algorithm has hit maximum iterations $(max_iter)!\n")
    print_with_color(:red, io, "Current Loss: $(loss)\n")
end

# default IO for print_maxiter is STDOUT
print_maxiter{T <: Float}(max_iter::Int, loss::T) = print_maxiter(STDOUT, max_iter, loss)

# verbose printing of cv results
function print_cv_results{T <: Float}(io::IO, errors::Vector{T}, path::DenseVector{Int}, k::Int)
    println(io, "\n\nCrossvalidation Results:")
    println(io, "k\tMSE")
    for i = 1:length(errors)
        println(io, path[i], "\t", errors[i])
    end
    println(io, "\nThe lowest MSE is achieved at k = ", k)
end

# default IO for print_cv_results is STDOUT
print_cv_results{T <: Float}(errors::Vector{T}, path::DenseVector{Int}, k::Int) = print_cv_results(STDOUT, errors, path, k)

# -------------------------------------------- #
# functions related to logistic regression

# container object for temporary arrays
type IHTLogVariables{T <: Float, V <: DenseVector}
    xk       :: Matrix{T}
    xk2      :: Matrix{T}
    d2b      :: Matrix{T}
    b        :: V
    b0       :: Vector{T}
    df       :: V
    xb       :: V
    lxb      :: V
    l2xb     :: Vector{T}
    bk       :: Vector{T}
    bk2      :: Vector{T}
    bk0      :: Vector{T}
    ntb      :: Vector{T}
    db       :: Vector{T}
    dfk      :: Vector{T}
    active   :: Vector{Int}
    bidxs    :: Vector{Int}
    dfidxs   :: Vector{Int}
    idxs     :: BitArray{1}
    idxs2    :: BitArray{1}
    idxs0    :: BitArray{1}

    IHTLogVariables(xk::Matrix{T}, xk2::Matrix{T}, d2b::Matrix{T}, b::V, b0::Vector{T}, df::V, xb::V, lxb::V, l2xb::Vector{T}, bk::Vector{T}, bk2::Vector{T}, bk0::Vector{T}, ntb::Vector{T}, db::Vector{T}, dfk::Vector{T}, active::Vector{Int}, bidxs::Vector{Int}, dfidxs::Vector{Int}, idxs::BitArray{1}, idxs2::BitArray{1}, idxs0::BitArray{1}) where {T <: Float, V <: DenseVector{T}} = new{T,V}(xk, xk2, d2b, b, b0, df, xb, lxb, l2xb, bk, bk2, bk0, ntb, db, dfk, active, bidxs, dfidxs, idxs, idxs2, idxs0)
    #IHTLogVariables(xk, xk2, d2b, b, b0, df, xb, lxb, l2xb, bk, bk2, bk0, ntb, db, dfk, active, bidxs, dfidxs, idxs, idxs2, idxs0) where {T <: Float, V <: DenseVector{T}} = new{T,V}(xk, xk2, d2b, b, b0, df, xb, lxb, l2xb, bk, bk2, bk0, ntb, db, dfk, active, bidxs, dfidxs, idxs, idxs2, idxs0)
end

#strongly typed constructor of IHTLogVariables
function IHTLogVariables{T <: Float}(
    xk       :: Matrix{T},
    xk2      :: Matrix{T},
    d2b      :: Matrix{T},
    b        :: DenseVector{T},
    b0       :: Vector{T},
    df       :: DenseVector{T},
    xb       :: DenseVector{T},
    lxb      :: DenseVector{T},
    l2xb     :: Vector{T},
    bk       :: Vector{T},
    bk2      :: Vector{T},
    bk0      :: Vector{T},
    ntb      :: Vector{T},
    db       :: Vector{T},
    dfk      :: Vector{T},
    active   :: Vector{Int},
    bidxs    :: Vector{Int},
    dfidxs   :: Vector{Int},
    idxs     :: BitArray{1},
    idxs2    :: BitArray{1},
    idxs0    :: BitArray{1}
)
    IHTLogVariables{T, typeof(b)}(xk, xk2, d2b, b, b0, df, xb, lxb, l2xb, bk, bk2, bk0, ntb, db, dfk, active, bidxs, dfidxs, idxs, idxs2, idxs0)
end

# construct IHTLogVariables from data x, y, k
function IHTLogVariables{T <: Float}(
    x :: DenseMatrix{T},
    y :: DenseVector{T},
    k :: Int
)
    (n,p)  = size(x)
    xk     = zeros(T, n, k)
    xk2    = zeros(T, n, k)
    d2b    = zeros(T, k, k)
    b      = zeros(T, p)
    b0     = zeros(T, p)
    df     = zeros(T, p)
    xb     = zeros(T, n)
    lxb    = zeros(T, n)
    l2xb   = zeros(T, n)
    bk     = zeros(T, k)
    bk2    = zeros(T, k)
    bk0    = zeros(T, k)
    ntb    = zeros(T, k)
    db     = zeros(T, k)
    dfk    = zeros(T, k)
    active = collect(1:p)
    bidxs  = collect(1:p)
    dfidxs = collect(1:p)
    idxs   = falses(p)
    idxs0  = falses(p)
    idxs2  = falses(p)

    IHTLogVariables{T, typeof(b)}(xk, xk2, d2b, b, b0, df, xb, lxb, l2xb, bk, bk2, bk0, ntb, db, dfk, active, bidxs, dfidxs, idxs, idxs2, idxs0)
end

# function to modify fields of IHTVariables that depend on k
function update_variables!{T <: Float}(
    v :: IHTLogVariables{T},
    x :: DenseMatrix{T},
    k :: Int
)
    n     = size(x,1)
    v.xk  = zeros(T, n, k)
    v.bk  = zeros(T, k)
    v.xk2 = zeros(T, n, k)
    v.d2b = zeros(T, k, k)
    v.bk0 = zeros(T, k)
    v.bk2 = zeros(T, k)
    v.ntb = zeros(T, k)
    v.db  = zeros(T, k)
    v.dfk = zeros(T, k)
    return nothing
end

###
### TODO: constructors for SharedArrays, BEDFiles!
###

# new return object for results file
immutable IHTLogResults{T <: Float, V <: DenseVector}
    time   :: T
    iter   :: Int
    loss   :: T
    beta   :: V
    active :: Vector{Int}

    IHTLogResults(time::T, iter::Int, loss::T, β::V, active::Vector{Int}) where {T <: Float, V <: DenseVector{T}} = new{T,V}(time, iter, loss, β, active)
end

## strongly typed external constructor for IHTLogResults
#function IHTLogResults{T <: Float}(
#    time   :: T,
#    iter   :: Int,
#    loss   :: T,
#    β      :: DenseVector{T},
#    active :: Vector{Int}
#)
#    IHTLogResults{T, typeof(β)}(time::T, iter::Int, loss::T, β::DenseVector{T}, active::Vector{Int})
#end

# function to display IHTLogResults object
function Base.show(io::IO, x::IHTLogResults)
    println(io, "IHT results:")
    @printf(io, "\nCompute time (sec):   %3.4f\n", x.time)
    @printf(io, "Final loss:           %3.7f\n", x.loss)
    @printf(io, "Iterations:           %d\n", x.iter)
    println(io, "IHT estimated ", countnz(x.beta), " nonzero coefficients.")
    print(io, DataFrame(Predictor=find(x.beta), Estimated_β=x.beta[find(x.beta)]))
    return nothing
end

# announce algo convergence
function print_log_convergence{T <: Float}(io::IO, iter::Int, loss::T, ctime::T, nrmdf::T)
    println("\nL0_log has converged successfully.")
    @printf("Results:\nIterations: %d\n", iter)
    @printf("Final Loss: %3.7f\n", loss)
    @printf("Norm of active gradient: %3.7f\n", nrmdf)
    @printf("Total Compute Time: %3.3f sec\n", ctime)
end

# default IO for print_convergence is STDOUT
print_log_convergence{T <: Float}(iter::Int, loss::T, ctime::T, nrmdf::T) = print_log_convergence(STDOUT, iter, loss, ctime, nrmdf)

# this prints the start of the algo
function print_header_log(io::IO)
     println(io, "\nBegin IHT algorithm\n")
     println(io, "Iter\tSteps\tHalves\t\tf(β)\t\t||df(β)||")
     println(io, "0\t0\t0\t\tInf\t\tInf")
end

# default IO for print_header is STDOUT
print_header_log() = print_header_log(STDOUT)


text:
using FileIO, JLD2
using CoexistenceHoles
using Test

betti = FileIO.load(joinpath(@__DIR__, "test_keys/betti_checks.jld2"), "betti")

@testset "ripserer_tests.jl" begin
    for f in keys(betti)
        H = read_hypergraph(joinpath(@__DIR__, "test_data", f, "Hp.dat"))
        test_betti = betti_hypergraph_ripscomplex(H; max_dim = 4)
        @test test_betti==betti[f]
    end
end

ext: jl

max_issues_repo_name: mtangulo/coexistenceholes

max_forks_repo_name: mtangulo/coexistenceholes


text:
using DomainSets: ×

"Does `z` agree with `val` up to the given `threshold`?"
isaccurate(z, val, threshold = 1e-5) = abs(z-val) < threshold

function test_readme_examples()
    # Examples of README.md
    threshold = 1e-5
    @test isaccurate(integral(cos, 0..1.0), 0.8414709848078965)
    @test isaccurate(integral(x -> exp(x[1]+x[2]), (0..1.0)^2), 2.9524924420120535)
    @test isaccurate(quadrature(cos, UnionDomain(0..1, 2..3))[1], 0.07329356604208204)
    @test isaccurate(integral(t -> sin(log(abs(t))), -1..1, LogPointSingularity(0.0)), -1.0000000021051316)
    @test isaccurate(integral( x -> exp(log(abs(x[1]-x[2]))), (2..3) × (1..4), DiagonalSingularity()),
2.333333333333333)
    @test isaccurate(integral(cos, ChebyshevTMeasure()), 2.403939430634413)
    @test isaccurate(integral(t -> cos(t)*1/sqrt(1-t^2), -1.0..1.0), 2.403939410869398)
    @test isaccurate(quadrature(QuadAdaptive(atol=1e-3, rtol = 1e-3), t->cos(t^2), 0..10)[1], 0.6011251848111901, 1e-3)
    @test isaccurate(integral(Q_GaussLaguerre(10), cos), 0.5000005097999486)
    @test isaccurate(integral(t -> cos(t)*exp(-t), HalfLine()), 0.5)
end

function test_some_integrals()
    z1 = sin(1.0)
    @test abs(integral(cos, 0..1.0) - z1) < 1e-8
    I, E = quadrature(cos, 0..1.0)
    @test abs(I-z1) < 1e-8
    @test E < 1e-8

    f2 = x -> exp(x[1]+x[2])
    z2 = hcubature(f2, (0,0), (1,1))[1]
    @test abs(integral(f2, (0..1)^2) - z2) < 1e-6
end

function test_integrals()
    test_some_integrals()
    test_readme_examples()
end


text:
# Copyright 2018 IBM Corp.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

module TestOntology
using Test

@testset "JSON" begin
  include("OntologyJSON.jl")
end

@testset "Database" begin
  include("OntologyDBs.jl")
end

@testset "RDF" begin
  include("rdf/OntologyRDF.jl")
end

end

max_stars_repo_name: cclauss/semanticflowgraph

max_issues_repo_name: cclauss/semanticflowgraph

max_forks_repo_name: cclauss/semanticflowgraph


text:
begin
    u = collect(-5:0.5:5)
    ω = ones(21) ./ 21
    prim = [1.0, 0.0, 1.0]
    w = KitML.KitBase.prim_conserve(prim, 5 / 3)
    M = KitML.KitBase.maxwellian(u, prim)
    τ = 0.01
    f = deepcopy(M)
    nn = Chain(Dense(21, 21, tanh), Dense(21, 21))
    p = initial_params(nn)
    nn1 = FastChain(FastDense(21, 21, tanh), FastDense(21, 21))
    p1 = initial_params(nn1)
end

df = KitML.ube_dfdt(f, (M, τ, (nn, p)), 0.1)
df2 = similar(df)
KitML.ube_dfdt!(df2, f, (M, τ, (nn, p)), 0.1)
@test df ≈ df2 atol = 0.01

KitML.ube_dfdt(f, (M, τ, (nn1, p1)), 0.1)
KitML.ube_dfdt!(similar(f), f, (M, τ, (nn1, p1)), 0.1)


text:
########## Common definitions ##########

const DiscreteTime = Integer
const ContinuousTime = Real
promote_rule(ContinuousTime, DiscreteTime) = ContinuousTime

"""
Type of time index of a stochastic process: discrete time (Integer) or continuous time (Real).
"""
const TimeStyle = Union{DiscreteTime, ContinuousTime}  # == Real

max_stars_repo_name: yanncalec/FracFin

max_issues_repo_name: yanncalec/FracFin

max_forks_repo_name: yanncalec/FracFin


text:
module Hubbard

using LinearMaps: FunctionMap
using StaticArrays
export TiltedSquare, hamiltonian, hops, index, mask, symmetries, ρ

const ρ = SA[0 -1; 1 0]  # Rotation by `½π`
const σ = SA[-1 0; 0 1]  # Reflection about `x` axis
const sum2 = Base.Fix1(sum, abs2)

"""
An integer greater than one can be written as a sum of two squares if and only if
its prime decomposition contains no factor pᵏ, where prime p ≡ 3 (mod 4) and k is odd.
"""
function istwosquares(n::Int)::Bool
    p = 2
    k = 0
    i = false
    while p ≤ n
        div, rem = divrem(n, p)
        if rem ≠ 0
            isodd(k) && return false
            p += 1
            k = 0
            i = p % 4 == 3
        else
            n = div
            k += i
        end
    end
    return iseven(k)
end

"""
    restrict(tilt)

Apply https://en.wikipedia.org/wiki/Periodic_boundary_conditions to restrict point `x`
within the `tilt × ρ * tilt` titled square
"""
restrict(tilt::SVector{2,Int}) =
    Base.Fix2(tilt) do x, tilt
        n = sum2(tilt)
        reduce([tilt, ρ * tilt], init = x) do u, v
            u - v * fld(u'v, n)
        end
    end

"""
    TiltedSquare

represents a compact subset of the Euclidean square lattice with `n` sites.
"""
struct TiltedSquare{N}
    tilt::SVector{2,Int}
    sites::SVector{N,SVector{2,Int}}

    function TiltedSquare{N}() where {N}
        istwosquares(N) || throw(DomainError("$n is not a sum of two squares"))
        tilt = first(Iterators.filter(==(N) ∘ sum2, SA[isqrt(N - v^2), v] for v ∈ 0:isqrt(N ÷ 2)))
        pred = isequal ∘ restrict(tilt)
        sites = [site for site ∈ (SA[x, y] for x ∈ -tilt[2]:tilt[1] for y ∈ 0:sum(tilt)) if pred(site)(site)]
        return new(tilt, SVector{N,SVector{2,Int}}(sites))
    end
end

"""
    Base.findall(r::Matrix{Int}, s::TiltedSquare)

Convert application of rotation matrix `r` to all sites in tilted square `s` to a
permutation `1:length(s.sites)` assuming periodic boundary conditions.

See the application of a rotation by `½π` to the 8-site tilted square for example.
The sites of `TiltedSquare{8}()` are labeled as follows:

      6
    2 5 8
    1 4 7
      3

Rotating these sites by `½π` results in

      8 7
    6 5 4 3
      2 1

Translating indices which landed outside the original shape (i.e. all but `3` and `7`)
by `a = [2, 2]` gives

      8
    6 5 4
    7 2 1
      3

Interpreting this new arrangement as a permutation, site `1` moved to position `7`, `2`
moved to position `4`, `3` stayed in place, ... → `(7 4 3 8 5 2 1 6)`

```jldoctest
julia> findall(ρ, TiltedSquare{8}())
8-element Vector{Int64}:
 7
 4
 3
 8
 5
 2
 1
 6

```
"""
function Base.findall(r::SMatrix{2,2,Int,4}, s::TiltedSquare)::Vector{Int}
    invperm(collect(sortperm(s.sites, by = restrict(s.tilt) ∘ Base.Fix1(*, r))))
end

"""
    findall(d::Vector{T}, s::TiltedSquare)

Convert application of displacement `d` to all sites in tilted square `s` to a
permutation of `1:length(s.sites)` assuming periodic boundary conditions.

      6→1'
    2→5→8→3'
    1→4→7→6'
      3→2'

```jldoctest
julia> findall(SA[1, 0], TiltedSquare{8}())
8-element Vector{Int64}:
 4
 5
 2
 7
 8
 1
 6
 3

"""
function Base.findall(d::SVector{2,Int}, s::TiltedSquare)::Vector{Int}
    invperm(collect(sortperm(s.sites, by = restrict(s.tilt) ∘ Base.Fix1(+, d))))
end

"""
    hops(tilted_square)

Computes neighbor bit masks for the `n` site square lattice with
periodic boundary conditions.

# Examples
```jldoctest
julia> [digits(h, base = 2, pad = 8) for h ∈ hops(TiltedSquare{8}())]
16-element Vector{Vector{Int64}}:
 [1, 0, 0, 1, 0, 0, 0, 0]
 [0, 1, 0, 0, 1, 0, 0, 0]
 [0, 1, 1, 0, 0, 0, 0, 0]
 [0, 0, 0, 1, 0, 0, 1, 0]
 [0, 0, 0, 0, 1, 0, 0, 1]
 [1, 0, 0, 0, 0, 1, 0, 0]
 [0, 0, 0, 0, 0, 1, 1, 0]
 [0, 0, 1, 0, 0, 0, 0, 1]
 [1, 1, 0, 0, 0, 0, 0, 0]
 [0, 1, 0, 0, 0, 0, 1, 0]
 [0, 0, 1, 1, 0, 0, 0, 0]
 [0, 0, 0, 1, 1, 0, 0, 0]
 [0, 0, 0, 0, 1, 1, 0, 0]
 [0, 0, 1, 0, 0, 1, 0, 0]
 [0, 0, 0, 0, 0, 0, 1, 1]
 [1, 0, 0, 0, 0, 0, 0, 1]

```
"""
function hops(s::TiltedSquare)::Vector{Int}
    masks(y) = [(1 << (i - 1)) | (1 << (j - 1)) for (i, j) ∈ enumerate(findall(y, s))]
    return [masks(SA[1, 0]); masks(SA[0, 1])]
end

"""
    symmetries(s)

Elements of dihedral group D₄ specialized to the `n` site tilted square (minus
translations by `tilt` and `ρ * tilt` as those are already taken care of by `restrict`)
as binary masks.  The group is generated by 2 elements:

1. Rotation by `½π`: `ρ`
1. Reflection about `x` axis: ``

Remaining elements can be represented as products of these generators:

1. Inversion: `ρ²`
1. Rotation by `-½π`: `ρ³ ≡ ρ⁻¹`
1. Reflection about anti-diagonal: `ρ σ`
1. Reflection about `y` axis: `ρ² σ ≡ ρ σ ρ⁻¹ ≡ τ`
1. Reflection about diagonal: `ρ³ σ ≡ ρ τ`

We already know that rotation by `½π` is equivalent to the `(7 4 3 8 5 2 1 6)` permutation.
Because fermions are indistinghishable, we convert this permutation to a bit mask where 1
indicates the site moves.  `(7 4 3 8 5 2 1 6)` becomes `0b11010111`

Reflection about the `x` axis does

      6
    8 5 2
    7 4 1
      3

or `(7 8 3 4 5 6 1 2)` as a permutation and `0b01011010` as a bitmask.

Let us confirm D₄ by working out one repeated application: `ρ σ`

                 2
      2 1      6 5 4
    6 5 4 3    1 8 7
      8 7        3

or `(1 6 3 8 5 2 7 4)` as a permutation.  We can easily verify that this equals the product of
`(7 4 3 8 5 2 1 6)(7 8 3 4 5 6 1 2)` by following the orbit of each site:

    1 → 7 → 1
    2 → 8 → 6
    3 → 3 → 3
    4 → 4 → 8
    5 → 5 → 5
    6 → 6 → 2
    7 → 1 → 7
    8 → 2 → 4

```jldoctest
julia> [digits(s, base = 2, pad = 8) for s ∈ symmetries(TiltedSquare{8}())]
8-element Vector{Vector{Int64}}:
 [0, 0, 0, 0, 0, 0, 0, 0]
 [1, 1, 0, 1, 0, 1, 1, 1]
 [0, 1, 0, 1, 0, 1, 0, 1]
 [1, 1, 0, 1, 0, 1, 1, 1]
 [1, 1, 0, 0, 0, 0, 1, 1]
 [0, 1, 0, 1, 0, 1, 0, 1]
 [1, 0, 0, 1, 0, 1, 1, 0]
 [0, 1, 0, 1, 0, 1, 0, 1]

```
"""
function symmetries(s::TiltedSquare)::Vector{Int}
    mask(r) = evalpoly(2, map(Base.splat(≠), enumerate(findall(r, s))))
    return [mask(r) for r ∈ [SA[1 0; 0 1], ρ, ρ^2, ρ^3, σ, ρ * σ, ρ^2 * σ, ρ^3 * σ]]
end

"""
    index(m, c)

Compute the index of the combination of `count_ones(mask)` items
in lexicographic order

# Examples
```jldoctest
julia> c = [binomial(i, j) for i=0:3, j=1:2];

julia> [index(parse(Int, b, base = 2), c) for b ∈ ["0011", "0101", "0110", "1001", "1010", "1100"]]
6-element Vector{Int64}:
 1
 2
 3
 4
 5
 6

```
"""
function index(m::Integer, c::Matrix)::Integer
    i = 1
    @inbounds for k ∈ 1:size(c, 2)
        i += c[trailing_zeros(m)+1, k]
        m &= m - 1
    end
    return i
end

index(n::Tuple{Int,Int}, c::Matrix) = CartesianIndex(Tuple(index(m, c) for m ∈ n))

"""
    mask(i, c)

Compute the binary mask of the combination of `k` items out of `n`
at position `index` in lexicographic order

# Examples
```jldoctest
julia> c = [binomial(i, j) for i=0:3, j=1:2];

julia> [digits(mask(i, c), base = 2, pad = 4) for i=1:6]
6-element Vector{Vector{Int64}}:
 [1, 1, 0, 0]
 [1, 0, 1, 0]
 [0, 1, 1, 0]
 [1, 0, 0, 1]
 [0, 1, 0, 1]
 [0, 0, 1, 1]

```
"""
function mask(i::Integer, c::Matrix)::Integer
    k = size(c, 2)
    m = 0
    @inbounds for n ∈ size(c, 1):-1:1
        j = i - c[n, k]
        if j ≥ 1
            m |= 1
            k == 1 && return m << (n - 1)
            i = j
            k -= 1
        end
        m <<= 1
    end
end

"""
    hamiltonian(n, k, t, U)

Implements the [Hubbard hamiltonian](https://en.wikipedia.org/wiki/Hubbard_model)
as a [mutating linear map](https://github.com/Jutho/LinearMaps.jl).

# Arguments
- `n::Integer`: sites in the square lattice with periodic boundary conditions
- `k::Integer`: spin up and spin down fermions
- `t::Real`: hopping integral
- `U::Real`: strength of on-site interaction

# Examples
```jldoctest
julia> h = hamiltonian(8, 4, 1., 4.)
3024×3024 LinearMaps.FunctionMap{Float64}(mult!; ismutating=true, issymmetric=true, ishermitian=true, isposdef=false)

julia> using Arpack: eigs

julia> λ, ϕ = eigs(h, nev=1, which=:SR);

julia> λ
1-element Vector{Float64}:
 -11.652313479748669

julia> size(ϕ)
(3024, 1)

```
"""
function hamiltonian(n::Integer, k::Integer, t::T, U::T) where {T<:Real}
    c = [binomial(i, j) for i = 0:(n-1), j = 1:k]
    nCk = binomial(n, k)
    s = TiltedSquare{n}()
    hop = hops(s)
    ok = isequal(k) ∘ count_ones

    # Apply all symmetries to all combinations
    # TODO: separate nup and ndn (now k = nup = ndn)
    masks = Base.Fix2(mask, c).(1:nCk) .⊻ unique(symmetries(s))'

    minpairs = [
        minimum(t for t ∈ zip(m, n) if all(ok, t))
        for m ∈ eachrow(masks), n ∈ eachrow(masks)
    ]

    uniquepairs = sort(vec(minpairs))
    unique!(uniquepairs)

    minindices = Base.Fix1(searchsortedfirst, uniquepairs).(minpairs)
    function mult!(y, x)
        fill!(y, 0.0)
        @inbounds for ((i, w), (u, d)) ∈ zip(enumerate(x), uniquepairs)
            w == 0 && continue

            y[i] += U * w * count_ones(u & d)

            iu = index(u, c)
            id = index(d, c)

            tw = t * w
            for p ∈ hop
                a = u ⊻ p
                if a ≠ u && ok(a)
                    y[minindices[index(a, c), id]] -= tw
                end
                a = d ⊻ p
                if a ≠ d && ok(a)
                    y[minindices[iu, index(a, c)]] -= tw
                end
            end
        end
    end
    return FunctionMap{T}(mult!, length(uniquepairs), ismutating = true, issymmetric = true)
end

end

max_stars_repo_name: jburgy/hubbard

max_issues_repo_name: jburgy/hubbard

max_forks_repo_name: jburgy/hubbard


text:
const scores = Dict("eggs" => 1, "peanuts" => 2, "shellfish" => 4, "strawberries" => 8,
    "tomatoes" => 16, "chocolate" => 32, "pollen" => 64, "cats" => 128)

function allergic_to(score, allergen)
    score & scores[allergen] != 0
end

function allergy_list(score)
    keys(filter(p -> score & p.second != 0, scores))
end

max_stars_repo_name: koladilip/julia-exercism

max_issues_repo_name: koladilip/julia-exercism

max_forks_repo_name: koladilip/julia-exercism


text:
using LinearAlgebraicRepresentation
using Plasm, LinearAlgebra
Lar = LinearAlgebraicRepresentation
using SparseArrays

store = []
scaling = 1.75
V,(VV,EV,FV,CV) = Lar.cuboid([0.25,0.25,0.25],true,[-0.25,-0.25,-0.25])
mybox = (V,CV,FV,EV)
for k=1:5
	size = rand()*scaling
	scale = Lar.s(size,size,size)
	transl = Lar.t(rand(3)...)
	alpha = 2*pi*rand()
	rx = Lar.r(alpha,0,0); ry = Lar.r(0,alpha,0); rz = Lar.r(0,0,alpha)
	rot = rx * ry * rz
	s = Lar.Struct([ transl, scale, rot, mybox ])
	push!(store, Lar.struct2lar(s))
end

s = Lar.Struct(store)
V,CV,FV,EV = Lar.struct2lar(s)
# V = Plasm.normalize3D(V) TODO:  solve MethodError bug
Plasm.view(V,CV)

cop_EV = Lar.coboundary_0(EV::Lar.Cells);
cop_EW = convert(Lar.ChainOp, cop_EV);
cop_FE = Lar.coboundary_1(V, FV::Lar.Cells, EV::Lar.Cells);
W = convert(Lar.Points, V');

V, copEV, copFE, copCF = Lar.Arrangement.spatial_arrangement(
	W::Lar.Points, cop_EW::Lar.ChainOp, cop_FE::Lar.ChainOp)

cc = [copEV, copFE, copCF]
output = Lar.lar2obj(V::Lar.Points, cc::Lar.ChainComplex)

f = open("test/out3d.obj", "w")
print(f, output)
close(f)

V,EVs,FVs = Lar.obj2lar("test/out3d.obj")

for k=1:length(FVs)
	Plasm.view(V, FVs[k])
end

Plasm.view(V, FVs[1])


text:
struct SimulationInputs
    base_power::Float64
    base_frequency::Float64
    injectors_data::Vector{<:PSY.StaticInjection}
    static_injection_data::Vector{<:PSY.StaticInjection}
    dynamic_branches::Vector{PSY.DynamicBranch}
    counts::Base.ImmutableDict{Symbol, Int}
    Ybus::SparseArrays.SparseMatrixCSC{Complex{Float64}, Int}
    dyn_lines::Bool
    voltage_buses_ix::Vector{Int}
    current_buses_ix::Vector{Int}
    global_index::Dict{String, Dict{Symbol, Int}}
    total_shunts::SparseArrays.SparseMatrixCSC{Complex{Float64}, Int}
    global_vars::Dict{Symbol, Real}
    lookup::Dict{Int, Int}
    DAE_vector::Vector{Bool}
    mass_matrix::SparseArrays.SparseMatrixCSC{Float64, Int}
    aux_arrays::Dict{Int, Vector}
    tspan::NTuple{2, Float64}
end

function SimulationInputs(sys::PSY.System, tspan::NTuple{2, Float64} = (0.0, 0.0))
    injector_data = PSY.get_components(
        PSY.StaticInjection,
        sys,
        x -> PSY.get_dynamic_injector(x) !== nothing && PSY.get_available(x),
    )

    if isempty(injector_data)
        error("System doesn't contain any DynamicInjection devices")
    end

    Ybus, lookup = _get_Ybus(sys)
    n_buses = length(PSY.get_bus_numbers(sys))

    dynamic_branches =
        collect(PSY.get_components(PSY.DynamicBranch, sys, x -> PSY.get_available(x)))

    branch_state_counts = 2 * length(dynamic_branches)
    injector_state_count = sum(PSY.get_n_states.(PSY.get_dynamic_injector.(injector_data)))
    var_count = injector_state_count + 2 * n_buses + branch_state_counts

    mass_matrix = SparseArrays.sparse(LinearAlgebra.I, var_count, var_count)
    mass_matrix[1:(2 * n_buses), 1:(2 * n_buses)] .= 0.0

    static_injection_data = PSY.get_components(
        PSY.StaticInjection,
        sys,
        x -> PSY.get_dynamic_injector(x) === nothing && PSY.get_available(x),
    )

    counts = Base.ImmutableDict(
        :injection_n_states => injector_state_count,
        :branches_n_states => branch_state_counts,
        :first_dyn_injection_pointer => 2 * n_buses + branch_state_counts + 1,
        :first_dyn_branch_point => isempty(dynamic_branches) ? 0 : 2 * n_buses + 1,
        :total_variables => var_count,
        :bus_count => n_buses,
    )

    return SimulationInputs(
        PSY.get_base_power(sys),
        PSY.get_frequency(sys),
        collect(injector_data),
        collect(static_injection_data),
        dynamic_branches,
        counts,
        Ybus,
        !isempty(dynamic_branches),
        Vector{Int}(),
        collect(1:n_buses),
        MAPPING_DICT(),
        SparseArrays.spzeros(Complex{Float64}, n_buses, n_buses),
        GLOBAL_VARS_IX(),
        lookup,
        _init_DAE_vector!(var_count, n_buses),
        mass_matrix,
        Dict{Int, Vector{Real}}(),
        tspan,
    )
end

get_base_power(inputs::SimulationInputs) = inputs.base_power
get_base_frequency(inputs::SimulationInputs) = inputs.base_frequency
get_injectors_data(inputs::SimulationInputs) = inputs.injectors_data
get_dynamic_branches(inputs::SimulationInputs) = inputs.dynamic_branches
get_static_injections_data(inputs::SimulationInputs) = inputs.static_injection_data
get_counts(inputs::SimulationInputs) = inputs.counts
get_voltage_buses_ix(inputs::SimulationInputs) = inputs.voltage_buses_ix
get_current_buses_ix(inputs::SimulationInputs) = inputs.current_buses_ix
get_global_index(inputs::SimulationInputs) = inputs.global_index
get_Ybus(inputs::SimulationInputs) = inputs.Ybus
get_total_shunts(inputs::SimulationInputs) = inputs.total_shunts
get_global_vars(inputs::SimulationInputs) = inputs.global_vars
get_lookup(inputs::SimulationInputs) = inputs.lookup
get_DAE_vector(inputs::SimulationInputs) = inputs.DAE_vector
get_mass_matrix(inputs::SimulationInputs) = inputs.mass_matrix
get_aux_arrays(inputs::SimulationInputs) = inputs.aux_arrays
get_tspan(inputs::SimulationInputs) = inputs.tspan
has_dyn_lines(inputs::SimulationInputs) = inputs.dyn_lines

get_injection_pointer(inputs::SimulationInputs) =
    get_counts(inputs)[:first_dyn_injection_pointer]
get_branches_pointer(inputs::SimulationInputs) = get_counts(inputs)[:first_dyn_branch_point]
get_n_injection_states(inputs::SimulationInputs) = get_counts(inputs)[:injection_n_states]
get_n_branches_states(inputs::SimulationInputs) = get_counts(inputs)[:branches_n_states]
get_variable_count(inputs::SimulationInputs) = get_counts(inputs)[:total_variables]
get_device_index(inputs::SimulationInputs, device::D) where {D <: PSY.DynamicInjection} =
    get_global_index(inputs)[device.name]
get_bus_count(inputs::SimulationInputs) = get_counts(inputs)[:bus_count]
get_ω_sys(inputs::SimulationInputs) = get_global_vars(inputs)[:ω_sys]

function change_vector_type!(inputs::SimulationInputs, ::Type{T}) where {T <: Real}
    for d in get_injectors_data(inputs)
        attach_inner_vars!(PSY.get_dynamic_injector(d), T)
    end
    return
end

function simulation_pre_step!(inputs::SimulationInputs, ::Type{T}) where {T <: Real}
    add_aux_arrays!(inputs, T)
    change_vector_type!(inputs, T)
    return
end

function add_aux_arrays!(inputs::SimulationInputs, ::Type{T}) where {T <: Real}
    @debug "Auxiliary Arrays created with Type $(T)"
    bus_count = get_bus_count(inputs)
    get_aux_arrays(inputs)[1] = collect(zeros(T, bus_count))                       #I_injections_r
    get_aux_arrays(inputs)[2] = collect(zeros(T, bus_count))                       #I_injections_i
    get_aux_arrays(inputs)[3] = collect(zeros(T, get_n_injection_states(inputs)))  #injection_ode
    get_aux_arrays(inputs)[4] = collect(zeros(T, get_n_branches_states(inputs)))   #branches_ode
    get_aux_arrays(inputs)[5] = collect(zeros(Complex{T}, bus_count))              #I_bus
    get_aux_arrays(inputs)[6] = collect(zeros(T, 2 * bus_count))                   #I_balance
    return
end

function _get_Ybus(sys::PSY.System)
    n_buses = length(PSY.get_components(PSY.Bus, sys))
    dyn_lines = PSY.get_components(PSY.DynamicBranch, sys)
    if !isempty(PSY.get_components(PSY.ACBranch, sys))
        Ybus_ = PSY.Ybus(sys)
        Ybus = Ybus_[:, :]
        lookup = Ybus_.lookup[1]
        for br in dyn_lines
            ybus_update!(Ybus, br, lookup, -1.0)
        end
    else
        Ybus = SparseArrays.SparseMatrixCSC{Complex{Float64}, Int}(zeros(n_buses, n_buses))
        lookup = Dict{Int.Int}()
    end
    return Ybus, lookup
end

function _init_DAE_vector!(n_vars::Int, n_buses::Int)
    DAE_vector = trues(n_vars)
    DAE_vector[1:(n_buses * 2)] .= false
    return DAE_vector
end

function _dynamic_lines_inputs!(
    inputs::SimulationInputs,
    state_space_ix::Vector{Int},
    n_buses::Int,
)
    static_bus_var_count = 2 * n_buses
    branches_n_states = 0
    if inputs.dyn_lines
        global_state_index = inputs.global_index
        for br in get_dynamic_branches(inputs)
            index_dynamic_lines!(inputs, br, n_buses)
            add_states_to_global!(global_state_index, state_space_ix, br)
            branches_n_states += PSY.get_n_states(br)
        end

        for (ix, val) in enumerate(inputs.DAE_vector[1:n_buses])
            if val
                #This V_ix should be V_number.
                global_state_index["V_$(ix)"] = Dict(:R => ix, :I => ix + n_buses)
                static_bus_var_count -= 2
                push!(inputs.voltage_buses_ix, ix)
                @assert static_bus_var_count >= 0
            end
        end
        unique!(inputs.voltage_buses_ix)
        setdiff!(inputs.current_buses_ix, inputs.voltage_buses_ix)
    else
        @debug("System doesn't contain Dynamic Branches")
    end
    return branches_n_states, static_bus_var_count
end

function _static_injection_inputs!(inputs::SimulationInputs, ::Vector{Int}, sys::PSY.System)
    for s in PSY.get_components(PSY.StaticInjection, sys)
        index_static_injection(inputs, s)
    end
    return
end

function _dynamic_injection_inputs!(inputs::SimulationInputs, state_space_ix::Vector{Int})
    dynamic_injection_states = 0
    for d in get_injectors_data(inputs)
        @debug PSY.get_name(d)
        _attach_control_refs!(d)
        dynamic_injector = PSY.get_dynamic_injector(d)
        dynamic_injection_states += PSY.get_n_states(dynamic_injector)
        index_dynamic_injection(inputs, dynamic_injector, state_space_ix)
    end
    return dynamic_injection_states
end

function _mass_matrix_inputs!(inputs::SimulationInputs)
    mass_matrix = get_mass_matrix(inputs)
    injection_range = PSID.get_injection_pointer(inputs):PSID.get_variable_count(inputs)
    mass_matrix_injectors = view(mass_matrix, injection_range, injection_range)
    for d in PSID.get_injectors_data(inputs)
        dynamic_injector = PSY.get_dynamic_injector(d)
        device_mass_matrix_entries!(mass_matrix_injectors, dynamic_injector)
    end

    if has_dyn_lines(inputs)
        sys_f = get_base_frequency(inputs)
        shunts = get_total_shunts(inputs)
        n_buses = get_bus_count(inputs)
        for i in 1:n_buses
            val = imag(shunts[i, i])
            if val > 0
                mass_matrix[i, i] =
                    mass_matrix[i + n_buses, i + n_buses] = val * (1 / (2.0 * π * sys_f))
            end
        end
    end
end

# Default implementation for both models. This implementation is to future proof if there is
# a divergence between the required build methods
function _build!(inputs::SimulationInputs, sys::PSY.System)
    n_buses = get_bus_count(inputs)
    state_space_ix = Int[n_buses * 2]
    branches_n_states, static_bus_var_count =
        _dynamic_lines_inputs!(inputs, state_space_ix, n_buses)

    _static_injection_inputs!(inputs, state_space_ix, sys)
    injection_n_states = _dynamic_injection_inputs!(inputs, state_space_ix)

    set_frequency_reference!(inputs, sys)
    IS.@assert_op get_ω_sys(inputs) != -1

    _mass_matrix_inputs!(inputs)

    IS.@assert_op injection_n_states == state_space_ix[1] - branches_n_states - n_buses * 2
    IS.@assert_op n_buses * 2 - static_bus_var_count >= 0
    IS.@assert_op length(inputs.DAE_vector) == state_space_ix[1]

    @debug inputs.counts
    return
end

"""
SimulationInputs build function for MassMatrixModels
"""
function build!(inputs::SimulationInputs, ::Type{MassMatrixModel}, sys::PSY.System)
    return _build!(inputs, sys)
end

"""
SimulationInputs build function for ImplicitModels
"""
function build!(inputs::SimulationInputs, ::Type{ImplicitModel}, sys::PSY.System)
    return _build!(inputs, sys)
end


text:
## changingnames.jl  methods for NamedArray that change some of the names

## (c) 2013, 2017 David A. van Leeuwen

## This code is licensed under the MIT license
## See the file LICENSE.md in this distribution

## sum etc: keep names in one dimension
for f in (:(Base.sum),
          :(Base.prod),
          :(Base.maximum),
          :(Base.minimum),
          :(Statistics.mean),
          :(Statistics.std),
          :(Statistics.var))
    @eval function ($f)(a::NamedArray{T,N}; dims=:) where {T,N}
        s = ($f)(a.array, dims=dims)
        if !isa(s, AbstractArray)
            return s
        else
            dicts = tuple([issubset(i,dims) ? OrderedDict(string($f,"(",a.dimnames[i],")") => 1) : a.dicts[i] for i=1:ndims(a)]...)
            NamedArray{T,N,typeof(a.array), typeof(dicts)}(s, dicts, a.dimnames)
        end
    end
end

## I forgot what `fan` stands for.  Function Apply Name?
function fan(f::Function, fname::AbstractString, a::NamedArray{T,N}, dim::Int) where {T,N}
    DimNamType = eltype(a.dimnames)
    dimnames = Array{DimNamType}(undef, N)
    for i=1:N
        if i==dim
            dimnames[i] = DimNamType(string(fname, "(", a.dimnames[dim], ")"))
        else
            dimnames[i] = a.dimnames[i]
        end
    end
     NamedArray(f(a.array,dim), a.dicts, tuple(dimnames...))
end

## rename a dimension
for f in (:cumprod, :cumsum)
    @eval Base.$f(a::NamedArray, d::Integer) = fan($f, string($f), a, d)
end


text:
######
##FEM Heat Nonlinear Test
######
using FiniteElementDiffEq

#Define a parabolic problem
prob = prob_femheat_birthdeathinteractingsystem
#@code_warntype solve(prob,alg=:Euler)
sol = solve(prob,FEMDiffEqHeatEuler())

#plot(sol,plot_analytic=false,zlim=(0,2),cbar=false)

prob = prob_femheat_birthdeathsystem
sol = solve(prob,FEMDiffEqHeatImplicitEuler(),iterations=1000,autodiff=true)

#plot(sol,plot_analytic=false,zlim=(0,2),cbar=false)

true


text:
# ---
# title: 354. Russian Doll Envelopes
# id: problem354
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Binary Search, Dynamic Programming
# link: <https://leetcode.com/problems/russian-doll-envelopes/description/>
# hidden: true
# ---
#
# You have a number of envelopes with widths and heights given as a pair of
# integers `(w, h)`. One envelope can fit into another if and only if both the
# width and height of one envelope is greater than the width and height of the
# other envelope.
#
# What is the maximum number of envelopes can you Russian doll? (put one inside
# other)
#
# **Note:**
# Rotation is not allowed.
#
# **Example:**
#
#
#
#     Input: [[5,4],[6,4],[6,7],[2,3]]
#     Output: 3
#     **Explanation: T** he maximum number of envelopes you can Russian doll is 3 ([2,3] => [5,4] => [6,7]).
#
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


text:
############################################################################################
abstract type ProgressLog end
struct SilentLog <: ProgressLog end
struct ConsoleLog <: ProgressLog end

"""
$(TYPEDEF)

Default arguments for logging sampling output during sampling process.

# Fields
$(TYPEDFIELDS)
"""
struct ProgressReport{P<:ProgressLog,K}
    "Boolean if progressbar enabled during sampling run"
    bar::Bool
    "Boolean if output diagnostics should be printed."
    log::P
    "Additional ProgressMeter keywords to be used to return output"
    kwargs::K
    function ProgressReport(;
        bar=true,
        log=SilentLog(),
        kwargs=(dt=0.1, desc="Computing...", color=:blue, showspeed=true),
    )
        if isa(log, ConsoleLog)
            ArgCheck.@argcheck bar "Log can only be printed if bar == true."
        end
        return new{typeof(log),typeof(kwargs)}(bar, log, kwargs)
    end
end

############################################################################################
"""
$(TYPEDEF)

Default arguments for sampling routine.

# Fields
$(TYPEDFIELDS)
"""
struct SampleDefault{
    D<:DataStructure,
    M<:TemperingMethod,
    P<:ProgressReport
    }
    "Format to split data before proposal step."
    dataformat::D
    "Tempering values for target distributions of each chain, per default = 1*logtarget(θ | data) for each chain"
    tempering::M
    "Number of chains run in sampling step."
    chains::Int64
    "Number of MCMC iterations. May be overwritten in SMC."
    iterations::Int64
    "Burn in samples."
    burnin::Int64
    "Number of consecutive samples taken for diagnostics output."
    thinning::Int64
    "Boolean if trace and algorithm should be safed to working directory."
    safeoutput::Bool
    "Boolean if summary statistics should be calculated."
    printoutput::Bool
    "Default arguments for printing summary statistics"
    printdefault::PrintDefault
    "Default arguments for printing output during sampling runs"
    report::P
    function SampleDefault(;
        dataformat=Batch(),
        tempering=IterationTempering(Float64, UpdateFalse(), 1.0, 1000),
        chains=4,
        iterations=2000,
        burnin=max(1, Int64(floor(iterations / 10))),
        thinning = 1,
        safeoutput=false,
        printoutput=true,
        printdefault=PrintDefault(),
        report=ProgressReport(),
    )
        ArgCheck.@argcheck 0 < chains
        ArgCheck.@argcheck 0 < thinning
        ArgCheck.@argcheck 0 <= burnin < iterations
        return new{
        typeof(dataformat), typeof(tempering), typeof(report)
        }(
            dataformat,
            tempering,
            chains,
            iterations,
            burnin,
            thinning,
            safeoutput,
            printoutput,
            printdefault,
            report,
        )
    end
end

############################################################################################
# Export
export
    ProgressLog,
    SilentLog,
    ConsoleLog,
    ProgressReport,
    SampleDefault


text:
using StatsBase
using Test
using DelimitedFiles

##### Location

## geomean

@test geomean([1, 2, 3])    ≈ cbrt(6.0)
@test geomean(1:3)          ≈ cbrt(6.0)
@test geomean([2, 8])       ≈ 4.0
@test geomean([4, 1, 1/32]) ≈ 0.5
@test geomean([1, 0, 2]) == 0.0

## harmmean

@test harmmean([1, 2, 3]) ≈ 3 / (1 + 1/2 + 1/3)
@test harmmean(1:3)       ≈ 3 / (1 + 1/2 + 1/3)
@test harmmean([1, 2, 4]) ≈ 12 / 7

## genmean
@test genmean([1,1,2,3], 1)         ≈ 7/4
@test genmean([1,4,2], -1)          ≈ 12/7
@test genmean([1,1,2,3], 0)         ≈ (6.0)^(1/4)
@test genmean([1.2,-0.5,0], 2)      ≈ sqrt(169/300)
@test genmean([16/9,0.25,1.0], 1.5) ≈ (755/648)^(2/3)
# Test numerical stability for `p` close to 0 (genmean should be close to geometric mean).
@test isapprox(genmean([1,1,2,3], -1e-8), (6.0)^(1/4), atol=1e-8)
# Test numerical stability for large `p` (genmean should be close to max).
@test isapprox(genmean([0.98,1.02], 1e4), 1.02, atol=1e-4)

## mode & modes

@test mode([1, 2, 3, 3, 2, 2, 1], 1:3) == 2
@test modes([1, 2, 3, 3, 2, 2, 1], 1:3) == [2]
@test modes([1, 3, 2, 3, 3, 2, 2, 1], 1:3) == [2, 3]

@test mode([1, 2, 3, 3, 2, 2, 1]) == 2
@test modes([1, 2, 3, 3, 2, 2, 1]) == [2]
@test sort(modes([1, 3, 2, 3, 3, 2, 2, 1])) == [2, 3]

## zscores

@test zscore([-3:3;], 1.5, 0.5) == [-9.0:2.0:3.0;]

a = [3 4 5 6; 7 8 1 2; 6 9 3 0]
z1 = [4. 6. 8. 10.; 5. 6. -1. 0.; 1.5 3.0 0.0 -1.5]
z2 = [8. 2. 3. 1.; 24. 10. -1. -1.; 20. 12. 1. -2.]

@test zscore(a, [1, 2, 3], [0.5, 1.0, 2.0])    ≈ z1
@test zscore(a, [1 3 2 4], [0.25 0.5 1.0 2.0]) ≈ z2

@test zscore!(collect(-3.0:3.0), 1.5, 0.5) == [-9.0:2.0:3.0;]
@test zscore!(float(a), [1, 2, 3], [0.5, 1.0, 2.0])    ≈ z1
@test zscore!(float(a), [1 3 2 4], [0.25 0.5 1.0 2.0]) ≈ z2

@test zscore!(zeros(7), [-3:3;], 1.5, 0.5) == [-9.0:2.0:3.0;]
@test zscore!(zeros(size(a)), a, [1, 2, 3], [0.5, 1.0, 2.0])    ≈ z1
@test zscore!(zeros(size(a)), a, [1 3 2 4], [0.25 0.5 1.0 2.0]) ≈ z2

@test zscore(a)    ≈ zscore(a, mean(a), std(a))
@test zscore(a, 1) ≈ zscore(a, mean(a, dims=1), std(a, dims=1))
@test zscore(a, 2) ≈ zscore(a, mean(a, dims=2), std(a, dims=2))

###### quantile & friends

@test nquantile(1:5, 2) ≈ [1, 3, 5]
@test nquantile(1:5, 4) ≈ [1:5;]

@test percentile([1:5;], 25)           ≈  2.0
@test percentile([1:5;], [25, 50, 75]) ≈ [2.0, 3.0, 4.0]

##### Dispersion

@test span([3, 4, 5, 6, 2]) == (2:6)

@test variation([1:5;]) ≈ 0.527046276694730

@test sem([1:5;]) ≈ 0.707106781186548

@test StatsBase.mad(1:5; center=3, normalize=true) ≈ 1.4826022185056018
@test StatsBase.mad!([1:5;]; center=3, normalize=true) ≈ 1.4826022185056018
@test mad(1:5, normalize=true) ≈ 1.4826022185056018
@test StatsBase.mad(1:5, normalize=false) ≈ 1.0
@test StatsBase.mad!([1:5;], normalize=false) ≈ 1.0
@test StatsBase.mad(1:5, center=3, normalize=false) ≈ 1.0
@test StatsBase.mad!([1:5;], center=3, normalize=false) ≈ 1.0
@test_throws ArgumentError mad(Int[])

# Issue 197
@test mad(1:2, normalize=true) ≈ 0.7413011092528009

@test iqr(1:5) ≈ 2.0

nutrient = readdlm(joinpath(@__DIR__, "data", "nutrient.txt"))[:,2:end]
@test @inferred(genvar(nutrient)) ≈ 2.8310418e19 rtol=1e-6
@test @inferred(totalvar(nutrient)) ≈ 2.83266877e6 rtol=1e-6

X = [1 2 5
     4 1 6
     4 0 4]
@test @inferred(genvar(X)) ≈ 0.0
@test @inferred(totalvar(X)) ≈ 5.0

x = rand(Float32, 10)
@test genvar(x) == totalvar(x) == var(x)

it = (xᵢ for xᵢ in x)
@test genvar(it) == totalvar(it) == var(it)

##### entropy

@test entropy([0.5, 0.5])      ≈ 0.6931471805599453
@test entropy([0.2, 0.3, 0.5]) ≈ 1.0296530140645737

@test entropy([0.5, 0.5],2)       ≈ 1.0
@test entropy([0.2, 0.3, 0.5], 2) ≈ 1.4854752972273344
@test entropy([1.0, 0.0]) ≈ 0.0

##### Renyi entropies
# Generate a random probability distribution
nindiv = 50
dist = rand(nindiv)
dist /= sum(dist)

# Check Shannon entropy against Renyi entropy of order 1
@test entropy(dist)         ≈ renyientropy(dist, 1)
@test renyientropy(dist, 1) ≈ renyientropy(dist, 1.0)

# Check Renyi entropy of order 0 is the natural log of the count of non-zeros
@test renyientropy(dist, 0) ≈ log(mapreduce(x -> x > 0 ? 1 : 0, +, dist))

# And is therefore not affected by the addition of non-zeros
zdist = dist
zdist = append!(dist, zeros(50))
@test renyientropy(dist, 0) ≈ renyientropy(zdist, 0)

# Indeed no Renyi entropy should be
loworder = rand() # Between 0 and 1
@test renyientropy(dist, loworder) ≈ renyientropy(zdist, loworder)
order = rand() * 49 + 1 # And over 1
@test renyientropy(dist, order) ≈ renyientropy(zdist, order)

# Renyi entropy of order infinity is -log(maximum(dist))
@test renyientropy(dist, Inf) ≈ -log(maximum(dist))

# And all Renyi entropies of uniform probabilities should be log n, where n is the length
udist = ones(nindiv) / nindiv
@test renyientropy(udist, order) ≈ log(nindiv)

# And test generalised probability distributions (sum(p) != 1)
scale = rand()
@test renyientropy(udist * scale, 0)     ≈ renyientropy(udist, 0) - log(scale)
@test renyientropy(udist * scale, 1)     ≈ renyientropy(udist, 1) - log(scale)
@test renyientropy(udist * scale, Inf)   ≈ renyientropy(udist, Inf) - log(scale)
@test renyientropy(udist * scale, order) ≈ renyientropy(udist, order) - log(scale)

##### Cross entropy
@test crossentropy([0.2, 0.3, 0.5], [0.3, 0.4, 0.3])    ≈ 1.1176681825904018
@test crossentropy([0.2, 0.3, 0.5], [0.3, 0.4, 0.3], 2) ≈ 1.6124543443825532

##### KL divergence
@test kldivergence([0.2, 0.3, 0.5], [0.3, 0.4, 0.3])    ≈ 0.08801516852582819
@test kldivergence([0.2, 0.3, 0.5], [0.3, 0.4, 0.3], 2) ≈ 0.12697904715521868

##### summarystats

s = summarystats(1:5)
@test isa(s, StatsBase.SummaryStats)
@test s.min == 1.0
@test s.max == 5.0
@test s.mean   ≈ 3.0
@test s.median ≈ 3.0
@test s.q25    ≈ 2.0
@test s.q75    ≈ 4.0


text:
#=
Copyright (c) 2018-2022 Chris Coey, Lea Kapelevich, and contributors

This Julia package Hypatia.jl is released under the MIT license; see LICENSE
file in the root directory or at https://github.com/chriscoey/Hypatia.jl

problem 1 (linear_dynamics = true)
eigenvalue problem related to Lyapunov stability example from sec 2.2.2, 6.3.2
"Linear Matrix Inequalities in System and Control Theory" by
S. Boyd, L. El Ghaoui, E. Feron, and V. Balakrishnan:
minimize    t
subject to  P in S_+
            [-A'*P - P*A - C'C, P*B;
            B'*P, tI] in S_+
for the system with linear dynamics x_dot = A*x

problem 2 (linear_dynamics = false)
Lyapunov stability example
see https://stanford.edu/class/ee363/sessions/s4notes.pdf:
minimize    t
subject to  P - I in S_+
            [-A'*P - P*A - alpha*P - t*gamma^2*I, -P;
            -P, tI] in S_+
originally a feasibility problem, a feasible P and t prove the existence of a
Lyapunov function for the system x_dot = A*x+g(x), norm(g(x)) <= gamma*norm(x)
=#

struct LyapunovStabilityJuMP{T <: Real} <: ExampleInstanceJuMP{T}
    num_rows::Int
    num_cols::Int
    linear_dynamics::Bool # solve problem 1 in the description, else problem 2
    use_matrixepipersquare::Bool # use matrixepipersquare cone, else PSD cone
end

function build(inst::LyapunovStabilityJuMP{T}) where {T <: Float64}
    (num_rows, num_cols) = (inst.num_rows, inst.num_cols)

    model = JuMP.Model()
    JuMP.@variable(model, t)
    JuMP.@objective(model, Min, t)

    if inst.linear_dynamics
        A = randn(num_rows, num_rows)
        A = -A * A'
        B = randn(num_rows, num_cols)
        C = randn(num_rows, num_rows)
        JuMP.@variable(model, P[1:num_rows, 1:num_rows], PSD)
        U = -A' * P - P * A - C' * C / 100
        W = P * B
    else
        @assert num_rows == num_cols
        # P = -A is a feasible solution, with alpha and gamma sufficiently small
        A = randn(num_rows, num_rows)
        A = -A * A' - I
        alpha = 0.01
        gamma = 0.01
        JuMP.@variable(model, P[1:num_rows, 1:num_rows], Symmetric)
        JuMP.@constraint(model, Symmetric(P - I) in JuMP.PSDCone())
        U = -A' * P - P * A - alpha * P - (t * gamma ^ 2) .*
            Matrix(I, num_rows, num_rows)
        W = -P
    end

    if inst.use_matrixepipersquare
        U_svec = Cones.smat_to_svec!(zeros(eltype(U),
            Cones.svec_length(num_rows)), U, sqrt(2))
        coneT = Hypatia.MatrixEpiPerSquareCone{T, T}
        JuMP.@constraint(model, vcat(U_svec, t / 2, vec(W)) in
            coneT(num_rows, num_cols))
    else
        JuMP.@constraint(model, Symmetric(
            [t .* Matrix(I, num_cols, num_cols) W'; W U]) in JuMP.PSDCone())
    end

    return model
end


text:
using RCall
using Distributions
using StatsBase
using IterTools
using Gurobi
using JLD2

using EmpiricalBayes

nreps = 4
k_max = 10 #in total 25*16 replications
ks_par = 1:10#1:10

true_dists = [EmpiricalBayes.easy_dist, EmpiricalBayes.hard_dist]
xs = collect(-3.0:0.5:3.0)
minimax_σs = [0.2]
ms = [10_000]
targets =[ [PosteriorTarget(PosteriorMeanNumerator(x)) for x in xs],
           [PosteriorTarget(LFSRNumerator(x)) for x in xs]]
combs = []

for (true_dist, target, σ, m) in IterTools.product(true_dists, targets, minimax_σs, ms)
    push!(combs, Dict(:dist=>true_dist, :target=>target, :σ=>σ, :m=>m ))
end

comb_reps = []
for (l, k_par) in IterTools.product(1:length(combs), ks_par)
    push!(comb_reps, (l,k_par))
end

j = parse(Int64, ARGS[1])
comb_rep = comb_reps[j]
i = comb_rep[1]
k_rep = comb_rep[2]
comb = combs[i]

m = comb[:m]
σ = comb[:σ]
true_dist = comb[:dist]
target = comb[:target]

marginal_grid = collect(range(-6,stop=6,length=1001));
prior_grid = collect(range(-3,stop=3,length=121));
marginal_h = marginal_grid[2]-marginal_grid[1];

d_true = NormalConvolutionProblem(true_dist, marginal_grid)
ds = MixingNormalConvolutionProblem(Normal, σ, prior_grid, marginal_grid);

for k=1:k_max
    sim_array = []
    for l = 1:nreps
        @show l
        # get the sample
        Xs = rand(d_true,m)
        m_total = length(Xs)
        m_test= ceil(Int, m_total/2)
        m_train = m_total - m_test
        idx_test = sample(1:m_total, m_test, replace=false)
        idx_train = setdiff(1:m_total, idx_test)
        Xs_train = Xs[idx_train]
        Xs_test = Xs[idx_test]

        brad_fit =  fit(SmoothedBradDeconvolveR, Xs, 0.05; c0=0.01,prior_grid = collect(-3.6:0.05:3.6))
        f_nb = fit(BinnedMarginalDensityNeighborhood, Xs_train, ds; nboot=400)

        calibrator_res = []
        for t in target
            @show t
            CB_num = ComteButucea(t.num, m_train, marginal_grid)
            CB_denom = ComteButucea(t.denom, m_train, marginal_grid)
            ceb_res =  CEB_ci(Xs_train, Xs_test,
                        ds,
                        t,
                        CB_num, #
                        CB_denom; f_nb=f_nb)
            push!(calibrator_res, ceb_res)
        end
        push!(sim_array, (calibrator_res, brad_fit))
    end
    res = (i, comb, sim_array)
    @save "/scratch/users/ignat/sims/base_sim_Jan10/mysim_$(i)_$(k)_$(k_rep).jld" res
end


text:
@doc raw"""
    NodaHilbert{T<:AbstractFloat}

Noda-Hilbert transform.

The Noda-Hilbert transformation is a square, skew symmetric,
[Toeplitz matrix](https://en.wikipedia.org/wiki/Toeplitz_matrix)
with zero diagonal defined as follows:

```math
N_{ij}
    = \begin{cases}
        0                     & i \eq j
        \frac{1}{\pi (j - i)} & i \neq i \\
      \end{cases}
```

# Examples
```jldoctest
julia> N === NodaHilbert{Float64}()
true

julia> NodaHilbert{Float32}()
NodaHilbert{Float32}()

julia> N(3)
3×3 ToeplitzMatrices.Toeplitz{Float64,Complex{Float64}}:
 -0.0        0.31831   0.159155
 -0.31831   -0.0       0.31831
 -0.159155  -0.31831  -0.0
```
"""
# TODO: should we subtype AbstractMatrix?
struct NodaHilbert{T<:AbstractFloat} end
NodaHilbert() = NodaHilbert{Float64}()
const N = NodaHilbert()

# Elements progressively decrease in magnitude.
@inline function calcfirstcolumn(T, n, nmax=21_361_417)
    c = zeros(T, n)
    @inbounds for i in 1:min(n-1, nmax)
        c[i+1] = 1 / (π * i)
    end
    return c
end

"""
    (N::NodaHilbert)(n::Int)

Realize the Noda-Hilbert transform matrix of size `n`.

`n` is the number of rows (and columns) in resulting matrix.
"""
function (N::NodaHilbert{T})(n::Int) where {T}
    c = calcfirstcolumn(T, n)
    return Toeplitz(-c, c)
end

# TODO: check eltype(A) when realizing the transform?
# TODO: probably improve the fact that it realizes an object every time?
Base.:*(N::NodaHilbert, A::AbstractArray) = N(size(A, 1)) * A


text:
#
# Just returns a fine mesh for the fit
#
function finexy(X,fine,model,fit)
  Xmin = minimum(X)
  Xmax = maximum(X)
  x = collect(Xmin:(Xmax-Xmin)/fine:Xmax)
  y = model(x,fit.param)
  ypred = model(X,fit.param)
  return x, y, ypred
end


text:
module Metis

using SparseArrays
using LinearAlgebra
import LightGraphs
using METIS_jll: libmetis

# Metis C API
include("metis_h.jl")
const options = fill(idx_t(-1), METIS_NOPTIONS)
options[METIS_OPTION_NUMBERING] = 1

# Julia interface
"""
    Metis.Graph

1-based CSR representation of a graph as defined in
section 5.5 "Graph data structure" in the Metis manual.
"""
struct Graph
    nvtxs::idx_t
    xadj::Vector{idx_t}
    adjncy::Vector{idx_t}
    vwgt::Vector{idx_t}
    Graph(nvtxs, xadj, adjncy) = new(nvtxs, xadj, adjncy)
    Graph(nvtxs, xadj, adjncy, vwgt) = new(nvtxs, xadj, adjncy, vwgt)
end

"""
    Metis.graph(G::SparseMatrixCSC; check_hermitian=true)

Construct the 1-based CSR representation of the sparse matrix `G`.
If `check_hermitian` is `false` the matrix is not checked for being hermitian
before constructing the graph.
"""
function graph(G::SparseMatrixCSC; check_hermitian=true)
    if check_hermitian
        ishermitian(G) || throw(ArgumentError("matrix must be Hermitian"))
    end
    N = size(G, 1)
    xadj = Vector{idx_t}(undef, N+1)
    xadj[1] = 1
    adjncy = Vector{idx_t}(undef, nnz(G))
    adjncy_i = 0
    @inbounds for j in 1:N
        n_rows = 0
        for k in G.colptr[j] : (G.colptr[j+1] - 1)
            i = G.rowval[k]
            if i != j # don't include diagonal elements
                n_rows += 1
                adjncy_i += 1
                adjncy[adjncy_i] = i
            end
        end
        xadj[j+1] = xadj[j] + n_rows
    end
    resize!(adjncy, adjncy_i)
    return Graph(idx_t(N), xadj, adjncy)
end

"""
    graph(G::LightGraphs.AbstractSimpleGraph)

Construct the 1-based CSR representation of the `LightGraphs` graph `G`.
"""
function graph(G::LightGraphs.AbstractSimpleGraph)
    N = LightGraphs.nv(G)
    xadj = Vector{idx_t}(undef, N+1)
    xadj[1] = 1
    adjncy = Vector{idx_t}(undef, 2*LightGraphs.ne(G))
    adjncy_i = 0
    for j in 1:N
        ne = 0
        for i in LightGraphs.outneighbors(G, j)
            ne += 1
            adjncy_i += 1
            adjncy[adjncy_i] = i
        end
        xadj[j+1] = xadj[j] + ne
    end
    resize!(adjncy, adjncy_i)
    return Graph(idx_t(N), xadj, adjncy)
end

"""
    graph(elements::AbstractMatrix{<:Integer}, ncommon=1)

Construct the 1-based CSR representation of the mesh described by the
connectivity matrix `elements`.

`elements[i,j]` contains the `i`-th node of the `j`-th element. Thus the size of
the first dimension `elements` is equal to the number of nodes per element, and
the size of the second dimension is the number of elements.

The nodes are assumed to be numbered consecutively, starting from 1.
"""
function graph(elements::AbstractMatrix{<:Integer}, ncommon=1)
    # Assume the number of nodes in the mesh is equal to the highest node
    # number.
    nnodes = maximum(elements)

    # Number of nodes per element.
    n_npe = size(elements, 1)

    # Number of elements in the mesh.
    ne = size(elements, 2)

    # Get the element connectivity data in the form METIS needs.
    eptr = convert(Array{Metis.idx_t}, collect(0:ne) .* n_npe)
    eind = convert(Array{Metis.idx_t}, elements[:])

    # With 1-based indexing, need to add 1 to eptr.
    eptr .+= 1

    # Create some pointers and stuff that will be needed for the METIS library call.
    ne = Metis.idx_t(ne)
    nn = Metis.idx_t(nnodes)
    ncommon = Metis.idx_t(ncommon)
    num_flag = Metis.options[Metis.METIS_OPTION_NUMBERING]
    xadj = [Ptr{Metis.idx_t}()]
    adjncy = [Ptr{Metis.idx_t}()]

    # Do it!
    Metis.METIS_MeshToDual(ne, nn, eptr, eind, ncommon, num_flag, xadj, adjncy)

    # Create a Julian copy of the xadj array. The length of xadj is one more
    # than the number of graph verticies (here, the number of mesh elements).
    xadj_out = copy(unsafe_wrap(Vector{Metis.idx_t}, xadj[1], ne+1))

    # Free the METIS memory.
    Metis.METIS_Free(xadj[1])

    # Create a Julian copy of the adjncy array.  The length of adjncy is twice
    # the number of graph edges. Turns out that's the last entry in xadj minus
    # 1 if we're using one-based indices.
    adjncy_out = copy(unsafe_wrap(Vector{Metis.idx_t}, adjncy[1], xadj_out[end]-1))

    # Free the METIS memory.
    Metis.METIS_Free(adjncy[1])

    return Graph(ne, xadj_out, adjncy_out)
end

"""
    perm, iperm = Metis.permutation(G)

Compute the fill reducing permutation `perm`
and its inverse `iperm` of `G`.
"""
permutation(G) = permutation(graph(G))

function permutation(G::Graph)
    perm = Vector{idx_t}(undef, G.nvtxs)
    iperm = Vector{idx_t}(undef, G.nvtxs)
    vwgt = isdefined(G, :vwgt) ? G.vwgt : C_NULL
    METIS_NodeND(G.nvtxs, G.xadj, G.adjncy, vwgt, options, perm, iperm)
    return perm, iperm
end

"""
    Metis.partition(G, n; alg = :KWAY)

Partition the graph `G` in `n` parts.
The partition algorithm is defined by the `alg` keyword:
 - :KWAY: multilevel k-way partitioning
 - :RECURSIVE: multilevel recursive bisection
"""
partition(G, nparts; alg = :KWAY) = partition(graph(G), nparts, alg = alg)

function partition(G::Graph, nparts::Integer; alg = :KWAY)
    part = Vector{idx_t}(undef, G.nvtxs)
    vwgt = isdefined(G, :vwgt) ? G.vwgt : C_NULL
    edgecut = fill(idx_t(0), 1)
    if alg === :RECURSIVE
        METIS_PartGraphRecursive(G.nvtxs, idx_t(1), G.xadj, G.adjncy, vwgt, C_NULL, C_NULL,
                                 idx_t(nparts), C_NULL, C_NULL, options, edgecut, part)
    elseif alg === :KWAY
        METIS_PartGraphKway(G.nvtxs, idx_t(1), G.xadj, G.adjncy, vwgt, C_NULL, C_NULL,
                            idx_t(nparts), C_NULL, C_NULL, options, edgecut, part)
    else
        throw(ArgumentError("unknown algorithm $(repr(alg))"))
    end
    return part
end

"""
    Metis.separator(G)

Compute a vertex separator of the graph `G`.
"""
separator(G) = separator(graph(G))

function separator(G::Graph)
    part = Vector{idx_t}(undef, G.nvtxs)
    sepsize = fill(idx_t(0), 1)
    vwgt = isdefined(G, :vwgt) ? G.vwgt : C_NULL
    # METIS_ComputeVertexSeparator segfaults with 1-based indexing
    xadj = G.xadj .- idx_t(1)
    adjncy = G.adjncy .- idx_t(1)
    METIS_ComputeVertexSeparator(G.nvtxs, xadj, adjncy, vwgt, options, sepsize, part)
    part .+= 1
    return part
end

end # module


text:
using NBConvolution
using Distributions
using Base.Test

# compare the two methods to verify accuracy
d1 = FurmanNegativeBinomialConvolution([4.0, 1.0], [0.01, 0.5], 1000)
d2 = NegativeBinomialConvolution([4, 1], [0.01, 0.5])
@test abs(pdf(d1, 10) - pdf(d2, 10)) < 1e-10

d1 = FurmanNegativeBinomialConvolution([4.0, 10.0], [0.5, 0.5], 1000)
d2 = NegativeBinomialConvolution([4, 10], [0.5, 0.5])
@test abs(pdf(d1, 30) - pdf(d2, 30)) < 1e-10

d1 = FurmanNegativeBinomialConvolution([4.0, 10.0, 3.0, 6.0], [0.4, 0.5, 0.3, 0.6], 1000)
d2 = NegativeBinomialConvolution([4, 10, 3, 6], [0.4, 0.5, 0.3, 0.6])
@test abs(pdf(d1, 30) - pdf(d2, 30)) < 1e-10

# cdf
d2 = NegativeBinomialConvolution([4, 10], [0.7, 0.5])
@test abs(cdf(d2, 1000) - 1.0) < 1e-10

# non-integer rs
d3 = NegativeBinomialConvolution([4.2, 10], [0.7, 0.5])
@test abs(cdf(d3, 1000) - 1.0) < 1e-10
@test abs(pdf(d2, 10) - pdf(d3, 10)) > 1e-10
d4 = FurmanNegativeBinomialConvolution([4.2, 10.0], [0.7, 0.5], 1000)
@test abs(pdf(d3, 4) - pdf(d4, 4)) < 1e-10

# integrate with the distributions package
modelInner = MixtureModel([NegativeBinomialConvolution([3, 4], [0.1, 0.5]), NegativeBinomial(5, 0.4)], [0.6, 0.4])
pdf(modelInner, 0)
pdf(modelInner, 10)


text:
using DSGE, ModelConstructors, HDF5, Random, JLD2, FileIO, SMC, Test, Distributions, StateSpaceRoutines, Dates

path = dirname(@__FILE__)
writing_output = false

if VERSION < v"1.5"
    ver = "111"
else
    ver = "150"
end

regenerate_sim_data = false
#########################
# Regime-switching test #
# w/ An-Schorfheide     #
#########################
m = AnSchorfheide()

# Set up SMC
m <= Setting(:sampling_method, :SMC)
m <= Setting(:data_vintage, "210101")
m <= Setting(:cond_vintage, "210101")
m <= Setting(:n_Φ, 100)
m <= Setting(:λ, 3.0)
m <= Setting(:n_particles, 500)
m <= Setting(:n_smc_blocks, 1)
m <= Setting(:use_parallel_workers, false)
m <= Setting(:step_size_smc, 0.5)
m <= Setting(:n_mh_steps_smc, 1)
m <= Setting(:resampler_smc, :polyalgo)
m <= Setting(:target_accept, 0.25)

m <= Setting(:mixture_proportion, .9)
m <= Setting(:tempering_target, 0.95)
m <= Setting(:resampling_threshold, .5)
m <= Setting(:use_fixed_schedule, true)

sim_filepath = joinpath(path, "..", "..", "reference/sim_data_regswitch_anschorfheide.h5")
if regenerate_sim_data
    simulate_data(m; filepath = sim_filepath)
end
data = h5read(sim_filepath, "data")

true_lik = DSGE.likelihood(m, data)

m <= Setting(:saveroot, normpath(joinpath(dirname(@__FILE__), "save")))

# With regime-switching
m <= Setting(:regime_switching, true, true, "rs2", "") # For file output purposes
m <= Setting(:n_regimes, 3)
m <= Setting(:regime_dates, Dict(1 => date_presample_start(m), 2 => DSGE.iterate_quarters(Date(1960, 3, 31), 50),
                                 3 => DSGE.iterate_quarters(Date(1960, 3, 31), 100)))
setup_regime_switching_inds!(m)
m <= Setting(:model2para_regimes, Dict{Int, Int}(1 => 1, 2 => 1, 3 => 2))

for i in 1:length(m.parameters)
    for k in 1:2
        ModelConstructors.set_regime_val!(m.parameters[i], k, m.parameters[i].value)
    end
end

param_mat = repeat([1 1 2], length(m.parameters))

DSGE.setup_param_regimes!(m, param_mat)

sys = compute_system(m)

regswitch_lik = DSGE.likelihood(m, data)

@testset "Estimate regime-switching AnSchorfheide with SMC (approx. 3 min.)" begin

    @test true_lik == regswitch_lik

    # Regime switching estimation
    true_para = ModelConstructors.get_values(m.parameters)
    Random.seed!(1793)
    DSGE.smc2(m, data, regime_switching = true, run_csminwel = false,
              verbose = :none)

    posterior_means = vec(mean(load_draws(m, :full), dims = 1))

    @test length(posterior_means) == length(m.parameters) * 2
    inds = vcat(1:4, 7:16, 18:20, 24:27, 29:length(posterior_means))
    @test maximum(abs.(true_para[inds] - posterior_means[inds])) < .6
    inds2 = [5, 17, 21, 23]
    @test maximum(abs.(true_para[inds] - posterior_means[inds])) < 1.4
end


text:
@testset "Metadata" begin
    sample_text1 = "This is a string"
    sample_text2 = "This is also a string"

    sd1 = StringDocument(sample_text1)
    sd2 = StringDocument(sample_text2)
    crps = Corpus([sd1, sd2])

    # Single document metadata getters
    @test isequal(name(sd1), "Unnamed Document")
    @test isequal(language(sd1), Languages.English())
    @test isequal(author(sd1), "Unknown Author")
    @test isequal(timestamp(sd1), "Unknown Time")

    # Single document metadata setters
    name!(sd1, "Document")
    language!(sd1, Languages.German())
    author!(sd1, "Author")
    timestamp!(sd1, "Time")
    @test isequal(name(sd1), "Document")
    @test isequal(language(sd1), Languages.German())
    @test isequal(author(sd1), "Author")
    @test isequal(timestamp(sd1), "Time")

    # Metadata getters for an entire corpus
    @test isequal(TextAnalysis.names(crps), ["Document", "Unnamed Document"])
    @test isequal(languages(crps), [Languages.German(), Languages.English()])
    @test isequal(authors(crps), ["Author", "Unknown Author"])
    @test isequal(timestamps(crps), ["Time", "Unknown Time"])

    # Metadata setters for an entire corpus
    names!(crps, "Document")
    languages!(crps, Languages.Spanish())
    authors!(crps, "Author")
    timestamps!(crps, "Time")
    @test isequal(TextAnalysis.names(crps), ["Document", "Document"])
    @test isequal(languages(crps), [Languages.Spanish(), Languages.Spanish()])
    @test isequal(authors(crps), ["Author", "Author"])
    @test isequal(timestamps(crps), ["Time", "Time"])
    names!(crps, ["Unnamed Document", "Unnamed Document"])
    languages!(crps, [Languages.English(), Languages.English()])
    authors!(crps, ["Unknown Author", "Unknown Author"])
    timestamps!(crps, ["Unknown Time", "Unknown Time"])
    @test isequal(TextAnalysis.names(crps), ["Unnamed Document", "Unnamed Document"])
    @test isequal(languages(crps), [Languages.English(), Languages.English()])
    @test isequal(authors(crps), ["Unknown Author", "Unknown Author"])
    @test isequal(timestamps(crps), ["Unknown Time", "Unknown Time"])
end


text:
using ContinuousTransformations
using FastGaussQuadrature
using ArgCheck

export Quadrature, quadrature_standard_normal, quadrature_standard_uniform

"""
Univariate numerical quadrature. A lightweight container for the
`nodes` and the `weights`, calculated from some quadrature rule. The
`domain` is provided for information purposes.
"""
immutable Quadrature{TD <: AbstractInterval, T}
    domain::TD
    nodes::Vector{T}
    weights::Vector{T}
    function Quadrature(domain, nodes, weights)
        @argcheck length(nodes) == length(weights)
        new(domain, nodes, weights)
    end
end

function Quadrature{TD, T}(domain::TD, nodes::Vector{T}, weights::Vector{T})
    Quadrature{TD, T}(domain, nodes, weights)
end

ContinuousTransformations.domain(q::Quadrature) = q.domain

Base.length(q::Quadrature) = length(q.nodes)

function Base.show(io::IO, q::Quadrature)
    print(io, "Quadrature of $(length(q)) nodes on $(q.domain)")
end

(q::Quadrature)(f) = dot(f.(q.nodes), q.weights)

"""
Return a quadrature that integrates a function under the standard
normal distribution using `N` nodes.
"""
function quadrature_standard_normal(N::Int)
    nodes, weights = gausshermite(N)
    Quadrature(ℝ, nodes*√2, weights/√π)
end

"""
Return a quadrature that integrates a function on the unit domain
[0,1] under the standard uniform distribution using `N` nodes.
"""
function quadrature_standard_uniform(N::Int)
    nodes, weights = gausslegendre(N)
    Quadrature(𝕀, fma.(nodes, 0.5, 0.5), weights .* 0.5)
end


text:
struct Local <: HamiltonianSolver end

min_path(T, stops, ::Local; kwargs...) = _min_length_path_local(T, stops; kwargs...)

function _min_length_path_local(T::TravelTimes, stops::Vector{<:AbstractStop}; kwargs...)
    _min_length_path_local(T.travel_Times, stops; kwargs...)
end

function _min_length_path_local(T::AbstractMatrix, stops::Vector{<:AbstractStop}; initial_stop=nothing)
    # initial stop
    x = isnothing(initial_stop) ? stops[1] : initial_stop
    x = BareStop(name(x))
    num_stops = length(stops)

    # prellocate output
    out = Vector{BareStop}()

    # store first stop in the output sequence
    push!(out, x)

    while length(out) < num_stops
        s = sort(view(T, name(x), :))

        i = 2 # ignore current position
        y = s.dicts[1].keys[i] |> BareStop
        while y ∈ out # iterate until finding a stop that hasn't been visited
            i += 1
            i > length(s.dicts[1]) && return out
            y = s.dicts[1].keys[i] |> BareStop
        end
        push!(out, y)
        x = y
    end
    return out
end


text:
"""Algorithm for computing H, as given by arxiv:1403.7698

H is related to Wigner's (small) d via

    dₗⁿᵐ = ϵₙ ϵ₋ₘ Hₗⁿᵐ,

where

         ⎧ 1 for k≤0
    ϵₖ = ⎨
         ⎩ (-1)ᵏ for k>0

H has various advantages over d, including the fact that it can be efficiently
and robustly valculated via recurrence relations, and the following symmetry
relations:

    H^{m', m}_n(β) = H^{m, m'}_n(β)
    H^{m', m}_n(β) = H^{-m', -m}_n(β)
    H^{m', m}_n(β) = (-1)^{n+m+m'} H^{-m', m}_n(π - β)
    H^{m', m}_n(β) = (-1)^{m+m'} H^{m', m}_n(-β)

Because of these symmetries, we only need to evaluate at most 1/4 of all the
elements.

"""

"""Return flat index into arrray of [n, m] pairs.

Assumes array is ordered as

    [
        [n, m]
        for n in range(n_max+1)
        for m in range(-n, n+1)
    ]

"""
nm_index(n, m) = m + n * (n + 1) + 1

"""Return flat index into arrray of [n, abs(m)] pairs

Assumes array is ordered as

    [
        [n, m]
        for n in range(n_max+1)
        for m in range(n+1)
    ]

"""
nabsm_index(n, absm) = absm + (n * (n + 1)) ÷ 2 + 1

"""Return flat index into arrray of [n, mp, m]

Assumes array is ordered as

    [
        [n, mp, m]
        for n in range(n_max+1)
        for mp in range(-n, n+1)
        for m in range(-n, n+1)
    ]

"""
nmpm_index(n, mp, m) = (((4n + 6) * n + 6mp + 5) * n + 3(m + mp)) ÷ 3 + 1

@inbounds function _step_1!(w::WignerMatrixCalculator)
    """If n=0 set H_{0}^{0,0}=1."""
    w.Hwedge[1] = 1
end

@inbounds function _step_2!(w::WignerMatrixCalculator, expiβ::Complex)
    """Compute values H^{0,m}_{n}(β)for m=0,...,n and H^{0,m}_{n+1}(β) for m=0,...,n+1

     Uses Eq. (32) of Gumerov-Duraiswami (2014) [arxiv:1403.7698]:

        H^{0,m}_{n}(β) = (-1)^m √((n-|m|)! / (n+|m|)!) P^{|m|}_{n}(cos β)
                       = (-1)^m P̄^{|m|}_{n}(cos β) / √(k (2n+1))

    Here, k=1 for m=0, and k=2 for m>0, and

        P̄ = √{k(2n+1)(n-m)!/(n+m)!} P

    We use the "fully normalized" associated Legendre functions (fnALF) P̄ because,
    as explained by Xing et al. (2020) [https://doi.org/10.1007/s00190-019-01331-0],
    it is possible to compute these values very efficiently and accurately, while
    also delaying the onset of overflow and underflow.

    NOTE: Though not specified in arxiv:1403.7698, there is not enough information
    for step 4 unless we also use symmetry to set H^{1,0}_{n} here.  Similarly,
    step 5 needs additional information, which depends on setting H^{0, -1}_{n}
    from its symmetric equivalent H^{0, 1}_{n} in this step.

    """
    n_max, mp_max, TW = ℓₘₐₓ(w), m′ₘₐₓ(w), T(w)
    Hwedge, Hextra, Hv = w.Hwedge, w.Hextra, w.Hv
    cosβ = expiβ.re
    sinβ = expiβ.im
    sqrt3 = √TW(3)

    # The general expressions for the constants are Eq. (13) of Xing et al.:
    #
    #   aₙ = √((2n+1)/TW(2n-1))
    #   bₙ = √(2*(n-1)*(2n+1)/TW(n*(2n-1)))
    #   cₙₘ = √(((n+m)*(n-m)*(2n+1)) / TW(2n-1)) / n
    #   dₙₘ = √(((n-m)*(n-m-1)*(2n+1)) / TW(2n-1)) / (2n)
    #   eₙₘ = √(((n+m)*(n+m-1)*(2n+1)) / TW(2n-1)) / (2n)
    #
    # Below, I factor aₙ out of each of these expressions, along with 1/2n where
    # relevant, to avoid divisions.
    #
    # We initialize with Eq. (14), then step through with Eq. (12), to compute all
    # values of P̄.  Finally, we normalize everything again to compute the H
    # values.

    if n_max > 0
        # n = 1
        n0n_index = WignerHindex(1, 0, 1, mp_max)
        Hwedge[n0n_index] = sqrt3 * sinβ
        Hwedge[n0n_index-1] = sqrt3 * cosβ
        # n = 2, ..., n_max+1
        for n in 2:n_max+1
            if n <= n_max
                n0n_index = WignerHindex(n, 0, n, mp_max)
                H = Hwedge
            else
                n0n_index = n + 1
                H = Hextra
            end
            nm10nm1_index = WignerHindex(n-1, 0, n-1, mp_max)
            inv2n = inv(TW(2n))
            aₙ = √((2n+1)/TW(2n-1))
            bₙ = aₙ * √((2*(n-1))/TW(n))

            # m = n
            eₙₘ = inv2n * √TW((2n)*(2n+1))
            H[n0n_index] = sinβ * eₙₘ * Hwedge[nm10nm1_index]

            # m = n-1
            eₙₘ = inv2n * √TW((2n-2)*(2n+1))
            cₙₘ = 2inv2n * √TW(2n+1)
            H[n0n_index-1] = cosβ * cₙₘ * Hwedge[nm10nm1_index] + sinβ * eₙₘ * Hwedge[nm10nm1_index-1]

            # m = n-2, ..., 2
            for i in 2:n-2
                # m = n-i
                cₙₘ = 2inv2n * aₙ * √TW((2n-i)*i)
                dₙₘ = inv2n * aₙ * √TW(i*(i-1))
                eₙₘ = inv2n * aₙ * √TW((2n-i)*(2n-i-1))
                H[n0n_index-i] = (
                    cosβ * cₙₘ * Hwedge[nm10nm1_index-i+1]
                    - sinβ * (
                        dₙₘ * Hwedge[nm10nm1_index-i+2]
                        - eₙₘ * Hwedge[nm10nm1_index-i]
                    )
                )
            end

            # m = 1
            cₙₘ = 2inv2n * aₙ * √TW((n+1)*(n-1))
            dₙₘ = inv2n * aₙ * √TW((n-1)*(n-2))
            eₙₘ = inv2n * aₙ * √TW(2n*(n+1))
            H[n0n_index-n+1] = (
                cosβ * cₙₘ * Hwedge[nm10nm1_index-n+2]
                - sinβ * (
                    dₙₘ * Hwedge[nm10nm1_index-n+3]
                    - eₙₘ * Hwedge[nm10nm1_index-n+1]
                )
            )

            # m = 0
            cₙₘ = aₙ
            dₙₘ = inv2n * aₙ * √TW(n*(n-1))
            eₙₘ = dₙₘ
            H[n0n_index-n] = (
                aₙ * cosβ * Hwedge[nm10nm1_index-n+1]
                - bₙ * sinβ * Hwedge[nm10nm1_index-n+2] / 2
            )

            # Supply extra edge cases as noted in docstring
            if n <= n_max
                Hv[nm_index(n, 1)] = Hwedge[WignerHindex(n, 0, 1, mp_max)]
                Hv[nm_index(n, 0)] = Hwedge[WignerHindex(n, 0, 1, mp_max)]
            end
        end

        # Supply extra edge cases as noted in docstring
        Hv[nm_index(1, 1)] = Hwedge[WignerHindex(1, 0, 1, mp_max)]
        Hv[nm_index(1, 0)] = Hwedge[WignerHindex(1, 0, 1, mp_max)]

        # Normalize, changing P̄ to H values
        for n in 1:n_max+1
            if n <= n_max
                n00_index = WignerHindex(n, 0, 0, mp_max)
                H = Hwedge
            else
                n00_index = 1
                H = Hextra
            end
            const0 = inv(√TW(2n+1))
            const1 = inv(√TW(4n+2))
            H[n00_index] *= const0
            for m in 1:n
                H[n00_index+m] *= const1
            end
            if n <= n_max
                Hv[nm_index(n, 1)] *= -const1
                Hv[nm_index(n, 0)] *= -const1
            end
        end
    end
end

@inbounds function _step_3!(w::WignerMatrixCalculator, expiβ::Complex)
    """Use relation (41) to compute H^{1,m}_{n}(β) for m=1,...,n.  Using symmetry and shift
    of the indices this relation can be written as

        b^{0}_{n+1} H^{1, m}_{n} =   (b^{−m−1}_{n+1} (1−cosβ))/2 H^{0, m+1}_{n+1}
                                   − (b^{ m−1}_{n+1} (1+cosβ))/2 H^{0, m−1}_{n+1}
                                   − a^{m}_{n} sinβ H^{0, m}_{n+1}

    """
    avalues = a(w)
    bvalues = b(w)
    n_max = ℓₘₐₓ(w)
    mp_max = m′ₘₐₓ(w)
    Hwedge = w.Hwedge
    Hextra = w.Hextra
    cosβ = expiβ.re
    sinβ = expiβ.im
    cosβ₊ = (1+cosβ)/2
    cosβ₋ = (1-cosβ)/2

    if n_max > 0 && mp_max > 0
        for n in 1:n_max
            # m = 1, ..., n
            i1 = WignerHindex(n, 1, 1, mp_max)
            if n+1 <= n_max
                i2 = WignerHindex(n+1, 0, 0, mp_max)
                H2 = Hwedge
            else
                i2 = 1
                H2 = Hextra
            end
            i3 = nm_index(n+1, 0)
            i4 = nabsm_index(n, 1)
            inverse_b5 = inv(bvalues[i3])
            for i in 0:n-1
                b6 = bvalues[-i+i3-2]
                b7 = bvalues[i+i3]
                a8 = avalues[i+i4]
                Hwedge[i+i1] = inverse_b5 * (
                    (
                          b6 * cosβ₋ * H2[i+i2+2]
                        - b7 * cosβ₊ * H2[i+i2]
                    )
                    - a8 * sinβ * H2[i+i2+1]
                )
            end
        end
    end
end

@inbounds function _step_4!(w::WignerMatrixCalculator)
    """Recursively compute H^{m'+1, m}_{n}(β) for m'=1,...,n−1, m=m',...,n using relation (50) resolved
    with respect to H^{m'+1, m}_{n}:

      d^{m'}_{n} H^{m'+1, m}_{n} =   d^{m'−1}_{n} H^{m'−1, m}_{n}
                                   − d^{m−1}_{n} H^{m', m−1}_{n}
                                   + d^{m}_{n} H^{m', m+1}_{n}

    (where the last term drops out for m=n).

    """
    dvalues, n_max, mp_max = d(w), ℓₘₐₓ(w), m′ₘₐₓ(w)
    Hwedge, Hv = w.Hwedge, w.Hv

    if n_max > 0 && mp_max > 0
        for n in 2:n_max
            for mp in 1:min(n, mp_max)-1
                # m = m', ..., n-1
                # i1 = WignerHindex(n, mp+1, mp, mp_max)
                i1 = WignerHindex(n, mp+1, mp+1, mp_max) - 1
                i2 = WignerHindex(n, mp-1, mp, mp_max)
                # i3 = WignerHindex(n, mp, mp-1, mp_max)
                i3 = WignerHindex(n, mp, mp, mp_max) - 1
                i4 = WignerHindex(n, mp, mp+1, mp_max)
                i5 = nm_index(n, mp)
                i6 = nm_index(n, mp-1)
                inverse_d5 = inv(dvalues[i5])
                d6 = dvalues[i6]
                let i=0
                    d7 = dvalues[i+i6]
                    d8 = dvalues[i+i5]
                    Hv[i+nm_index(n, mp+1)] = inverse_d5 * (
                          d6 * Hwedge[i+i2]
                        - d7 * Hv[i+nm_index(n, mp)]
                        + d8 * Hwedge[i+i4]
                    )
                end
                for i in 1:n-mp-1
                    d7 = dvalues[i+i6]
                    d8 = dvalues[i+i5]
                    Hwedge[i+i1] = inverse_d5 * (
                          d6 * Hwedge[i+i2]
                        - d7 * Hwedge[i+i3]
                        + d8 * Hwedge[i+i4]
                    )
                end
                # m = n
                let i=n-mp
                    Hwedge[i+i1] = inverse_d5 * (
                          d6 * Hwedge[i+i2]
                        - dvalues[i+i6] * Hwedge[i+i3]
                    )
                end
            end
        end
    end
end

@inbounds function _step_5!(w::WignerMatrixCalculator)
    """Recursively compute H^{m'−1, m}_{n}(β) for m'=−1,...,−n+1, m=−m',...,n using relation (50)
    resolved with respect to H^{m'−1, m}_{n}:

      d^{m'−1}_{n} H^{m'−1, m}_{n} = d^{m'}_{n} H^{m'+1, m}_{n}
                                     + d^{m−1}_{n} H^{m', m−1}_{n}
                                     − d^{m}_{n} H^{m', m+1}_{n}

    (where the last term drops out for m=n).

    NOTE: Although arxiv:1403.7698 specifies the loop over mp to start at -1, I
    find it necessary to start at 0, or there will be missing information.  This
    also requires setting the (m',m)=(0,-1) components before beginning this loop.

    """
    dvalues, n_max, mp_max = d(w), ℓₘₐₓ(w), m′ₘₐₓ(w)
    Hwedge, Hv = w.Hwedge, w.Hv

    if n_max > 0 && mp_max > 0
        for n in 0:n_max
            for mp in 0:-1:1-min(n, mp_max)
                # m = -m', ..., n-1
                # i1 = WignerHindex(n, mp-1, -mp, mp_max)
                i1 = WignerHindex(n, mp-1, -mp+1, mp_max) - 1
                # i2 = WignerHindex(n, mp+1, -mp, mp_max)
                i2 = WignerHindex(n, mp+1, -mp+1, mp_max) - 1
                # i3 = WignerHindex(n, mp, -mp-1, mp_max)
                i3 = WignerHindex(n, mp, -mp, mp_max) - 1
                i4 = WignerHindex(n, mp, -mp+1, mp_max)
                i5 = nm_index(n, mp-1)
                i6 = nm_index(n, mp)
                i7 = nm_index(n, -mp-1)
                i8 = nm_index(n, -mp)
                inverse_d5 = inv(dvalues[i5])
                d6 = dvalues[i6]
                let i=0
                    d7 = dvalues[i+i7]
                    d8 = dvalues[i+i8]
                    if mp == 0
                        Hv[i+nm_index(n, mp-1)] = inverse_d5 * (
                              d6 * Hv[i+nm_index(n, mp+1)]
                            + d7 * Hv[i+nm_index(n, mp)]
                            - d8 * Hwedge[i+i4]
                        )
                    else
                        Hv[i+nm_index(n, mp-1)] = inverse_d5 * (
                              d6 * Hwedge[i+i2]
                            + d7 * Hv[i+nm_index(n, mp)]
                            - d8 * Hwedge[i+i4]
                        )
                    end
                end
                for i in 1:n+mp-1
                    d7 = dvalues[i+i7]
                    d8 = dvalues[i+i8]
                    Hwedge[i+i1] = inverse_d5 * (
                          d6 * Hwedge[i+i2]
                        + d7 * Hwedge[i+i3]
                        - d8 * Hwedge[i+i4]
                    )
                end
                # m = n
                let i=n+mp
                    Hwedge[i+i1] = inverse_d5 * (
                          d6 * Hwedge[i+i2]
                        + dvalues[i+i7] * Hwedge[i+i3]
                    )
                end
            end
        end
    end
end

"""
    H!(w, expiβ)

Compute (a quarter of) the H matrix

WARNING: The returned array will be a view into the `workspace` variable (see
below for an explanation of that).  If you need to call this function again
using the same workspace before extracting all information from the first call,
you should use `numpy.copy` to make a separate copy of the result.

Parameters
----------
expiβ : array_like
    Value of exp(i*β) on which to evaluate the H matrix.

Returns
-------
Hwedge : array
    This is a 1-dimensional array of floats; see below.
workspace : array_like, optional
    A working array like the one returned by Wigner.new_workspace().  If not
    present, this object's default workspace will be used.  Note that it is not
    safe to use the same workspace on multiple threads.  Also see the WARNING
    above.

See Also
--------
d : Compute the full Wigner d matrix
D : Compute the full Wigner 𝔇 matrix
rotate : Avoid computing the full 𝔇 matrix and rotate modes directly
evaluate : Avoid computing the full 𝔇 matrix and evaluate modes directly

Notes
-----
H is related to Wigner's (small) d via

    dₗⁿᵐ = ϵₙ ϵ₋ₘ Hₗⁿᵐ,

where

         ⎧ 1 for k≤0
    ϵₖ = ⎨
         ⎩ (-1)ᵏ for k>0

H has various advantages over d, including the fact that it can be efficiently
and robustly valculated via recurrence relations, and the following symmetry
relations:

    H^{m', m}_n(β) = H^{m, m'}_n(β)
    H^{m', m}_n(β) = H^{-m', -m}_n(β)
    H^{m', m}_n(β) = (-1)^{n+m+m'} H^{-m', m}_n(π - β)
    H^{m', m}_n(β) = (-1)^{m+m'} H^{m', m}_n(-β)

Because of these symmetries, we only need to evaluate at most 1/4 of all the
elements.

"""
function H!(w::WignerMatrixCalculator, expiβ::Complex)
    _step_1!(w)
    _step_2!(w, expiβ)
    _step_3!(w, expiβ)
    _step_4!(w)
    _step_5!(w)
    w.Hwedge
end