text:
## index_set.jl : representation of an index set
#
# Representation of Multilevel and Multi-Index sets. All indices in the index set for a
# given size parameter `sz` can be computed using `get_index_set()`. This is implemented
# using the `filter`-function from `Base.Iterators`.
#
# This file is part of MultilevelEstimators.jl - A Julia toolbox for
# Multilevel Monte Carlo Methods (c) Pieterjan Robbe, 2019

## Indexset ##
abstract type AbstractIndexSet{d} end

## SL ##
"""
    SL()

Return a single-level index set.

# Examples
```jldoctest; setup = :(using MultilevelEstimators)
julia> SL()
SL
```
See also: [`ML`](@ref), [`FT`](@ref), [`TD`](@ref), [`HC`](@ref), [`ZC`](@ref), [`AD`](@ref), [`U`](@ref)
"""
struct SL{d} <: AbstractIndexSet{d} end

SL() = SL{1}()

filter(::SL, sz) = itr -> sum(itr) == sz

## ML ##
"""
    ML()

Return a multi-level index set.

# Examples
```jldoctest; setup = :(using MultilevelEstimators)
julia> ML()
ML
```
See also: [`SL`](@ref), [`FT`](@ref), [`TD`](@ref), [`HC`](@ref), [`ZC`](@ref), [`AD`](@ref), [`U`](@ref)
"""
struct ML{d} <: AbstractIndexSet{d} end

ML() = ML{1}()

filter(::ML, sz) = itr -> sum(itr) ≤ sz

## FT ##
"""
    FT(d::Integer)
    FT(δ::Real...)
	FT(δ::AbstractVector{<:Real})
    FT(δ::NTuple{N, <:Real})

Return a full tensor index set in `d` dimenions with optional weights `δ`. Default weights are all 1's.

# Examples
```jldoctest; setup = :(using MultilevelEstimators)
julia> FT(2)
FT{2}

julia> FT([1, 1/2, 1/2])
FT{3}

julia> print(FT(2), 4)
  ◼ ◼ ◼ ◼ ◼
  ◼ ◼ ◼ ◼ ◼
  ◼ ◼ ◼ ◼ ◼
  ◼ ◼ ◼ ◼ ◼
  ◼ ◼ ◼ ◼ ◼
```
See also: [`SL`](@ref), [`ML`](@ref), [`TD`](@ref), [`HC`](@ref), [`ZC`](@ref), [`AD`](@ref), [`U`](@ref)
"""
struct FT{d, T} <: AbstractIndexSet{d}
    δ::T
end

filter(idxset::FT, sz) = itr -> all(itr.I ./ idxset.δ .≤ sz)

## TD ##
"""
    TD(d::Integer)
    TD(δ::Real...)
	TD(δ::AbstractVector{<:Real})
    TD(δ::NTuple{N, <:Real})

Return a total degree index set in `d` dimenions with optional weights `δ`. Default weights are all 1's.

# Examples
```jldoctest; setup = :(using MultilevelEstimators)
julia> TD(2)
TD{2}

julia> TD([1, 1/2, 1/2])
TD{3}

julia> print(TD(2), 4)
  ◼
  ◼ ◼
  ◼ ◼ ◼
  ◼ ◼ ◼ ◼
  ◼ ◼ ◼ ◼ ◼
```

See also: [`SL`](@ref), [`ML`](@ref), [`FT`](@ref), [`HC`](@ref), [`ZC`](@ref), [`AD`](@ref), [`U`](@ref)
"""
struct TD{d, T} <: AbstractIndexSet{d}
    δ::T
end

filter(idxset::TD, sz) = itr -> sum(itr.I ./ idxset.δ) ≤ sz

## HC ##
"""
    HC(d::Integer)
    HC(δ::Real...)
	HC(δ::AbstractVector{<:Real})
    HC(δ::NTuple{N, <:Real})

Return a hyperbolic cross index set in `d` dimenions with optional weights `δ`. Default weights are all 1's.

# Examples
```jldoctest; setup = :(using MultilevelEstimators)
julia> HC(2)
HC{2}

julia> HC([1, 1/2, 1/2])
HC{3}

julia> print(HC(2), 4)
  ◼
  ◼
  ◼
  ◼ ◼
  ◼ ◼ ◼ ◼ ◼
```

See also: [`SL`](@ref), [`ML`](@ref), [`FT`](@ref), [`TD`](@ref), [`ZC`](@ref), [`AD`](@ref), [`U`](@ref)
"""
struct HC{d, T} <: AbstractIndexSet{d}
    δ::T
end

filter(idxset::HC, sz) = itr -> prod(itr.I ./ idxset.δ .+ 1) - 1 ≤ sz

## ZC ##
"""
    ZC(d::Integer)
    ZC(δ::Real...)
	ZC(δ::AbstractVector{<:Real})
	ZC(δ::NTuple{N, <:Real})

Return a Zaremba cross index set in `d` dimenions with optional weights `δ`. Default weights are all 1's.

# Examples
```jldoctest; setup = :(using MultilevelEstimators)
julia> ZC(2)
ZC{2}

julia> ZC([1, 1/2, 1/2])
ZC{3}

julia> print(ZC(2), 4)
  ◼ ◼
  ◼ ◼
  ◼ ◼ ◼
  ◼ ◼ ◼ ◼ ◼
  ◼ ◼ ◼ ◼ ◼
```

See also: [`SL`](@ref), [`ML`](@ref), [`FT`](@ref), [`TD`](@ref), [`HC`](@ref), [`AD`](@ref), [`U`](@ref)
"""
struct ZC{d, T} <: AbstractIndexSet{d}
    δ::T
end

filter(idxset::ZC, sz) = itr -> sz == 0 ? all(itr.I .== 0) : prod(max.(1, itr.I ./ idxset.δ)) ≤ sz

## constructors ##
for T in ["FT" "TD" "HC" "ZC"]
    @eval begin
        $(Symbol(T))(d::Integer) = $(Symbol(T))(fill(1, d))

        $(Symbol(T))(δ::Real...) = $(Symbol(T))(promote(δ...))

        $(Symbol(T))(δ::AbstractVector{<:Real}) = $(Symbol(T))(Tuple(δ))

        $(Symbol(T))(δ::NTuple{d, <:Real} where d) = check_args(δ, $(T)) && $(Symbol(T)){length(δ), typeof(δ)}(δ ./ maximum(δ))
    end
end

## AD ##
"""
    AD(d::Integer)

Return an adaptive index set in `d` dimenions.

# Examples
```jldoctest; setup = :(using MultilevelEstimators)
julia> AD(2)
AD{2}
```
See also: [`SL`](@ref), [`ML`](@ref), [`FT`](@ref), [`TD`](@ref), [`HC`](@ref), [`ZC`](@ref), [`U`](@ref)
"""
struct AD{d} <: AbstractIndexSet{d} end

AD(d::Integer) = check_larger_than(AD, d, "dimension", 1) && AD{d}()

## U ##
"""
    U(d::Integer)

Return an unbiased  multi-level or multi-index index set in `d` dimensions.

# Examples
```jldoctest; setup = :(using MultilevelEstimators)
julia> U(2)
U{2}
```
See also: [`SL`](@ref), [`ML`](@ref), [`FT`](@ref), [`TD`](@ref), [`HC`](@ref), [`ZC`](@ref), [`AD`](@ref)
"""
struct U{d} <: AbstractIndexSet{d} end

U(d::Integer) = check_larger_than(U, d, "dimension", 0) && U{d}()

## utilities ##
AbstractMI = Union{TD, FT, HC, ZC, AD, U}

AbstractML = Union{ML, U{1}}

ndims(::AbstractIndexSet{d}) where {d} = d

length(idxset::AbstractIndexSet, sz::Integer) = length(collect(get_index_set(idxset, sz)))

eltype(::Type{<:AbstractIndexSet{d}}) where d = Index{d}

"""
    get_index_set(idxset::AbstractIndexSet, sz::Integer)

Return an iterator over all indices in the index set `idxset` for a given size parameter `sz`.

# Examples
```jldoctest; setup = :(using MultilevelEstimators)
julia> collect(get_index_set(TD(2), 2))
6-element Array{CartesianIndex{2},1}:
 (0, 0)
 (1, 0)
 (2, 0)
 (0, 1)
 (1, 1)
 (0, 2)
```
"""
get_index_set(idxset::AbstractIndexSet{d}, sz::Integer) where d = Base.Iterators.filter(filter(idxset, sz), CartesianIndices(ntuple(i -> 0:sz + 1, d)))

get_index_set(idxset::AD, sz::Integer) = get_index_set_not_implemented("AD")

get_index_set(idxset::U, sz::Integer) = get_index_set_not_implemented("U")

get_index_set_not_implemented(name) = throw(ArgumentError(string("get_index_set not implemented for index sets of type ", name)))

function is_admissable(idxset::Set{Index{d}}, index::Index{d})  where d
    if index ∈ idxset
        return false
    else
        for i in 1:d
            bw_neighbour = index - CartesianIndex(ntuple(j -> j == i, d)) # backward neighbour
            if !( bw_neighbour ∈ idxset || bw_neighbour[i] < 0 )
                return false
            end
        end
        return true
    end
end

## output formatting ##
"""
    print(idxset::AbstractIndexSet{2}, sz::Integer)

Print the indices in the index set `idxset` for a given size parameter `sz`. Only implemented for two-dimensional index sets.

# Examples
```jldoctest; setup = :(using MultilevelEstimators)
julia> print(TD(2), 3)
  ◼
  ◼ ◼
  ◼ ◼ ◼
  ◼ ◼ ◼ ◼
```
"""
print(idxset::AbstractIndexSet, sz::Integer) = print(collect(get_index_set(idxset, sz)))

function print(idxset::Vector{<:Index{2}}, idxset2::Vector{Index{2}}=Index{2}[])
	if !(isempty(idxset) && isempty(idxset2))
		char = "\u25FC"
		char2 = "\u25A3"
		n = maximum(union(idxset, idxset2)).I
		R = CartesianIndices(ntuple(i -> 0:n[i], 2))
		A = map(i -> i ∈ idxset, R)
		B = map(i -> i ∈ idxset2, R)
		str = ""
		for j in n[2]+1:-1:1
			str = string(str, "  ")
			for i in 1:n[1]+1
				str = A[i,j] ? string(str, char, " ") : str
				str = B[i,j] ? string(str, char2, " ") : str
			end
			str = string(str,"\n")
		end
		print(str)
	end
end

print(::Vector{<:Index}) = throw(ArgumentError("print only available for index sets with d = 2"))

show(io::IO, idxset::AbstractIndexSet) = print(io, shortname(idxset))

for T in ["SL" "ML"]
    @eval shortname(::$(Symbol(T))) = $(T)
end

for T in ["TD" "FT" "HC" "ZC" "AD" "U"]
    @eval shortname(idxset::$(Symbol(T))) = string($(T), "{", ndims(idxset), "}")
end

## input checking ##
function check_args(δ::NTuple{d, <:Real}, name) where d
    check_larger_than(name, d, "dimension", 1)
    check_larger_than(string("weighted ", name), δ, "weights δ", 0)
    check_smaller_than_or_equal_to(string("weighted ", name), δ, "weights δ", 1)
end

ext: jl


text:
export extract_con_headers

"""
extract_con_headers(con::SeisCon, keys::Array{String,1}, blocks::Array{Int,1};
                                prealloc_traces::Int = 50000)

Read `keys` from `blocks` in `con`. Skips all data and only reads and returns headers in an array.
This method is similar to `read_con_headers`, except the data is returned

# Examples

h = extract_con_headers(s, ["SourceX"; "SourceY"], Array(1:11))
"""
function extract_con_headers(con::SeisCon, keys::Array{String,1}, blocks::Vector;
                                                    f_map::Function = pmap)
    nblocks = length(blocks)
    h = f_map(x -> extractor(con, keys, x), blocks)
    return vcat(h...)
end

function extract_con_headers(con::SeisCon, keys::Array{String,1}, blocks::Vector, reducer::Function;
                                                    f_map::Function = pmap)
    nblocks = length(blocks)
    h = f_map(x -> reducer(extractor(con, keys, x)), blocks)
    return vcat(h...)
end

function extractor(con, keys, block)

    println("Extracting ... $block")
    headers = read_con_headers(con, copy(keys), block, prealloc_traces = 10000000)
    ntraces = length(headers)
    nkeys = length(keys)
    header_array = Array{Int32,2}(undef, ntraces, nkeys)

    for cur_key in 1:nkeys
        header_array[:,cur_key] = get_header(headers, Symbol(keys[cur_key]), scale=false)
    end

    return header_array
end


text:
"""Julia program to print the n'th term in the Golomb sequence
Golomb sequence is a non-decreasing integer sequence where n'th
term is  equal to the number of times n appears in the sequence"""

function golomb_sequence(n)
    dp = zeros(Int, n + 1)
    dp[1] = 1

    for i in 2:n
        dp[i] = 1 + dp[i - dp[dp[i - 1]]]
    end
    return dp[n]
end

