text:
import Base: getindex, getproperty
using Base.Threads

function getindex(g::StorageGraph, v::Integer)
    get_prop(g, v)
end

function getindex(g::StorageGraph, dep::Pair)
    (lastn, cpaths), t = @timed walkdep(g, dep)
    @debug "Finding last node on path took $t seconds."
    if lastn == endof(dep) && length(cpaths) ≠ 0
        neighbors = outneighbors(g, lastn)
        mask = Vector{Bool}(undef, length(neighbors))
        @threads for i in eachindex(neighbors)
            mask[i] = on_path(g, neighbors[i], cpaths)
        end
        return neighbors[mask]
    else
        return eltype(g)[]
    end
end

function getindex(g::StorageGraph, name::Symbol, dep::Pair)
    get.(get_prop.(Ref(g), g[dep]), name, nothing)
end

function getindex(g::StorageGraph, conditions::Dict{Symbol,F}, nodes::Vararg{NamedTuple}) where {F<:Function}
    paths = filter_paths(g, nodes, conditions)
    vs = Iterators.filter(v->on_path(g,v,paths), outneighbors(g, nodes[end]))
    get_prop.(Ref(g), vs)
end

function getindex(g::StorageGraph, name::Symbol, conditions::Dict{Symbol,F}, nodes::Vararg{NamedTuple}) where {F<:Function}
    paths = filter_paths(g, nodes, conditions)
    outn = outneighbors(g, g[nodes[end]])
    i = findfirst(n->has_prop(g, n, name), outn)
    if i === nothing
        vs = walkpath(g, paths, g[nodes[1]], stopcond=(g,v)->has_prop(g,v,name))
        # filter out the cases where it stopped before reaching stopcond
        filter!(v->has_prop(g,v,name), vs)
    else
        # there is no need to traverse the graph if the last given node has the
        # desired nodes as outneighbors
        vs = filter(v->on_path(g,v,paths), outn)
    end
    get.(get_prop.(Ref(g), vs), name, nothing)
end

function getindex(g::StorageGraph, nodes::Vararg{NamedTuple})
    return getindex(g, Dict{Symbol,Function}(), nodes...)
end

function getindex(g::StorageGraph, data::NamedTuple)
    return g.index[data]
end

function getindex(g::StorageGraph, name::Symbol, nodes::Vararg{NamedTuple})
    getindex(g, name, Dict{Symbol,Function}(), nodes...)
end

function getindex(g::StorageGraph, names::NTuple{N, Symbol}, nodes::Vararg{NamedTuple}) where {N}
    (getindex(g, n, nodes...) for n in names)
end

function getindex(g::StorageGraph, name::Symbol)
    extractvals(findnodes(g, name), name)
end

getindex(g::StorageGraph) = []

"""
    paths_through(g, v::Integer; dir=:out)

Return a vector of the paths going through the given vertex. If `dir` is specified,
use the corresponding edge direction (`:in` and `:out` are acceptable values).
"""
paths_through(g, v; dir=:out) = paths_through!(Set{eltype(g)}(), g, v, dir=dir)

function paths_through!(paths, g, v::Integer; dir=:out)
    v == 0 && return paths
    if dir == :out
        out = outneighbors(g, v)
        if isempty(out)
            return paths
        else
            es = [Edge(v, i) for i in out]
        end
    elseif dir == :in
        in = inneighbors(g, v)
        if isempty(in)
            return paths
        else
            es = [Edge(i, v) for i in in]
        end
    else
        in = inneighbors(g, v)
        out = outneighbors(g, v)
        l = length(in)
        n = l + length(out)
        n == 0 && return paths
        es = Vector{Edge}(undef, n)
        for i in eachindex(in)
            es[i] = Edge(in[i], v)
        end
        for i in eachindex(out)
            es[i+l] = Edge(v, out[i])
        end
    end
    for e in es
        union!(paths, g.paths[e])
    end
    return paths
end

function paths_through!(paths, g, dep::Pair; dir=:out)
    intersect!(paths_through!(Set{eltype(g)}(), g, dep[2], dir=dir), paths_through!(paths, g, dep[1], dir=dir))
end

function paths_through!(paths, g, node::NamedTuple; dir=:out)
    !haskey(g.index, node) && return paths
    paths_through!(paths, g, g[node], dir=dir)
end

"""
    final_neighborhs(g, dep::Pair; dir=:out)

Return the vertex indices for the neighbors at the end of the dependency chain.
Note: this assumes that the dependency chain is valid (all the nodes exist).
"""
function final_neighborhs(g, dep::Pair; dir=:out)
    v = g[endof(dep)]
    dir == :out ? outneighbors(g, v) : inneighbors(g, v)
end

"""
    findnodes(g, name::Symbol)

Finds the nodes containing `name`.
"""
function findnodes(g, name::Symbol)
    findall(v -> has_prop(g, v, name), g.index)
end

"""
    extractvals(nodes, name::Symbol)

Return an array of values corresponding to `name` form the array of `NamedTuple`s
`nodes`.
"""
function extractvals(nodes, name::Symbol)
    [n[name] for n in nodes]
end

function filter_paths!(paths, g, conditions)
    for (name, cond) in conditions
        target = findnodes(g, name)
        valid_vals = target[cond.(target)]
        all_valid_paths = paths_through.(Ref(g), valid_vals)
        valid_paths = Set{eltype(g)}()
        for p in all_valid_paths
            union!(valid_paths, p)
        end
        intersect!(paths, valid_paths)
    end
end

function filter_paths(g, nodes, conditions)
    possible_paths = paths_through.(Ref(g), nodes)
    paths = possible_paths[1]
    for p in possible_paths
        intersect!(paths, p)
    end
    filter_paths!(paths, g, conditions)

    return paths
end


text:
ENV["GKSwstype"] = 322 # workaround for gr segfault on GH actions
# ENV["GKS_WSTYPE"]=100 # try this if above does not work
using Documenter, RobustAndOptimalControl

using Plots
gr()

makedocs(
      sitename = "RobustAndOptimalControl Documentation",
      doctest = true,
      modules = [RobustAndOptimalControl],
      pages = [
            "Home" => "index.md",
            "Uncertainty modeling" => "uncertainty.md",
      ],
      format = Documenter.HTML(prettyurls = haskey(ENV, "CI")),
)

deploydocs(
      deps = Deps.pip("pygments", "mkdocs", "python-markdown-math", "mkdocs-cinder"),
      repo = "github.com/JuliaControl/RobustAndOptimalControl.jl.git",
)

ext: jl