print("Enter the value of n?, where you need the n'th number in the golomb sequence. ")
n = readline()
n = parse(Int, n)
if (n <= 0)
    println("The given value of n is invalid.")
    exit()
end
res = golomb_sequence(n)
println("The $n'th term in the golomb sequence is $res.")

"""
Time Complexity: O(n), where 'n' is the given number
Space Complexity: O(n)

SAMPLE INPUT AND OUTPUT

SAMPLE 1
Enter the value of n?, where you need the n'th number in the golomb sequence. 5
The 5'th term in the golomb sequence is 3.

SAMPLE 2
Enter the value of n?, where you need the n'th number in the golomb sequence. 867
The 867'th term in the golomb sequence is 79.

"""


text:
# Note that this script can accept some limited command-line arguments, run
# `julia build_tarballs.jl --help` to see a usage message.
using BinaryBuilder

name = "Qhull"
version = v"2019.1"

# Collection of sources required to build
sources = [
    ArchiveSource(
        "https://github.com/qhull/qhull/archive/$(version.major).$(version.minor).tar.gz", # URL
        "cf7235b76244595a86b9407b906e3259502b744528318f2178155e5899d6cf9f"                 # sha256 hash
    ),
]

# Bash recipe for building across all platforms
script = raw"""
# initial setup
cd $WORKSPACE/srcdir/qhull*/
export CMAKE_ARGS="-DCMAKE_INSTALL_PREFIX=${prefix} -DCMAKE_TOOLCHAIN_FILE=${CMAKE_TARGET_TOOLCHAIN} -DCMAKE_BUILD_TYPE=Release"

# begin the build process
cd build

# Generate makefiles
cmake $CMAKE_ARGS -G "Unix Makefiles" .. && cmake $CMAKE_ARGS ..

# Run the build script
make -j${nproc}

# Install Qhull
make install
"""

# These are the platforms we will build for by default, unless further
# platforms are passed in on the command line
platforms = supported_platforms()

# The products that we will ensure are always built
products = [
    # Executables
    # compute convex hulls and related structures
    ExecutableProduct("qhull", :qhull),
    # generate various point distributions
    ExecutableProduct("rbox", :rbox),
    # compute the convex hull
    ExecutableProduct("qconvex", :qconvex),
    # compute the Delaunay triangulation
    ExecutableProduct("qdelaunay", :qdelaunay),
    # compute the Voronoi diagram
    ExecutableProduct("qvoronoi", :qvoronoi),
    # compute the halfspace intersection about a point
    ExecutableProduct("qhalf", :qhalf),

    # Libraries
    # reentrant Qhull
    LibraryProduct(["libqhull_r", "qhull_r"], :libqhull_r),

    # Files
    # Reentrant Qhull header files
    FileProduct("include/libqhull_r/geom_r.h", :geom_r_h),
    FileProduct("include/libqhull_r/io_r.h", :io_r_h),
    FileProduct("include/libqhull_r/libqhull_r.h", :libqhull_r_h),
    FileProduct("include/libqhull_r/mem_r.h", :mem_r_h),
    FileProduct("include/libqhull_r/merge_r.h", :merge_r_h),
    FileProduct("include/libqhull_r/poly_r.h", :poly_r_h),
    FileProduct("include/libqhull_r/qhull_ra.h", :qhull_ra_h),
    FileProduct("include/libqhull_r/qset_r.h", :qset_r_h),
    FileProduct("include/libqhull_r/random_r.h", :random_r_h),
    FileProduct("include/libqhull_r/stat_r.h", :stat_r_h),
    FileProduct("include/libqhull_r/user_r.h", :user_r_h),
]

# Dependencies that must be installed before this package can be built
dependencies = Dependency[]

# Build the tarballs, and possibly a `build.jl` as well.
build_tarballs(ARGS, name, version, sources, script, platforms, products, dependencies)


text:
"""
Acts as a store for build information. Passed around to the important setps
"""
type MarbleDoc
    docname::AbstractString # For loggin and tex output
    settings::SettingsBundle
    raw # The file. The stream. The whatever
    tree::Markdown.MD # Parsed markdown tree
    content::AbstractString # Rendered markdown tree w/o main template
    final::AbstractString # Final document, ready for handoff to another processor

    cache::States # Storage of cached files
    templates::LazyTemplateLoader # Templateing environment
    scratch # Place for individual elements to store and share data while rendering
    # log

    function MarbleDoc(docname, raw, cache, settings)
        templates = LazyTemplateLoader([
                settings["paths"]["template"],
                mrbldir("templates"),
                "$(Pkg.dir("Marble"))/templates"
            ];
            block_start_string=settings["JINJA_block_start_string"],
            block_end_string=settings["JINJA_block_end_string"],
            variable_start_string=settings["JINJA_variable_start_string"],
            variable_end_string=settings["JINJA_variable_end_string"],
            comment_start_string=settings["JINJA_comment_start_string"],
            comment_end_string=settings["JINJA_comment_end_string"],
            keep_trailing_newline=true,
        )

        return new(
            docname,
            settings,
            raw,
            Markdown.MD(),
            "",
            "",
            cache,
            templates,
            Dict())
    end
end

"""
Simplification of MarbleDoc constructor
"""
function Marble.MarbleDoc(contents, path, docname; cache=nothing)
    # Configure the build env
    settings = get_settings(path)
    create_paths(settings)
    cache = cache == nothing ? States("$(settings["paths"]["cache"])/hashes.json") : cache

    return Marble.MarbleDoc(
        docname,
        contents,
        cache,
        settings)
end

############## The Main Build Process ##############

"""
Takes in a MarbleDoc and parses the Markdown in the primary_file setting
"""
function Base.parse(env::MarbleDoc)
    if env.settings["debug"]
        println("PARSING") # LOGGING
    end

    env.scratch[:analysis] = get_analysis(env)
    env.tree = Markdown.parse(env.raw, flavor=:mrbl)
end

"""
Takes and manipulates the freshly parsed Tree
"""
function process(env::MarbleDoc)
    if env.settings["debug"]
        println("PROCESSING") # LOGGING
    end

    # Make any Document tags part of the ENV
    for e in env.tree.content
        if isa(e, Document)
            add!(env.settings, e.data)
        end
    end

    #=
    I don't like this implementation. I would prefere something more like

    # References

    {references}

    where the header does not have any magic meaning
    =#
    ref_ind = 0
    addrefs = false
    walk(env) do e, c, env
        try
            if isa(e, Markdown.Header) &&
                strip(e.text[1]) == env.settings["references_header"]
                ref_ind = c[1]
                addrefs = true
            end
        catch
        end
    end
    # replace header with actual element
    if ref_ind != 0
        # deleteat!(env.tree.content, ref_ind)
        insert!(env.tree.content, ref_ind + 1, Tex("\\printbibliography[heading=none]\n"))
    end

    # probably side load the processors, e.g. for processor in processors...
    # skip
end

"""
Takes a marble and converst the abstract tree object, and renders
each object to it's corrisponding Latex (could be HTML. Who knows)
"""
function render(env::MarbleDoc)
    if env.settings["debug"]
        println("RENDERING") # LOGGING
    end

    env.content = jinja(env, env.tree)
end

"""
Takes the finished document and puts it into template
"""
function template(env::MarbleDoc)
    if env.settings["debug"]
        println("TEMLATING") # LOGGING
    end

    env.final = JinjaTemplates.render(env.templates, "documents/$(env.settings["template"]).tex";
        settings=flatten(env.settings),
        content=env.content
    )

    # Write to file
    texfile = "$(env.settings["paths"]["base"])/$(get_basename(env)).tex"
    open(texfile, "w") do f
        write(f, env.final)
    end
    # show(env.final)
end

"""
Takes the finished document, and calls a build script (probably lualatex)
"""
function build(env::MarbleDoc)
    if env.settings["debug"]
        println("BUILDING") # LOGGING
    end

    print("Building $(env.docname)... ")

    basename = get_basename(env)
    texfile = "$(env.settings["paths"]["base"])/$(get_basename(env)).tex"
    builddir = relpath(env.settings["paths"]["build"])

    logf = open(joinpath(env.settings["paths"]["log"], "$(get_basename(env))_build.log"), "w")
    write(logf, "\n===== Build at $(now()) =====\n")

    if env.settings["topdf"]
        try
            runindir(builddir) do
                run(pipeline(`latexmk -$(env.settings["texcmd"]) -shell-escape -halt-on-error $(env.settings["paths"]["base"])/$(basename).tex`; stdout=logf, stderr=logf))
            end
            println("DONE")
        catch y
            if !isa(y, ErrorException)
                rethrow(y)
            end

            print_with_color(:red, "BUILD FAILED: ")
            println("See $(abspath(env.settings["paths"]["log"]))/$(basename)_build.log for details.")
            return nothing
        finally
            close(logf)
        end
    end
end

############## The Supporting Functions ##############
"""
Run the analysis CMD specified in the `analysis` command
"""
function get_analysis(env)
    # Execute the analysis script.
    if "analysis" in keys(env.settings)
        a_file = env.settings["analysis"]

        # Use cached settings if analysis.jl is up to date
        if changed(env.cache, a_file)
            # Get the executable command
            exec =  map(split(env.settings["analysiscmd"], ' ')) do command
                if command == "\$filename"
                    return a_file
                else
                    return command
                end
            end

            try
                out = JSON.parse(readstring(`$exec`))
                open(f->JSON.print(f, out),
                    "$(env.settings["paths"]["cache"])/analysis.json",
                    "w")
            catch y
                out = Dict()
                warn("Analysis script `$(env.settings["analysis"])` failed to run.")
                println(y)
            end
            pin!(env.cache, a_file)
        else
            println("Using cached analysis file")
            out = JSON.parsefile("$(env.settings["paths"]["cache"])/analysis.json")
        end
    else
        out = Dict()
    end

    return out
end

"""
Needs to add support for inlines and nested block elements
NEEDS TO MAKE A COPY OF THE TREE, OR ELSE WALKING GETS WIERD

f should accept e, c, env::MarbleEnv
e → The element itself
c → The coordinates for this element (array of indicies)
env → The tree itself
"""
function walk(f::Function, env::MarbleDoc)
    for i in 1:length(env.tree.content)
        n = length(env.tree.content) + 1 - i
        f(env.tree.content[n], (n,), env)
    end
end


text:
using Distributions
using StatsBase: var

########## GDEMO ##################
# Example partially based on https://stats.stackexchange.com/a/266672/234110 and
# http://www2.stat.duke.edu/~rcs46/modern_bayes17/lecturesModernBayes17/lecture-7/07-gibbs.pdf.
# We follow the model used in https://github.com/TuringLang/Turing.jl/blob/60724e22a9831066fc2e0e82d428bd4922bda6e8/test/test_utils/models.jl#L2,
# with normally distributed observations and conjugate priors for both parameters:
# λ ~ Gamma(α₀, θ₀)
# m ~ Normal(0, √(1 / λ))
# Xᵢ ~ Normal(m, √(1 / λ)) (iid)

function gdemo_forward(N, α₀, θ₀)
    λ = rand(Gamma(α₀, θ₀))
    σ = √(1 / λ)
    m = rand(Normal(0, σ))
    x = rand(Normal(m, σ), N)
    return (m=m, λ=λ), x
end

function gdemo_logjoint(α₀, θ₀, x, m, λ)
    σ = √(1 / λ)
    return logpdf(Gamma(α₀, θ₀), λ) +
        logpdf(Normal(0, σ), m) +
        sum(logpdf.(Normal(m, σ), x))
end

function gdemo_statistics(x)
    # The conditionals and posterior can be formulated in terms of the following statistics:
    N = length(x) # number of samples
    x̄ = mean(x) # sample mean
    s² = var(x; mean=x̄, corrected=false) # sample variance
    return N, x̄, s²
end

function gdemo_cond_m(α₀, θ₀, x, λ)
    N, x̄, s² = gdemo_statistics(x)
    mₙ = N * x̄ / (N + 1)
    λₙ = λ * (N + 1)
    σₙ = √(1 / λₙ)
    return Normal(mₙ, σₙ)
    # return Normal(7/6, 0.00001)
end

function gdemo_cond_λ(α₀, θ₀, x, m)
    N, x̄, s² = gdemo_statistics(x)
    αₙ = α₀ + (N - 1) / 2 + 1
    βₙ = (s² * N / 2 + m^2 / 2 + inv(θ₀))
    return Gamma(αₙ, inv(βₙ))
end

function gdemo(x; α₀=2.0, θ₀=inv(3.0))
    cond_m((λ,)) = gdemo_cond_m(α₀, θ₀, x, λ)
    cond_λ((m,)) = gdemo_cond_λ(α₀, θ₀, x, m)
    conditionals = (m=cond_m, λ=cond_λ)
    return conditionals
end

############### Mixture of two Gaussians ####################
# see http://www.cs.columbia.edu/~blei/fogm/2015F/notes/mixtures-and-gibbs.pdf
# μ ~ arraydist(Normal.(fill(m, K), fill(λ, K)))
# z ~ arraydist(Categorical.(fill(π, N)))
# x ~ arraydist(Normal.(μ[z], σ))