text:
function calculate_ideal_deltas(Q::Matrix{WEIGHT}, y::OPINIONS, α::PSYCHOLOGICAL_FACTOR, budget::WEIGHT)

    (n, m) = size(Q)
    T = Q * diagm(α)
    new_vals = fill(WEIGHT(0.0), n)
    counts = fill(WEIGHT(0.0), n)
    for i = 1:n
        if y[i] != 0.0
            @simd for j = 1:m
                if (T[i,j] > 0) && (i != j)
                    counts[j] += 1
                    numer = -((T[i, :]' * y) - (T[i, j] * y[j]))
                    denom = T[i,j]
                    new_vals[j] += (numer / denom)
                    @show j
                    @show new_vals[j]
                end
            end
        end
    end

    @simd for j = 1:n
        if counts[j] > 0.0
            new_vals[j] = new_vals[j] / counts[j]
        else
            new_vals[j] = y[j]
        end
    end

    δ = new_vals - y
    cost = sum(δ .* δ)

end

max_stars_repo_name: InzamamRahaman/SignedDegroot

max_issues_repo_name: InzamamRahaman/SignedDegroot

max_forks_repo_name: InzamamRahaman/SignedDegroot


text:
_default_onerr(path, err) = (err isa InterruptException) ? rethrow(err) : nothing
_default_keepout(dir) = false


text:
###########################################################
#
#   (non-weighted) sampling
#
###########################################################

### Algorithms for sampling with replacement

# Direct sampling
#
#   for each i, draw an index from 1:length(a), and
#   pick the corresponding element from a.
#
#   x[i] and x[j] are independent, when i != j
#
function direct_sample!(a::UnitRange, x::AbstractArray)
    s = RandIntSampler(length(a))
    b = a[1] - 1
    if b == 0
        for i = 1:length(x)
            @inbounds x[i] = rand(s)
        end
    else
        for i = 1:length(x)
            @inbounds x[i] = b + rand(s)
        end
    end
    return x
end

function direct_sample!(a::AbstractArray, x::AbstractArray)
    s = RandIntSampler(length(a))
    for i = 1:length(x)
        @inbounds x[i] = a[rand(s)]
    end
    return x
end

# Expanded Multinomial sampling
#
#   for each element in a, we draw the number of its
#   occurrences, and fill this element to x for
#   this number of times.
#
function xmultinom_sample!(a::AbstractArray, x::AbstractArray)
    n = length(a)
    k = length(x)
    offset = 0
    i = 1
    while offset < k
        rk = k - offset
        if i == n
            @inbounds ai = a[i]
            for j = 1:rk
                @inbounds x[offset + j] = ai
            end
            offset = k
        else
            m = rand_binom(rk, 1.0 / (n - i + 1))
            if m > 0
                @inbounds ai = a[i]
                for j = 1:m
                    @inbounds x[offset + j] = a[i]
                end
                offset += m
            end
            i += 1
        end
    end
    x
end

### draw a pair of distinct integers in [1:n]

function samplepair(n::Int)
    i1 = randi(n)
    i2 = randi(n-1)
    return (i1, ifelse(i2 == i1, n, i2))
end

function samplepair(a::AbstractArray)
    i1, i2 = samplepair(length(a))
    return a[i1], a[i2]
end

### Algorithm for sampling without replacement

# Knuth's Algorithm S
#
#   Reference: The Art of Computer Programming, Vol 2, 3.4.2, p.142
#
function knuths_sample!(a::AbstractArray, x::AbstractArray; initshuffle::Bool=true)
    n = length(a)
    k = length(x)
    k <= n || error("length(x) should not exceed length(a)")

    # initialize
    for i = 1:k
        @inbounds x[i] = a[i]
    end
    if initshuffle
        @inbounds for j = 1:k
            l = randi(j, k)
            if l != j
                t = x[j]
                x[j] = x[l]
                x[l] = t
            end
        end
    end

    # scan remaining
    s = RandIntSampler(k)
    for i = k+1:n
        if rand() * i < k  # keep it with probability k / i
            @inbounds x[rand(s)] = a[i]
        end
    end
    return x
end

# Fisher-Yates sampling
#
#   create an array of index inds = [1:n]
#
#   for i = 1:k
#       swap inds[i] with a random one in inds[i:n]
#       set x[i] = a[inds[i]]
#   end
#
#   O(n) for initialization + O(k) for random shuffling
#
function fisher_yates_sample!(a::AbstractArray, x::AbstractArray)
    n = length(a)
    k = length(x)
    k <= n || error("length(x) should not exceed length(a)")

    inds = Array(Int, n)
    for i = 1:n
        @inbounds inds[i] = i
    end

    @inbounds for i = 1:k
        j = randi(i, n)
        t = inds[j]
        inds[j] = inds[i]
        inds[i] = t
        x[i] = a[t]
    end
    return x
end

# Self-avoid sampling
#
#   Use a set to maintain the index that has been sampled,
#
#   each time draw a new index, if the index has already
#   been sampled, redraw until it draws an unsampled one.
#
function self_avoid_sample!(a::AbstractArray, x::AbstractArray)
    n = length(a)
    k = length(x)
    k <= n || error("length(x) should not exceed length(a)")

    s = Set{Int}()
    sizehint(s, k)
    rgen = RandIntSampler(n)

    # first one
    idx = rand(rgen)
    x[1] = a[idx]
    push!(s, idx)

    # remaining
    for i = 2:k
        idx = rand(rgen)
        while idx in s
            idx = rand(rgen)
        end
        x[i] = a[idx]
        push!(s, idx)
    end
    return x
end

# Random subsequence sampling
#
#  References:
#
#   Jeffrey Scott Vitter.
#   "Faster Methods for Random Sampling"
#   Communications of the ACM, 27 (7), July 1984
#
#
# This paper presents three algorithms, respectively
# named Algorithm A, C, and D.
#
# These algorithms are for sequential sampling, and the
# outputs are ordered. They are implemented below
#

## Algorithm A (page 714)
#
#  Require O(n) random numbers.
#
function seqsample_a!(a::AbstractArray, x::AbstractArray)
    n = length(a)
    k = length(x)
    k <= n || error("length(x) should not exceed length(a)")

    i = 0
    j = 0
    while k > 1
        u = rand()
        q = (n - k) / n
        while q > u  # skip
            i += 1
            n -= 1
            q *= (n - k) / n
        end
        @inbounds x[j+=1] = a[i+=1]
        n -= 1
        k -= 1
    end

    if k > 0  # checking k > 0 is necessary: x can be empty
        s = itrunc(n * rand())
        x[j+1] = a[i+(s+1)]
    end
    return x
end

## Algorithm C (page 715)
#
#  Require O(k^2) random numbers
#
function seqsample_c!(a::AbstractArray, x::AbstractArray)
    n = length(a)
    k = length(x)
    k <= n || error("length(x) should not exceed length(a)")

    i = 0
    j = 0
    while k > 1
        l = n - k + 1
        minv = l
        u = n
        while u >= l
            v = u * rand()
            if v < minv
                minv = v
            end
            u -= 1
        end
        s = itrunc(minv) + 1
        x[j+=1] = a[i+=s]
        n -= s
        k -= 1
    end

    if k > 0
        s = itrunc(n * rand())
        x[j+1] = a[i+(s+1)]
    end
    return x
end

## TODO: implement Algorithm D (page 716 - 717)

### Interface functions (poly-algorithms)

sample(a::AbstractArray) = a[randi(length(a))]

function sample!(a::AbstractArray, x::AbstractArray; replace::Bool=true, ordered::Bool=false)
    n = length(a)
    k = length(x)
    k == 0 && return x

    if replace  # with replacement
        if ordered
            if k > 10 * n
                xmultinom_sample!(a, x)
            else
                sort!(direct_sample!(a, x))
            end
        else
            direct_sample!(a, x)
        end

    else  # without replacement
        k <= n || error("Cannot draw more samples without replacement.")

        if ordered
            if n > 10 * k * k
                seqsample_c!(a, x)
            else
                seqsample_a!(a, x)
            end
        else
            if k == 1
                @inbounds x[1] = sample(a)
            elseif k == 2
                @inbounds (x[1], x[2]) = samplepair(a)
            elseif n < k * 24
                fisher_yates_sample!(a, x)
            else
                self_avoid_sample!(a, x)
            end
        end
    end
    return x
end

function sample{T}(a::AbstractArray{T}, n::Integer; replace::Bool=true, ordered::Bool=false)
    sample!(a, Array(T, n); replace=replace, ordered=ordered)
end

function sample{T}(a::AbstractArray{T}, dims::Dims; replace::Bool=true, ordered::Bool=false)
    sample!(a, Array(T, dims); replace=replace, ordered=ordered)
end

################################################################
#
#  Weighted sampling
#
################################################################

function sample(wv::WeightVec)
    t = rand() * sum(wv)
    w = values(wv)
    n = length(w)
    i = 1
    cw = w[1]
    while cw < t && i < n
        i += 1
        @inbounds cw += w[i]
    end
    return i
end

sample(a::AbstractArray, wv::WeightVec) = a[sample(wv)]

function direct_sample!(a::AbstractArray, wv::WeightVec, x::AbstractArray)
    n = length(a)
    length(wv) == n || throw(DimensionMismatch("Inconsistent lengths."))
    for i = 1:length(x)
        x[i] = a[sample(wv)]
    end
    return x
end

function make_alias_table!(w::AbstractVector{Float64}, wsum::Float64,
                           a::AbstractVector{Float64},
                           alias::AbstractVector{Int})
    # Arguments:
    #
    #   w [in]:         input weights
    #   wsum [in]:      pre-computed sum(w)
    #
    #   a [out]:        acceptance probabilities
    #   alias [out]:    alias table
    #
    # Note: a and w can be the same way, then that away will be
    #       overriden inplace by acceptance probabilities
    #
    # Returns nothing
    #

    n = length(w)
    length(a) == length(alias) == n ||
        throw(DimensionMismatch("Inconsistent array lengths."))

    ac = n / wsum
    for i = 1:n
        @inbounds a[i] = w[i] * ac
    end

    larges = Array(Int, n)
    smalls = Array(Int, n)
    kl = 0  # actual number of larges
    ks = 0  # actual number of smalls

    for i = 1:n
        @inbounds ai = a[i]
        if ai > 1.0
            larges[kl+=1] = i  # push to larges
        elseif ai < 1.0
            smalls[ks+=1] = i  # push to smalls
        end
    end

    while kl > 0 && ks > 0
        s = smalls[ks]; ks -= 1  # pop from smalls
        l = larges[kl]; kl -= 1  # pop from larges
        @inbounds alias[s] = l
        @inbounds al = a[l] = (a[l] - 1.0) + a[s]
        if al > 1.0
            larges[kl+=1] = l  # push to larges
        else
            smalls[ks+=1] = l  # push to smalls
        end
    end

    # this loop should be redundant, except for rounding
    for i = 1:ks
        @inbounds a[smalls[i]] = 1.0
    end
    nothing
end

function alias_sample!(a::AbstractArray, wv::WeightVec, x::AbstractArray)
    n = length(a)
    length(wv) == n || throw(DimensionMismatch("Inconsistent lengths."))

    # create alias table
    ap = Array(Float64, n)
    alias = Array(Int, n)
    make_alias_table!(values(wv), sum(wv), ap, alias)

    # sampling
    s = RandIntSampler(n)
    for i = 1:length(x)
        j = rand(s)
        x[i] = rand() < ap[j] ? a[j] : a[alias[j]]
    end
    return x
end

function xmultinom_sample!(a::AbstractArray, wv::WeightVec, x::AbstractArray)
    n = length(a)
    length(wv) == n || throw(DimensionMismatch("Inconsistent lengths."))
    k = length(x)
    offset = 0
    i = 1
    wsum = sum(wv)
    w = values(wv)

    while offset < k
        rk = k - offset
        wi = w[i]

        if i == n || wi >= wsum
            for j = 1 : rk
                @inbounds x[offset + j] = a[i]
            end
            offset = k
        else
            m = rand_binom(rk, wi / wsum)
            for j = 1 : m
                @inbounds x[offset + j] = a[i]
            end
            i += 1
            wsum -= wi
            offset += m
        end
    end
    return x
end

function naive_wsample_norep!(a::AbstractArray, wv::WeightVec, x::AbstractArray)
    n = length(a)
    length(wv) == n || throw(DimensionMismatch("Inconsistent lengths."))
    k = length(x)

    w = Array(Float64, n)
    copy!(w, values(wv))
    wsum = sum(wv)

    for i = 1:k
        u = rand() * wsum
        j = 1
        c = w[1]
        while c < u && j < n
            @inbounds c += w[j+=1]
        end
        @inbounds x[i] = a[j]

        @inbounds wsum -= w[j]
        @inbounds w[j] = 0.0
    end
    return x
end

function sample!(a::AbstractArray, wv::WeightVec, x::AbstractArray;
                 replace::Bool=true, ordered::Bool=false)
    n = length(a)
    k = length(x)

    if replace
        if ordered
            if k < n && k < 256
                sort!(direct_sample!(a, wv, x))
            else
                xmultinom_sample!(a, wv, x)
            end
        else
            if k > 3 * n
                shuffle!(xmultinom_sample!(a, wv, x))
            else
                if n < 40
                    direct_sample!(a, wv, x)
                else
                    t = ifelse(n < 500, 64, 32)
                    if k < t
                        direct_sample!(a, wv, x)
                    else
                        alias_sample!(a, wv, x)
                    end
                end
            end
        end
    else
        naive_wsample_norep!(a, wv, x)
        if ordered
            sort!(x)
        end
    end
    return x
end

sample{T}(a::AbstractArray{T}, wv::WeightVec, n::Integer; replace::Bool=true, ordered::Bool=false) =
    sample!(a, wv, Array(T, n); replace=replace, ordered=ordered)

sample{T}(a::AbstractArray{T}, wv::WeightVec, dims::Dims; replace::Bool=true, ordered::Bool=false) =
    sample!(a, wv, Array(T, dims); replace=replace, ordered=ordered)

# wsample interface

wsample(w::RealVector) = sample(weights(w))
wsample(a::AbstractArray, w::RealVector) = sample(a, weights(w))

wsample!(a::AbstractArray, w::RealVector, x::AbstractArray; replace::Bool=true, ordered::Bool=false) =
    sample!(a, weights(w), x; replace=replace, ordered=ordered)

wsample{T}(a::AbstractArray{T}, w::RealVector, n::Integer; replace::Bool=true, ordered::Bool=false) =
    wsample!(a, w, Array(T, n); replace=replace, ordered=ordered)

wsample{T}(a::AbstractArray{T}, w::RealVector, dims::Dims; replace::Bool=true, ordered::Bool=false) =
    wsample!(a, w, Array(T, dims); replace=replace, ordered=ordered)


text:
using jInv.Mesh
using jInv.LinearSolvers
using Test
using KrylovMethods
using LinearAlgebra
println("===  Example 2D DivSigGrad ====");

domain = [0.0, 1.0, 0.0, 1.0];
n      = [50,50];
Mr     = getRegularMesh(domain,n)
G      = getNodalGradientMatrix(Mr);
m      = sparse(Diagonal(exp.(randn(size(G,1)))));
Ar     = G'*m*G;
Ar     = Ar + 1e-1*norm(Ar,1)*sparse(1.0I,size(Ar,2),size(Ar,2));
N      = size(Ar,2);
B      = Ar*rand(N,4);

sPCG   = getBlockIterativeSolver(KrylovMethods.blockCG,out=1,sym=1);
X,  = solveLinearSystem(Ar,B,sPCG);
@test norm(Ar*X-B)/norm(B) < sPCG.tol

IterMethod = (A,B;M=M,X=X,tol=1e-1,maxIter=10,out=-1) ->
                  blockBiCGSTB(A,B,M1=M,tol=tol,x=X,maxIter=maxIter,out=out)
sBiCG = getBlockIterativeSolver(IterMethod,out=1);
X, = solveLinearSystem(Ar,B,sBiCG);
@test norm(Ar*X-B)/norm(B) < sBiCG.tol

println("===  Test for nonsymmetric matrices ====");
n = 100
m = 4
sBiCG.PC = :jac
sBiCG.doClear=true
sBiCG.out=0
A = sprandn(n,n,5/n) + 10*sparse(1.0I,n,n)
B = randn(n,m)
X, = solveLinearSystem(A,B,sBiCG)
@test norm(A*X - B)/norm(B) < sBiCG.tol

sBiCG.isTranspose=true
Xt, = solveLinearSystem(sparse(A'),B,sBiCG)
@test norm(A*X - B)/norm(B) < sBiCG.tol
@test norm(Xt - X)/norm(X) < 1e-13


text:
# functions for getting data into H2O

function upload_file(path;
                     destination_frame = "",
                     parse = true,
                     header = null)

    this
    that
    other

end

h2o.uploadFile(path, destination_frame = "", parse = TRUE, header = NA,
sep = "", col.names = NULL, col.types = NULL, na.strings = NULL,
progressBar = FALSE, parse_type = NULL)

function create_frame(;rows = 10000,
                     cols = 1000,
                     randomize = true,
                     value = 0,
                     real_range = 100,
                     categorical_fraction = 0.2,
                     factors = 100,
                     integer_range = 100,
                     binary_fraction = 0.1,
                     binary_ones_fraction = 0.02,
                     missing_fraction = 0.01,
                     response_factors = 2,
                     has_response = false,
                     seed = 0,
                     key = 0)

  results = h2o_post("CreateFrame",
                    "rows=$(rows)",
                    "cols=$(cols)",
                    "randomize=$(randomize)",
                    "value=$(value)",
                    "real_range=$(real_range)",
                    "categorical_fraction=$(categorical_fraction)",
                    "factors=$(factors)",
                    "integer_range=$(integer_range)",
                    "binary_fraction=$(binary_fraction)",
                    "binary_ones_fraction=$(binary_ones_fraction)",
                    "missing_fraction=$(missing_fraction)",
                    "response_factors=$(response_factors)",
                    "has_response=$(has_response)")
                    # "seed=$(seed)",
                    # "key=$(key)")
  return results
end


text:
## VISIBLE SECTION
#

function flock!(Ba, Bb; W, nfun=neighbors_n, tau=0.1)
  Ba0 = deepcopy(Ba)
  Bb0 = deepcopy(Bb)
  @simd for bi in Ba
    deli = neighbors(bi, Ba0)

    si = seperate(bi, deli)
    ai = align(bi, deli)
    ci = cohere(bi, deli)
    ev = evade(bi, Bb0, c=tau)
    bp = bound_position(bi)

    bi.vel += bi.ws*si + bi.wa*ai + bi.wc*ci + bi.we*ev + bi.wb*bp
    limit_velocity!(bi)
    bi.pos += tau*bi.vel
  end
  @simd for bi in Bb
    si = seperate(bi, Bb0)
    pu = pursuit(bi, Ba0, c=tau)
    bp = bound_position(bi)

    bi.vel += bi.ws*si + bi.wp*pu + bi.wb*bp
    limit_velocity!(bi)
    bi.pos += tau*bi.vel
  end
end

function pursuit(bi, O; c)
  j, cg = 0, zeros(dim(bi))
  for bj in O
    if dis(bi, bj) < bi.dp
      cg += bj.pos + c*dis(bi, bj)*bj.vel
      j += 1
    end
  end
  if j != 0
    pit = cg/j - bi.pos
    dv = (pit/norm(pit))*bi.vmax - bi.vel
    return dv
  else
    dv = cg #zeros(dim(bi))
    return dv
  end
end

function evade(bi, O; c)
  bj = partialsort(collect(O), 1, by = bj -> dis(bi, bj))
  if dis(bi, bj) < bi.de
    pit = (bj.pos + c*dis(bi, bj)*bj.vel) - bi.pos
    dv = -(pit/norm(pit))*bi.vmax - bi.vel
    return dv
  else
    dv = zeros(dim(bi))
    return dv
  end
end

## HIDDEN SECTION
#

function initialize_pureva(;
  dim=2, nbb=2,
  nba=100, dv=5., av=3.05, ds=dv/2., nc=4, de=3*dv, dp=3*dv,
  w0=1., ws=.1, wa=.05, wc=.01, ww=0., wb=.1, wp=.01, we=.1,
  vmax=5., vmin=.4*vmax,
  bmin=-(200/dim^2)*ones(dim), bmax=(200/dim^2)*ones(dim)
)
  Random.seed!(0)
  pfun = x -> bmin + x .* (bmax .- bmin)
  vfun = x -> (x ./ norm(x)) .* vmax / 2
  Bfun = nb -> [Boid(i, pfun(rand(dim)), vfun(rand(dim)*2 - ones(dim)),
    dv, av, ds, w0, ws, wa, wc, vmax, vmin, bmin, bmax, ww, wb, wp, we, nc, dp, de) for i = 1:nb]
  if dim == 2
    Wxy = B -> Node(collect(Point2f0(b.pos) for b in B))
  else
    Wxy = B -> Node(collect(Point3f0(b.pos) for b in B))
  end
  B = map(Bfun, [nba, nbb])
  W = World()
  W.xy = map(Wxy, B)
  W.sc = scatter(W.xy[1])
  scatter!(W.xy[2], color = "red")
  return (B[1], B[2], W)
end

function draw!(W, Ba, Bb; t, dim=2, fps=60)
  if dim == 2
    Wxy = B -> collect(Point2f0(b.pos) for b in B)
  else
    Wxy = B -> collect(Point3f0(b.pos) for b in B)
  end
  W.xy[1][] = Wxy(Ba)
  W.xy[2][] = Wxy(Bb)
  dt = time() - t
  if dt > 1/fps
    sleep(1/fps)
    return time()
  else
    sleep(1/fps - dt)
    return time()
  end
  #sleep(1/fps)
  #return time()
end

function run_pureva(; nba=100, nbb=2, nfun=neighbors_n, nmax=1000, dim=2)
  Ba, Bb, W = initialize_pureva(nba=nba, nbb=nbb, dim=dim)
  t = time()
  display(W.sc)

  try
    @fastmath for _ = 1:nmax
      t = draw!(W, Ba, Bb, t=t, dim=dim)
      #stat!(W, Ba, Bb)
      flock!(Ba, Bb, W=W, nfun=nfun)
    end
  catch InterruptException
    println("\nSimulation interrupted.")
  end
end


text:
using Lsode
@static if VERSION < v"0.7.0-DEV.2005"
    using Base.Test
else
    using Test
end

# write your own tests here
@testset "Basic Tests" begin
    include("basic_tests.jl")
end
@testset "chem2diff Tests" begin
    include("chem2diff_test.jl")
end


text:
using Test, InteractiveDynamics, Observables

# include("od_test.jl")
# include("psos_test.jl")


text:
using ClimateMachine
using ClimateMachine.SystemSolvers
using ClimateMachine.ODESolvers
using ClimateMachine.MPIStateArrays
using ClimateMachine.SingleStackUtils
using ClimateMachine.Checkpoint
using ClimateMachine.BalanceLaws: vars_state
using JLD2, FileIO
const clima_dir = dirname(dirname(pathof(ClimateMachine)));

include(joinpath(clima_dir, "experiments", "AtmosLES", "stable_bl_model.jl"))
include("edmf_model.jl")
include("edmf_kernels.jl")

function main(::Type{FT}, cl_args) where {FT}

    surface_flux = cl_args["surface_flux"]

    # DG polynomial order
    N = 1
    nelem_vert = 50

    # Prescribe domain parameters
    zmax = FT(400)

    t0 = FT(0)

    # Simulation time
    timeend = FT(3600 * 6)
    CFLmax = FT(40.0)

    config_type = SingleStackConfigType

    ode_solver_type = ClimateMachine.ExplicitSolverType(
        solver_method = LSRK144NiegemannDiehlBusch,
    )

    N_updrafts = 1
    N_quad = 3 # Using N_quad = 1 leads to norm(Q) = NaN at init.
    turbconv = NoTurbConv()

    C_smag = FT(0.23)

    model = stable_bl_model(
        FT,
        config_type,
        zmax,
        surface_flux;
        turbulence = SmagorinskyLilly{FT}(C_smag),
        turbconv = turbconv,
    )

    # Assemble configuration
    driver_config = ClimateMachine.SingleStackConfiguration(
        "SBL_EDMF",
        N,
        nelem_vert,
        zmax,
        param_set,
        model;
        hmax = FT(40),
        solver_type = ode_solver_type,
    )

    solver_config = ClimateMachine.SolverConfiguration(
        t0,
        timeend,
        driver_config,
        init_on_cpu = true,
        Courant_number = CFLmax,
    )

    #################### Change the ode_solver to implicit solver

    dg = solver_config.dg
    Q = solver_config.Q

    vdg = DGModel(
        driver_config;
        state_auxiliary = dg.state_auxiliary,
        direction = VerticalDirection(),
    )

    # linear solver relative tolerance rtol which should be slightly smaller than the nonlinear solver tol
    linearsolver = BatchedGeneralizedMinimalResidual(
        dg,
        Q;
        max_subspace_size = 30,
        atol = -1.0,
        rtol = 5e-5,
    )

    """
    N(q)(Q) = Qhat  => F(Q) = N(q)(Q) - Qhat

    F(Q) == 0
    ||F(Q^i) || / ||F(Q^0) || < tol

    """
    # ϵ is a sensity parameter for this problem, it determines the finite difference Jacobian dF = (F(Q + ϵdQ) - F(Q))/ϵ
    # I have also try larger tol, but tol = 1e-3 does not work
    nonlinearsolver =
        JacobianFreeNewtonKrylovSolver(Q, linearsolver; tol = 1e-4, ϵ = 1.e-10)

    # this is a second order time integrator, to change it to a first order time integrator
    # change it ARK1ForwardBackwardEuler, which can reduce the cost by half at the cost of accuracy
    # and stability
    # preconditioner_update_freq = 50 means updating the preconditioner every 50 Newton solves,
    # update it more freqent will accelerate the convergence of linear solves, but updating it
    # is very expensive
    ode_solver = ARK2ImplicitExplicitMidpoint(
        dg,
        vdg,
        NonLinearBackwardEulerSolver(
            nonlinearsolver;
            isadjustable = true,
            preconditioner_update_freq = 50,
        ),
        Q;
        dt = solver_config.dt,
        t0 = 0,
        split_explicit_implicit = false,
        variant = NaiveVariant(),
    )

    solver_config.solver = ode_solver

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

    # --- Zero-out horizontal variations:
    vsp = vars_state(model, Prognostic(), FT)
    horizontally_average!(
        driver_config.grid,
        solver_config.Q,
        varsindex(vsp, :turbconv),
    )
    horizontally_average!(
        driver_config.grid,
        solver_config.Q,
        varsindex(vsp, :energy, :ρe),
    )
    vsa = vars_state(model, Auxiliary(), FT)
    horizontally_average!(
        driver_config.grid,
        solver_config.dg.state_auxiliary,
        varsindex(vsa, :turbconv),
    )
    # ---

    dgn_config = config_diagnostics(driver_config)

    cbtmarfilter = GenericCallbacks.EveryXSimulationSteps(1) do
        Filters.apply!(
            solver_config.Q,
            (turbconv_filters(turbconv)...,),
            solver_config.dg.grid,
            TMARFilter(),
        )
        nothing
    end

    diag_arr = [single_stack_diagnostics(solver_config)]
    time_data = FT[0]

    # Define the number of outputs from `t0` to `timeend`
    n_outputs = 5
    # This equates to exports every ceil(Int, timeend/n_outputs) time-step:
    every_x_simulation_time = ceil(Int, timeend / n_outputs)

    cb_data_vs_time =
        GenericCallbacks.EveryXSimulationTime(every_x_simulation_time) do
            diag_vs_z = single_stack_diagnostics(solver_config)

            nstep = getsteps(solver_config.solver)
            # Save to disc (for debugging):
            # @save "bomex_edmf_nstep=$nstep.jld2" diag_vs_z

            push!(diag_arr, diag_vs_z)
            push!(time_data, gettime(solver_config.solver))
            nothing
        end

    check_cons =
        (ClimateMachine.ConservationCheck("ρ", "3000steps", FT(0.001)),)

    cb_print_step = GenericCallbacks.EveryXSimulationSteps(100) do
        @show getsteps(solver_config.solver)
        nothing
    end

    result = ClimateMachine.invoke!(
        solver_config;
        diagnostics_config = dgn_config,
        check_cons = check_cons,
        user_callbacks = (cbtmarfilter, cb_data_vs_time, cb_print_step),
        check_euclidean_distance = true,
    )

    diag_vs_z = single_stack_diagnostics(solver_config)
    push!(diag_arr, diag_vs_z)
    push!(time_data, gettime(solver_config.solver))

    return solver_config, diag_arr, time_data
end

# add a command line argument to specify the kind of surface flux
# TODO: this will move to the future namelist functionality
sbl_args = ArgParseSettings(autofix_names = true)
add_arg_group!(sbl_args, "StableBoundaryLayer")
@add_arg_table! sbl_args begin
    "--surface-flux"
    help = "specify surface flux for energy and moisture"
    metavar = "prescribed|bulk|custom_sbl"
    arg_type = String
    default = "custom_sbl"
end

cl_args = ClimateMachine.init(
    parse_clargs = true,
    custom_clargs = sbl_args,
    output_dir = get(ENV, "CLIMATEMACHINE_SETTINGS_OUTPUT_DIR", "output"),
    fix_rng_seed = true,
)

solver_config, diag_arr, time_data = main(Float64, cl_args)

include(joinpath(@__DIR__, "report_mse_sbl_ss_implicit.jl"))

nothing


text:
"""
    optimaldesign(
      n1::Integer,
      parameters::Parameters,
      solver::MathProgBase.AbstractMathProgSolver
    )

Find the optimal two-stage design for given parameters and fixed stage-one
sample size.

# Parameters

| Parameter    | Description |
| -----------: | :---------- |
| n1           | stage-one sample size |
| nparameters  | paramters object defining the optimality criterion |
| solver       | MathProgBase solver used for optimization |
| VERBOSE      | control verbosity during optimization |

# Return Value

Design object optimized for given parameter set.

# Examples
```julia-repl
julia> ss     = SimpleSampleSpace(10:25, 100, n2min = 5)
julia> design = optimaldesign(15, params, solver = Gurobi.GurobiSolver())
```
"""
function optimaldesign(
    n1::Integer,
    parameters::Parameters,
    solver::MathProgBase.AbstractMathProgSolver
)

    !possible(n1, samplespace(parameters)) ? warn("n1 not compatible with sample space") : nothing
    ipm = IPModel(samplespace(parameters), n1, UNIMODAL = isunimodal(parameters))
    completemodel(ipm, parameters, n1)
    JuMP.setsolver(ipm.m, solver)
    JuMP.solve(ipm.m) in (:Optimal, :UserLimit) ? nothing : error("no feasible solution reached")
    design = extractsolution(ipm, parameters)
    _isfeasible(design, parameters) ? nothing : error("solution infeasible")
    return design

end # optimaldesign

"""
    getoptimaldesign{TS<:MathProgBase.AbstractMathProgSolver}(
        parameters::Parameters,
        solver::TS;
        VERBOSE::Integer = 1
    )

Find the optimal two-stage design for given parameters (optimizes over `n1` as well).

# Parameters

| Parameter    | Description |
| -----------: | :---------- |
| nparameters  | paramters object defining the optimality criterion |
| solver       | MathProgBase solver used for optimization |
| VERBOSE      | control verbosity during optimization |

# Return Value

Tuple, first element is the Design object optimized for given parameter set
and the second element is a dictionary with n1, scores and respective optimal designs.

# Examples
```julia-repl
julia> ss = SimpleSampleSpace(10:25, 100, n2min = 5)
julia> interimsamplesize(ss)
julia> design, res = getoptimaldesign(params, solver = Gurobi.GurobiSolver())
```
"""
function optimaldesign(
  parameters::Parameters,
  solver::MathProgBase.AbstractMathProgSolver;
  VERBOSE::Integer = 1,
  EARLYTERMINATION::Bool = false
)

  n1range   = samplespace(parameters) |> interimsamplesizerange
  designs   = Array{Design}(length(n1range))
  scores    = ones(length(n1range))*Inf

  if VERBOSE > 0
    t0 = now()
    println()

    println(@sprintf(
      "optimizing design for parameters '%s'", label(parameters)
    ))
    print(@sprintf(
      "\rconsidering %i stage-one sample sizes between %i and %i", length(n1range), minimum(n1range), maximum(n1range)
    ))
    print(@sprintf(
      " using %s as solver\n\r", typeof(solver)
    ))
  end

  function reldiff(newval, best)

    diff = best - newval
    if diff == 0
      return 0.0
    else
      return -sign(diff) * abs(diff/best)
    end

  end

  for i in 1:length(n1range)
    t0i = now()
    try
      designs[i] = optimaldesign(n1range[i], parameters, solver)
      scores[i]  = score(designs[i])
      if VERBOSE > 0
        i == 1 ? println("    time    n1   % done   sol. time [s]   cum. time [min]       score        best   % diff to best") : nothing
        print(@sprintf("\r%s   %3i    %5.1f   %13i   %15.1f   %+.2e   %+.2e   %14.1f\n\r",
          Dates.format(t0i, "HH:MM:SS"),
          n1range[i],
          i / length(n1range) * 100,
          Int(round(Int(Dates.value(now() - t0i)) / 1000)),
          Int(Dates.value(now() - t0)) / 60 / 1000,
          scores[i],
          scores[findmin(scores)[2]],
          reldiff(scores[i], scores[findmin(scores)[2]]) * 100
        ))
      end
    catch e
      print(@sprintf("\r%s   %3i   error: %s, continuing...\n\r", Dates.format(t0i, "HH:MM:SS"), n1range[i], e))
      scores[i] = Inf
    end
    if VERBOSE > 1
      try
        println()
        print(UnicodePlots.stairs(
          collect(0:n1range[i]) / n1range[i],
          designs[i].n,
          title = "n vs. x1/n1",
          ylim = [0, Int(round(maximum(samplesize(designs[i])*1.1)))],
          canvas = UnicodePlots.AsciiCanvas,
          style = :pre
        ))
        print("\n\n")
      catch e
      end
      println(convert(DataFrame(design[i])))
    end
    if (EARLYTERMINATION == true) & (i > 3)
      try
        if all(reldiff.(scores[[i - 2; i - 1; i]], scores[findmin(scores)[2]]) .>= .05)
          print("\n")
          info("terminating optimization early")
          break
        end
      catch e
      end
    end
  end
  print("\n\n")
  if VERBOSE > 1
    valid = scores .!= Inf
    try
      print(UnicodePlots.stairs(n1range[valid], scores[valid], title = "score vs. n1", xlim = [minimum(n1range), maximum(n1range)], canvas = UnicodePlots.AsciiCanvas, style = :pre))
      print("\n\n")
    catch e
    end
  end
  if VERBOSE > 0
    print(@sprintf("\rdone after %i minutes.\n\r", Int(round(Int(Dates.value(now() - t0)) / 60 / 1000))))
    print(@sprintf("\rOptimal stage-one sample size is %i resulting in a minimal score of %.2e\n\n\r", n1range[findmin(scores)[2]], scores[findmin(scores)[2]]))
  end
  return designs[findmin(scores)[2]], Dict("n1" => n1range, "scores" => scores, "designs" => designs)

end

max_stars_repo_name: kkmann/binaryTwoStageDesigns

max_issues_repo_name: kkmann/binaryTwoStageDesigns

max_forks_repo_name: kkmann/binaryTwoStageDesigns


text:
using HorizonSideRobots
inverse(side::HorizonSide) = HorizonSide(mod(Int(side) + 2, 4))

function counter(r::Robot)
    side=Ost
    num=0
    while isborder(r,Nord)==false
        while isborder(r,side)==false
            if borders(r,side)==true
                num+=1
                walk_by(r,side)
            end
        end
        move!(r,Nord)
        side=inverse(side)
    end
    return num
end

function borders(r::Robot,side::HorizonSide)
    while isborder(r,Nord)==false
        if isborder(r,side)==true
            return false
        end
        move!(r,side)
    end
    return true
end

function walk_by(r::Robot,side::HorizonSide)
    while isborder(r,Nord)==true
        move!(r,side)
    end
end

counter(r)


text:
# Automatically generated using Clang.jl

const DWFINC = TRUE

# Skipping MacroDefinition: IsBitSet ( fs , bit ) ( ( fs & ( 1 << bit ) ) != 0 )

const BOOL = Cint
const BYTE = Cuchar
const HDWF = Cint
const ENUMFILTER = Cint
const DEVID = Cint
const DEVVER = Cint
const TRIGSRC = BYTE
const DwfState = BYTE
const DwfEnumConfigInfo = Cint
const ACQMODE = Cint
const FILTER = Cint
const TRIGTYPE = Cint
const DwfTriggerSlope = Cint
const TRIGLEN = Cint
const DWFERC = Cint
const FUNC = BYTE
const ANALOGIO = BYTE
const AnalogOutNode = Cint
const DwfAnalogOutMode = Cint
const DwfAnalogOutIdle = Cint
const DwfDigitalInClockSource = Cint
const DwfDigitalInSampleMode = Cint
const DwfDigitalOutOutput = Cint
const DwfDigitalOutType = Cint
const DwfDigitalOutIdle = Cint
const DwfAnalogImpedance = Cint
const DwfParam = Cint
const TRIGCOND = Cint
const STS = BYTE


text:
using ASCIIPlots
using Base.Test

x = randn(20)
y = 21.0 * x + randn(20)
scatterplot(x, y)


text:
function Power(arg0::jdouble)
    return Power((jdouble,), arg0)
end

function value(obj::Power, arg0::DerivativeStructure)
    return jcall(obj, "value", DerivativeStructure, (DerivativeStructure,), arg0)
end

function value(obj::Power, arg0::jdouble)
    return jcall(obj, "value", jdouble, (jdouble,), arg0)
end


text:
"""
    matrixA(match::ConeFund = conefund(SPECENV.cmf), type::Symbol = :standard)

Creates matrix A from a cone fundamental for Cohen’s matrix R calculations.

Available types:

- :protanopia
- :deuteranopia
- :tritanopia
- :standard

# Examples:

```jldoctest
julia> matrixA(conefund(SPECENV.cmf))
441×3 Matrix{Float64}:
 0.0108549   0.00218811  0.028312
 0.0131472   0.00265254  0.0343045
 0.0159195   0.00321355  0.0415584
 0.0192642   0.00388931  0.0503177
[...]
 2.6301e-6   1.56603e-6  0.0
 2.48078e-6  1.47742e-6  0.0
 2.34056e-6  1.39418e-6  0.0
 2.20894e-6  1.31601e-6  0.0
```
"""
function matrixA(match::ConeFund = conefund(SPECENV.cmf), type::Symbol = :standard)
    type == :protanopia ? [match.m match.s] :
    type == :deuteranopia ? [match.l match.s] :
    type == :tritanopia ? [match.l match.m] :
    type == :standard ? [match.l match.m match.s] :
    throw(DomainError(type, "Type does not exist."))
end

"""
    matrixR(match::ConeFund = conefund(SPECENV.cmf), type::Symbol = :standard)

Creates Cohen’s matrix R from a given CMF. For available types, see `matrixA`.

# Examples:

```jldoctest
julia> matrixR(conefund(SPECENV.cmf), :deuteranopia)
441×441 Matrix{Float64}:
 9.97123e-7   1.20808e-6   1.4634e-6    1.77166e-6   …  4.61275e-11  4.35087e-11  4.10494e-11  3.8741e-11
 1.20808e-6   1.46367e-6   1.77301e-6   2.14648e-6      5.57733e-11  5.26069e-11  4.96333e-11  4.68422e-11
 1.4634e-6    1.77301e-6   2.14773e-6   2.60013e-6      6.73948e-11  6.35686e-11  5.99755e-11  5.66027e-11
[...]
 4.35087e-11  5.26069e-11  6.35686e-11  7.6741e-11      5.56939e-14  5.2532e-14   4.95627e-14  4.67755e-14
 4.10494e-11  4.96333e-11  5.99755e-11  7.24034e-11     5.25459e-14  4.95627e-14  4.67613e-14  4.41316e-14
 3.8741e-11   4.68422e-11  5.66027e-11  6.83317e-11  …  4.95909e-14  4.67755e-14  4.41316e-14  4.16498e-14
```
"""
function matrixR(match::ConeFund = conefund(SPECENV.cmf), type::Symbol = :standard)
    A = matrixA(match, type)
    A * ((A' * A)^-1) * A'
end

"""
    fundamental_metamer(illum::LSpec, rad::Spectrum, match::ConeFund = conefund(SPECENV.cmf),type::Symbol = :standard)

Creates fundamental metamer of a spectrum according to Cohen’s matrix R method.

# Examples:

```jldoctest
```
"""
function fundamental_metamer(illum::LSpec, rad::Spectrum, match::ConeFund = conefund(SPECENV.cmf),type::Symbol = :standard)
    # radiometric function N*
    N = illum * rad
    Nstar = matrixR(match, type) * N.l
    # fundamental metamer
    luminance_spec(illum.λ, Nstar)
end

"""
    metameric_black(illum::LSpec, rad::Spectrum, match::ConeFund = conefund(SPECENV.cmf), type::Symbol = :standard)

Creates metameric black  according to Cohen’s matrix R method.

# Examples

```jldoctest
```
"""
function metameric_black(illum::LSpec, rad::Spectrum, match::ConeFund = conefund(SPECENV.cmf), type::Symbol = :standard)
    N = illum * rad
    Nstar = matrixR(match, type) * N.l
    luminance_spec(N.λ, N.l .- Nstar)
end

"""
    cohen_pair(illum::LSpec, rad::Spectrum, match::ConeFund = conefund(SPECENV.cmf), type::Symbol = :standard)

Returns a pair of Cohen’s fundamental metamer and metameric black according to a light stimulus.

# Examples

```jldoctest
```
"""
function cohen_pair(illum::LSpec, rad::Spectrum, match::ConeFund = conefund(SPECENV.cmf), type::Symbol = :standard)
    fundamental_metamer(illum, rad, match, type), metameric_black(illum, rad, match, type)
end

"""
    tristimulus(illum::LSpec, rad::Spectrum, match::ConeFund = conefund(SPECENV.cmf), type::Symbol = :standard)

Creates tristimulus values resulting from radiometric function/stimulus reaching the eye, according to Cohen’s matrix R method.

# Examples:

```
julia> tristimulus(D65_illuminant(), adapt_spec(agfait(100), SPECENV, :zero, :cubic, :natural))
3-element Vector{Any}:
 2827.9824977740263
 4605.154579468038
 6719.982399962246
 ````
"""
function tristimulus(illum::LSpec, rad::Spectrum, match::ConeFund = conefund(SPECENV.cmf), type::Symbol = :standard)
    matrixA(match, type)' * (illum * rad).l
end


text:
function multinetwork_data(sn_data::Dict{String,Any}, extradata::Dict{String,Any}, global_keys::Set{String})
    count = extradata["dim"]
    if InfrastructureModels.ismultinetwork(sn_data)
        error("replicate can only be used on single networks")
    end

    name = get(sn_data, "name", "anonymous")

    mn_data = Dict{String,Any}(
        "nw" => Dict{String,Any}()
    )

    sn_data_tmp = deepcopy(sn_data)
    for k in global_keys
        if haskey(sn_data_tmp, k)
            mn_data[k] = sn_data_tmp[k]
        end

        # note this is robust to cases where k is not present in sn_data_tmp
        delete!(sn_data_tmp, k)
    end
    mn_data["multinetwork"] = true
    mn_data["name"] = "$(count) replicates of $(name)"

    for k in global_keys
        delete!(sn_data_tmp, k)
    end
    sn_data_tmp["dcpol"] = sn_data["dcpol"]
    sn_data_tmp["dcline"] = sn_data["dcline"]

    for n in 1:count
        mn_data["nw"]["$n"] = copy(sn_data_tmp)
        for (key, element) in extradata
            if key == "dim"
            else
                if haskey(mn_data["nw"]["$n"], key)
                    mn_data["nw"]["$n"][key] = copy(sn_data_tmp[key])
                    for (l, element) in extradata[key]
                        if haskey(mn_data["nw"]["$n"][key], l)
                            mn_data["nw"]["$n"][key][l] = deepcopy(sn_data_tmp[key][l])
                            for (m, property) in extradata[key][l]
                                if haskey(mn_data["nw"]["$n"][key][l], m)
                                    mn_data["nw"]["$n"][key][l][m] = property[n]
                                else
                                    warn(_PM.LOGGER, ["Property ", m ," for , ", key, " " , l, " not found, will be ingnored"])
                                end
                            end
                        else
                            warn(_PM.LOGGER, [key, " " , l,  " not found, will be ingnored"])
                        end
                    end
                else
                    warn(_PM.LOGGER, ["Key ", key, " not found, will be ingnored"])
                end
            end
        end

    end

    return mn_data
end


text:
@testset "VoxelwiseConstant" begin
    N = (10,10)
    val = rand(N...)
    bnd_min, bnd_max = -5, 5
    vxlbound = (bnd_min:bnd_max, bnd_min:bnd_max)
    vc = VoxelwiseConstant(val, vxlbound)

    @test size(vc) == N
    @test bounds(vc) == ([bnd_min,bnd_min], [bnd_max,bnd_max])
    @test vc([-0.5, -0.5]) == val[5, 5]
    @test vc([0, -0.5]) == mean(val[5:6, 5])
    @test vc([-0.5, 0]) == mean(val[5, 5:6])
    @test vc([0, 0]) == mean(val[5:6, 5:6])

end  # @testset "VoxelwiseConstant"


text:
# The core functionality in the manuscript.
#
# The functions in this file could be applied to real data.

using LinearAlgebra: qr, Diagonal, tr, det, svd, eigen, dot
using Random: MersenneTwister, randexp, GLOBAL_RNG, AbstractRNG
using Random
using EllipsisNotation

#################################################
################ Random Sampling ################
#################################################

export randomPSD, samplePSD, sampleStateSpace

"""
    randomPSD([rng = GLOBAL_RNG,] n, λ)

Generate a random `n×n` positive definite matrix with i.i.d.
exponentially distributed eigenvalues.

### Arguments
 - `n`: Dimensionality of the matrix
 - `λ`: Exponential distribution parameter

### Returns
A random `n×n` matrix
"""
function randomPSD(rng::AbstractRNG, n, λ)
    mat = randn(rng,n,n)
    ortho, ~ = qr(mat)
    return ortho' * Diagonal(randexp(rng,n)/λ) * ortho
end

randomPSD(n, λ) = randomPSD(GLOBAL_RNG, n, λ)

"""
    samplePSD([rng = GLOBAL_RNG,] K, n, λ)

Generate `k` random `n×n` positive definite matrices with
i.i.d. exponentially distributed eigenvalues.

### Arguments
 - `K` - Number of samples
 - `n` - Dimensionality of the matrix
 - `λ` - Exponential distribution parameter

### Returns
An `n×n×K` array representing `K` random matrices
"""
function samplePSD(rng::AbstractRNG, K, n, λ)
    out = zeros(n,n,K)
    for i in 1:K
        out[:,:,i] = randomPSD(rng, n, λ)
    end
    return out
end

samplePSD(K, n, λ) = samplePSD(GLOBAL_RNG, K, n, λ)

"""
    sampleStateSpace(system::AbstractSystem, trajectorySampleCount, timestepCount, burnin, τ)

Generates `stateSampleCount` random samples of the statespace.
For now, initializes according to a high-variance Gaussian distribution,
then steps the system forward in time to shape the density according to the system.

A small i.i.d. Gaussian random vector is added after each step to avoid degeneracy in value iteration.

### Arguments
 - `system`                 - Chosen dynamical system
 - `trajectorySampleCount`  - Number of trajectory samples
 - `timestepCount`          - Number of timesteps per trajectory
 - 'burnin`                 - Initial time advancement
 - `τ`                      - timestep size

### Returns
A 2D array representing `trajectorySampleCount * timestepCount` samples, where columns represent the state.
"""
function sampleStateSpace(rng::AbstractRNG, system::AbstractSystem, trajectorySampleCount, timestepCount, burnin, τ)
    out = zeros(dimension(system), trajectorySampleCount * timestepCount)
    state = zeros(dimension(system))
    for i in 1:trajectorySampleCount
        if burnin == 0
            state .= 5*randn(rng, dimension(system))
        else
            state = flow(5*randn(rng, dimension(system)), burnin, system)
        end
        for j in 1:timestepCount
            index = j + (i-1) * timestepCount
            out[:,i] .= state
            state .= flow(state, τ, system) + .25 * randn(rng, dimension(system))
        end
    end
    return out
end

function sampleStateSpace(system::AbstractSystem, trajectorySampleCount, timestepCount, burnin, τ)
    return sampleStateSpace(Random.GLOBAL_RNG, system, trajectorySampleCount, timestepCount, burnin, τ)
end

#################################################
################ Interpolation ##################
#################################################

export nearestNeighbor, buildNearestNeighbor, localAverage
export min_dist, localAverageWeights

"""
    min_dist(v, D)

Finds the minimum distance vector in `D` from `v`.
Allows multidimensional array representation,
assumes last axis of `D` indexes the vectors.

### Arguments
 - `v` - Vector being compared
 - `D` - Set of vectors

### Returns
    (out, out_ind)

Where `out` is the closest vector, and `out_ind` is the position.
"""
function min_dist(v,D)
    alloc = zeros(size(v))
    alloc .= v.- @view D[..,1]
    out = sum(abs2,alloc)
    out_ind = 1
    @inbounds for i in 2:size(D)[end]
        alloc .= v .- @view D[..,i]
        d = sum(abs2,alloc)
        if d < out
            out = d
            out_ind = i
        end
    end
    return out, out_ind
end

"""
    nearestNeighbor(target, samples, values)

Approximates the value at `target` with the value at
the closest known sample.

### Arguments
 - `target`  - Point being evaluated, domain of function
 - `samples` - Known sample locations, domain of function
 - `values`  - Values at known locations, range of function

### Returns
Approximate value at `target`
"""
function nearestNeighbor(target, samples, values)
    ~, i = min_dist(target, samples)
    return values[i]
end

"""
    buildNearestNeighbor(f, samples)

Evaluates a function `f` at a set of chosen locations
to build the nearest neighbor approximation.
Glorified broadcast function.

The last axis of `samples` is assumed to index the sample locations.

### Arguments
 - `f`      - The function to be approximated
 - `samples`- Set of representation points

### Returns
An array of values at the given sample locations.
"""
function buildNearestNeighbor(f, samples)
    values = zeros(size(samples)[end])
    Threads.@threads for i in 1:size(samples)[end]
        values[i] = f(samples[..,i])
    end
    return values
end

"""
    localAverage(targetPSD, targetState, psdSamples, stateSamples, values, d_max)

Computes the local average interpolation, defaulting to a
nearest neighbor interpolation if there are no points within `d_max`
of the desired `target` location.

The last axis of `samples` is assumed to index the sample locations.

### Arguments
 - `targetPSD`    - PSD point being evaluated, domain of function
 - `targetState`  - State point being evaluated, domain of function
 - `psdSamples`   - Known PSD sample locations, domain of function
 - `stateSamples` - Known state sample locations, domain of function
 - `values`       - Values at known locations, range of function
 - `d_max`        - Maximum distance to average

### Returns
A local average approximation at a given location

### Notes
values is indexed as ...
"""
function localAverage(targetPSD, targetState, psdSamples, stateSamples, values, d_max)
    if length(values) != size(psdSamples)[end] * size(stateSamples)[end]
        throw(ArgumentError("`values` length mismatch with sample counts"))
    end
    num   = 0
    denom = 0
    psdScratch   = zeros(size(targetPSD))
    stateScratch = zeros(size(targetState))

    # As a backup, track the nearest neighbor
    nearestIndex = 1
    minDist  = d_max
    minDist -= sqrt(sum(abs2,psdSamples[..,1] - targetPSD) + sum(abs2,stateSamples[..,1] - targetState))

    # Loop through all points performing local average
    # Note that through a hash approach, this loop could
    # be done in a much more efficient manner.
    psdSampleCount = size(psdSamples)[end]

    for i in 1:length(values)
        stateIndex = Int(floor((i-1)/psdSampleCount)) + 1
        psdIndex = ((i - 1) % psdSampleCount) + 1

        psdScratch   .= targetPSD   - @view psdSamples[..,psdIndex]
        stateScratch .= targetState - @view stateSamples[..,stateIndex]

        dist = d_max - sqrt(sum(abs2,psdScratch) + sum(abs2,stateScratch))

        if dist > 0
            num += dist * values[i]
            denom += dist
        elseif minDist < dist
            minDist = dist
            nearestIndex = i
        end

    end
    if denom == 0
        return values[nearestIndex]
    end

    return num/denom
end

"""
    localAverageWeights(targetPSD, targetState, psdSamples, stateSamples, values, d_max)

Computes the local average interpolation, defaulting to a
nearest neighbor interpolation if there are no points within `d_max`
of the desired `target` location.

The last axis of `samples` is assumed to index the sample locations.

### Arguments
 - `targetPSD`    - PSD point being evaluated, domain of function
 - `targetState`  - State point being evaluated, domain of function
 - `psdSamples`   - Known PSD sample locations, domain of function
 - `stateSamples` - Known state sample locations, domain of function
 - `d_max`        - Maximum distance to average

### Returns
    (indices, weights)

### Notes
values is indexed as ...
"""
function localAverageWeights(targetPSD, targetState, psdSamples, stateSamples, d_max)
    psdScratch   = zeros(size(targetPSD))
    stateScratch = zeros(size(targetState))

    # As a backup, track the nearest neighbor
    nearestIndex = 1
    minDist  = d_max
    minDist -= sqrt(sum(abs2,psdSamples[..,1] - targetPSD) + sum(abs2,stateSamples[..,1] - targetState))

    # Loop through all points performing local average
    # Note that through a hash approach, this loop could
    # be done in a much more efficient manner.
    psdSampleCount = size(psdSamples)[end]

    weights = zeros(Float64,0)
    indices = zeros(Int64,0)

    for i in 1:(size(psdSamples)[end] * size(stateSamples)[end])
        stateIndex = Int(floor((i-1)/psdSampleCount)) + 1
        psdIndex = ((i - 1) % psdSampleCount) + 1

        psdScratch   .= targetPSD   - @view psdSamples[..,psdIndex]
        stateScratch .= targetState - @view stateSamples[..,stateIndex]

        dist = d_max - sqrt(sum(abs2,psdScratch) + sum(abs2,stateScratch))

        if dist > 0
            push!(weights, dist)
            push!(indices, i)
        elseif minDist < dist
            minDist = dist
            nearestIndex = i
        end
    end
    if length(weights) == 0
        return ([nearestIndex], [1])
    end

    return indices, weights / sum(weights)
end

#################################################
################# POMDP Tools ###################
#################################################

export updateCRLB, updateCRLB_naive, optimalAction_NearestNeighbor, valueUpdate_NearestNeighbor, valueIterate_NearestNeighbor, valueIterate_NearestNeighbor_precompute!
export optimalAction_LocalAverage, valueUpdate_LocalAverage, valueIterate_LocalAverage

"""
    updateCRLB_naive(    crlb,
                         action,
                         jacobian,
                         σ²)

Updates the Fisher information based on the Jacobian of the
flow and the current measurement under Gaussian noise.

### Arguments
 - `crlb`- Previous Fisher information matrix
 - `action`     - Measurement vector
 - `jacobian`   - Jacobian of flow for current timestep
 - `σ²`         - Measurement variance

### Returns
Updated Fisher Information
"""
function updateCRLB_naive(crlb, action, jacobian, σ²)
    return jacobian * inv(inv(crlb) + action * action'/σ²) * jacobian'
end

"""
    updateCRLB(   crlb,
                  action,
                  jacobian,
                  σ²)

Updates the Fisher information based on the Jacobian of the
flow and the current measurement under Gaussian noise.

### Arguments
 - `crlb`       - Previous Fisher information matrix
 - `action`     - Measurement vector
 - `jacobian`   - Jacobian of flow for current timestep
 - `σ²`         - Measurement variance

### Returns
Updated CRLB
"""
function updateCRLB(crlb, action, jacobian, σ²)
    tmp = crlb * action
    return jacobian * (crlb - (tmp * tmp')/(σ² + dot(action, tmp))) * jacobian'
end

"""
    optimalAction_NearestNeighbor(  crlb, actionSpace,
                                    samples, values,
                                    jacobian, σ²)

Computes the optimal action for a given state based on the current
nearest neighbor approximation of the value function.
Assumes linear measurement under additive Gaussian noise.

### Arguments
 - `crlb`       - Current Cramér-Rao bound
 - `actionSpace`- Action Space (finite)
 - `samples`    - Sample locations in value function
 - `values`     - Sample evaluations of value function
 - `jacobian`   - Jacobian of flow for current timestep
 - `σ²`         - Measurement variance

### Returns
The action that minimizes the Cramér-Rao bound
"""
function optimalAction_NearestNeighbor( crlb, actionSpace,
                                        samples, values, jacobian, σ²)

    vals = zeros(length(actionSpace))

    for i in 1:length(actionSpace)
        new_crlb  = updateCRLB(crlb, actionSpace[i], jacobian, σ²)
                vals[i]   = nearestNeighbor(new_crlb, samples, values)
    end
    return findmin(vals)[2]
end

"""
    optimalAction_LocalAverage(     crlb, newState, actionSpace,
                                    psdSamples, stateSamples, values,
                                    jacobian, σ², d_max)

Computes the optimal action for a given state based on the current
nearest neighbor approximation of the value function.
Assumes linear measurement under additive Gaussian noise.

### Arguments
 - `crlb`         - Current Cramér-Rao bound
 - `newState`     - Updated system state
 - `actionSpace`  - Action Space (finite)
 - `psdSamples`   - Sample PSD locations in value function
 - `stateSamples` - Sample state locations in value function
 - `values`       - Sample evaluations of value function
 - `jacobian`     - Jacobian of flow for current timestep
 - `σ²`           - Measurement variance
 - `d_max`        - Maximum distance for local average

### Returns
The action that minimizes the Cramér-Rao bound
"""
function optimalAction_LocalAverage(    crlb, newState, actionSpace,
                                        psdSamples, stateSamples, values, jacobian, σ², d_max)

    vals = zeros(length(actionSpace))

    for i in 1:length(actionSpace)
        new_crlb  = updateCRLB(crlb, actionSpace[i], jacobian, σ²)
        vals[i]   = localAverage(new_crlb, newState, psdSamples, stateSamples, values, d_max)
    end
    return findmin(vals)[2]
end

"""
    valueUpdate_NearestNeighbor(   crlb,
                                   γ,
                                   jacobian,
                                   actionSpace,
                                   samples, values)

Iterates the discounted Bellman equation for minimizing the
CRLB at a given current Fisher information under the
nearest neighbor value approximation for a given point.

Assumes linear measurement under additive Gaussian noise.

### Arguments
 - `crlb`       - Current Cramér-Rao bound
 - `γ`          - Discount factor
 - `jacobian`   - Jacobian of flow for current timestep
 - `actionSpace`- Action Space (finite)
 - `samples`    - Sample locations in value function
 - `values`     - Sample evaluations of value function
 - `σ²`         - Noise variance

### Returns
The updated value function approximation evaluated at `crlb`
"""
function valueUpdate_NearestNeighbor(   crlb,
                                        γ,
                                        jacobian,
                                        actionSpace,
                                        samples, values,
                                        σ²)

    update_vals = zeros(length(actionSpace))
    for i in 1:length(actionSpace)
        new_crlb = updateCRLB(crlb, actionSpace[i], jacobian, σ²)
        update_vals[i] = nearestNeighbor(new_crlb, samples, values)
    end
    return γ * minimum(update_vals) + tr(crlb)
end

"""
    valueUpdate_LocalAverage(   crlb,
                                newState,
                                γ,
                                jacobian,
                                actionSpace,
                                psdSamples, stateSamples,
                                values,
                                σ²,
                                d_max)

Iterates the discounted Bellman equation for minimizing the
CRLB at a given current Fisher information under the
nearest neighbor value approximation for a given point.

Assumes linear measurement under additive Gaussian noise.

### Arguments
 - `crlb`         - Current Cramér-Rao bound
 - `newState`     - Updated state value
 - `γ`            - Discount factor
 - `jacobian`     - Jacobian of flow for current timestep
 - `actionSpace`  - Action Space (finite)
 - `psdSamples`   - Sample PSD locations in value function
 - `stateSamples` - Sample state locations in value function
 - `values`       - Sample evaluations of value function
 - `σ²`           - Noise variance
 - `d_max`        - Maximum distance for local average

### Returns
The updated value function approximation evaluated at `crlb`
"""
function valueUpdate_LocalAverage(      crlb,
                                        newState,
                                        γ,
                                        jacobian,
                                        actionSpace,
                                        psdSamples, stateSamples,
                                        values,
                                        σ²,
                                        d_max)

    update_vals = zeros(length(actionSpace))
    for i in 1:length(actionSpace)
        newCRLB = updateCRLB(crlb, actionSpace[i], jacobian, σ²)
        update_vals[i] = localAverage(newCRLB, newState, psdSamples, stateSamples, values, d_max)
    end
    return γ * minimum(update_vals) + tr(crlb)
end

"""
    valueIterate_NearestNeighbor(γ, jacobian, actionSpace, samples, values)

Iterates the discounted Bellman equation for minimizing the
CRLB at a given current Fisher information under the
nearest neighbor value approximation.

Assumes linear measurement under additive Gaussian noise.

### Arguments
 - `γ`          - Discount factor
 - `jacobian`   - Jacobian of flow for current timestep
 - `actionSpace`- Action Space (finite)
 - `samples`    - Sample locations in value function
 - `values`     - Sample evaluations of value function
 - `σ²`         - Measurement variance

### Returns
The updated value function approximation values for all points in `samples`.

### Notes
This function is multithreaded, remember to give Julia multiple threads when launching with
`julia -t NTHREADS`, where `NTHREADS` is the desired number of threads.
"""
function valueIterate_NearestNeighbor(γ, jacobian, actionSpace, samples, values, σ²)
    new_out = zeros(length(values))
    Threads.@threads for i in 1:length(values)
        new_out[i] = valueUpdate_NearestNeighbor(  view(samples,:,:,i),
                                                           γ,
                                                           jacobian,
                                                           actionSpace,
                                                           samples, values,
                                                           σ²)

    end
    return new_out
end

function valueIterate_LocalAverage_precompute!( out,
                                                γ,
                                                psdSamples,
                                                values,
                                                indices,
                                                weights)

    psdSampleCount = size(psdSamples)[end]
    Threads.@threads for i in 1:length(values)
        crlbIndex = ((i-1) % psdSampleCount) + 1
        out[i] = valueUpdate_precomputedWeights( values, psdSamples[:,:,crlbIndex], γ, indices[:,i], weights[:,i])
    end
end

function valueUpdate_precomputedWeights(values, crlb, γ, indices, weights)
    BestValue = dot(weights[1], values[indices[1]])
    for i in 2:length(weights)
        value = dot(weights[i], values[indices[i]])
        if value < BestValue
            BestValue = value
        end
    end
    return γ * BestValue + tr(crlb)
end

"""
    valueIterate_LocalAverage!( out,
                                γ,
                                jacobianList,
                                updateStateList,
                                actionSpace,
                                psdSamples,
                                stateSamples,
                                values,
                                σ²
                                d_max)

Iterates the discounted Bellman equation for minimizing the
CRLB at a given current Fisher information under the
nearest neighbor value approximation.

Assumes additive Gaussian noise.

### Arguments
 - `out`          - Output array
 - `γ`            - Discount factor
 - `jacobianList` - Jacobian of flow for current timestep
 - `updateStateList` -
 - `actionSpace`  - Action Space (finite)
 - `psdSamples`      - Sample locations in value function
 - `stateSamples`    -
 - `values`       - Sample evaluations of value function
 - `σ²`           - Measurement variance
 - `d_max`

### Returns
Nothing, but stores the updated value function approximation values for all points in `out`.

### Notes
This function is multithreaded, remember to give Julia multiple threads when launching with
`julia -t NTHREADS`, where `NTHREADS` is the desired number of threads.
"""
function valueIterate_LocalAverage!(    out,
                                        γ,
                                        jacobianList,
                                        updateStateList,
                                        actionSpace,
                                        psdSamples,
                                        stateSamples,
                                        values,
                                        σ²,
                                        d_max)

    psdSampleCount = Int(length(values)//size(jacobianList)[3])
    Threads.@threads for i in 1:length(values)
        stateIndex = Int(floor((i-1)/psdSampleCount)) + 1
        psdIndex = ((i - 1) % psdSampleCount) + 1

        out[i] = valueUpdate_LocalAverage(  view(psdSamples,:,:,psdIndex),
                                            view(updateStateList,:,stateIndex),
                                            γ,
                                            jacobianList[:,:,stateIndex],
                                            actionSpace,
                                            psdSamples, stateSamples,
                                            values,
                                            σ²,
                                            d_max)

    end
end

"""
    valueIterate_NearestNeighbor_precompute!(   out,
                                                values,
                                                transitionMap::Array{Int,2},
                                                psdSamples,
                                                γ,
                                                stateSampleCount)

Completes one step of value iteration using precomputed transition maps.

### Arguments
 - `out`              - Output Array (new value function approximation)
 - `values`           - Current value function approximation
 - `transitionMap`    - Transitions of state/crlb based on the action
 - `psdSamples`       - CRLB samples
 - `γ`                - Discount Factor
 - `stateSampleCount` - Number of state-space samples (for indexing psdSamples)

### Returns
Nothing

### Notes
This function is multithreaded
"""
function valueIterate_NearestNeighbor_precompute!(  out,
                                                    values,
                                                    transitionMap::Array{Int,2},
                                                    psdSamples,
                                                    γ,
                                                    psdSampleCount)

    Threads.@threads for i in 1:length(values)
        crlbIndex = ((i-1) % psdSampleCount) + 1
        crlb = @view psdSamples[:,:,crlbIndex]
        updatedStates = @view transitionMap[:,i]

        futureValue = minimum(values[updatedStates])
        out[i] = tr(crlb) + γ * futureValue
    end
end

#################################################
################# POMDP Tools ###################
#################################################
export ValueFunctionApproximation_NearestNeighbor_precompute, NearestNeighbor_OptimalPolicy
export ValueFunctionApproximation_LocalAverage, LocalAverage_OptimalPolicy, ValueFunctionApproximation_LocalAverage_precompute

"""
    ValueFunctionApproximation_NearestNeighbor_precompute(  system::AbstractSystem,
                                                            τ,
                                                            γ,
                                                            actionSpace,
                                                            λ = 1,
                                                            psdSampleCount = 1000,
                                                            trajectorySampleCount = 100
                                                            timestepSampleCount = 10,
                                                            σ² = 1,
                                                            max_iterations=50)

Approximates the value function in the discounted cost optimal sampling problem.
Requires a model of the system dynamics, timestep, discount-factor, and a discrete action space.

Precomputes all state transitions for all actions. Uses significant memory in exchange for
less expensive computation.

Computes `max_iterations` steps of value iteration.

### Arguments
 - `system` - Representation of the dynamical system
 - `τ`                      - Timestep size
 - `γ`                      - Discount Factor
 - `actionSpace`            - Vector of possible actions
 - `λ`                      - CRLB space sampling parameter
 - `psdSampleCount`         - Number of covariance-matrix samples
 - `trajectorySampleCount`  - Number of state-space trajectory samples
 - `timestepSampleCount`    - Number of state-space timestep samples
 - `σ²`                     - Measurement variance
 - `max_iterations`         - Maximum number of iterations

### Returns
    (values, psdSamples, stateSamples)

 - `values`         - Output of value function at sample points
 - `psdSamples`     - CRLB sample locations
 - `stateSamples`   - State sample locations

### Note

`values` is a 1D array indexed as crlbIndex + psdSampleCount * (stateIndex - 1)
"""
function ValueFunctionApproximation_NearestNeighbor_precompute( rng::AbstractRNG,
                                                                system::AbstractSystem,
                                                                τ,
                                                                γ,
                                                                actionSpace,
                                                                λ = 1,
                                                                psdSampleCount = 1000,
                                                                trajectorySampleCount = 100,
                                                                timestepSampleCount = 10,
                                                                σ² = 1,
                                                                max_iterations=50)

    # Discretize the space
    stateSampleCount = trajectorySampleCount * timestepSampleCount
    psdSamples   = samplePSD(rng, psdSampleCount, dimension(system), λ)
    stateSamples = sampleStateSpace(rng, system, trajectorySampleCount, timestepSampleCount, 0, τ)
    values = rand(rng, psdSampleCount*stateSampleCount)
    updateValues = copy(values)

    # Precompute the nearest-neighbor maps based on actions
    systemMap = zeros(Int,length(actionSpace), psdSampleCount*stateSampleCount)
    Threads.@threads for i in 1:stateSampleCount
        # Compute flow and jacobian
        nextState = flow(stateSamples[:,i], τ, system)
        jacobian = flowJacobian(stateSamples[:,i], τ, system)

        # Find nearest state
        ~, nearestState = min_dist(nextState, stateSamples)

        for j in 1:psdSampleCount
            crlbInv = inv(psdSamples[:,:,j])
            for k in 1:length(actionSpace)
                # Compute Updated CRLB
                nextCRLB = jacobian * inv(crlbInv + actionSpace[k] * (actionSpace[k]')) * jacobian'
                ~, nearestJacobian = min_dist(nextCRLB, psdSamples)

                systemMap[k, (i-1)*psdSampleCount + j] = (nearestState-1)*psdSampleCount + nearestJacobian
            end
        end
    end

    # Apply Value Iteration
    for i in 1:max_iterations
        valueIterate_NearestNeighbor_precompute!(updateValues, values, systemMap, psdSamples, γ, psdSampleCount)
        values .= updateValues
    end

    return values, psdSamples, stateSamples
end

function ValueFunctionApproximation_NearestNeighbor_precompute( system::AbstractSystem,
                                                                τ,
                                                                γ,
                                                                actionSpace,
                                                                λ = 1,
                                                                psdSampleCount = 1000,
                                                                trajectorySampleCount = 100,
                                                                timestepSampleCount = 10,
                                                                σ² = 1,
                                                                max_iterations=50)

    return  ValueFunctionApproximation_NearestNeighbor_precompute(Random.GLOBAL_RNG,
                                                                  system,
                                                                  τ,
                                                                  γ,
                                                                  actionSpace,
                                                                  λ,
                                                                  psdSampleCount,
                                                                  trajectorySampleCount,
                                                                  timestepSampleCount,
                                                                  σ²,
                                                                  max_iterations)
end
"""
    ValueFunctionApproximation_LocalAverage_precompute(    system::AbstractSystem,
                                                           τ,
                                                           γ,
                                                           actionSpace,
                                                           λ = 1,
                                                           psdSampleCount = 1000,
                                                           trajectorySampleCount = 100
                                                           timestepSampleCount = 10,
                                                           σ² = 1,
                                                           max_iterations=50,
                                                           d_max = 0.05)

Approximates the value function using a LocalAverage interpolation in the discounted cost optimal sampling problem.
Requires a model of the system dynamics, timestep, discount-factor, and a discrete action space.

Computes `max_iterations` steps of value iteration.

Precomputes the local averaging weights to accelerate the value iteration steps.

### Arguments
 - `system`                 - Representation of the dynamical system
 - `τ`                      - Timestep size
 - `γ`                      - Discount Factor
 - `actionSpace`            - Vector of possible actions
 - `λ`                      - CRLB space sampling parameter
 - `psdSampleCount`         - Number of covariance-matrix samples
 - `trajectorySampleCount`  - Number of state-space trajectory samples
 - `timestepSampleCount`    - Number of state-space timestep samples
 - `σ²`                     - Measurement variance
 - `max_iterations`         - Maximum number of iterations
 - `d_max`                  -

### Returns
    (values, psdSamples, stateSamples)

 - `values`         - Output of value function at sample points
 - `psdSamples`     - CRLB sample locations
 - `stateSamples`   - State sample locations

### Note

`values` is a 1D array indexed as crlbIndex + psdSampleCount * (stateIndex - 1)
"""
function ValueFunctionApproximation_LocalAverage_precompute(    rng::AbstractRNG,
                                                                system::AbstractSystem,
                                                                τ,
                                                                γ,
                                                                actionSpace,
                                                                λ = 1,
                                                                psdSampleCount = 1000,
                                                                trajectorySampleCount = 100,
                                                                timestepSampleCount = 10,
                                                                σ² = 1,
                                                                max_iterations=50,
                                                                d_max=.05)

    # Discretize the space
    stateSampleCount = trajectorySampleCount * timestepSampleCount
    psdSamples   = samplePSD(rng, psdSampleCount, dimension(system), λ)
    stateSamples = sampleStateSpace(rng, system, trajectorySampleCount, timestepSampleCount, 0, τ)
    values = rand(rng, psdSampleCount*stateSampleCount)
    updateValues = copy(values)

    # Precompute the list of Jacobians
    jacobianList = zeros(dimension(system),dimension(system), stateSampleCount)
    updateStateList = zeros(dimension(system), stateSampleCount)

    for i in 1:stateSampleCount
        jacobianList[:,:,i] .= flowJacobian(stateSamples[:,i], τ, system)
        updateStateList[:,i] .= flow(stateSamples[:,i], τ, system)
    end

    # Precompute Local Average weights
    weights = Array{Array{Float64,1}, 2}(undef, length(actionSpace), length(values))
    indices = Array{Array{Int64,1},2}(undef, length(actionSpace), length(values))

    Threads.@threads for i in 1:length(values)
        stateIndex = Int(floor((i-1)/psdSampleCount)) + 1
        psdIndex = ((i - 1) % psdSampleCount) + 1

        targetState = updateStateList[:,stateIndex]
        jacobian = jacobianList[:,:,stateIndex]

        for j in 1:length(actionSpace)
            targetPSD = updateCRLB(psdSamples[:,:,psdIndex], actionSpace[j], jacobian, σ²)

            chosen_indices, chosen_weights = localAverageWeights(targetPSD,
                                                                 targetState,
                                                                 psdSamples,
                                                                 stateSamples,
                                                                 d_max)
            weights[j,i] = chosen_weights
            indices[j,i] = chosen_indices
        end
    end

    # Apply Value Iteration
    for i in 1:max_iterations
        valueIterate_LocalAverage_precompute!( updateValues,
                                               γ,
                                               psdSamples,
                                               values,
                                               indices,
                                               weights)

        values .= updateValues
    end

    return values, psdSamples, stateSamples
end

function ValueFunctionApproximation_LocalAverage_precompute(    system::AbstractSystem,
                                                                τ,
                                                                γ,
                                                                actionSpace,
                                                                λ = 1,
                                                                psdSampleCount = 1000,
                                                                trajectorySampleCount = 100,
                                                                timestepSampleCount = 10,
                                                                σ² = 1,
                                                                max_iterations=50,
                                                                d_max=.05)

    return ValueFunctionApproximation_LocalAverage_precompute(  Random.GLOBAL_RNG,
                                                                system,
                                                                τ,
                                                                γ,
                                                                actionSpace,
                                                                λ,
                                                                psdSampleCount,
                                                                trajectorySampleCount,
                                                                timestepSampleCount,
                                                                σ²,
                                                                max_iterations,
                                                                d_max)
end

"""
    ValueFunctionApproximation_LocalAverage(    system::AbstractSystem,
                                                τ,
                                                γ,
                                                actionSpace,
                                                λ = 1,
                                                psdSampleCount = 1000,
                                                trajectorySampleCount = 100
                                                timestepSampleCount = 10,
                                                σ² = 1,
                                                max_iterations=50,
                                                d_max = 0.05)

Approximates the value function using a LocalAverage interpolation in the discounted cost optimal sampling problem.
Requires a model of the system dynamics, timestep, discount-factor, and a discrete action space.

Computes `max_iterations` steps of value iteration.

### Arguments
 - `system`                 - Representation of the dynamical system
 - `τ`                      - Timestep size
 - `γ`                      - Discount Factor
 - `actionSpace`            - Vector of possible actions
 - `λ`                      - CRLB space sampling parameter
 - `psdSampleCount`         - Number of covariance-matrix samples
 - `trajectorySampleCount`  - Number of state-space trajectory samples
 - `timestepSampleCount`    - Number of state-space timestep samples
 - `σ²`                     - Measurement variance
 - `max_iterations`         - Maximum number of iterations
 - `d_max`                  -

### Returns
    (values, psdSamples, stateSamples)

 - `values`         - Output of value function at sample points
 - `psdSamples`     - CRLB sample locations
 - `stateSamples`   - State sample locations

### Note

`values` is a 1D array indexed as crlbIndex + psdSampleCount * (stateIndex - 1)
"""
function ValueFunctionApproximation_LocalAverage(   system::AbstractSystem,
                                                    τ,
                                                    γ,
                                                    actionSpace,
                                                    λ = 1,
                                                    psdSampleCount = 1000,
                                                    trajectorySampleCount = 100,
                                                    timestepSampleCount = 10,
                                                    σ² = 1,
                                                    max_iterations=50,
                                                    d_max=.05)

    # Discretize the space
    stateSampleCount = trajectorySampleCount * timestepSampleCount
    psdSamples   = samplePSD(psdSampleCount, dimension(system), λ)
    stateSamples = sampleStateSpace(system, trajectorySampleCount, timestepSampleCount, 0, τ)
    values = rand(psdSampleCount*stateSampleCount)
    updateValues = copy(values)

    # Precompute the list of Jacobians
    jacobianList = zeros(dimension(system),dimension(system), stateSampleCount)
    updateStateList = zeros(dimension(system), stateSampleCount)

    for i in 1:stateSampleCount
        jacobianList[:,:,i] .= flowJacobian(stateSamples[:,i], τ, system)
        updateStateList[:,i] .= flow(stateSamples[:,i], τ, system)
    end

    # Apply Value Iteration
    for i in 1:max_iterations
        valueIterate_LocalAverage!( updateValues,
                                    γ,
                                    jacobianList,
                                    updateStateList,
                                    actionSpace,
                                    psdSamples,
                                    stateSamples,
                                    values,
                                    σ²,
                                    d_max)

        values .= updateValues
    end

    return values, psdSamples, stateSamples
end

"""
    NearestNeighbor_OptimalPolicy(  state,
                                    crlb,
                                    system::AbstractSystem,
                                    σ²,
                                    τ,
                                    values,
                                    psdSamples,
                                    stateSamples,
                                    actionSpace)

Given a nearest neighbor approximation method for the value function, as well as the current state
and crlb, returns the optimal action.

### Arguments
 - `state`        - Current state of the system
 - `crlb`         - Current CRLB
 - `system`       - Represntation of the system
 - `σ²`           - Measurement noise power
 - `τ`            - Timestep size
 - `values`       - Value function evaluated at samples
 - `psdSamples`   - Positive Semidefinite matrix samples
 - `stateSamples  - Samples in the state-space
 - `actionSpace`  - Set of possible actions

### Returns
    (action, index, new_state, new_crlb)

 - `action` - The optimal action
 - `index`  - The index of the action
"""
function NearestNeighbor_OptimalPolicy( state,
                                        crlb,
                                        system::AbstractSystem,
                                        σ²,
                                        τ,
                                        values,
                                        psdSamples,
                                        stateSamples,
                                        actionSpace)
    # If only one action, return
    if length(actionSpace) == 1
        return (actionSpace[1], 1)
    end

    new_state = flow(state, τ, system)
    jacobian = flowJacobian(state, τ, system)
    psdSampleCount = size(psdSamples)[3]

    # if not invertible, operate in invertible subspace
    # Project action to be orthogonal to non-invertible subspace
    new_crlb = updateCRLB(crlb, actionSpace[1], jacobian, σ²)
    chosen_crlb = copy(new_crlb)

    ~, crlb_index  = min_dist(new_crlb, psdSamples)
    ~, state_index = min_dist(new_state, stateSamples)
    base_index = (state_index-1)*psdSampleCount

    index = 1
    minvalue = values[base_index + crlb_index]

    for i in 2:length(actionSpace)
        new_crlb = updateCRLB(crlb, actionSpace[i], jacobian, σ²)
        ~, crlb_index  = min_dist(new_crlb, psdSamples)

        if minvalue > values[base_index + crlb_index]
            index = i
            minvalue = values[base_index + crlb_index]
            chosen_crlb = copy(new_crlb)
        end
    end
    return (actionSpace[index], index, new_state, chosen_crlb)
end

"""
    LocalAverage_OptimalPolicy( state,
                                crlb,
                                system::AbstractSystem,
                                σ²,
                                τ,
                                values,
                                psdSamples,
                                stateSamples,
                                actionSpace,
                                d_max)

Given a nearest neighbor approximation method for the value function, as well as the current state
and crlb, returns the optimal action.

### Arguments
 - `state`        - Current state of the system
 - `crlb`         - Current CRLB
 - `system`       - Represntation of the system
 - `σ²`           - Measurement noise power
 - `τ`            - Timestep size
 - `values`       - Value function evaluated at samples
 - `psdSamples`   - Positive Semidefinite matrix samples
 - `stateSamples  - Samples in the state-space
 - `actionSpace`  - Set of possible actions
 - `d_max`        - Maximum distance for local average

### Returns
    (action, index, new_state, new_crlb)

 - `action` - The optimal action
 - `index`  - The index of the action
"""
function LocalAverage_OptimalPolicy(    state,
                                        crlb,
                                        system::AbstractSystem,
                                        σ²,
                                        τ,
                                        values,
                                        psdSamples,
                                        stateSamples,
                                        actionSpace,
                                        d_max)
    # If only one action, return
    if length(actionSpace) == 1
        return (actionSpace[1], 1)
    end

    new_state = flow(state, τ, system)
    jacobian = flowJacobian(state, τ, system)
    psdSampleCount = size(psdSamples)[3]

    new_crlb = updateCRLB(crlb, actionSpace[1], jacobian, σ²)
    chosen_crlb = copy(new_crlb)

    index = 1
    minvalue = localAverage(new_crlb, new_state, psdSamples, stateSamples, values, d_max)

    for i in 2:length(actionSpace)
        new_crlb = updateCRLB(crlb, actionSpace[i], jacobian, σ²)
        val = localAverage(new_crlb, new_state, psdSamples, stateSamples, values, d_max)

        if minvalue > val
            index = i
            minvalue = val
            chosen_crlb = copy(new_crlb)
        end
    end
    return (actionSpace[index], index, new_state, chosen_crlb)
end

########################################
########## State Estimation ############
########################################
export stateupdate_EKF

"""
    stateupdate_EKF(prediction, covariance, observation, action, σ²)

Updates the state based on the measurement according to the Extended Kalman Filter (EKF).

### Arguments
 - `prediction`  - Predicted State
 - `covariance`  - Covariance approximation
 - `observation` - Observed value
 - `action`      - Measurement functional
 - `σ²`          - Noise variance

### Returns
The updated state vector
"""
function stateupdate_EKF(prediction, covariance, observation, action, σ²)
    update = copy(prediction)

    num   = covariance * action * (observation - dot(action, prediction))
    denom = dot(action, covariance * action)

    update += num/denom
    return update
end

########################################
########## Direct Solution #############
########################################
export exponentialmap_sphere, actiongradient_1DApprox
export optimalaction_1DApprox, measurementvalue_1DApprox
export actiongradientstep_1DApprox

"""
    measurementvalue_1DApprox(action, limitvector::Vector, crlb, σ²)

Gives a value for the measurement action based on the 1D approximation
of the limiting behavior.

### Arguments
 - `action`      - Measurement functional
 - `limitvector` - Non-zero right singular vector of limiting flow Jacobian
 - `crlb`        - Current CRLB for state estimation
 - `σ²`          - Measurement noise variance

### Returns
A value of the action which should be maximized to minimize the CRLB for prediction.
"""
function measurementvalue_1DApprox(action, limitvector::AbstractVector, crlb, σ²)
    u  = action
    v = limitvector
    Σu = crlb * u

    num   = dot(v, Σu)^2
    denom = dot(u, Σu)^2 + σ²

    return num/denom
end

"""
    actiongradient_1DApprox(action, limitvector::Vector, crlb, σ²)

Returns the gradient of the 1D approximation value in the N-sphere
"""
function actiongradient_1DApprox(action, limitvector::AbstractVector, crlb, σ²)
    u  = action
    v = limitvector
    Σu = crlb * u
    Σv = crlb * v

    num   = dot(v, Σu)^2
    denom = dot(u, Σu)^2 + σ²

    gradient_num = 2 * (denom * dot(v, Σu) * Σv - num * (σ² * u + 2 * Σu))
    gradient_denom = denom^2

    gradient = gradient_num ./ gradient_denom
    #f(a) = measurementvalue_1DApprox(a, limitvector, crlb, σ²)
    #gradient = ForwardDiff.gradient(f, action)

    # Project onto tangent space
    gradient .-= dot(gradient, action) * action

    return gradient
end

"""
    exponentialmap_sphere(state::Vector, gradient::Vector, τ)

Advances a state on an n-sphere along a particular gradient
"""
function exponentialmap_sphere(state::AbstractVector, gradient::AbstractVector, τ)
    if gradient == zeros(length(gradient))
        return state
    end

    speed = sqrt(sum(abs2,gradient))
    return cos(speed * τ) * state + sin(speed * τ) * gradient ./ speed
end

"""
    actiongradientstep_1DApprox(action, limitvector, crlb, σ², τ)

Gradient step along action space.
"""
function actiongradientstep_1DApprox(action, limitvector, crlb, σ², τ)
    gradient = actiongradient_1DApprox(action, limitvector, crlb, σ²)
    return exponentialmap_sphere(action, gradient, τ)
end

"""
    optimalaction_1DApprox(actionspace, limitvector::Vector, crlb, σ²)

Returns the optimal action out of a finite `actionspace`.

### Arguments
 - `actionspace` - List of measurement functionals
 - `limitvector` - Non-zero right singular vector of limiting flow Jacobian
 - `crlb`        - Current CRLB for state estimation
 - `σ²`          - Measurement noise variance

### Returns
    `index, action`

- `index`  - The index of the optimal action
- `action` - The optimal action
"""
function optimalaction_1DApprox(actionspace, limitvector::Vector, crlb, σ²)
    params = (limitvector, crlb, σ²)

    maxind = 1
    maxval = measurementvalue_1DApprox(actionspace[1], params...)

    for i = 2:length(actionspace)
        val = measurementvalue_1DApprox(actionspace[i], params...)
        if val > maxval
            maxval = val
            maxind = i
        end
    end

    return maxind, actionspace[maxind]
end

"""
    optimalaction_1DApprox(actionspace, system::AbstractSystem, state, crlb, σ², T)

Returns the optimal action out of a finite `actionspace`.

### Arguments
 - `actionspace` - List of measurement functionals
 - `system`      - Dynamical System
 - `state`       - Current system state
 - `crlb`        - Current CRLB for state estimation
 - `σ²`          - Measurement noise variance
 - `T`           - Time horizon to approximate limit

### Returns
    `index, action`

- `index`  - The index of the optimal action
- `action` - The optimal action
"""
function optimalaction_1DApprox(actionspace, system::AbstractSystem, state, crlb, σ², T)
    # Approximate limiting singular vector
    jacobian = flowJacobian(state, T, system)
    ~, Σ, V = svd(jacobian)
    limitvector = V[:,1]

    return optimalaction_1DApprox(actionspace, limitvector, crlb, σ²)
end