function mixture_cond_z(π, K, m, λ, σ, x, μ)
    function mixtureweight(x)
        p = π .* pdf.(Normal.(μ, σ), Ref(x))
        return p ./ sum(p)
    end
    return Product(Categorical.(mixtureweight.(x)))
end

function mixture_cond_μ(π, K, m, λ, σ, x, z)
    n = [count(z .== k) for k = 1:K]

    # If there were no observations assigned to center `k`, `n[k] == 0`, and
    # we use the prior instead.
    x_bar = [(n[k] != 0) ? (sum(x[z .== k]) / n[k]) : m for k = 1:K]
    λ_hat = [(n[k] != 0) ? inv(n[k] / σ^2 + 1/λ^2) : λ for k = 1:K]
    μ_hat = [(n[k] != 0) ? x_bar[k] * (n[k] / σ^2) * λ_hat[k] : m for k = 1:K]

    return Product(Normal.(μ_hat, λ_hat))
end

function mixture(π, K, m, λ, σ, x)
    cond_z((μ,)) = mixture_cond_z(π, K, m, λ, σ, x, μ)
    cond_μ((z,)) = mixture_cond_μ(π, K, m, λ, σ, x, z)
    return (z=cond_z, μ=cond_μ)
end

########### 1D ISING MODEL #########################
function ising_conditional(β, i, M)
    return function (neighbours)
        x_prev = neighbours[mod1(i - 1, M - 1)]
        x_next = neighbours[mod1(i, M - 1)]
        p₁ = exp(-β * (x_prev + x_next))
        p₂ = exp(-β * (abs(x_prev - 1) + abs(x_next - 1)))
        return Bernoulli(p₁ / (p₁ + p₂))
    end
end

function isingdemo(β, M)
    # One dimensional Ising model on a circle, without external field;
    # see https://stats.stackexchange.com/a/312044/234110
    conditional_dists = [Symbol(:x, i) => ising_conditional(β, i, M) for i = 1:M]
    conditionals = merge(NamedTuple(), conditional_dists)
end

function ising_conditional(β, i, j, M, N)
    return function (neighbours)
        x_prev = neighbours[mod1(i - 1, M - 1)]
        x_next = neighbours[mod1(i, M - 1)]
        x_down = neighbours[mod1(j - 1, N - 1)]
        x_up = neighbours[mod1(j, N - 1)]
        p₁ = exp(-β * (x_prev + x_next + x_up + x_down))
        p₂ = exp(-β * (abs(x_prev - 1) + abs(x_next - 1) + abs(x_up - 1) + abs(x_down - 1)))
        return Bernoulli(p₁ / (p₁ + p₂))
    end
end

function isingdemo(β, M, N)
    # One dimensional Ising model on a circle, without external field;
    # see https://stats.stackexchange.com/a/312044/234110
    conditional_dists = [Symbol(:x, i, j) => ising_conditional(β, i, j, M, N) for i = 1:M, j = 1:N]
    conditionals = merge(NamedTuple(), conditional_dists)
end

function visualize_ising(θ::AbstractVector{<:Real})
    join(v == 1 ? '\u2588' : ' ' for v in values(θ))
end

function visualize_ising(θ::AbstractMatrix{<:Real})
    join((join(v == 1 ? '\u2588' : ' ' for v in row) for row in eachrow(θ)), '\n')
end


text:
abstract type AbstractImage end
abstract type AbstractAnalogueImage <: AbstractImage end
abstract type AbstractDigitalImage <: AbstractImage end

struct AnalogueImage <: AbstractAnalogueImage end

# function get_data(image::AnalogueImage)
#     image.data
# end


text:
function open_and_read(path)
    bytes = read(path)
    if bytes[1:3] == [0xef, 0xbb, 0xbf]
        # Why skip the first three bytes of the response?

        # From https://docs.python.org/3/library/codecs.html
        # To increase the reliability with which a UTF-8 encoding can be detected,
        # Microsoft invented a variant of UTF-8 (that Python 2.5 calls "utf-8-sig")
        # for its Notepad program: Before any of the Unicode characters is written
        # to the file, a UTF-8 encoded BOM (which looks like this as a byte sequence:
        # 0xef, 0xbb, 0xbf) is written.
        return String(bytes[4:end])
    else
        return String(bytes)
    end
end

function getXML(open_file)

    return xml = XMLDict.xml_dict(open_file)

end

# get potentially missing value out of dict
function get_and_parse(dict, key)
    try
        return val = Parsers.parse(Float64, dict[key])
    catch y
        if isa(y, KeyError)
            return val = missing
        else
            throw(y)
        end
    end
end

struct XTbMLTable{S,U}
    select::S
    ultimate::U
    d::TableMetaData
end

function parseXTbMLTable(x, path)
    md = x["XTbML"]["ContentClassification"]
    name = get(md, "TableName", nothing) |> strip
    content_type = get(get(md, "ContentType", nothing), "", nothing) |> strip
    id = get(md, "TableIdentity", nothing) |> strip
    provider = get(md, "ProviderName", nothing) |> strip
    reference = get(md, "TableReference", nothing) |> strip
    description = get(md, "TableDescription", nothing) |> strip
    comments = get(md, "Comments", nothing) |> strip
    source_path = path
    d = TableMetaData(
        name=name,
        id=id,
        provider=provider,
        reference=reference,
        content_type=content_type,
        description=description,
        comments=comments,
        source_path=source_path,
    )

    if isa(x["XTbML"]["Table"], Vector)
        # for a select and ultimate table, will have multiple tables
        # parsed into a vector of tables
        sel = map(x["XTbML"]["Table"][1]["Values"]["Axis"]) do ai
            (issue_age = Parsers.parse(Int, ai[:t]),
                rates = [(duration = Parsers.parse(Int, aj[:t]), rate = get_and_parse(aj, "")) for aj in ai["Axis"]["Y"] if !ismissing(get_and_parse(aj, ""))])
        end

        ult = map(x["XTbML"]["Table"][2]["Values"]["Axis"]["Y"]) do ai
            (age  = Parsers.parse(Int, ai[:t]), rate = get_and_parse(ai, ""),)
        end

    else
        # a table without select period will just have one set of values

        ult = map(x["XTbML"]["Table"]["Values"]["Axis"]["Y"]) do ai
            (age = Parsers.parse(Int, ai[:t]),
                rate = get_and_parse(ai, ""))
        end

        sel = nothing

    end

    tbl = XTbMLTable(
        sel,
        ult,
        d
    )

    return tbl
end

function XTbML_Table_To_MortalityTable(tbl::XTbMLTable)
    ult = UltimateMortality(
                [v.rate for v in  tbl.ultimate],
                start_age=tbl.ultimate[1].age
            )

    ult_omega = lastindex(ult)

    if !isnothing(tbl.select)
        sel =   map(tbl.select) do (issue_age, rates)
            last_sel_age = issue_age + rates[end].duration - 1
            first_defined_select_age =  issue_age + rates[1].duration - 1
            last_age = max(last_sel_age, ult_omega)
            vec = map(issue_age:last_age) do attained_age
                if attained_age < first_defined_select_age
                    return missing
                else
                    if attained_age <= last_sel_age
                        return rates[attained_age - first_defined_select_age + 1].rate
                    else
                        return ult[attained_age]
                    end
                end
            end
            return mortality_vector(vec, start_age=issue_age)
        end
        sel = OffsetArray(sel, tbl.select[1].issue_age - 1)

        return MortalityTable(sel, ult, metadata=tbl.d)
    else
        return MortalityTable(ult, metadata=tbl.d)
    end
end

"""
    readXTbML(path)

Loads the [XtbML](https://mort.soa.org/About.aspx) (the SOA XML data format for mortality tables) stored at the given path and returns a `MortalityTable`.
"""
@memoize function readXTbML(path)
    x = open_and_read(path) |> getXML
    XTbML_Table_To_MortalityTable(parseXTbMLTable(x, path))
end

# Load Available Tables ###

"""
    read_tables(dir=nothing)

Loads the [XtbML](https://mort.soa.org/About.aspx) (the SOA XML data format for mortality tables) stored in the given path. If no path is specified, will load the packages in the MortalityTables package directory. To see where your system keeps packages, run `DEPOT_PATH` from a Julia REPL.
"""
function read_tables(dir=nothing)
    if isnothing(dir)
        table_dir = artifact"mort.soa.org"
    else
        table_dir = dir
    end
    tables = []
    @info "Loading built-in Mortality Tables..."
    for (root, dirs, files) in walkdir(table_dir)
        for file in files
            if endswith(file,".xml") && !startswith(file,".")
                tbl =  readXTbML(joinpath(root,file))
                push!(tables,tbl)
            end
        end
    end
    return Dict(tbl.metadata.name => tbl for tbl in tables if ~isnothing(tbl))
end

# this is used to generate the table mapping in table_source_map.jl
function _write_available_tables()
        table_dir = artifact"mort.soa.org"
    tables = []
    @info "Loading built-in Mortality Tables..."
    for (root, dirs, files) in walkdir(table_dir)
        for file in files
            if endswith(file,".xml") && !startswith(file,".")
            x = open_and_read(joinpath(root,file)) |> XMLDict.xml_dict
            md = x["XTbML"]["ContentClassification"]
            name = get(md, "TableName", nothing) |> strip
            content_type = get(get(md, "ContentType", nothing), "", nothing) |> strip
            id = get(md, "TableIdentity", nothing) |> strip
            push!(tables,(source="mort.soa.org",name=name,id=Parsers.parse(Int,id)))
            end
        end
    end
    return sort!(tables,by=last)
end


text:
"""interpolate missing values (nodata) in a grid
methods available:
nearest neighbor
ordinary kriging
inverste distance interpolation

inputs
img          the image to be interpolated
nodata       nodata value in img
mask         boolean array. only interpolate where true
mode         "nn", "kriging", "idw"
k            number of neighbors (default=10; for kriging and idw only)
p            idw power (default=10.0; for idw only)
extrapolate  check if point inside triangle of neighbors (default=false; for kriging and idw only)
"""
function interp_missing!(
    img::Array{T,2},
    nodata::T;
    mask=trues(size(img)),
    mode="nn",
    k=10, # number of neighboring points used in kringing and idw
    p=2.0, # idw power
    extrapolate=true) where T <: AbstractFloat

    # initiate interpolators or spatial index
    xy_p, z_p = xy_z_grid(img, nodata=nodata)

    # skip images with only no-data values
    if length(xy_p) > 0
        if mode in ["nn", "kriging", "idw"]
            tree = KDTree(xy_p)
        else
            error("mode $(mode) not implemented")
        end

        # interpolation
        ismissing = (img .== nodata) .& mask
        dims = size(img)
        for I in 1:length(img)
            ismissing[I] || continue
            x,y = ind2sub(dims, I)
            if mode == "nn"
                img[I] = interpolator_nn(x, y, xy_p, z_p, tree)
            elseif mode == "kriging"
                img[I] = interpolator_kriging(x, y, xy_p, z_p, tree; k=k, extrapolate=extrapolate)
            elseif mode == "idw"
                img[I] = interpolator_idw(x, y, xy_p, z_p, tree; p=p, k=k, extrapolate=extrapolate)
            end
        end
    end

    img
end

"""interpolate z at xy based on values z_p at xy_p"""
function interp2d(xy::Matrix{U}, xy_p::Matrix{U}, z_p::Vector{T};
    mode="nn",
    k=10, # number of neighboring points used in kringing and idw
    p=2.0, # idw power
    extrapolate=true,
    tree::NearestNeighbors.NNTree=KDTree(xy_p)) where T <: AbstractFloat where U <: Real

    # output
    n = size(xy,2)
    z = zeros(T,length(xy))

    # interpolation
    for I in 1:n
        x,y = xy[:,I]
        if mode == "nn"
            z[I] = interpolator_nn(x, y, xy_p, z_p, tree)
        elseif mode == "kriging"
            z[I] = interpolator_kriging(x, y, xy_p, z_p, tree; k=k, extrapolate=extrapolate)
        elseif mode == "idw"
            z[I] = interpolator_idw(x, y, xy_p, z_p, tree; p=p, k=k, extrapolate=extrapolate)
        end
    end

    z
end

function interpolator_kriging(x::Real, y::Real,
    xy_p::Array{U,2}, z_p::Array{T,1},
    tree::NearestNeighbors.NNTree=KDTree(xy_p);
    k=10, # number of neighboring points used in kriging and idw
    extrapolate=true,
    nodata=-9999.0) where T <: AbstractFloat where U <: Real

    idxs, dists = knn(tree, [x, y], k, true)
    if !extrapolate && !in_hull(x, y, [xy_p[:,i] for i in idxs])
        T(nodata)
    elseif length(idxs) > 0
        # define a covariance model
        γ = GaussianVariogram(sill=1.0f0, range=1.0f0, nugget=0.0f0)
        # define an estimator (i.e. build the Kriging system)
        ordkrig = OrdinaryKriging(xy_p[:,idxs], z_p[idxs], γ)
        # estimate at target location
        μ, _ = estimate(ordkrig, T[x,y]) # not using σ²
        μ
    else
        error("no points to interpolate")
    end
end

function interpolator_idw(x::Real, y::Real,
    xy_p::Array{U,2}, z_p::Array{T,1},
    tree::NearestNeighbors.NNTree=KDTree(xy_p);
    p=2.0, # idw power
    k=10, # number of neighboring points used in kringing and idw
    extrapolate=true,
    nodata=-9999.0) where T <: AbstractFloat where U <: Real

    idxs, dists = knn(tree, [x, y], k, true)
    if !extrapolate && !in_hull(x, y, [xy_p[:,i] for i in idxs])
        T(nodata)
    elseif length(idxs) > 0
        idw_interpolation(z_p[idxs], dists; power=p)
    else
        error("no points to interpolate")
    end
end

function interpolator_nn(x::Real, y::Real,
    xy_p::Array{U,2}, z_p::Array{T,1},
    tree::NearestNeighbors.NNTree=KDTree(xy_p)) where T <: AbstractFloat where U <: Real

    idx, dists = knn(tree, Float32[x, y], 1)
    z_p[idx[1]]
end

"""loop over image to find missings & make arrays with datapoints"""
function xy_z_grid(img::Array{T,2}; nodata::Real=-9999.0) where T
    n = count(x -> x != nodata, img)
    xy = zeros(Float32, (2, n))
    z = zeros(Float32, n)
    dims = size(img)
    i = 0
    for I in 1:length(img)
        if img[I] != nodata
            i += 1
            x,y = ind2sub(dims, I)
            # build xy and z arrays for KDTree
            xy[1, i] = x
            xy[2, i] = y
            z[i] = img[I]
        end
    end
    xy, z
end

"fill las dropouts, but keep no data areas next to strips open"
function create_mask(img::Array; nodata::Real=-9999.0, n=10)
    mask = img .!= nodata
    for i in 1:n
        ImageMorphology.dilate!(mask)
    end
    for i in 1:n
        ImageMorphology.erode!(mask)
    end
    Array{Bool,2}(mask)
end

"""calculate the Inverse Distance Weighting interpolation for one point"""
function idw_interpolation(z::Vector{U}, dis::Vector{U}; power=10.) where U <: AbstractFloat
    invdisp = 1.0 ./ (dis .^ power)
    weights = invdisp ./ sum(invdisp)
    dot(weights, z)
end

"""linear interpolation based fox xy withing triangle in 3D plane"""
function linear_interpolation(x::Real, y::Real, a::Vector{U},
        b::Vector{U}, c::Vector{U}) where U <: AbstractFloat

    t1 = ((b[1]-a[1])*(c[3]-a[3])-(c[1]-a[1])*(b[3]-a[3]))
    t2 = ((b[2]-a[2])*(c[3]-a[3])-(c[2]-a[2])*(b[3]-a[3]))
    n = ((b[1]-a[1])*(c[2]-a[2])-(c[1]-a[1])*(b[2]-a[2]))
    z = a[3] + t1/n*(y-a[2]) - t2/n*(x-a[1])
end


text:
using FunctionalCollections: PersistentSet, PersistentHashMap, dissoc, assoc, conj, disj

####################################################
# factor graph, variable elimination, and sampling #
####################################################

# TODO performance optimize?

struct VarNode{T,V} # T would be FactorNode, but for https://github.com/JuliaLang/julia/issues/269
    addr::Any
    factor_nodes::PersistentSet{T}
    idx_to_domain::Vector{V}
    domain_to_idx::Dict{V,Int}
end

addr(node::VarNode) = node.addr
factor_nodes(node::VarNode) = node.factor_nodes
num_values(node::VarNode) = length(node.idx_to_domain)
idx_to_value(node::VarNode{T,V}, idx::Int) where {T,V} = node.idx_to_domain[idx]::V
value_to_idx(node::VarNode{T,V}, value::V) where {T,V} = node.domain_to_idx[value]

function remove_factor_node(node::VarNode{T,V}, factor_node::T) where {T,V}
    return VarNode{T,V}(
        node.addr, disj(node.factor_nodes, factor_node),
        node.idx_to_domain, node.domain_to_idx)
end

function add_factor_node(node::VarNode{T,V}, factor_node::T) where {T,V}
    return VarNode{T,V}(
        node.addr, conj(node.factor_nodes, factor_node),
        node.idx_to_domain, node.domain_to_idx)
end

struct FactorNode{N} # N is the number of variables in the (original) factor graph
    id::Int
    vars::Vector{Int} # immutable
    log_factor::Array{Float64,N} # immutable
end

vars(node::FactorNode) = node.vars
get_log_factor(node::FactorNode) = node.log_factor

struct FactorGraph{N}
    num_factors::Int
    var_nodes::PersistentHashMap{Int,VarNode}

    # NOTE: when variables get eliminated from a factor graph, they don't get reindex]
    # (i.e. these fields are unchanged)
    addr_to_idx::Dict{Any,Int}
end

# just for testing purposes:
function factor_value(fg::FactorGraph, node::FactorNode{N}, values::Dict) where {N}
    idxs = (idx in vars(node) ? value_to_idx(idx_to_var_node(fg, idx), values[idx]) : 1 for idx in 1:N)
    return exp.(node.log_factor[idxs...])
end

function draw_factor_graph(fg::FactorGraph, graphviz, fname)
    dot = graphviz.Digraph()
    factor_idx = 1
    for node in values(fg.var_nodes)
        shape = "ellipse"
        color = "white"
        name = replace(replace(replace(string(addr(node)), "=>" => "_"), "(" => ""), ")" => "")
        dot[:node](name, name, shape=shape, fillcolor=color, style="filled")
        for factor_node in factor_nodes(node)
            shape = "box"
            color = "gray"
            factor_name = string(factor_node.id)
            dot[:node](factor_name, factor_name, shape=shape, fillcolor=color, style="filled")
            dot[:edge](name, factor_name)
        end
    end
    dot[:render](fname, view=true)
end

idx_to_var_node(fg::FactorGraph, idx::Int) = fg.var_nodes[idx]
addr_to_idx(fg::FactorGraph, addr) = fg.addr_to_idx[addr]

# variable elimination
# - generates a sequence of factor graphs
# - multiply all factors that mention the variable, generating a product factor, which replaces the other factors
# - then sum out the product factor, and remove the variable
# ( we could break these into two separate operations -- NO )

# all factors are of the same dimension, but with singleton dimensions for
# variables that are eliminated

function multiply_and_sum(log_factors::Vector{Array{Float64,N}}, idx_to_sum_over::Int) where {N}
    result = broadcast(+, log_factors...)
    m = maximum(result, dims=idx_to_sum_over)
    return m .+ log.(sum(exp.(result .- m), dims=idx_to_sum_over))
end

function eliminate(fg::FactorGraph{N}, addr::Any) where{N}
    eliminated_var_idx = addr_to_idx(fg, addr)
    eliminated_var_node = idx_to_var_node(fg, eliminated_var_idx)
    log_factors_to_combine = Vector{Array{Float64,N}}()
    other_involved_var_nodes = Dict{Int,VarNode{FactorNode{N}}}()
    for factor_node in factor_nodes(eliminated_var_node)
        push!(log_factors_to_combine, get_log_factor(factor_node))

        # remove the reference to this factor node from its variable nodes
        for other_var_idx::Int in vars(factor_node)
            if other_var_idx == eliminated_var_idx
                continue
            end
            if !haskey(other_involved_var_nodes, other_var_idx)
                other_var_node = idx_to_var_node(fg, other_var_idx)
                other_involved_var_nodes[other_var_idx] = other_var_node
            else
                other_var_node = other_involved_var_nodes[other_var_idx]
            end
            @assert factor_node in factor_nodes(other_var_node)
            other_var_node = remove_factor_node(other_var_node, factor_node)
            other_involved_var_nodes[other_var_idx] = other_var_node
        end
    end

    # compute the new factor
    new_log_factor = multiply_and_sum(log_factors_to_combine, eliminated_var_idx)

    # add the new factor node
    new_factor_node = FactorNode{N}(
        fg.num_factors+1, collect(keys(other_involved_var_nodes)), new_log_factor)
    for (other_var_idx, other_var_node) in other_involved_var_nodes
        other_involved_var_nodes[other_var_idx] = add_factor_node(other_var_node, new_factor_node)
    end

    # remove the eliminated var node
    new_var_nodes = dissoc(fg.var_nodes, eliminated_var_idx)

    # replace old other var nodes with new other var nodes
    for (other_var_idx, other_var_node) in other_involved_var_nodes
        new_var_nodes = assoc(new_var_nodes, other_var_idx, other_var_node)
    end

    return FactorGraph{N}(fg.num_factors+1, new_var_nodes, fg.addr_to_idx)
end

function conditional_dist(fg::FactorGraph{N}, values::Vector{Any}, addr::Any) where {N}

    # other_values must contain a value for all variables that have a factor in
    # common with variable addr in fg
    var_idx = addr_to_idx(fg, addr)
    var_node = idx_to_var_node(fg, var_idx)
    n = num_values(var_node)
    #println("conditional_dist, addr: $addr, var_idx: $var_idx, num_values: $n")
    log_probs = zeros(n)
    # TODO : writing the slow version first..
    # LATER: use generated function to generate a version that is specialized
    # to N? (unroll this loop, and inline the indices..)
    indices = Vector{Int}(undef, N)
    for i in 1:n
        for factor_node in factor_nodes(var_node)
            log_factor::Array{Float64,N} = get_log_factor(factor_node)
            fill!(indices, 1)
            for other_var_idx in vars(factor_node)
                if other_var_idx != var_idx
                    other_var_node = idx_to_var_node(fg, other_var_idx)
                    indices[other_var_idx] = value_to_idx(other_var_node, values[other_var_idx])
                end
            end
            indices[var_idx] = i
            log_probs[i] += log_factor[CartesianIndex{N}(indices...)]
        end
    end
    return exp.(log_probs .- logsumexp(log_probs))
end

struct VariableEliminationResult{N}
    elimination_order::Any
    intermediate_fgs::Vector{FactorGraph{N}}
end

function variable_elimination(fg::FactorGraph{N}, elimination_order) where {N}
    intermediate_fgs = Vector{FactorGraph{N}}(undef, N)
    for addr in elimination_order
        var_idx = addr_to_idx(fg, addr)
        intermediate_fgs[var_idx] = fg
        fg = eliminate(fg, addr)
    end
    return VariableEliminationResult(elimination_order, intermediate_fgs)
end


text:
module Digits

export
  reversedigits,
  reversedigits!,
  undigit,
  digithist,
  isanagram,
  ispalindrome,
  contains,
  startswith,
  endswith,
  crop,
  combine,
  crosssum,
  select,
  replace,
  replace!,
  digitroot

function undigit(l::Vector{Int}, base::Int = 10)
  acc = 0
  for (i, digit) in enumerate(l)
    acc += digit * base ^ (i-1)
  end
  acc::Integer
end

undigit(n::Integer) = n

reversedigits(n::Integer) = undigit(reverse!(digits(n)))

reversedigits(l::Array{Int,1}) = reverse(l)

reversedigits!(l::Array{Int,1}) = reverse!(l)

function digithist(l::Array{Int,1})
  return hist(abs(l),-0.5:10)[2]
end

digithist(n::Integer) = digithist(digits(n))

function isanagram(a::Array{Int,1}, b::Array{Int,1})
  if sign(a) != sign(b)
    return false
  end
  return digithist(a) == digithist(b)
end

isanagram(a::Integer, b::Integer) = isanagram(digits(a),digits(b))
isanagram(a::Integer, b::Vector{Int}) = isanagram(digits(a),b)
isanagram(a::Vector{Int}, b::Integer) = isanagram(a,digits(b))

function ispalindrome(l::Array{Int,1})
  m = div(length(l),2)
  for i = 1:m
    if l[i] != l[end-i+1]
      return false
    end
  end
  return true
end

ispalindrome(n::Integer) = ispalindrome(digits(n))

function contains(a::Array{Int,1},b::Array{Int,1})
  la = length(a)
  lb = length(b)
  if la > lb
    for i = 1:la-lb+1
      if a[i:i+lb-1] == b
        return true
      end
    end
  elseif lb > la
    for i = 1:lb-la+1
      if b[i:i+la-1] == a
        return true
      end
    end
  elseif a == b
    return true
  end
  return false
end

contains(a::Array{Int,1},b::Integer) = contains(a,digits(b))
contains(a::Integer,b::Array{Int,1}) = contains(digits(a),b)
contains(a::Integer,b::Integer) = contains(digits(a),digits(b))

function endswith(a::Array{Int,1},b::Array{Int,1})
  if a == b
    return true
  end
  la = length(a)
  lb = length(b)
  if la > lb
    if a[1:lb] == b
      return true
    end
  else
    if b[1:la] == a
      return true
    end
  end
  return false
end

endswith(a::Array{Int,1},b::Integer) = endswith(a,digits(b))
endswith(a::Integer,b::Array{Int,1}) = endswith(digits(a),b)
endswith(a::Integer,b::Integer) = endswith(digits(a),digits(b))

function startswith(a::Array{Int,1},b::Array{Int,1})
  if a == b
    return true
  end
  la = length(a)
  lb = length(b)
  if la > lb
    if a[end-lb+1:end] == b
      return true
    end
  else
    if b[end-la+1:end] == a
      return true
    end
  end
  return false
end

startswith(a::Array{Int,1},b::Integer) = startswith(a,digits(b))
startswith(a::Integer,b::Array{Int,1}) = startswith(digits(a),b)
startswith(a::Integer,b::Integer) = startswith(digits(a),digits(b))

function crop(l::Array{Int,1},i::Int)
  if abs(i) > length(l)
    error("Crop value too big!")
  elseif abs(i) == length(l)
    return [0]
  elseif i == 0
    return l
  elseif i > 0
    return l[1:end-i]
  else i < 0
    return l[1+abs(i):end]
  end
end

crop(n::Integer,i::Int) = undigit(crop(digits(n),i))

function replace!(l::Array{Int,1},idx::AbstractArray{Int,1},d::Array{Int,1})
  l[idx] = d
  return l
end

function replace!(l::Array{Int,1},old::Int,new::Int)
  l[l.==old] = new
  return l
end

replace(l::Array{Int,1},idx::AbstractArray{Int,1},d::Array{Int,1}) = replace!(deepcopy(l),idx,d)
replace(l::Array{Int,1},old::Int,new::Int) = replace!(deepcopy(l),old,new)

function replace(n::Integer,idx::AbstractArray{Int,1},d::Array{Int,1})
  return undigit(reverse(replace(reverse(digits(n)),idx,d)))
end

function replace(n::Integer,old::Int,new::Int)
  return undigit(replace(digits(n),old,new));
end

function select(l::Array{Int,1},idx::AbstractArray{Int,1})
  return l[idx]
end

function select(n::Integer,idx::AbstractArray{Int,1})
  return undigit(reverse(reverse(digits(n))[idx]))
end

function combine(a::Integer,b::Integer)
  if a == 0
    return b
  end
  f = 10^ndigits(b)
  return a*f+b
end

combine(a::Array{Int,1},b::Array{Int,1}) = [b;a]
combine(a::Integer,b::Vector{Int}) = combine(a,undigit(b))
combine(a::Vector{Int},b::Integer) = combine(a,digits(b))

crosssum(l::Array{Int,1}) = sum(l)
crosssum(n::Integer) = sum(digits(n))

function digitroot(n::Integer)
  x = abs(n)
  r = x-9*div(x-1,9)
  return copysign(r,n)
end

digitroot(l::Array{Int,1}) = digitroot(undigit(l))

end # module


text:
"""
set and get function overloads
"""
#MOI.get(model::Optimizer, ::MOI.ObjectiveSense) = model.inner.sense

"""
function MOI.set(model::Optimizer, ::MOI.ObjectiveSense, sense::MOI.OptimizationSense)
    model.inner.sense = sense
    return
end
"""

"""
    MOI.set(model::Optimizer, MOI.ObjectiveFunction, func<:AbstractScalarFunction)

Set the objective function of your model.
Support linear ones only at the moment put we will move to non-linear support soon.
"""

function MOI.set(model::Optimizer, ::MOI.ObjectiveFunction, svf::MOI.SingleVariable)
    model.moimodel.objective_identifier = svf.variable
end

function MOI.set(model::Optimizer, ::MOI.ObjectiveSense, sense::MOI.OptimizationSense)
    @assert sense == MOI.MIN_SENSE "SeaPearl does not support maximisation"
end

function MOI.set(model::Optimizer, ::MOI.ObjectiveFunction, saf::MOI.ScalarAffineFunction{T}) where {T<:Real}
    moiaff = MOIAffineFunction(nothing, saf)
    push!(model.moimodel.affines, moiaff)

    model.moimodel.objective_identifier = AffineIndex(length(model.moimodel.affines))
end


text:
import Base.+
import Base.-
import Base.*

export ℝ⁴

"""
    Represents a point in ℝ⁴.

field: a.
"""
struct ℝ⁴ <: VectorSpace
    a::Array{Float64}
    ℝ⁴(x₁::Real, x₂::Real, x₃::Real, x₄::Real) = new(float.([x₁; x₂; x₃; x₄]))
    ℝ⁴(a::Array) = begin
        @assert(length(a) == 4, "The input vector must contain exactly four elements.")
        ℝ⁴(a...)
    end
end

## Unary Operators ##

+(r::ℝ⁴) = r
-(r::ℝ⁴) = ℝ⁴(-vec(r))

## Binary Operators ##

+(r1::ℝ⁴, r2::ℝ⁴) = ℝ⁴(vec(r1) + vec(r2))
-(r1::ℝ⁴, r2::ℝ⁴) = ℝ⁴(vec(r1) - vec(r2))
*(r::ℝ⁴, λ::Real) = ℝ⁴(λ .* vec(r))
*(λ::Real, r::ℝ⁴) = r * λ

max_issues_repo_name: iamazadi/Porta

max_forks_repo_name: iamazadi/Porta


text:
using PyCall
ct = pyimport("cantera")

P1 = 1e5
T1 = 300
X1 = "H2:42, O2:21,N2:79"
mech = "gri30.xml"

gas = ct.Solution(mech)

T = T1
print("new\n")

function set_T(gas)
    global T *= 1.01
    gas.TPX = T,P1,X1
    print("$T\n")
    print("$(gas.T)\n")
end

for i = 1:100

    # global T *= 1.01
    set_T(gas)

end


text:
# Helper functions
function z_subproblem!{T}(z::Array{T, 2},
                          t::Array{T, 2}, mu_t::T, lambda_t::Array{T, 2},
                          DtD::Array{T, 2}, Dtx::Array{T, 2})
  dictSize = size(DtD, 1);
  # z = ( DtD + mu_t.*I ) \ ( Dtx + mu_t.*t - lambda_t )
  A = copy(DtD);
  for i=1:dictSize; A[i, i] += mu_t; end
  b = copy(t);
  scale!(b, mu_t);
  b += Dtx;
  b -= lambda_t;
  z[:, :] = A\b;
  return z::Array{T, 2};
end

function t_subproblem!{T}(t::Array{T, 2}, z::Array{T, 2}, mu_t::T, lambda_t::Array{T, 2}, beta::T)
  if mu_t>0;
    # shrink!(t, z+lambda_t/mu_t, beta/mu_t);
    for i=1:length(t);
      t[i] = lambda_t[i]/mu_t + z[i];
    end
    shrink!(t, t, beta/mu_t);
  else;
    fill!(t, 0);
  end
  return t::Array{T, 2};
end

function fsc_coding_fastmulti{T}(data::FastCodingData{T}, optParams::FastCodingOpts{T})
# Alg. parameters:
#   mu_max: 1e5
#   tau: [1.01 1.5] larger means faster convergence but sacrifices primal feasibility,
#        smaller means slower convergence

  const mu_max = 1e5;
  const tau = 1.01;

  const Dtx          = data.Dtx;
  const DtD          = data.DtD;
        z            = optParams.z;
        t            = optParams.t;
        lambda_t     = optParams.lambda_t;
  const beta::T      = optParams.beta;
        mu_t::T      = optParams.mu_t;
  const max_iter     = optParams.max_iter;
  const tol          = optParams.tol;
  const DEBUG        = optParams.DEBUG;
        debugData    = optParams.debugData;

  (filtLen, numImgs) = size(Dtx);
  dictSize = size(DtD, 1);

  assert(size(Dtx, 1)==dictSize, "Dtx does not have the correct dimensions.")
  z = initcheck(z, T, (dictSize, numImgs), "z does not have the correct dimensions.");
  t = initcheck(t, T, (dictSize, numImgs), "t does not have the correct dimensions.");
  lambda_t = initcheck(lambda_t, T, (dictSize, numImgs), "lambda_t does not have the correct dimensions.");

  if DEBUG;
    debugData["z"] = z;
    debugData["t"] = t;
    debugData["lambda_t"] = lambda_t;
    (prev_obj, rerr, spar) = compute_lagrangian(debugData);
  end

  iter = 0;
  converged = false;
  while ~converged;
    iter += 1;

    if max_iter>1; @printf("\b\b\b\b\b%05d", iter); end

    # Solve for z* in the frequency domain
    z_subproblem!(z, t, mu_t, lambda_t, DtD, Dtx);
    if DEBUG;
      debugData["z"] = z;
      (obj, rerr, spar) = compute_lagrangian(debugData);
      is_mono = (prev_obj>=obj) ? 'O' : 'X';
      @printf("            Z: %c obj=%.10f (rerr=%.10f, spar=%.10f)\n", is_mono, obj, rerr, spar);
      prev_obj = obj;
    end

    # Solve for t* in the spatial domain
    t_subproblem!(t, z, mu_t, lambda_t, beta);
    if DEBUG;
      debugData["t"] = t;
      (obj, rerr, spar) = compute_lagrangian(debugData);
      is_mono = (prev_obj>=obj) ? 'O' : 'X';
      @printf("            T: %c obj=%.10f (rerr=%.10f, spar=%.10f)\n", is_mono, obj, rerr, spar);
      prev_obj = obj;
    end

    if max_iter==1; return (z, t); end # Don't do unnecessary work

    # Update the Lagrangian multipliers
    for i=1:filtLen*dictSize; lambda_t[i] += mu_t*( z[i]-t[i] ); end
    if DEBUG;
      debugData["lambda_t"] = lambda_t;
      (obj, rerr, spar) = compute_lagrangian(debugData);
      is_mono = (prev_obj>=obj) ? 'O' : 'X';
      @printf("            L: %c obj=%.10f (rerr=%.10f, spar=%.10f)\n", is_mono, obj, rerr, spar);
      prev_obj = obj;
    end

    # Update penalty term
    mu_t = min(tau*mu_t, mu_max);
    if DEBUG;
      debugData["mu_t"] = mu_t;
      (obj, rerr, spar) = compute_lagrangian(debugData);
      is_mono = (prev_obj>=obj) ? 'O' : 'X';
      @printf("            M: %c obj=%.10f (rerr=%.10f, spar=%.10f)\n", is_mono, obj, rerr, spar);
      prev_obj = obj;
    end

    codesConverged = true;
    for j=1:numImgs;
      if mu_t>0 && mean((z[:, j]-t[:, j]).^2)>tol;
        codesConverged = false;
        break;
      end
    end
    if iter>=max_iter || codesConverged;
      converged = true;
    end
  end

  return (z, t);
end

max_stars_repo_name: bkong/FSC

max_issues_repo_name: bkong/FSC

max_forks_repo_name: bkong/FSC


text:
using OffsetArrays
using StaticArrays
using DataStructures
export pgbf, cgbf, contract, addbf!, PGBF, CGBF, build_basis,eri_fetcher, Shell, Basis
export m2ao, shell_indices,nao

"""
    PGBF(expn,xyz,I,J,K,norm)

Create a primitive Gaussian basis function
    g(x,y,z) = norm * (x-x0)^I (y-y0)^J (z-z0)^K exp(-expn*r^2)
The function parameters `xyz` correspond to [x0,y0,z0].
"""
mutable struct PGBF
    expn::Float64
    xyz::SVector{3,Float64}
    I::Int64
    J::Int64
    K::Int64
    norm::Float64
end

"""
    pgbf(expn,x=0,y=0,z=0,I=0,J=0,K=0,norm=1)

Helper function to create a normalized PGBF with some optional
defaults set.
"""
function pgbf(expn,x=0.0,y=0.0,z=0.0,I=0,J=0,K=0,norm=1)
    p = PGBF(expn,SA[x,y,z],I,J,K,norm)
    normalize!(p)
    return p
end

function amplitude(bf::PGBF,x,y,z)
    dx,dy,dz = dxyz = bf.xyz-SA[x,y,z]
    r2 = dist2(dxyz)
    return bf.norm*(dx^bf.I)*(dy^bf.J)*(dz^bf.K)*exp(-bf.expn*r2)
end
(bf::PGBF)(x,y,z) = amplitude(bf::PGBF,x,y,z)

function normalize!(pbf::PGBF)
    pbf.norm /= sqrt(overlap(pbf,pbf))
end

"""
    CGBF(xyz,I,J,K,norm,[pbgfs],[coefs])

Create a contracted Gaussian basis function made of
the functions in [pgbfs] with coefficients [coefs].
Also track the origin `xyz` and powers `I,J,K`.
"""
mutable struct CGBF
    xyz::SVector{3,Float64}
    I::Int64
    J::Int64
    K::Int64
    norm::Float64
    pgbfs::Vector{PGBF}
    coefs::Vector{Float64}
end

"""
    cgbf(expn,x=0,y=0,z=0,I=0,J=0,K=0,norm=1)

Helper function to create a CGBF with optional defaults.
"""
cgbf(x=0.0,y=0.0,z=0.0,I=0,J=0,K=0) = CGBF(SA[x,y,z],I,J,K,1.0,[],[])

amplitude(bf::CGBF,x,y,z) = bf.norm*sum(c*amplitude(pbf,x,y,z) for (c,pbf) in primitives(bf))
(bf::CGBF)(x,y,z) = amplitude(bf::CGBF,x,y,z)

function normalize!(bf::CGBF)
    bf.norm /= sqrt(overlap(bf,bf))
end

primitives(a::CGBF) = zip(a.coefs,a.pgbfs)

function addbf!(cbf::CGBF,expn,coef)
    Base.push!(cbf.pgbfs,pgbf(expn,cbf.xyz...,cbf.I,cbf.J,cbf.K))
    Base.push!(cbf.coefs,coef)
    normalize!(cbf)
end

contract(f,a::CGBF,b::CGBF) = a.norm*b.norm*sum(ca*cb*f(abf,bbf) for (ca,abf) in primitives(a) for (cb,bbf) in primitives(b))
function contract(f,a::CGBF,b::CGBF,c::CGBF,d::CGBF)
    s = 0
    for (ca,abf) in primitives(a)
        for (cb,bbf) in primitives(b)
            for (cc,cbf) in primitives(c)
                for (cd,dbf) in primitives(d)
                    s += ca*cb*cc*cd*f(abf,bbf,cbf,dbf)
                end
            end
        end
    end
    return a.norm*b.norm*c.norm*d.norm*s
end

function shells(atoms::Vector{Atom},name="sto3g")
    data = basis_data[name]
    shs = Shell[]
    for atom in atoms
        for (sym,primlist) in data[atom.atno]
            L = lvalue[sym]
            expns = [expn for (expn,coef) in primlist]
            coefs = [coef for (expn,coef) in primlist]
            bfs = CGBF[]
            for (I,J,K) in shell_indices[L]
                cbf = cgbf(atom.xyz...,I,J,K)
                push!(bfs,cbf)
                for (expn,coef) in zip(expns,coefs)
                    addbf!(cbf,expn,coef)
                end
            end
            push!(shs,Shell(atom.xyz,lvalue[sym],expns,coefs,bfs))
        end
    end
    return shs
end

function build_basis(atoms::Vector{Atom},name="sto3g")
    shs = shells(atoms,name)
    bfs = CGBF[]
    ishs = Int[]
    mshs = Int[]
    for (ish,sh) in enumerate(shs)
        for (msh,(I,J,K)) in enumerate(shell_indices[sh.L])
            cbf = cgbf(sh.xyz...,I,J,K)
            push!(bfs,cbf)
            push!(ishs,ish)
            push!(mshs,msh)
            for (expn,coef) in zip(sh.expns,sh.coefs)
                addbf!(cbf,expn,coef)
            end
        end
    end
    return Basis(bfs,shs,ishs,mshs)
end

"""
    Shell(xyz,L,expns,coeffs)

Structure for a basis function shell, containing multiple
CGBFs of different angular momenta.
"""
mutable struct Shell
    xyz::SVector{3,Float64}
    L::Int
    expns::Vector{Float64}
    coefs::Vector{Float64}
    cgbfs::Vector{CGBF}
end
nbf(sh::Shell) = length(shell_indices(sh.L))
function pgbfs(sh::Shell)
    x,y,z = sh.xyz
    pgbfs = PGBF[]
    for (I,J,K) in shell_indices[sh.L]
        for (expn,coef) in zip(sh.expns,sh.coefs)
            push!(pgbfs,pgbf(expn,x,y,z,I,J,K))
        end
    end
    return pgbfs
end

"""
    Basis(cgbfs,shells,ishell,mshell)

Structure to hold a basis set, with info about shells and other data
"""
mutable struct Basis # subset of AbstractVector{CGBF}?
    cgbfs::Vector{CGBF}
    shells::Vector{Shell}
    ishell::Vector{Int64} # Which shell corresponds to bf i
    mshell::Vector{Int64} # Which m-value (px,dxy, etc.) corresponds to bf i
end
Base.size(b::Basis) = (length(b.cgbfs),)
Base.length(b::Basis) = length(b.cgbfs)
Base.iterate(b::Basis,i::Int) = iterate(b.cgbfs,i)
Base.iterate(b::Basis) = iterate(b.cgbfs)
Base.getindex(b::Basis, i::Int) = b.cgbfs[i]
nbf(b::Basis) = length(b.cgbfs)
nshells(b::Basis) = length(b.shells)

"""
  eri_fetcher

Compute all of the required ijkl terms that go into an ERI
record, and compute the required calls to hrr, and how to
unpack the results into the record.

eri_fetcher returns a dictionary such that the integral structure
may be formed via:
```
    fetcher[ishell,jshell,kshell,lshell] = (ijkl,hi,hj,hk,hl)
    hrrs = hrr(ishell,jshell,kshell,lshell)
    ints[ijkl] = hrrs[hi,hj,hk,hl]
```
"""
function eri_fetcher(bfs::Basis)
    fetcher = DefaultOrderedDict{NTuple{4, Int}, Vector{NTuple{5, Int}}}(Vector{NTuple{5, Int}}[])
    for (index,ijkl) in enumerate(iiterator(length(bfs)))
        i,j,k,l = ijkl
        li,mi = bfs.ishell[i],bfs.mshell[i]
        lj,mj = bfs.ishell[j],bfs.mshell[j]
        lk,mk = bfs.ishell[k],bfs.mshell[k]
        ll,ml = bfs.ishell[l],bfs.mshell[l]
        push!(fetcher[li,lj,lk,ll],(index,mi,mj,mk,ml))
    end
    return fetcher
end

function eri_fetcher(shs::Vector{Shell})
    fetcher = DefaultOrderedDict{NTuple{4, Int}, Vector{NTuple{5, Int}}}(Vector{NTuple{5, Int}}[])
    nsh = length(shs)
    ijk2lm = make_ijk2lm()

    i = 0
    index = 0
    for (ish,shi) in enumerate(shs)
        shi = shs[ish]
        for ibf in shi.cgbfs
            i += 1
            mi = ijk2lm[ibf.I,ibf.J,ibf.K][2]

            j = 0
            for (jsh,shj) in enumerate(take(shs,ish))
                for jbf in shj.cgbfs
                    j += 1
                    if j > i break end
                    mj = ijk2lm[jbf.I,jbf.J,jbf.K][2]

                    k = 0
                    for (ksh,shk) in enumerate(shs)
                        for kbf in shk.cgbfs
                            k += 1
                            mk = ijk2lm[kbf.I,kbf.J,kbf.K][2]

                            l = 0
                            for (lsh,shl) in enumerate(take(shs,ksh))
                                @inbounds for lbf in shl.cgbfs
                                    l += 1
                                    if l > k break end
                                    if k*l > i*j break end
                                    ml = ijk2lm[lbf.I,lbf.J,lbf.K][2]
                                    index += 1
                                    push!(fetcher[ish,jsh,ksh,lsh],(index,mi,mj,mk,ml))
                                end # lbf
                            end # lsh

                        end # kbf
                    end # ksh

                end # jbf
            end # jsh

        end # ibf
    end # ish

    return fetcher
end

# shell_indices map from a shell l-value to the Cartesian version of m-values that are the
#   powers of the Cartesian Gaussian basis functions.
#
# If desired, we can also invert shell_indices to map IJK triplets to l,m pairs:
# IJK2lm = Dict(IJK =>(l,m) for l in 0:4 for (m,IJK) in enumerate(shell_indices[l]))
const shell_indices = Dict(
    0 => [MVector(0,0,0)], # 1
    1 => [MVector(1,0,0),MVector(0,1,0),MVector(0,0,1)], # 3
    2 => [MVector(2,0,0),MVector(1,1,0),MVector(1,0,1),MVector(0,2,0),MVector(0,1,1),MVector(0,0,2)],
    3 => [MVector(3,0,0),MVector(2,1,0),MVector(2,0,1),
            MVector(1,2,0),MVector(1,0,2),MVector(1,1,1),
            MVector(0,3,0),MVector(0,2,1),MVector(0,1,2),MVector(0,0,3)], # 10
    4 => [MVector(4,0,0),MVector(3,1,0),MVector(3,0,1),MVector(2,2,0),MVector(2,1,1),MVector(2,0,2),
            MVector(1,3,0),MVector(1,2,1),MVector(1,1,2),MVector(1,0,3),
            MVector(0,4,0),MVector(0,3,1),MVector(0,2,2),MVector(0,1,3),MVector(0,0,4)] # 15
)

llabel = Dict(0=>"s",1=>"p",2=>"d",3=>"f",4=>"g",5=>"h")
lvalue = merge(Dict((v,k) for (k,v) in llabel),Dict((uppercase(v),k) for (k,v) in llabel))

bflabel(bf) = llabel[bf.I+bf.J+bf.K]*bfpow("x",bf.I)*bfpow("y",bf.J)*bfpow("z",bf.K)
function bfpow(s,j)
	if j == 0
		return ""
	elseif j == 1
		return s
	end
	return "$s$j"
end

"make_m2ao - Make map between a sequential list of (mx,my,mz) values and ao indices"
function make_m2ao(lmax=4)
    m2ao = Dict{Vector{Int64}, Int64}()
    iao = 0
    for l in 0:lmax
        for ms in shell_indices[l]
            iao += 1
            m2ao[ms] = iao
        end
    end
    return m2ao
end
function make_ao2m(lmax=4)
    ao2m = MVector{3, Int64}[]
    for l in 0:lmax
        for ms in shell_indices[l]
            push!(ao2m,ms)
        end
    end
    return ao2m
end

const m2ao = make_m2ao()
const ao2m = make_ao2m()

"make_nao - Number of AOs for system with l shells"
make_nao(l) = sum(length(shell_indices[i]) for i in 0:l)
const nao = OffsetArray([make_nao(l) for l in 0:4],0:4)

# Data functions and tables over basis functions
# todo: keep a constant of maxao or something
# and use it to construct tables
function make_shift_index()
    n = length(ao2m)
    shift_index = zeros(Int,n,3)
    for a in 1:n
        m = ao2m[a]
        for i in 1:3
            if m[i] == 0
                shift_index[a,i] = 0
            else
                mm = copy(m)
                mm[i] -= 1
                am = m2ao[mm]
                shift_index[a,i] = am
            end
        end
    end
    return shift_index
end
const shift_index = make_shift_index()

function make_shift_direction()
    n = length(ao2m)
    shift_direction = zeros(Int,n)
    for a in 1:n
        shift_direction[a] = argmax(ao2m[a])
    end
    return shift_direction
end
const shift_direction = make_shift_direction()

function make_shell_number()
    n = length(ao2m)
    shell_number = zeros(Int,n)
    for a in 1:n
        shell_number[a] = sum(ao2m[a])
    end
    return shell_number
end
const shell_number = make_shell_number()

function make_shift_index_plus()
    n = length(ao2m)
    maxsh = shell_number[n]
    shift_index = zeros(Int,n,3)
    for a in 1:n
        m = ao2m[a]
        for i in 1:3
            if sum(m) == maxsh
                shift_index[a,i] = 0
            else
                mm = copy(m)
                mm[i] += 1
                am = m2ao[mm]
                shift_index[a,i] = am
            end
        end
    end
    return shift_index
end
const shift_index_plus = make_shift_index_plus()

function make_ijk2lm(maxl=4)
    ijk2lm = Dict{NTuple{3,Int},NTuple{2,Int}}()
    for l in 0:maxl
        for (m,(i,j,k)) in enumerate(shell_indices[l])
            ijk2lm[i,j,k] = (l,m)
        end
    end
    return ijk2lm
end


text:
using Test: @test, @testset
using Statistics: mean, cor
using LinearAlgebra: norm, Diagonal
using Random: bitrand
using LogExpFunctions: softmax
using RestrictedBoltzmannMachines: RBM, Spin, Binary, Potts, Gaussian, training_epochs, zerosum!
using RestrictedBoltzmannMachines: mean_h_from_v, var_h_from_v, batchmean, batchvar
using RestrictedBoltzmannMachines: mean_from_inputs, extensive_sample, zerosum
using RestrictedBoltzmannMachines: sample_v_from_h, sample_v_from_v, initialize!, pcd!, free_energy, wmean
import Random
import Flux

@testset "extensive_sample" begin
    @test extensive_sample(Binary(1)) == [0 1]
    @test extensive_sample(Binary(2)) == reduce(hcat, [σ1,σ2] for σ1 in 0:1, σ2 in 0:1)
    @test extensive_sample(Spin(1)) == Int8[-1 1]
    @test extensive_sample(Spin(2)) == reduce(hcat, [s1,s2] for s1 in (-1,1), s2 in (-1,1))
    @test extensive_sample(Potts(2,1)) == reshape([1,0,0,1], 2,1,2)
    @test all(sum(extensive_sample(Potts(3,4)); dims=1) .== 1)
end

@testset "pcd -- teacher/student, Binary, with weights, exact" begin
    N = 5
    batchsize = 2^N
    nupdates = 10000
    teacher = RBM(Binary(N), Binary(1), zeros(N,1))
    teacher.w[:,1] .= range(-2, 2, length=N)
    data = extensive_sample(teacher.visible)
    wts = softmax(-free_energy(teacher, data))
    @test sum(wts) ≈ 1
    nsamples = size(data)[end]
    epochs = training_epochs(; nsamples, batchsize, nupdates)
    student = RBM(Binary(N), Binary(1), zeros(N,1))
    initialize!(student, data; wts)
    student.w .= cos.(range(-10, 10, length=N))
    @test mean_from_inputs(student.visible) ≈ wmean(data; wts, dims=2)
    pcd!(student, data; wts, epochs, batchsize, mode=:exact, optim=Flux.AdaBelief())
    @info @test cor(free_energy(teacher, data), free_energy(student, data)) > 0.9999

    # moment matching conditions
    wts_student = softmax(-free_energy(student, data))

    v_student = batchmean(student.visible, data; wts = wts_student)
    v_teacher = batchmean(teacher.visible, data; wts)
    @info @test norm(v_student - v_teacher) < 1e-10

    h_student = batchmean(student.hidden, mean_h_from_v(student, data); wts=wts_student)
    h_teacher = batchmean(teacher.hidden, mean_h_from_v(student, data); wts)
    @info @test norm(h_student - h_teacher) < 1e-10

    vh_student = data * Diagonal(wts_student) * mean_h_from_v(student, data)' / sum(wts_student)
    vh_teacher = data * Diagonal(wts) * mean_h_from_v(student, data)' / sum(wts)
    @info @test norm(vh_student - vh_teacher) < 1e-10
end

@testset "pcd -- teacher/student, Spin, with weights, exact" begin
    N = 5
    batchsize = 2^N
    nupdates = 10000
    teacher = RBM(Spin(N), Spin(1), zeros(N,1))
    teacher.w[:,1] .= range(-1, 1, length=N)
    data = extensive_sample(teacher.visible)
    wts = softmax(-free_energy(teacher, data))
    @test sum(wts) ≈ 1
    nsamples = size(data)[end]
    epochs = training_epochs(; nsamples, batchsize, nupdates)
    student = RBM(Spin(N), Spin(1), zeros(N,1))
    initialize!(student, data; wts)
    student.w .= cos.(1:N)
    @test mean_from_inputs(student.visible) ≈ wmean(data; wts, dims=2)
    pcd!(student, data; wts, epochs, batchsize, mode=:exact, shuffle=false, optim=Flux.RADAM())
    @info @test cor(free_energy(teacher, data), free_energy(student, data)) > 0.9999
end

@testset "pcd -- teacher/student, Gaussian, with weights, exact" begin
    N = 5
    batchsize = 2^N
    nupdates = 10000
    teacher = RBM(Binary(N), Gaussian(1), zeros(N,1))
    teacher.w[:,1] .= range(-1, 1, length=N)
    data = extensive_sample(teacher.visible)
    wts = softmax(-free_energy(teacher, data))
    @test sum(wts) ≈ 1
    nsamples = size(data)[end]
    epochs = training_epochs(; nsamples, batchsize, nupdates)
    student = RBM(Binary(N), Gaussian(1), zeros(N,1))
    initialize!(student, data; wts)
    student.w .= cos.(1:N)
    @test mean_from_inputs(student.visible) ≈ wmean(data; wts, dims=2)
    pcd!(student, data; wts, epochs, batchsize, ϵh=1e-2, shuffle=false, mode=:exact, optim=Flux.ADAM())
    @info @test cor(free_energy(teacher, data), free_energy(student, data)) > 0.99
    wts_student = softmax(-free_energy(student, data))
    ν_int = batchmean(student.hidden, var_h_from_v(student, data); wts = wts_student)
    ν_ext = batchvar(student.hidden, mean_h_from_v(student, data); wts = wts_student)
    @test only(ν_int + ν_ext) ≈ 1 - 1e-2 # not exactly 1 because of ϵh
end

@testset "pcd -- teacher/student, Potts, with weights, exact" begin
    q = 2
    N = 5
    batchsize = q^N
    nupdates = 2000
    teacher = RBM(Potts(q, N), Spin(1), zeros(q,N,1))
    teacher.w[1,:,1] .= range(-1, 1, length=N)
    teacher.w[2,:,1] .= -teacher.w[1,:,1]
    data = extensive_sample(teacher.visible)
    wts = softmax(-free_energy(teacher, data))
    @test sum(wts) ≈ 1
    nsamples = size(data)[end]
    epochs = training_epochs(; nsamples, batchsize, nupdates)
    student = RBM(Potts(q, N), Spin(1), zeros(q, N, 1))
    initialize!(student, data; wts)
    student.w[1,:,1] .= -student.w[2,:,1] .= cos.(1:N)
    @test mean_from_inputs(student.visible) ≈ wmean(data; wts, dims=3)
    pcd!(student, data; wts, epochs, batchsize, mode=:exact, optim=Flux.AdaBelief())
    @info @test cor(free_energy(teacher, data), free_energy(student, data)) > 0.9999
end

@testset "pcd -- teacher/student, Spin" begin
    N = 12
    batchsize = 8
    nupdates = 10000
    teacher = RBM(Spin(N), Spin(1), randn(N,1) ./ √N)
    # since h = ±1 are equally likely, the following gives an unbiased sample
    data = sample_v_from_h(teacher, repeat(Int8[1 -1], 1, 10000))
    student = RBM(Spin(N), Spin(1), zeros(N,1))
    nsamples = size(data)[end]
    epochs = training_epochs(; nsamples, batchsize, nupdates)
    initialize!(student, data)
    @test mean_from_inputs(student.visible) ≈ mean(data; dims=2)
    pcd!(student, data; epochs, batchsize)
    @info @test cor(free_energy(teacher, data), free_energy(student, data)) > 0.9
end

@testset "pcd -- teacher/student, Spin, with weights" begin
    N = 12
    batchsize = 8
    nupdates = 10000
    teacher = RBM(Spin(N), Spin(1), zeros(N,1))
    teacher.w .= range(-1, 1, length=N)
    data = extensive_sample(teacher.visible)
    wts = softmax(-free_energy(teacher, data))
    @test sum(wts) ≈ 1
    nsamples = size(data)[end]
    epochs = training_epochs(; nsamples, batchsize, nupdates)
    student = RBM(Spin(N), Spin(1), zeros(N,1))
    initialize!(student, data; wts)
    @test mean_from_inputs(student.visible) ≈ wmean(data; wts, dims=2)
    pcd!(student, data; wts, epochs, batchsize)
    @info @test cor(free_energy(teacher, data), free_energy(student, data)) > 0.95
end

@testset "pcd -- teacher/student, Potts" begin
    q = 2
    N = 12
    batchsize = 8
    nupdates = 10000
    teacher = RBM(Potts(q, N), Spin(1), zerosum(randn(q, N, 1)))
    # since h = ±1 are equally likely, the following gives an unbiased sample
    data = sample_v_from_h(teacher, repeat(Int8[1 -1], 1, 10000))
    student = RBM(Potts(q, N), Spin(1), zeros(q, N, 1))
    nsamples = size(data)[end]
    epochs = training_epochs(; nsamples, batchsize, nupdates)
    initialize!(student, data)
    @test mean_from_inputs(student.visible) ≈ mean(data; dims=3)
    pcd!(student, data; epochs, batchsize)
    @info @test cor(free_energy(teacher, data), free_energy(student, data)) > 0.9
end

@testset "pcd -- teacher/student, Potts, with weights" begin
    q = 2
    N = 12
    batchsize = 8
    nupdates = 10000
    teacher = RBM(Potts(q,N), Spin(1), zeros(q,N,1))
    teacher.w[1,:,1] .= range(-1, 1, length=N)
    teacher.w[2,:,1] .= -teacher.w[1,:,1]
    data = extensive_sample(teacher.visible)
    wts = softmax(-free_energy(teacher, data))
    @test sum(wts) ≈ 1
    nsamples = size(data)[end]
    epochs = training_epochs(; nsamples, batchsize, nupdates)
    student = RBM(Potts(q,N), Spin(1), zeros(q,N,1))
    initialize!(student, data; wts)
    @test mean_from_inputs(student.visible) ≈ wmean(data; wts, dims=3)
    pcd!(student, data; wts, epochs, batchsize)
    @info @test cor(free_energy(teacher, data), free_energy(student, data)) > 0.95
end

@testset "pcd -- teacher/student, Gaussian hidden" begin
    q = 2
    N = 12
    batchsize = 8
    nupdates = 10000
    teacher = RBM(Potts(q, N), Gaussian(1), zerosum(randn(q, N, 1) / 2))
    data = sample_v_from_v(teacher, falses(q, N, 10000); steps=1000)
    student = RBM(Potts(q, N), Gaussian(1), zeros(q, N, 1))
    nsamples = size(data)[end]
    epochs = training_epochs(; nsamples, batchsize, nupdates)
    initialize!(student, data)
    @test mean_from_inputs(student.visible) ≈ mean(data; dims=3)
    pcd!(student, data; epochs, batchsize)
    @info @test cor(free_energy(teacher, data), free_energy(student, data)) > 0.9
end

@testset "pcd -- teacher/student, Potts visible, Gaussian hidden, with weights" begin
    q = 2
    N = 12
    batchsize = 16
    nupdates = 10000
    teacher = RBM(Potts(q, N), Gaussian(1), zeros(q, N, 1))
    teacher.w[1,:,1] .= range(-0.4, 0.5, length=N)
    teacher.w[2,:,1] .= -teacher.w[1,:,1]
    data = extensive_sample(teacher.visible)
    wts = softmax(-free_energy(teacher, data))
    student = RBM(Potts(q, N), Gaussian(1), zeros(q, N, 1))
    nsamples = size(data)[end]
    epochs = training_epochs(; nsamples, batchsize, nupdates)
    initialize!(student, data)
    @test mean_from_inputs(student.visible) ≈ mean(data; dims=3)
    pcd!(student, data; wts, epochs, batchsize)
    @info @test cor(free_energy(teacher, data), free_energy(student, data)) > 0.9
end

@testset "pcd -- teacher/student, Potts, with weights and regularization, exact" begin
    Random.seed!(4)
    q = 2
    N = 5
    teacher = RBM(Potts(q, N), Spin(1), zerosum(1.5randn(q, N, 1)))
    data = extensive_sample(teacher.visible)
    wts_teacher = softmax(-free_energy(teacher, data))
    @test sum(wts_teacher) ≈ 1

    nupdates = 20000
    nsamples = size(data)[end]
    batchsize = nsamples
    epochs = training_epochs(; nsamples, batchsize, nupdates)
    student = RBM(Potts(q, N), Spin(1), zeros(q, N, 1))
    initialize!(student, data; wts = wts_teacher)
    @test mean_from_inputs(student.visible) ≈ wmean(data; wts = wts_teacher, dims=3)
    l2_weights = 0.05
    function callback(; rbm, kw...)
        rbm.visible.θ .= 0
        rbm.hidden.θ .= 0
    end
    pcd!(
        student, data;
        wts = wts_teacher, epochs, batchsize, mode=:exact, optim=Flux.AdaBelief(),
        l2_weights, zerosum=false, center=true, rescale=false, callback
    )

    wts_student = softmax(-free_energy(student, data))

    v_teacher = batchmean(student.visible, data; wts = wts_teacher)
    v_student = batchmean(student.visible, data; wts = wts_student)
    @info @test v_student ≈ v_teacher rtol=1e-3

    h_student = batchmean(student.hidden, mean_h_from_v(student, data); wts=wts_student)
    h_teacher = batchmean(student.hidden, mean_h_from_v(student, data); wts=wts_teacher)
    @info @test h_student ≈ h_teacher atol=1e-10

    data_flat = reshape(data, q*N, :)
    vh_teacher = data_flat * Diagonal(wts_teacher) * mean_h_from_v(student, data)' / sum(wts_teacher)
    vh_student = data_flat * Diagonal(wts_student) * mean_h_from_v(student, data)' / sum(wts_student)
    vh_student = reshape(vh_student, q, N, :)
    vh_teacher = reshape(vh_teacher, q, N, :)
    @info @test vh_student ≈ vh_teacher - l2_weights * student.w rtol=1e-4
end


text:
"""
Case 1:
This case study defines a classical machine against an infinite bus. The fault
drop a circuit on the (double circuit) line connecting the two buses, doubling its impedance
"""

##################################################
############### LOAD DATA ########################
##################################################

############### Data Network ########################

nodes_case1 = [PSY.Bus(1 , #number
                   "Bus 1", #Name
                   "REF" , #BusType (REF, PV, PQ)
                   0, #Angle in radians
                   1.05, #Voltage in pu
                   (min=0.94, max=1.06), #Voltage limits in pu
                   69), #Base voltage in kV
                   Bus(2 , "Bus 2"  , "PV" , 0 , 1.0 , (min=0.94, max=1.06), 69)]

branch_case1 = [PSY.Line("Line1", #name
                     true, #available
                     0.0, #active power flow initial condition (from-to)
                     0.0, #reactive power flow initial condition (from-to)
                     Arc(from=nodes_case1[1], to=nodes_case1[2]), #Connection between buses
                     0.01, #resistance in pu
                     0.05, #reactance in pu
                     (from=0.0, to=0.0), #susceptance in pu
                     18.046, #rate in MW
                     1.04)]  #angle limits (-min and max)

#Trip of a single circuit of Line 1 -> Resistance and Reactance doubled.
branch_case1_fault = [PSY.Line("Line1", #name
                           true, #available
                           0.0, #active power flow initial condition (from-to)
                           0.0, #reactive power flow initial condition (from-to)
                           Arc(from=nodes_case1[1], to=nodes_case1[2]), #Connection between buses
                           0.02, #resistance in pu
                           0.1, #reactance in pu
                           (from=0.0, to=0.0), #susceptance in pu
                           18.046, #rate in MW
                           1.04)]  #angle limits (-min and max)

loads_case1 = [PowerLoad("Bus1", true, nodes_case1[2], PowerSystems.ConstantPower, 0.3, 0.01, 0.3, 0.01)]

############### Data devices ########################

inf_gen_case1 = StaticSource(1, #number
                :InfBus, #name
                nodes_case1[1], #bus
                1.05, #VR
                0.0, #VI
                0.000001) #Xth

######## Machine Data #########

### Case 1: Classical machine against infinite bus ###
case1_machine = BaseMachine(0.0, #R
                            0.2995, #Xd_p
                            0.7087, #eq_p
                            100.0)  #MVABase

######## Shaft Data #########

### Shaft for Case 1 ###
case1_shaft = SingleMass(3.148, #H
                     2.0) #D

######## PSS Data #########
cases_no_pss = PSSFixed(0.0)

######## TG Data #########

### No TG for Cases 1, 2, 3, 4 ###
case1234_no_tg = TGFixed(1.0) #eff

########  AVR Data #########
case1_avr = AVRFixed(0.0) #Vf not applicable in Classic Machines

### Case 1 Generator ###
case1_gen = DynGenerator(1, #Number
                      :Case1Gen,
                      nodes_case1[2], #bus
                      1.0, # ω_ref,
                      1.0, #V_ref
                      0.5, #P_ref
                      case1_machine, #machine
                      case1_shaft, #shaft
                      case1_avr, #avr
                      case1234_no_tg, #tg
                      cases_no_pss) #pss

######################### Dynamical System ########################

case1_DynSystem = DynamicSystem(nodes_case1,
                              branch_case1,
                              [case1_gen],
                              vcat(inf_gen_case1,loads_case1),
                              100.0,
                              60.0)

##################################################
############### SOLVE PROBLEM ####################
##################################################

#Compute Y_bus after fault
Ybus_fault = PSY.Ybus(branch_case1_fault, nodes_case1)[:,:]

#Initialize variables
dx0 = zeros(LITS.get_total_rows(case1_DynSystem))
x0 = [1.05, #VR_1
      1.0, #VR_2
      0.0, #VI_1
      0.01, #VI_2
      0.2, #δ
      1.0] #ω
tspan = (0.0, 30.0);

#Find initial condition
inif! = (out,x) -> LITS.system_model!(out, dx0 ,x, (Ybus_fault,case1_DynSystem), 0.0)
sys_solve = nlsolve(inif!, x0)
x0_init = sys_solve.zero

#Define Fault using Callbacks
cb = DiffEqBase.DiscreteCallback(LITS.change_t_one, LITS.Y_change!)

#Define Simulation Problem
sim = DynamicSimulation(case1_DynSystem, tspan, Ybus_fault, cb, x0_init)

#Solve problem in equilibrium
run_simulation!(sim, IDA(), dtmax=0.02);

#Obtain data for angles
series = get_state_series(sim, (:Case1Gen, :δ));

@test sim.solution.retcode == :Success


text:
#!/usr/bin/env julia

using HTTP
using YAML
using LibGit2

function GetFileExtension(filename)
    index=findlast(isequal('.'),filename)
    if isnothing(index)
        return filename, ""
    elseif index == 1
        return "", filename
    else
        return filename[begin:index-1], filename[index:end]
    end
end

function ProcessScheme(tempdir, nix, scheme_name, scheme_repo)
    println("$scheme_name => $scheme_repo :")
    write(nix, "  $scheme_name = {\n")

    dir = Base.Filesystem.joinpath(tempdir, scheme_name)
    LibGit2.clone(scheme_repo, dir)
    names = Dict{String,Int}()
    for yaml_file in Base.Filesystem.readdir(dir; join=false, sort=false)
        base, ext = GetFileExtension(yaml_file)
        if !(ext == ".yaml" || ext == ".yml")
            continue
        end
        scheme = YAML.load_file(Base.Filesystem.joinpath(dir, yaml_file))
        name = scheme["scheme"]
        if haskey(names, name)
            count = names[name]
            names[name] = names[name] + 1
            name = "$name#$count"
        else
            names[name] = 1
        end
        println("  - $name : $yaml_file")

        write(nix, "    \"$name\" = {\n");
        for color in [ "base00" "base01" "base02" "base03" "base04" "base05" "base06" "base07"
                        "base08" "base09" "base0A" "base0B" "base0C" "base0D" "base0E" "base0F" ]
            r = scheme[color][1:2]
            g = scheme[color][3:4]
            b = scheme[color][5:6]
            write(nix, "      $color.hex = { r = \"$r\"; g = \"$g\"; b = \"$b\"; };\n")
        end
        write(nix, "    };\n")
    end

    Base.Filesystem.rm(dir; force=true, recursive=true)
    write(nix, "  };\n");
end

schemes = Dict{String,Dict{String,String}}()

open("schemes.nix", "w") do nix
    write(nix, "{\n");
    tempdir = Base.Filesystem.mktempdir(; prefix="nix_base16_", cleanup=true)
    print("Storing temporary files in ")
    println(tempdir)

    schemes_url = "https://raw.githubusercontent.com/chriskempson/base16-schemes-source/master/list.yaml"
    schemes_yaml = String(HTTP.request("GET", schemes_url).body)

    for (scheme_name,scheme_repo) in YAML.load(schemes_yaml)
        ProcessScheme(tempdir, nix, scheme_name, scheme_repo)
    end
    write(nix, "}\n");
end


text:
include("typestest.jl")

kf0 = copy(kf)

kf1 = predict(kf)
predict!(kf)

@test kf1.x.x == kf.x.x
@test_approx_eq kf.x.x [1.0,-0.01]
@test kf1.x.p == kf.x.p

y = Observation([1.01])

kf2 = update(kf,y)
update!(kf,y)

@test kf2.x.x == kf.x.x
@test kf2.x.p == kf.x.p


text:
include("blosc1.jl")
include("context.jl")
include("superchunk.jl")
include("metalayers.jl")


text:
using DataFrames

"""
    _load_dataframe(m::AbstractModel, comp_name::Symbol, item_name::Symbol), df::Union{Nothing,DataFrame}=nothing)

Load a DataFrame from the variable or parameter `item_name` in component `comp_name`. If `df` is
nothing, a new DataFrame is allocated. Returns the populated DataFrame.
"""
function _load_dataframe(m::AbstractModel, comp_name::Symbol, item_name::Symbol, df::Union{Nothing,DataFrame}=nothing)
    md, mi = m isa MarginalModel ? (m.base.md, m.base.mi) : (m.md, m.mi)

    dims = dim_names(m, comp_name, item_name)

    # Create a new df if one was not passed in
    df = df === nothing ? DataFrame() : df

    if hasproperty(df, item_name)
        error("An item named $item_name already exists in this DataFrame")
    end

    num_dims = length(dims)
    if num_dims == 0
        error("Cannot create a dataframe for a scalar parameter :$item_name")
    end

    data = deepcopy(m[comp_name, item_name])

    if num_dims == 1
        dim1name = dims[1]
        dim1 = dimension(md, dim1name)
        df[!, dim1name] = collect(keys(dim1))
        # @info "dim: $dim1name size(df): $(size(df))"

        # df[item_name] = data

        if dim1name == :time && size(data)[1] != length(time_labels(md))
            ci = compinstance(mi, comp_name)
            first = dim1[ci.first]  # Dimension converts year key to index in array
            last  = dim1[ci.last]

            # Pad the array with NaNs outside this component's bounds
            shifted_data = vcat(repeat([missing], inner=first - 1),
                                data[1:(last-first+1)], # ignore padding after these values
                                repeat([missing], inner=length(dim1) - last))
            # @info "len shifted: $(length(shifted_data))"
            df[!, item_name] = shifted_data
        else
            df[!, item_name] = data
        end
    else
        df = _df_helper(m, comp_name, item_name, dims, data)
    end

    return df
end

function _df_helper(m::AbstractModel, comp_name::Symbol, item_name::Symbol, dims::Vector{Symbol}, data::AbstractArray)
    md, mi = m isa MarginalModel ? (m.base.md, m.base.mi) : (m.md, m.mi)
    num_dims = length(dims)

    dim1name = dims[1]
    dim1 = dimension(md, dim1name)
    keys1 = collect(keys(dim1))
    len_dim1 = length(dim1)

    df = DataFrame()

    if num_dims == 2
        dim2name = dims[2]
        dim2 = dimension(md, dim2name)
        keys2 = collect(keys(dim2))
        len_dim2 = length(dim2)

        df[!, dim1name] = repeat(keys1, inner = [len_dim2])
        df[!, dim2name] = repeat(keys2, outer = [len_dim1])

        if dim1name == :time && size(data)[1] != len_dim1
            ci = compinstance(mi, comp_name)
            t = dimension(m, :time)
            first = t[ci.first]
            last  = t[ci.last]

            top    = fill(missing, first - 1, len_dim2)
            bottom = fill(missing, len_dim1 - last, len_dim2)
            data = vcat(top, data, bottom)
        end

        df[!, item_name] = cat([vec(data[i, :]) for i = 1:len_dim1]...; dims=1)
    else

        # shift the data to be padded with missings if this data is shorter than the model
        if dim1name == :time && size(data)[1] != len_dim1
            ci = compinstance(mi, comp_name)
            t = dimension(m, :time)
            first = t[ci.first]
            last  = t[ci.last]

            rest_dim_lens = [length(dimension(md, dimname)) for dimname in dims[2:end]]
            top = fill(missing, first - 1, rest_dim_lens...)
            bottom = fill(missing, len_dim1 - last, rest_dim_lens...)

            data = vcat(top, data, bottom)
        end

        # Indexes is #, :, :, ... for each index of first dimension
        indexes = repeat(Any[Colon()], num_dims)

        for i in 1:size(data)[1]
            indexes[1] = i
            subdf = _df_helper(m, comp_name, item_name, dims[2:end], data[indexes...])
            subdf[!, dims[1]] .= keys1[i]

            if i == 1
                # add required columns in the first iteration
                df_names = names(df)
                for name in names(subdf)
                    if ! (name in df_names)
                        df[!, name] = []
                    end
                end
            end
            df = vcat(df, subdf)
        end
    end

    return df
end

"""
    getdataframe(m::AbstractModel, comp_name::Symbol, pairs::Pair{Symbol, Symbol}...)

Return a DataFrame with values for the given variables or parameters of model `m`
indicated by `pairs`, where each pair is of the form `comp_name => item_name`.
If more than one pair is provided, all must refer to items with the same
dimensions, which are used to join the respective item values.
"""
function getdataframe(m::AbstractModel, pairs::Pair{Symbol, Symbol}...)
    (comp_name1, item_name1) = pairs[1]
    dims = dim_names(m, comp_name1, item_name1)
    df = getdataframe(m, comp_name1, item_name1)

    for (comp_name, item_name) in pairs[2:end]
        next_dims = dim_names(m, comp_name, item_name)
        if dims != next_dims
            error("Can't create DataFrame from items with different dimensions ($comp_name1.$item_name1: $dims vs $comp_name.$item_name: $next_dims)")
        end
        result = getdataframe(m, comp_name, item_name)
        df = hcat(df, result[!, [item_name]])      # [[xx]] retrieves a 1 column DataFrame
    end

    return df
end

"""
    getdataframe(m::AbstractModel, pair::Pair{Symbol, NTuple{N, Symbol}})

Return a DataFrame with values for the given variables or parameters
indicated by `pairs`, where each pair is of the form `comp_name => item_name`.
If more than one pair is provided, all must refer to items with the same
dimensions, which are used to join the respective item values.
"""
function getdataframe(m::AbstractModel, pair::Pair{Symbol, NTuple{N, Symbol}}) where N
    comp_name = pair.first
    expanded = [comp_name => param_name for param_name in pair.second]
    return getdataframe(m, expanded...)
end

"""
    getdataframe(m::AbstractModel, comp_name::Symbol, item_name::Symbol)

Return the values for variable or parameter `item_name` in `comp_name` of
model `m` as a DataFrame.
"""
function getdataframe(m::AbstractModel, comp_name::Symbol, item_name::Symbol)
    if ! is_built(m)
        error("Cannot get DataFrame: model has not been built yet.")
    end

    df = _load_dataframe(m, comp_name, item_name)
    return df
end