"""
output index format:
x[i][j][k][l]
i-th spin system, k-th magnetically equivalent nuclei, l-th local spin index.
"""
function getmageqcompound(g,
    H_inds_sys,
    dict_ind_to_H_ID,
    dict_H_inds_to_css,
    dict_H_IDs_to_css,
    dict_J_ID_to_val,
    J_IDs;
    atol = 1e-6)

    C_g = Graphs.maximal_cliques(g)

    N_spin_systems = length(H_inds_sys)

    mag_eq_sys_inds_local = Vector{Vector{Vector{Int}}}(undef, N_spin_systems)
    mag_eq_sys_inds_global = Vector{Vector{Vector{Int}}}(undef, N_spin_systems)
    mag_eq_sys_IDs = Vector{Vector{Vector{Int}}}(undef, N_spin_systems)

    for i = 1:N_spin_systems
        C = keeptargetintegers(C_g, H_inds_sys[i])

        mag_eq_IDs0, mag_eq_inds0 = getmageqIDs(C,
            dict_ind_to_H_ID,
            dict_H_inds_to_css,
            dict_H_IDs_to_css, dict_J_ID_to_val, J_IDs;
            atol = atol)

        mag_eq_inds = combinetransitiveeqgroups(mag_eq_inds0)
        mag_eq_IDs = combinetransitiveeqgroups(mag_eq_IDs0)

        # mag_eq_inds = mag_eq_inds0
        # mag_eq_IDs = mag_eq_IDs0

        mag_eq_sys_inds_local[i] = getmageqlocalinds(H_inds_sys[i], mag_eq_inds)
        mag_eq_sys_IDs[i] = mag_eq_IDs
        mag_eq_sys_inds_global[i] = mag_eq_inds
    end

    return mag_eq_sys_inds_local, mag_eq_sys_IDs, mag_eq_sys_inds_global
end

"""
convertlabelsglobaltolocal(H_inds_sys::Vector{Vector{Int}},
    mag_eq_sys_inds::Vector{Vector{Int}})

The return type is Vector{Vector{Int}}.

Local: Each spin system will have its own spin nucleui numbering that starts at 1.
Global: The spins systems will together have one single spin nucleui numbering that starts at 1.
"""
function getmageqlocalinds(H_inds::Vector{Int},
    mag_eq_sys_inds::Vector{Vector{Int}})

    N_eqs = length(mag_eq_sys_inds)
    labels_sys_local = Vector{Vector{Int}}(undef, N_eqs)

    for i = 1:N_eqs

        x = H_inds
        conversion_dict = Dict(x .=> collect(1:length(x)))

        z = mag_eq_sys_inds[i]
        labels_sys_local[i] = collect( conversion_dict[z[k]] for k = 1:length(z) )
    end

    return labels_sys_local
end

"""
getmageqIDs(C::Vector{Vector{Int}},
    dict_ind_to_H_ID,
    dict_H_inds_to_css;
    atol::Float64 = 1e-6)

Given a list of node indices `C` of cliques of an undirected graph, with the indices of a clique being C[i],
Returns the list of node indices for each clique that are magnetically equivalent.
"""
function getmageqIDs(C::Vector{Vector{Int}},
    dict_ind_to_H_ID,
    dict_H_inds_to_css,
    dict_H_IDs_to_css,
    dict_J_ID_to_val,
    J_IDs;
    atol::Float64 = 1e-6)

    # traverse each maximally connected cliques.
    C_IDs_mag_eq = Vector{Vector{Vector{Int}}}(undef, length(C))
    C_inds_mag_eq = Vector{Vector{Vector{Int}}}(undef, length(C))
    for i = 1:length(C)

        ### partition the nucleus by unique cs.
        cs = collect( dict_H_inds_to_css[C[i][k]] for k = 1:length(C[i]) )
        cs_IDs = collect( dict_ind_to_H_ID[C[i][k]] for k = 1:length(C[i]) )
        cs_inds = C[i]

        # unique_cs = unique(round.(cs, digits = cs_round_digits))
        # unique_cs_inds = collect( inds = findall(xx->isapprox(unique_cs[k], xx; atol = atol), cs) for k = 1:length(unique_cs) )

        unique_cs, unique_cs_inds = uniqueinds(cs; atol = atol)

        # check magnetic equivalence for the chemically equivalent entries in `unique_cs_inds`.
        pass_flags = collect( checkmageq(unique_cs_inds[k], cs_IDs, dict_J_ID_to_val,
            dict_H_IDs_to_css, J_IDs; atol = atol) for k = 1:length(unique_cs_inds) )

        # store result.
        common_cs_IDs = collect( cs_IDs[unique_cs_inds[k]] for k = 1:length(unique_cs_inds) )
        C_IDs_mag_eq[i] = common_cs_IDs[pass_flags]

        common_cs_inds = collect( cs_inds[unique_cs_inds[k]] for k = 1:length(unique_cs_inds) )
        C_inds_mag_eq[i] = common_cs_inds[pass_flags]

        if length(common_cs_IDs[pass_flags]) > 2
            println("Warning for clique $(i): more than two groups of magnetically equivalent nuclei!")
        end
    end

    # remove empty entries.
    keep_flags = trues(length(C))
    for i = 1:length(C)
        if isempty(C_IDs_mag_eq[i])
            keep_flags[i] = false
        end
    end
    C_IDs_mag_eq = C_IDs_mag_eq[keep_flags]
    C_inds_mag_eq = C_inds_mag_eq[keep_flags]

    mag_eq_IDs::Vector{Vector{Int}} = combinevectors(C_IDs_mag_eq)
    mag_eq_inds::Vector{Vector{Int}} = combinevectors(C_inds_mag_eq)

    return mag_eq_IDs, mag_eq_inds
end

### routines related to checking whether the J-coupling value of A-C and B-C
# are the same, for all C's connected to either A or B. A and B are
# chemically equivalent nuclei.
"""
matchJlabels((label_pair_list::Vector{Tuple{Int,Int}}, search_list::Vector{Int})::Vector{Int}

Returns the indices of entries of `label_pair_list` that has any of the labels in the `search_list`.
"""
function matchanyJlabels(label_pair_list::Vector{Tuple{Int,Int}}, search_list::Vector{Int})::Vector{Int}

    inds = Vector{Int}(undef, 0)
    for k = 1:length(label_pair_list)
        i, j = label_pair_list[k]

        if any(i .== search_list) || any(j .== search_list)
            push!(inds, k)
        end
    end

    return inds
end

"""
getJIDstest(J_IDs::Vector{Tuple{Int,Int}},
    common_cs_IDs::Vector{Int},
    dict_H_IDs_to_css; atol = 1e-6)

Get the list of A-C, B-C, A-D, B-D, etc pairs to test.
"""
function getJIDstest(J_IDs::Vector{Tuple{Int,Int}},
    common_cs_IDs::Vector{Int},
    dict_H_IDs_to_css; atol = 1e-6)

    inds_any = matchanyJlabels(J_IDs, common_cs_IDs)
    inds_both = matchJlabels(J_IDs, common_cs_IDs)
    inds = setdiff(inds_any, inds_both)

    #inds_both =
    J_IDs_tmp = J_IDs[inds]

    target_IDs = Vector{Int}(undef, 0)
    for l = 1:length(J_IDs_tmp)
        i, j = J_IDs_tmp[l]

        # look for nuclei pairs that have the same chemical shift.
        if !isapprox(dict_H_IDs_to_css[i], dict_H_IDs_to_css[j], atol = atol)

            # find the value between i and j that isn't in common_cs_IDs.
            target_ID = i
            if any(i .== common_cs_IDs)
                target_ID = j
            end

            push!(target_IDs, target_ID)
        end
    end

    return target_IDs
end

"""
checkmageq(test_inds::Vector{Vector{Int}},
    cs_IDs::Vector{Int},
    dict_J_ID_to_val;
    atol = 1e-6)::Bool

`test_inds` is a list of indices in cs_IDs. This function checks for magnetic equivalence for the nuclei in `test_inds`.
"""
function checkmageq(test_inds::Vector{Int},
    cs_IDs::Vector{Int},
    dict_J_ID_to_val,
    dict_H_IDs_to_css,
    J_IDs;
    atol = 1e-6)::Bool

    if length(test_inds) == 1
        # need to have at least two nuclei to be magnetically equivalent.
        return false
    end

    common_cs_IDs = cs_IDs[test_inds]
    p_IDs = getpairs(common_cs_IDs)

    ## test common J within pairs.
    J_p = collect( getJfromdict(p_IDs[l][1], p_IDs[l][2], dict_J_ID_to_val) for l = 1:length(p_IDs) )
    pass_common_J_flag = isallsame(J_p; atol = atol)

    ## test common J with other IDs.

    # find all connections to first pair.
    t_IDs = getJIDstest(J_IDs, common_cs_IDs, dict_H_IDs_to_css; atol = atol)

    # check J-values.
    pass_test_J_flag = true
    for k = 1:length(t_IDs)

        J_t = collect( getJfromdict(t_IDs[k], base_ID, dict_J_ID_to_val) for base_ID in common_cs_IDs )

        pass_test_J_flag = pass_test_J_flag & isallsame(J_t; atol = atol)
    end

    return pass_common_J_flag & pass_test_J_flag
end

"""
getJfromdict(i::Int, j::Int, dict::Dict{Tuple{Int,Int},T} ) where T

query `dict` for the entry `(i,j)` and `(j,i)`. Returns the entry for `(i,j)`. Returns the entry for `(j,i)` if `(i,j)` is not found. Return zero if both entries are not found.
"""
function getJfromdict(i::Int, j::Int, dict::Dict{Tuple{Int,Int},T} ) where T
    #
    if haskey(dict, (i,j))
        return dict[(i,j)]
    end

    if haskey(dict, (j,i))
        return dict[(j,i)]
    end

    return zero(T)
end

"""
combinetransitiveeqgroups(eq_inds::Vector{Vector{Int}})

`eq_inds` contains theequivalent nuclei indices.
This function combine the groups of equivalent nuclei via transitivity.

Assumes the equivalence relation cannot be true for a group containg only one nucleus.
"""
function combinetransitiveeqgroups(eq_inds::Vector{Vector{Int}})

    if isempty(eq_inds)
        return Vector{Vector{Int}}(undef, 0)
    end

    eq_inds_flat_unique = unique(combinevectors(eq_inds))
    V = 1:length(eq_inds_flat_unique)
    dict_eq_to_V = Dict(eq_inds_flat_unique .=> V)
    dict_V_to_eq = Dict(V .=> eq_inds_flat_unique)

    eq_inds_V = collect( collect( dict_eq_to_V[eq_inds[i][k]] for k = 1:length(eq_inds[i])) for i = 1:length(eq_inds) )

    # get the edges of fully connected
    E = Vector{Tuple{Int,Int}}(undef, 0)
    for k = 1:length(eq_inds)
        E_k = getconnectpath(eq_inds_V[k])

        push!(E, E_k...)
    end

    #
    h = Graphs.SimpleGraph(V[end])
    for k = 1:length(E)
        Graphs.add_edge!(h, E[k][1], E[k][2])
    end

    H = Graphs.connected_components(h)

    # convert back to the nuclei indexing of `eq_inds`.
    H_out = collect( collect( dict_V_to_eq[H[i][k]] for k = 1:length(H[i])) for i = 1:length(H) )

    return H_out
end

### final package for a compound.
function getmageqinfo(H_IDs, H_css::Vector{T}, J_IDs, J_vals;
    unique_cs_atol = 1e-6)::Tuple{Vector{Vector{Vector{Int}}},Vector{Vector{Vector{Int}}},Vector{Vector{Vector{Int}}}} where T <: Real

    # parse the spin systems and singlets.
    J_inds_sys, J_inds_sys_local, J_IDs_sys, J_vals_sys, H_inds_sys,
    cs_sys, H_inds_singlets, cs_singlets, H_inds, J_inds,
    g = setupcsJ(H_IDs, H_css, J_IDs, J_vals)

    # set up look-up.
    dict_H_IDs_to_css = Dict(H_IDs .=> H_css)
    dict_H_inds_to_css = Dict(H_inds .=> H_css)
    dict_H_ID_to_ind = Dict(H_IDs .=> collect(1:length(H_IDs)))
    dict_ind_to_H_ID = Dict( collect(1:length(H_IDs)) .=> H_IDs)

    dict_J_ID_to_val = Dict(J_IDs .=> J_vals)
    dict_J_ind_to_val = Dict(J_inds .=> J_vals)

    dict_J_ID_to_ind = Dict(J_IDs .=> J_inds)
    dict_J_ind_to_ID = Dict(J_inds .=> J_IDs)

    #
    mag_eq_sys_inds_local, mag_eq_sys_IDs,
    mag_eq_sys_inds_global = getmageqcompound(g,
    H_inds_sys, dict_ind_to_H_ID, dict_H_inds_to_css,
    dict_H_IDs_to_css, dict_J_ID_to_val, J_IDs; atol = unique_cs_atol)

    return mag_eq_sys_inds_local, mag_eq_sys_IDs,
    mag_eq_sys_inds_global
end

function getmageqinfomixture(target_names::Vector{String},
    base_path::String,
    dict_compound_to_filename;
    unique_cs_atol = 1e-6)

    N_compounds = length(target_names)
    MEs = Vector{Vector{Vector{Vector{Int}}}}(undef, N_compounds)

    for n = 1:N_compounds

        # TODO add error-handling if name is not found in the dictionary, or filename does not exist.
        load_path = joinpath(base_path, dict_compound_to_filename[target_names[n]]["file name"])
        H_IDs, H_css, J_IDs, J_vals = loadcouplinginfojson(load_path)

        MEs[n], _ = getmageqinfo(H_IDs, H_css, J_IDs, J_vals; unique_cs_atol = unique_cs_atol)
    end

    return MEs
end


export
    default_initializer

using ..Operators

"""
Deprecated
"""
function default_initializer(runtime::InstantRuntime, default_range_size::Int=10,
        placeholder_beliefs :: Dict = Dict())
    ensure_all!(runtime)
    for (pname, bel) in placeholder_beliefs
        node = get_node(runtime, pname)
        inst = current_instance(runtime, node)
        post_belief!(runtime, inst, bel)
    end
    order = topsort(get_initial_graph(get_network(runtime)))
    set_ranges!(runtime, Dict{Symbol, Score}(), default_range_size, 1, order, placeholder_beliefs)
end

#=
function default_initializer(runtime::DynamicRuntime{T}, default_time::T, default_range_size::Int=10) where {T}
    ensure_all!(runtime, default_time)
    net = get_network(runtime)
    ord = topsort(get_transition_graph(net))
    ranges = Dict{Symbol, Array{Any, 1}}()
    for var in ord
        inst = current_instance(runtime, var)
        sf = get_sfunc(inst)
        parranges = collect(map(p -> ranges[p.name], get_parents(net, var)))
        rng = support(sf, Tuple(parranges), default_range_size, Vector{output_type(sf)}())
        set_range!(runtime, inst, rng)
        ranges[var.name] = rng
    end
end

function default_initializer(runtime::DynamicRuntime{Int})
    default_initializer(runtime, 0)
end
=#

function  preserve_evidence(evid_i, typeOfEvidence)
    preserve_values = Vector{typeOfEvidence}()
    if isa(evid_i, Dict) # soft evidence
        push!(preserve_values, collect(keys(evid_i)))
    elseif isa(evid_i, NoiseEvidence) # noise evidence
        push!(preserve_values, evid_i.mean)
    elseif isa(evid_i, Function)# hard evidence
        # do nothing
    else # hard evidence
        push!(preserve_values, evid_i)
    end
    return preserve_values
end


export predict
using TupleVectors
using SampleChains

predict(args...; kwargs...) = predict(Random.GLOBAL_RNG, args...; kwargs...)

# TODO: Fix this hack
predict(d::AbstractMeasure, x) = x
predict(d::Dists.Distribution, x) = x
predict(d::AbstractModel, args...; kwargs...) = predict(Random.GLOBAL_RNG, d, args...; kwargs...)

@inline function predict(rng::AbstractRNG, m::AbstractModel, nt::NamedTuple{N}) where {N}
    pred = predictive(Model(m), N...)
    rand(rng, pred(merge(argvals(m), nt)))
end

predict(rng::AbstractRNG, m::AbstractModel; kwargs...) = predict(rng, m, (;kwargs...))

@inline function predict(rng::AbstractRNG, d::AbstractModel, nt::LazyMerge)
    predict(rng, d, convert(NamedTuple, nt))
end

function predict(rng::AbstractRNG, d::AbstractModel, post::AbstractVector{<:NamedTuple{N}}) where {N}
    m = Model(d)
    pred = predictive(m, N...)
    args = argvals(d)
    y1 = rand(rng, pred(merge(args,post[1])))
    n = length(post)
    v = TupleVectors.chainvec(y1, n)
    @inbounds for j in 2:n
        newargs = merge(args,post[j])
        v[j] = rand(rng, pred(newargs))
    end

    v
end

function predict(rng::AbstractRNG, d::ConditionalModel, post::MultiChain)
    [predict(rng, d, c) for c in getchains(post)]
end


struct priorInfo{T1, T2, T3 <:Function}
    f::T1
    df::T2
    d2f::T3
end

function initialize_prior(range, type="Uniform")
    if type=="Uniform"
        return priorInfo(x -> 1/(range[2]-range[1]), x-> 0, x->0)
    else
        throw(ArgumentError("Prior type not supported."))
    end
end


using Trilinos
using Test
using MPI
import Trilinos.Testing.arrayview_sum

a = collect(1.0:10.0)
s = sum(a)
s_half = sum(a[1:5])

@test arrayview_sum(Teuchos.ArrayView(a)) == s
@test arrayview_sum(a) == s
@test arrayview_sum(a[1:5]) == s_half
@test arrayview_sum(@view a[1:5]) == s_half
@test arrayview_sum(Teuchos.ArrayView(a,5)) == s_half

function time_av(a)
  cumsum = 0.0
  for i in 1:10000
    cumsum += arrayview_sum(Teuchos.ArrayView(a,5))
  end
  return cumsum
end

@time time_av(a)
@time time_av(a)
@time time_av(a)


######################################################################
# policies.jl
# Does basic policy stuff.
# Policy's don't need to remember vehicle's max_step,
#  as the action function ensures the step is normalized to max_step.
######################################################################

abstract type Policy end

"""
`action(m, x, o, f, p)`

Returns an action given SearchDomain `m`, vehicle `x`, observation `o`, filter `f`, and policy `p`.
"""
function action(m::SearchDomain, x::Vehicle, o, f::AbstractFilter,p::Policy)
	return error(typeof(p), " does not yet implement action.")
end

# TODO: should ensure the heading is not too much either
function normalize(a::Action, x::Vehicle)
	ax = a[1]; ay = a[2]
	den = sqrt(ax*ax + ay*ay)
	return x.max_step * ax / den, x.max_step * ay / den, a[3]
end

function reset!(p::Policy) end

# to make action lists
include("action_list.jl")

######################################################################
# greedy policies
######################################################################
include("greedy.jl")
include("greedyp.jl")
include("greedyp2.jl")
include("greedy2.jl")

######################################################################
# GaussianMPC
######################################################################
include("gaussian_mpc.jl")

######################################################################
# Waypoint policy
######################################################################
include("waypoint.jl")

######################################################################
# Random policy
######################################################################
type RandomPolicy <: Policy end

function action(m::SearchDomain, x::Vehicle, o, f::AbstractFilter, p::RandomPolicy)
	ax = rand() - 0.5
	ay = rand() - 0.5
	ah = 10.0
	if rand() < 0.5
		ah = -10.0
	end
	return normalize((ax,ay,ah), x)
end

# random policy for rotate for bearing method
include("random2.jl")

######################################################################
# Sit policy: Just sit and don't move
######################################################################
type SitPolicy <: Policy end

action(m::SearchDomain, x::Vehicle, o, f::AbstractFilter, p::SitPolicy) = (0.0,0.0,0.0)

######################################################################
# Spin Policy: Just sit there and spin
######################################################################
type SpinPolicy <: Policy end

function action(m::SearchDomain, x::Vehicle, o, f::AbstractFilter, p::SpinPolicy)
	return (0.0, 0.0, 10.0)
end

######################################################################
# CirclePolicy: Move orthogonally to last measurement

# This ends up tracing a circle around the jammer
######################################################################
type CirclePolicy <: Policy
	last::Action

	CirclePolicy() = new( (0.0, 0.0, 0.0) )
end

# Remembers last action to ensure we follow same direction around circle
# Otherwise, it will "chatter" back and forth
# TODO: take into account distance from edge
function action(m::SearchDomain, x::Vehicle, o, f::AbstractFilter, p::CirclePolicy)
	#ax = -1.0 / sind(o)
	#ay = 1.0 / cosd(o)

	ax = -cosd(o);
	ay = sind(o);

	if ax*p.last[1] + ay*p.last[2] < 0.
		ax = -ax
		ay = -ay
	end

	temp_a = normalize((ax,ay,0.0), x)
	temp_x = x.x + temp_a[1]
	temp_y = x.y + temp_a[2]
	ax = temp_a[1]
	ay = temp_a[2]
	# If this moves you out of bounds, go the other way
	if temp_x < 0. || temp_x > m.length || temp_y < 0. || temp_y > m.length
		ax *= -1.
		ay *= -1.
	end

	p.last = (ax, ay, 0.0)
	return (ax, ay, 0.0)
	#return normalize((ax,ay,0.0), x)
end

######################################################################
# SpiralPolicy: Like CirclePolicy, but also move towards jammer
#
# TODO: actually implement this
######################################################################
# Move orthogonally to last measurement
type SpiralPolicy <: Policy
	last::Action

	SpiralPolicy() = new( (0.0, 0.0) )
end

# Remembers last action to ensure we follow same direction around circle
# Otherwise, it will "chatter" back and forth
# TODO: take into account distance from edge
function action(m::SearchDomain, x::Vehicle, o, f::DF, p::SpiralPolicy)
	ax = -1.0 / sind(o)
	ay = 1.0 / cosd(o)

	if ax*p.last[1] + ay*p.last[2] < 0.
		ax = -ax
		ay = -ay
	end

	p.last = (ax, ay)

	return normalize((ax,ay), x)
end


using JLD
using ProgressMeter
using Distributions
using StatsBase
include("make_SQ_model.jl")
include("utilities.jl")
include("self_confidence.jl")

function get_net_input_vals(net::Dict{Symbol,Any},inpts::Array{Symbol},sch::ML.Schema;return_scaled::Bool=false)
    info("# Getting inputs")
    ins = Array{Float64}(length(inpts),1)
    scaled_ins = Array{Float64}(length(inpts),1)
    for (i,n) in zip(collect(1:length(inpts)),inpts)
        ins[i] = net[:training_data][n]
        if return_scaled
            scaled_ins[i] = (ins[i]-mean(sch[n]))/std(sch[n])
        end
    end
    if return_scaled
        return ins, scaled_ins
    else
        return ins
    end
end

function get_surrogate(fname::String,epocs::Int64,fldr::String,net_type::String,
                       inpts::Array{Symbol},trusted_fname::String)
    # make an sq model if it doesn't exist already
    if !(any([contains(x,fname*"_") for x in readdir(fldr)]))
        info("No nn file exists, making one now")
        make_sq_model(net_type,inpts,num_epoc=epocs,trusted_fname=trusted_fname)
    end

    info("Loading xQ Model")
    param_files = searchdir(fldr,fname,".params")
    num_epocs = parse(split(match(r"-\d+",param_files[1]).match,"-")[2])

    fn = joinpath(fldr,fname)
    SQmodel = load_network(fn,num_epocs,fn*"_SQmodel.jld")

    return SQmodel
end

function calculate_sq(experiment_dict)
    trusted_cond = 0
    for c in keys(experiment_dict[:conditions])
        if experiment_dict[:conditions][c][:role] == :trusted
            trusted_cond = c
        end
    end
    for expr in keys(experiment_dict[:conditions])
        net_type = experiment_dict[:name]
        inpts = experiment_dict[:xQ][:inpts]
        epocs = experiment_dict[:xQ][:epocs]
        println("Processing: $(experiment_dict[:conditions][expr])")

        println("running '$expr' data")

        ###### Define the variables for the problem
        inputs = Dict()
        for i in inpts
            inputs[i] = "ML.Continuous"
        end

        log_fname = "$(net_type)_$(make_label_from_keys(inputs))"
        log_loc = experiment_dict[:xQ][:nn_loc]
        tfname = experiment_dict[:conditions][trusted_cond][:fname]
        tfold = experiment_dict[:conditions][trusted_cond][:fldr]
        cfname = experiment_dict[:conditions][expr][:fname]
        cfold = experiment_dict[:conditions][expr][:fldr]
        trusted_path = joinpath(tfold,tfname*".csv")

        ###### get surrogate model
        SQmodel = get_surrogate(log_fname,epocs,log_loc,net_type,inpts,trusted_path)
        info("restoring limits")
        limits = restore_eng_units(SQmodel.range,SQmodel.output_sch)
        println("limits: $limits")
        println("out_sch: $(SQmodel.output_sch)")

        ###### calculate xQ for each net
        df = JLD.load(joinpath(cfold,cfname*".jld"))
        d = df["problem_dict"]
        g = df["global_rwd_range"]
        info("Processing Networks")
        p = Progress(length(d),barglyphs=BarGlyphs("[=> ]"),color=:white)
        j = 0
        for net in d
            net_num = net[1]
            net = net[2]
            e_start = net[:evader_start]
            p_start = net[:pursuer_start]
            ins, scaled_ins = get_net_input_vals(net,inpts,SQmodel.input_sch,return_scaled=true)

            info("getting predictions")
            println("ins: $ins, $(typeof(ins))")
            _notused, R_star = SQ_predict(SQmodel,scaled_ins,use_eng_units=true)
            R_star_μ = R_star[:X3_1][1]
            R_star_σ = R_star[:X3_2][1]

            R_μ = net[:training_data][:X3_1] # mean of rwd dist
            R_σ = net[:training_data][:X3_2] # std of rwd dist

            info("R_star_μ = $R_star_μ, R_star_σ = $R_star_σ")
            info("R_μ = $R_μ, R_σ = $R_σ")
            # make sure sigma isn't zero or negative (i.e. model could have bad prediction, or data could be all an identical number)
            # is R_star <= 0. ? if yes R_star_σ = 1. : otherwise do nothing
            R_star_σ <= 0. ? R_star_σ = 1. : nothing
            R_σ <= 0. ? R_σ = 1. : nothing

            x_Q, x_Q_Dict = X3(Normal(R_μ,R_σ),Normal(R_star_μ,R_star_σ),global_rwd_range=g,return_raw_sq=true)

            net[:solver_quality] = Dict{Symbol,Any}()
            net[:solver_quality][:x_QDict] = deepcopy(x_Q_Dict)
            net[:solver_quality][:R_star] = R_star
            net[:solver_quality][:X_Q] = x_Q
            j += 1
            ProgressMeter.next!(p; showvalues = [(:network,"$j of $(length(d))")])
        end
        # re-write original file with xQ added to each network
        out_fname = joinpath("logs",net_type*"_$(expr)_"*make_label_from_keys(inputs)*"_solver.jld")
        info("# writing to $out_fname")
        jldopen(out_fname, "w") do file
            write(file, "problem_dict", d)
            write(file, "global_rwd_range", g)
        end
    end
end


function R_tree(tree::Tree ; file = "tree.Rdata")
    edges, branch_lengths, temp_rates, tip_labels = make_ape(tree)
    ntip = (tree.n_nodes + 1)/2

    @rput edges
    @rput branch_lengths
    @rput temp_rates
    @rput ntip
    @rput file
    @rput tip_labels

    reval("""
        temp_tree = list(edge = edges, Nnode = ntip - 1, edge.length = branch_lengths, tip.label = tip_labels)
        class(temp_tree) = "phylo"
    """)
end

"""
    save_ClaDS_in_R(co::CladsOutput, path::String ; maxit = Inf, ...)
sss
Save the output of a ClaDS run as a Rdata file.

# Arguments
- `co::CladsOutput`: the result of a run of `infer_ClaDS`.
- `path::String`: the path the file should be saved to.

# .Rdata file
The .Rdata file contains an object called `CladsOutput` that is a list with the following fields :
- `tree`: the phylogeny, saved as a `phylo` object
- `chains`: the MCMC chains
- `rtt_chains`: the MCMCs with the rate through time information
- `sig_map`: the MAP estimate for the σ parameter
- `al_map`: the MAP estimate for the α parameter
- `eps_map`: the MAP estimate for the ε parameter
- `lambda0_map`: the MAP estimate for the λ0 parameter
- `lambdai_map`: the MAP estimate for the branch-specific speciation rates
- `lambdatip_map`: the MAP estimate for the tip speciation rates
- `DTT_mean`: the point estimate for the diversity through time
- `RTT_map` : the rate through time estimates
- `time_points` : the times at which `DTT_mean` and `RTT_map` are computed
- `enhanced_trees` : a sample of the complete tree distribution, as a list of `phylo` objects
- `gelm` : the gelman parameter
"""
function save_ClaDS_in_R(co::CladsOutput, path::String)

    chains = co.chains
    rtt_chains = co.rtt_chains
    sig_map = co.σ_map
    al_map = co.α_map
    eps_map = co.ε_map
    lambda0_map = co.λ0_map
    lambdai_map = co.λi_map
    lambdatip_map = co.λtip_map
    DTT_mean = co.DTT_mean
    RTT_map = co.RTT_map
    time_points = co.time_points
    enhanced_trees = co.enhanced_trees
    gelm_id, gelm_stat  = co.gelm

    R_tree(co.tree)
    reval("""tree = temp_tree""")

    reval("""enhanced_trees = list()""")
    for i in 1:length(co.enhanced_trees)
        @rput i
        R_tree(co.enhanced_trees[i])
        reval("""enhanced_trees[[i]] = list(tree = temp_tree, rates = temp_rates)""")
    end

    @rput chains
    @rput rtt_chains
    @rput sig_map
    @rput al_map
    @rput eps_map
    @rput lambda0_map
    @rput lambdai_map
    @rput lambdatip_map
    @rput DTT_mean
    @rput RTT_map
    @rput time_points
    @rput gelm_id
    @rput gelm_stat

    @rput path
    reval("""
        CladsOutput = list(
            tree = tree,
            rtt_chains = rtt_chains,
            chains = chains,
            sig_map = sig_map,
            al_map = al_map,
            eps_map = eps_map,
            lambda0_map = lambda0_map,
            lambdai_map = lambdai_map,
            lambdatip_map = lambdatip_map,
            DTT_mean = DTT_mean,
            RTT_map = RTT_map,
            time_points = time_points,
            enhanced_trees = enhanced_trees,
            gelm = c(gelm_id, gelm_stat)
        )

        save(CladsOutput, file = path)
    """)
end


using BinDeps
using Compat.Libdl

@BinDeps.setup

fmmlib2d = library_dependency("fmmlib2d")

provides(Sources, URI("https://github.com/ludvigak/fmmlib2d/archive/master.zip"),
         fmmlib2d, unpacked_dir = "fmmlib2d-master")

fmmsrcdir = joinpath(BinDeps.srcdir(fmmlib2d), "fmmlib2d-master/src")
libname = "fmmlib2d." * Libdl.dlext
libfile = joinpath(BinDeps.libdir(fmmlib2d),libname)
buildfile = joinpath(fmmsrcdir, libname)

fsrcs = "hfmm2dpart.f hfmm2drouts.f d2tstrcr_omp.f d2mtreeplot.f h2dterms.f helmrouts2d.f cdjseval2d.f hank103.f prini.f cfmm2dpart.f zfmm2dpart.f rfmm2dpart.f lfmm2dpart.f lfmm2drouts.f l2dterms.f laprouts2d.f"

fsrcs = split(fsrcs, " ")
buildcmd = `gfortran -O2 -fPIC -shared $fsrcs -o $libname`

provides(BuildProcess,
         (@build_steps begin
          GetSources(fmmlib2d)
          @build_steps begin
          ChangeDirectory(fmmsrcdir)
          FileRule(libfile, @build_steps begin
                   buildcmd
                   CreateDirectory(libdir(fmmlib2d))
                   `cp $libname $libfile`
                   end)
          end
          end),
         fmmlib2d)

@BinDeps.install Dict(:fmmlib2d => :fmmlib2d)


# Support for @simd for

module SimdLoop

export @simd

# Error thrown from ill-formed uses of @simd
type SimdError <: Exception
    msg::ASCIIString
end

# Parse iteration space expression
#       symbol '=' range
#       symbol 'in' range
function parse_iteration_space(x)
    Meta.isexpr(x, [:(=), :in]) || throw(SimdError("= or in expected"))
    length(x.args) == 2 || throw(SimdError("simd range syntax is wrong"))
    isa(x.args[1], Symbol) || throw(SimdError("simd loop index must be a symbol"))
    x.args # symbol, range
end

# Compile Expr x in context of @simd.
function compile(x)
    Meta.isexpr(x, :for) || throw(SimdError("for loop expected"))
    length(x.args) == 2 || throw(SimdError("1D for loop expected"))

    var,range = parse_iteration_space(x.args[1])
    r = gensym("r") # Range value
    n = gensym("n") # Trip count
    s = gensym("s") # Step
    i = gensym("i") # Index variable
    quote
        # Evaluate range value once, to enhance type and data flow analysis by optimizers.
        let $r = $range, $n = length($r)
            if zero($n) < $n
                let $var = first($r)
                    # LLVM vectorizer needs to compute a trip count, so make it obvious.
                    local $s = step($r)
                    local $i = zero($n)
                    while $i < $n
                        $(x.args[2])
                        $var += $s
                        $i += 1
                        $(Expr(:simdloop))  # Mark loop as SIMD loop
                    end
                end
                # Set index to last value just like a regular for loop would
                $var = last($r)
            end
        end
        nothing
    end
end

macro simd(forloop)
    esc(compile(forloop))
end

end # module SimdLoop


"""
    Cell

A type describing the location at the center of a grid cell.
"""
struct Cell end

"""
	Face

A type describing the location at the face of a grid cell.
"""
struct Face end

"""
    Field{LX, LY, LZ, A, G} <: AbstractField{A, G}

A field defined at the location (`LX`, `LY`, `LZ`) which can be either `Cell` or `Face`.
"""
struct Field{Lx, Ly, Lz, A, G} <: AbstractField{A, G}
    data :: A
    grid :: G
end

"""
	Field(L::Tuple, arch::AbstractArchitecture, grid)

Construct a `Field` on architecture `arch` and `grid` with location defined by the tuple `L` of length 3 whose elements
are `Cell` or `Face`.
"""
function Field(L::Tuple, arch::AbstractArchitecture, grid)
    data = zeros(arch, grid)
    return Field(L, data, grid)
end

"""
	Field(L::Tuple, data::AbstractArray, grid)

Construct a `Field` on `grid` using the array `data` with location defined by the tuple `L` of length 3 whose elements
are `Cell` or `Face`.
"""
Field(L::Tuple, data::AbstractArray, grid) = Field{L[1], L[2], L[3], typeof(data), typeof(grid)}(data, grid)

"""
	Field(Lx, Ly, Lz, data, grid)

Construct a `Field` on `grid` using the array `data` with location defined by `Lx`, `Ly`, and `Lz` which are `Cell`
or `Face`.
"""
Field(Lx, Ly, Lz, data, grid) = Field((Lx, Ly, Lz), data, grid)

"""
	CellField

A field defined at the cell centers. Used for pressure and tracers.
"""
const CellField  = Field{Cell, Cell, Cell}

"""
	FaceFieldX

A field defined at the faces along the x-direction. Used for horizontal velocity u.
"""
const FaceFieldX = Field{Face, Cell, Cell}

"""
	FaceFieldY

A field defined at the faces along the y-direction. Used for horizontal velocity v.
"""
const FaceFieldY = Field{Cell, Face, Cell}

"""
	FaceFieldY

A field defined at the faces along the z-direction. Used for vertical velocity w.
"""
const FaceFieldZ = Field{Cell, Cell, Face}

"""
    CellField([T=eltype(grid)], arch, grid)

Return a `CellField` on architecture `arch` and `grid`.
"""
CellField(T, arch, grid) = Field(Cell, Cell, Cell, zeros(T, arch, grid), grid)

"""
    FaceFieldX([T=eltype(grid)], arch, grid)

Return a `FaceFieldX` on architecture `arch` and `grid`.
"""
FaceFieldX(T, arch, grid) = Field(Face, Cell, Cell, zeros(T, arch, grid), grid)

"""
    FaceFieldY([T=eltype(grid)], arch, grid)

Return a `FaceFieldY` on architecture `arch` and `grid`.
"""
FaceFieldY(T, arch, grid) = Field(Cell, Face, Cell, zeros(T, arch, grid), grid)

"""
    FaceFieldZ([T=eltype(grid)], arch, grid)

Return a `FaceFieldZ` on architecture `arch` and `grid`.
"""
FaceFieldZ(T, arch, grid) = Field(Cell, Cell, Face, zeros(T, arch, grid), grid)

 CellField(arch, grid) = Field(Cell, Cell, Cell, arch, grid)
FaceFieldX(arch, grid) = Field(Face, Cell, Cell, arch, grid)
FaceFieldY(arch, grid) = Field(Cell, Face, Cell, arch, grid)
FaceFieldZ(arch, grid) = Field(Cell, Cell, Face, arch, grid)

fieldtype(f::AbstractField) = typeof(f).name.wrapper
location(::Field{Lx, Ly, Lz}) where {Lx, Ly, Lz} = (Lx, Ly, Lz)

@inline size(f::AbstractField) = size(f.grid)
@inline length(f::AbstractField) = length(f.data)

@propagate_inbounds getindex(f::AbstractField, inds...) = getindex(f.data, inds...)
@propagate_inbounds setindex!(f::AbstractField, v, inds...) = setindex!(f.data, v, inds...)
@inline lastindex(f::AbstractField) = lastindex(f.data)
@inline lastindex(f::AbstractField, dim) = lastindex(f.data, dim)

"Returns a view over the interior points of the `field.data`."
@inline data(f::AbstractField) = view(f.data, 1:f.grid.Nx, 1:f.grid.Ny, 1:f.grid.Nz)

"Returns a reference to the interior points of `field.data.parent.`"
@inline parentdata(f::AbstractField) = @inbounds f.data.parent[1+f.grid.Hx:f.grid.Nx+f.grid.Hx,
                                                               1+f.grid.Hy:f.grid.Ny+f.grid.Hy,
                                                               1+f.grid.Hz:f.grid.Nz+f.grid.Hz]

show(io::IO, f::AbstractField) = show(io, f.data)
iterate(f::AbstractField, state=1) = iterate(f.data, state)

# Define +, -, and * on fields as element-wise calculations on their data. This
# is only true for fields of the same type, e.g. when adding a FaceFieldY to
# another FaceFieldY, otherwise some interpolation or averaging must be done so
# that the two fields are defined at the same point, so the operation which
# will not be commutative anymore.
for ft in (:CellField, :FaceFieldX, :FaceFieldY, :FaceFieldZ)
    for op in (:+, :-, :*)
        @eval begin
            # +, -, * a Field by a Number on the left.
            function $op(num::Number, f::$ft)
                ff = similar(f)
                @. ff.data = $op(num, f.data)
                ff
            end

            # +, -, * a Field by a Number on the right.
            $op(f::$ft, num::Number) = $op(num, f)

            # Multiplying two fields together
            function $op(f1::$ft, f2::$ft)
                f3 = similar(f1)
                @. f3.data = $op(f1.data, f2.data)
                f3
            end
        end
    end
end

xnodes(ϕ::AbstractField) = reshape(ϕ.grid.xC, ϕ.grid.Nx, 1, 1)
ynodes(ϕ::AbstractField) = reshape(ϕ.grid.yC, 1, ϕ.grid.Ny, 1)
znodes(ϕ::AbstractField) = reshape(ϕ.grid.zC, 1, 1, ϕ.grid.Nz)

xnodes(ϕ::Field{Face}) = reshape(ϕ.grid.xF[1:end-1], ϕ.grid.Nx, 1, 1)
ynodes(ϕ::Field{X, Face}) where X = reshape(ϕ.grid.yF[1:end-1], 1, ϕ.grid.Ny, 1)
znodes(ϕ::Field{X, Y, Face}) where {X, Y} = reshape(ϕ.grid.zF[1:end-1], 1, 1, ϕ.grid.Nz)

nodes(ϕ) = (xnodes(ϕ), ynodes(ϕ), znodes(ϕ))

# Niceties
const AbstractCPUField = AbstractField{A, G} where {A<:OffsetArray{T, D, <:Array} where {T, D}, G}
@hascuda const AbstractGPUField = AbstractField{A, G} where {A<:OffsetArray{T, D, <:CuArray} where {T, D}, G}

set!(u::AbstractField, v::Number) = @. u.data = v
set!(u::AbstractField{A}, v::AbstractField{A}) where A = @. u.data.parent = v.data.parent

"Set the CPU field `u` to the array `v`."
function set!(u::AbstractCPUField, v::Array)
    for k in 1:u.grid.Nz, j in 1:u.grid.Ny, i in 1:u.grid.Nx
        u[i, j, k] = v[i, j, k]
    end
    return nothing
end

# Set the GPU field `u` to the array `v`.
@hascuda function set!(u::AbstractGPUField, v::Array)
    FieldType = fieldtype(u)
    v_field = FieldType(location(u), CPU(), u.grid)
    set!(v_field, v)
    set!(u, v_field)
    return nothing
end

# Set the GPU field `u` to the CuArray `v`.
@hascuda function set!(u::AbstractGPUField, v::CuArray)
    @launch device(GPU()) config=launch_config(u.grid, 3) _set_gpu!(u.data, v, u.grid)
    return nothing
end

function _set_gpu!(u, v, grid)
    @loop for k in (1:grid.Nz; (blockIdx().z - 1) * blockDim().z + threadIdx().z)
        @loop for j in (1:grid.Ny; (blockIdx().y - 1) * blockDim().y + threadIdx().y)
            @loop for i in (1:grid.Nx; (blockIdx().x - 1) * blockDim().x + threadIdx().x)
	            u[i, j, k] = v[i, j, k]
	        end
	    end
    end
    return nothing
end

# Set the GPU field `u` data to the CPU field data of `v`.
@hascuda set!(u::AbstractGPUField, v::AbstractCPUField) = copyto!(u.data.parent, v.data.parent)

# Set the CPU field `u` data to the GPU field data of `v`.
@hascuda set!(u::AbstractCPUField, v::AbstractGPUField) = u.data.parent .= Array(v.data.parent)

"Set the CPU field `u` data to the function `f(x, y, z)`."
set!(u::AbstractField, f::Function) = data(u) .= f.(nodes(u)...)

# Set the GPU field `u` data to the function `f(x, y, z)`.
@hascuda function set!(u::AbstractGPUField, f::Function)
    FieldType = fieldtype(u)
    u_cpu = FieldType(location(u), CPU(), u.grid)
    set!(u_cpu, f)
    set!(u, u_cpu)
    return nothing
end

"""
    set!(model; kwargs...)

Set velocity and tracer fields of `model`. The keyword arguments
`kwargs...` take the form `name=data`, where `name` refers to one of the
fields of `model.velocities` or `model.tracers`, and the `data` may be an array,
a function with arguments `(x, y, z)`, or any data type for which a
`set!(ϕ::AbstractField, data)` function exists.

Example
=======

model = Model(grid=RegularCartesianGrid(N=(32, 32, 32), L=(1, 1, 1))

# Set u to a parabolic function of z, v to random numbers damped
# at top and bottom, and T to some silly array of half zeros,
# half random numbers.

u₀(x, y, z) = z/model.grid.Lz * (1 + z/model.grid.Lz)
v₀(x, y, z) = 1e-3 * rand() * u₀(x, y, z)

T₀ = rand(size(model.grid)...)
T₀[T₀ .< 0.5] .= 0

set!(model, u=u₀, v=v₀, T=T₀)
"""
function set!(model::AbstractModel; kwargs...)
    for (fldname, value) in kwargs
        if fldname ∈ propertynames(model.velocities)
            ϕ = getproperty(model.velocities, fldname)
        elseif fldname ∈ propertynames(model.tracers)
            ϕ = getproperty(model.tracers, fldname)
        else
            throw(ArgumentError("name $fldname not found in model.velocities or model.tracers."))
        end
        set!(ϕ, value)
    end
    return nothing
end

function set!(Φ::NamedTuple; kwargs...)
    for (fldname, value) in kwargs
        ϕ = getproperty(Φ, fldname)
        set!(ϕ, value)
    end
    return nothing
end

set_ic!(model; kwargs...) = set!(model; kwargs...) # legacy wrapper


#Small utility library for the definition of the State Vector
import Base.size, Base.getindex, Base.setindex!, Base.IndexStyle
import Base.*, Base.+, Base.-, Base./

!isdefined(:Quaternion) ? include("Quatlib.jl") : nothing

mutable struct State <: AbstractArray{Float64,1}
    vx::Float64
    vy::Float64
    vz::Float64
    x::Float64
    y::Float64
    z::Float64
    wx::Float64
    wy::Float64
    wz::Float64
    q1::Float64
    q2::Float64
    q3::Float64
    q4::Float64
end

Base.size(::State) = (13,)
Base.IndexStyle(::Type{State}) = IndexLinear()
function Base.getindex(state::State, i::Int)
    if (i > 0 && i < 14)
        if i == 1
            return state.vx
        elseif i == 2
            return state.vy
        elseif i == 3
            return state.vz
        elseif i == 4
            return state.x
        elseif i == 5
            return state.y
        elseif i == 6
            return state.z
        elseif i == 7
            return state.wx
        elseif i == 8
            return state.wy
        elseif i == 9
            return state.wz
        elseif i == 10
            return state.q1
        elseif i == 11
            return state.q2
        elseif i == 12
            return state.q3
        elseif i == 13
            return state.q4
        end
    else
        error("index out of range")
    end
end

function setindex!(state::State, v::T, i::Int) where {T <: Real}
    if (i > 0 && i < 14)
        if i == 1
            state.vx = Float64(v)
        elseif i == 2
            state.vy = Float64(v)
        elseif i == 3
            state.vz = Float64(v)
        elseif i == 4
            state.x = Float64(v)
        elseif i == 5
            state.y = Float64(v)
        elseif i == 6
            state.z = Float64(v)
        elseif i == 7
            state.wx = Float64(v)
        elseif i == 8
            state.wy = Float64(v)
        elseif i == 9
            state.wz = Float64(v)
        elseif i == 10
            state.q1 = Float64(v)
        elseif i == 11
            state.q2 = Float64(v)
        elseif i == 12
            state.q3 = Float64(v)
        elseif i == 13
            state.q4 = Float64(v)
        end
    else
        error("index out of range")
    end
end

velocity(s::State) = s[1:3]
position(s::State) = s[4:6]
angularvelocity(s::State) = s[7:9]
orientation(s::State) = Quaternion(s[10], s[11:13])
function State(vel::Vector{T} = zeros(3),pos::Vector{S} = zeros(3),
        angvel::Vector{R} = zeros(3), quat::Quaternion = Quaternion(0., [0., 0., 0.])) where {T,S,R <: Real}
    State(vel[1], vel[2], vel[3], pos[1], pos[2], pos[3],
        angvel[1], angvel[2], angvel[3], quat[1], quat[2], quat[3], quat[4])
end


@testset "isapprox" begin
    @testset "Composite{Tuple}" begin
        @testset "basic" begin
            x_tup = (1.5, 2.5, 3.5)
            x_comp = Composite{typeof(x_tup)}(x_tup...)
            @test x_comp ≈ x_tup
            @test x_tup ≈ x_comp
            @test x_comp ≈ x_comp

            @test_throws Exception x_comp ≈  collect(x_tup)
        end

        @testset "different types" begin
            # both of these are reasonable diffentials for the `Tuple{Int, Int}` primal
            @test Composite{Tuple{Int, Int}}(1f0, 2f0) ≈ Composite{Tuple{Int, Int}}(1.0, 2.0)

            D = Diagonal(randn(5))
            @test Composite{typeof(D)}(diag=D.diag) ≈ Composite{typeof(D)}(diag=D.diag)

            # But these have different primals so should not be equal
            @test !(Composite{Tuple{Int, Int}}(1.0, 2.0) ≈ Composite{Tuple{Float64, Float64}}(1.0, 2.0))
        end
    end
end


#///////////////////////////////////////
#// File Name: type_definition.jl
#// Author: Haruki Nishimura (hnishimura@stanford.edu)
#// Date Created: 2021/01/15
#// Description: Type definitions for Risk Sensitive Stochastic SAC
#///////////////////////////////////////

import Base: ==, isequal, isapprox
import RobotOS: Time

using DataStructures
using LinearAlgebra
using RobotOS

abstract type Parameter end;

abstract type State end;

# WayPoint
struct WayPoint2D
    pos::Vector{Float64}
    t::Time
    function WayPoint2D(pos::Vector{Float64}, t::Time)
        @assert length(pos) == 2 "Invalid waypoint dimension!"
        return new(pos, t)
    end
end

# Trajectory
const Trajectory2D = OrderedDict{Time, Vector{Float64}};

function Trajectory2D(points::Vector{WayPoint2D})
    trajectory = Trajectory2D();
    for waypoint in points
        trajectory[waypoint.t] = waypoint.pos;
    end
    sort!(trajectory)
    return trajectory
end

function get_position(trajectory::Trajectory2D, t::Time)
    if haskey(trajectory, t)
        return trajectory[t]
    else
        traj_array = collect(trajectory);
        index2 = findfirst(x -> x[1] > t, traj_array);
        @assert !(index2 == 1 || index2 == nothing) "Trajectory Interpolation failed for t = $(to_sec(t)) [s]!"
        t2 = traj_array[index2][1];
        index1 = index2 - 1;
        t1 = traj_array[index1][1];
        # linear interpolation
        pos = traj_array[index1][2] + to_nsec(t - t1)/to_nsec(t2 - t1)*
              (traj_array[index2][2] - traj_array[index1][2])
        return pos
    end
end

# Ego Robot State
struct RobotState <: State
    x::Vector{Float64} # [pos_x, pos_y, vel_x, vel_y]
    t::Time
    function RobotState(x::Vector{Float64}, t::Time)
        @assert length(x) == 4 "Invalid ego state dimension!"
        return new(x, t)
    end
end

RobotState(x::Vector{Float64}) = RobotState(x, Time());

function isequal(e1::RobotState, e2::RobotState)
    return e1.x == e2.x && e1.t == e2.t
end

==(e1::RobotState, e2::RobotState) = isequal(e1, e2)

function isapprox(e1::RobotState, e2::RobotState; kwargs...)
    return isapprox(e1.x, e2.x; kwargs...) && e1.t == e2.t
end

get_position(e::RobotState) = e.x[1:2];
get_velocity(e::RobotState) = e.x[3:4];

# Joint State
mutable struct WorldState <: State
    e_state::RobotState
    ap_dict::Dict{String, Vector{Float64}}
    t::Time
    t_last_m::Union{Time, Nothing}
    function WorldState(e_state::RobotState,
                        ap_dict::Dict{String, Vector{Float64}},
                        t::Time, t_last_m::Union{Time, Nothing})
        @assert all(length.(values(ap_dict)) .== 2) "Invalid ado state dimension!"
        @assert e_state.t == t "Inconsistent time stamps!";
        if t_last_m != nothing
            @assert t_last_m <= t "Last measurement time cannot be in future!"
        end
        return new(e_state, ap_dict, t, t_last_m)
    end
end

function WorldState(e_state::RobotState,
                    ap_dict::Dict{String, Vector{Float64}},
                    t_last_m=nothing)
    return WorldState(e_state, ap_dict, e_state.t, t_last_m)
end

function WorldState(x::Vector{Float64},
                    ap_dict::Dict{String, Vector{Float64}},
                    t::Time=Time(), t_last_m=nothing)
    return WorldState(RobotState(x, t), ap_dict, t, t_last_m)
end
function isequal(w1::WorldState, w2::WorldState)
    return begin w1.e_state == w2.e_state && w1.ap_dict == w2.ap_dict &&
                 w1.t == w2.t && w1.t_last_m == w2.t_last_m
           end
end

==(w1::WorldState, w2::WorldState) = isequal(w1, w2);

function isapprox(w1::WorldState, w2::WorldState; kwargs...)
    e_approx = isapprox(w1.e_state, w2.e_state; kwargs...);
    ap_approx = keys(w1.ap_dict) == keys(w2.ap_dict) &&
                all([isapprox(w1.ap_dict[key], w2.ap_dict[key]; kwargs...)
                     for key in keys(w1.ap_dict)]);
    return e_approx && ap_approx && w1.t == w2.t && w1.t_last_m == w2.t_last_m
end


module DataGenerators

#generator macro
export @generator

# exported functions that are used to generate objects from a generator
export choose, robustchoose, generate, meta, choicepointinfo, choicemodel, setchoicemodel!, GenerationTerminatedException

# exported functions that are used to register and match generators
export register, generatorfor

# exported choice models application function
export setsimplechoicemodel!, setsamplerchoicemodel!, setnmcschoicemodel!, mctschoicemodel!
export SimpleChoiceModel, SamplerChoiceModel, NMCSChoiceModel, MCTSChoiceModel, GEChoiceModel, SwitchChoiceModel

# exported Choice Model functions
export paramranges, setparams!, getparams, numparams, estimateparams!, estimateconditionalmodel!

using DataGeneratorTranslators
using Distributions
import Base.show

const THIS_MODULE = current_module() # used when creating calls to this module in the macro

# list of number types that are supported directly by generator macro
# (could do this automatically as leaf subtypes of Integer and AbstractFloat, but some of these we can't really handle directly yet
# - e.g. BigInt, BigFloat - and it is possible that custom subtypes could have been added)
const GENERATOR_SUPPORTED_CHOOSE_NUMBER_TYPES = Type[Bool, Int8, Int16, Int32, Int64, UInt8, UInt16, UInt32, UInt64, Float16, Float32, Float64,]

# list of number types that are supported directly by generator macro
# (could do this automatically as leaf subtypes of AbstractString, but some of these we can't really handle directly yet)
const GENERATOR_SUPPORTED_CHOOSE_STRING_TYPES = Type[String,]

# all choose types supported directly by the generator maco
const GENERATOR_SUPPORTED_CHOOSE_TYPES = Type[GENERATOR_SUPPORTED_CHOOSE_NUMBER_TYPES; GENERATOR_SUPPORTED_CHOOSE_STRING_TYPES]

# generation and associated functions
include("generation.jl")

# @generator macro
include("generator_macro.jl")

# updating and querying registry metadata
include("registry.jl")

# choice models
include(joinpath("choice_model","choice_model.jl"))

end


Base.:(==)(x::nodeTypeDln, y::Real) = ==(Int(x), y)
Base.:(==)(x::Real, y::nodeTypeDln) = ==(x, Int(y))
Base.convert(::Type{nodeTypeDln}, x::Real) = nodeTypeDln(Int(x))
Base.zero(::Type{nodeTypeDln}) = nodeTypeDln(0)

Base.getindex(x::DislocationLoop, i) = i == 1 ? x : throw(BoundsError())
Base.eachindex(x::DislocationLoop) = 1

function Base.zero(::Type{DislocationLoop})
    return DislocationLoop(;
        loopType = loopDln(),
        numSides = convert(Int, 0),
        nodeSide = convert(Int, 0),
        numLoops = convert(Int, 0),
        segLen = convert(Float64, 0),
        slipSystemIdx = convert(Int, 0),
        slipSystem = SlipSystem(nothing, zeros(3), zeros(3)),
        label = zeros(nodeTypeDln, 0),
        buffer = convert(Float64, 0),
        range = zeros(3, 0),
        dist = Zeros(),
    )
end

function Base.zero(::Type{DislocationNetwork})
    return DislocationNetwork(;
        links = zeros(Int, 2, 0),
        slipPlane = zeros(3, 0),
        bVec = zeros(3, 0),
        coord = zeros(3, 0),
        label = zeros(nodeTypeDln, 0),
        nodeVel = zeros(3, 0),
        nodeForce = zeros(3, 0),
        numNode = [0],
        numSeg = [0],
        maxConnect = 0,
        segForce = zeros(3, 2, 0),
        linksConnect = zeros(Int, 2, 0),
        extSeg = zeros(Bool, 0),
        segIdx = zeros(Int, 0, 3),
    )
end
function Base.push!(network::DislocationNetwork, n)
    elemT = eltype(network.coord)
    network = DislocationNetwork(;
        links = hcat(network.links, zeros(Int, 2, n)),
        slipPlane = hcat(network.slipPlane, zeros(elemT, 3, n)),
        bVec = hcat(network.bVec, zeros(elemT, 3, n)),
        coord = hcat(network.coord, zeros(elemT, 3, n)),
        label = vcat(network.label, zeros(nodeTypeDln, n)),
        nodeVel = hcat(network.nodeVel, zeros(elemT, 3, n)),
        nodeForce = hcat(network.nodeForce, zeros(elemT, 3, n)),
        numNode = copy(network.numNode),
        numSeg = copy(network.numSeg),
        maxConnect = copy(network.maxConnect),
        connectivity = hcat(
            network.connectivity,
            zeros(Int, size(network.connectivity, 1), n),
        ),
        linksConnect = hcat(network.linksConnect, zeros(Int, 2, n)),
        extSeg = vcat(network.extSeg, zeros(Bool, n)),
        segIdx = vcat(network.segIdx, zeros(Int, n, 3)),
        segForce = cat(network.segForce, zeros(elemT, 3, 2, n), dims = 3),
    )

    return network
end
function Base.getindex(network::DislocationNetwork, i)
    return @views network.links[:, i],
    network.slipPlane[:, i],
    network.bVec[:, i],
    network.coord[:, i],
    network.label[i],
    network.nodeVel[:, i],
    network.nodeForce[:, i],
    network.connectivity[:, i],
    network.linksConnect[:, i],
    network.extSeg[i],
    network.segIdx[i, :],
    network.segForce[:, :, i]
end
function Base.iszero(network::DislocationNetwork)
    (network.numNode[1] == 0 || network.numSeg[1] == 0) && return true
    return false
end
function Base.getindex(slipSystem::SlipSystem, i)
    return @views slipSystem.slipPlane[:, i], slipSystem.bVec[:, i]
end

"""
```
loopDistribution(::Rand, n, args...; kw...)
```
Returns a `3 × n` matrix whose points follow a random uniform distribution. Used by [`DislocationNetwork`](@ref) when the `dist` parameter of [`DislocationLoop`](@ref) is `Rand()`.
"""
loopDistribution(::Rand, n, args...; kw...) = rand(3, n)
"""
```
loopDistribution(::Zeros, n, args...; kw...)
```
Returns a `3 × n` zeros matrix. Used by [`DislocationNetwork`](@ref) when the `dist` parameter of [`DislocationLoop`](@ref) is `Zeros()`.
"""
loopDistribution(::Zeros, n, args...; kw...) = zeros(3, n)
"""
```
loopDistribution(::Rand, n, args...; kw...)
```
Returns a `3 × n` matrix whose points follow a random normal distribution. Used by [`DislocationNetwork`](@ref) when the `dist` parameter of [`DislocationLoop`](@ref) is `Randn()`.
"""
loopDistribution(::Randn, n, args...; kw...) = randn(3, n)
"""
```
loopDistribution(::Rand, n, args...; kw...)
```
Returns a `3 × n` matrix whose points follow are regularly placed. Used by [`DislocationNetwork`](@ref) when the `dist` parameter of [`DislocationLoop`](@ref) is `Regular()`.

!!! note
    Not yet implemented.
"""
function loopDistribution(::Regular, n, args...; kw...)
    return error("loopDistribution: regular distribution yet not implemented")
end

"""
```
limits!(lims, segLen, range, buffer)
```
Compute the bounding limits within which a [`DislocationLoop`](@ref) will be distributed in a [`DislocationNetwork`](@ref).
"""
function limits!(lims, segLen, range, buffer)
    @inbounds for i in 1:size(lims, 2)
        @simd for j in 1:size(lims, 1)
            lims[j, i] = range[j, i] + buffer * segLen
        end
    end
    return lims
end

"""
```
translatePoints!(coord, lims, disp)
```
Translates coordinates using the limits and displacements calculated by [`limits!`](@ref) and [`loopDistribution`](@ref).
"""
function translatePoints!(coord, lims, disp)
    @inbounds for i in 1:size(coord, 2)
        @simd for j in 1:size(coord, 1)
            coord[j, i] += lims[j, 1] + (lims[j, 2] - lims[j, 1]) * disp[j]
        end
    end
    return coord
end

"""
```
makeSegment(::segEdge, slipPlane, bVec)
```
Return the edge segment `(slipPlane × bVec) / || slipPlane × bVec ||`. Used to create segments for [`DislocationNetwork`](@ref).
"""
function makeSegment(::segEdge, slipPlane, bVec)
    edge = slipPlane × bVec
    return edge / norm(edge)
end
"""
```
makeSegment(::segEdge, slipPlane, bVec)
```
Return the edge segment `slipPlane / || slipPlane ||`. Used to create segments for [`DislocationNetwork`](@ref).
"""
function makeSegment(::segEdgeN, slipPlane, bVec)
    return slipPlane / norm(slipPlane)
end
"""
```
makeSegment(::segScrew, slipPlane, bVec)
```
Return the screw segment `bVec / || bVec ||`. Used to create segments for [`DislocationNetwork`](@ref).
"""
function makeSegment(::segScrew, slipPlane, bVec)
    return bVec / norm(bVec)
end

"""
```
makeConnect(links, maxConnect)
makeConnect!(network::DislocationNetwork)
```
Creates `connectivity` and `linksConnect` matrices for [`DislocationNetwork`](@ref).

* `connectivity` contains the number of other other nodes each node is connected to, up to `maxConnect` other nodes. Every `(2i, j)` entry contains a node connected to node `j`. Every `(2i+1, j)` coordinate contains whether that node is the first or second node in the link.
* `linksConnect` relates connections enabled by a link. Analogous to the connectivity of a link.
"""
function makeConnect(links, maxConnect)
    lenLinks = size(links, 2)
    connectivity = zeros(Int, 1 + 2 * maxConnect, lenLinks)
    linksConnect = zeros(Int, 2, lenLinks)

    _makeConnect!(links, connectivity, linksConnect)

    return connectivity, linksConnect
end
function makeConnect!(network::DislocationNetwork)
    links = network.links
    connectivity = network.connectivity
    linksConnect = network.linksConnect
    connectivity .= 0
    linksConnect .= 0

    _makeConnect!(links, connectivity, linksConnect)

    return nothing
end
function _makeConnect!(links, connectivity, linksConnect)
    lenLinks = size(links, 2)
    # Indices of defined links, undefined links are always at the end so we only need to know the first undefined entry.
    idx = findfirst(x -> x == 0, @view links[1, :])
    isnothing(idx) ? idx = lenLinks : idx -= 1
    # Loop through indices.
    @inbounds @simd for i in 1:idx
        # Node 1, it is the row of the coord matrix.
        n1 = links[1, i]
        n2 = links[2, i]
        # Num of other nodes node n1 is connected to.
        connectivity[1, n1] += 1
        connectivity[1, n2] += 1
        # Next column for n1 in connectivity.
        tmp1 = 2 * connectivity[1, n1]
        tmp2 = 2 * connectivity[1, n2]
        # i = linkID. 1, 2 are the first and second nodes in link with linkID
        connectivity[tmp1:(tmp1 + 1), n1] .= (i, 1)
        connectivity[tmp2:(tmp2 + 1), n2] .= (i, 2)
        # Connectivity of the nodes in link i.
        linksConnect[1, i] = connectivity[1, n1]
        linksConnect[2, i] = connectivity[1, n2]
    end
    return nothing
end

# TODO #15
"""
```
getSegmentIdx(links, label)
```
Finds the indices of a link and corresponding nodes.

## Returns

`n × 3` matrix is of the form `[i, node1, node2]`.

* `i` can be used to find the Burgers vector, slip plane and segment forces of segment `i`, eg `bVec[:, i]`.
* `node1` and `node2` can be used to find the coordinate and velocity of the nodes, eg `l = coord[:, node2] - coord[:, node1]`.
"""
function getSegmentIdx(links, label)
    lenLinks = size(links, 2)
    extSeg = zeros(Bool, lenLinks)
    segIdx = zeros(Int, lenLinks, 3)  # Indexing matrix.

    _getSegmentIdx!(segIdx, extSeg, lenLinks, links, label)

    return segIdx, extSeg
end
"""
```
getSegmentIdx!(network::DislocationNetwork)
```
Mutates the `segIdx` matrix in `network`. Works the same way as [`getSegmentIdx`](@ref).
"""
function getSegmentIdx!(network::DislocationNetwork)
    links = network.links
    label = network.label
    extSeg = network.extSeg
    segIdx = network.segIdx

    segIdx .= 0
    extSeg .= false
    lenLinks = size(links, 2)

    _getSegmentIdx!(segIdx, extSeg, lenLinks, links, label)

    return nothing
end

function _getSegmentIdx!(segIdx, extSeg, lenLinks, links, label)
    idx = findfirst(x -> x == 0, @view links[1, :])
    isnothing(idx) ? idx = lenLinks : idx -= 1

    @inbounds for i in 1:idx
        n1 = links[1, i]
        n2 = links[2, i]
        segIdx[i, :] .= (i, n1, n2)
        (label[n1] == extDln || label[n2] == extDln) ? extSeg[i] = true : continue
    end

    return nothing
end

"""
```
checkNetwork(network::DislocationNetwork)
```
Checks the validity of the dislocation network. It ensures the following conditions are met by the member variables of `network`:
1. `connectivity` and `links` have the same number of non-zero entries;
1. all entries in `bVec` are non-zero;
1. only the trailing columns of `connectivity` are zeros;
1. consistency between `connectivity` and `links`;
1. `bVec` is conserved among connected nodes;
1. entries in `links` are unique;
1. consistency betwen `connectivity` and `linksConnect`

"""
function checkNetwork(network::DislocationNetwork)
    links = network.links
    label = network.label
    connectivity = network.connectivity
    linksConnect = network.linksConnect

    idx = findfirst(x -> x == 0, @view links[1, :])
    isnothing(idx) ? idx = size(links, 2) : idx -= 1

    # Max value of connectivity should be the last link.
    maximum(connectivity) == idx ? nothing :
    error(
        "Non-empty entries of connectivity should be the same as the non-empty entries of links.",
    )

    bVec = network.bVec
    elemT = eltype(network.bVec)
    bSum = zeros(3)
    @inbounds for i in 1:idx
        iLinkBuffer = zeros(Int, 0)
        col = connectivity[1, i]

        # Check for zero Burgers vectors.
        staticBVec = SVector{3, elemT}(bVec[1, i], bVec[2, i], bVec[3, i])
        norm(staticBVec) < eps(elemT) ? error("Burgers vector must be non-zero.
                                       norm(bVec) = $(norm(bVec[i,:]))") : nothing

        # Trailing columns must be zero.
        sum(connectivity[(2 * (col + 1)):end, i]) != 0 ?
        error(
            "Trailing columns of connectivity must be zero.
      connectivity[$(i), $(2 * (col + 1)):end] = $(connectivity[i, (2 * (col + 1)):end]))",
        ) : nothing

        # Burgers vector conservation.
        bSum .= 0
        # No repeated links.
        neighbours = -ones(Int, col)
        for j in 1:col
            j2 = 2 * j
            iLink = connectivity[j2, i]     # Link ID.
            colLink = connectivity[j2 + 1, i] # Link position in links.
            colOppLink = 3 - connectivity[2 * j + 1, i] # Opposite column in links
            staticBVec = SVector{3, elemT}(bVec[1, iLink], bVec[2, iLink], bVec[3, iLink])
            bSum .+= (3 - 2 * colLink) * staticBVec # Running total of burgers vectors.

            neighbours[j] = links[colOppLink, iLink] # Find neighbouring nodes.
            push!(iLinkBuffer, iLink) # Push to link buffer for error printing.

            # Check that node does not appear twice in connectivity.
            @simd for k in (j + 1):col
                connectivity[j2, i] == connectivity[2 * k, i] ?
                error(
                    "Node $(i) cannot appear twice in connectivity, $(connectivity[i, j2])",
                ) : nothing
            end

            # Check links and connectivity are consistent.
            testNode = links[connectivity[j2 + 1, i], connectivity[j2, i]]
            testNode != i ?
            error(
                "Connectivity and links are not consistent.
connectivity[:, $(i)] = $(connectivity[:, i])
links[:, $(connectivity[j2, i])] = [$(links[1, connectivity[j2, i]]), $(links[2, connectivity[j2, i]])]",
            ) : nothing
        end

        # Check for Burgers vector conservation.
        abs(bSum ⋅ bSum) > eps(elemT) ?
        error("Burgers vector is not conserved, bSum = $(bSum).") : nothing

        # Find unique neighbours.
        uniqueNeighbours = unique(neighbours)
        length(uniqueNeighbours) != length(neighbours) ?
        error(
            "There must be no duplicate links. Each entry in links must correspond to a unique pair of adjacent entries in neighbours.
links = $(iLinkBuffer[1:end - 1])
uniqueNeighbours = $(uniqueNeighbours)
neighbours = $(neighbours)",
        ) : nothing

        # Check connectivity and linksConnect are consistent.
        check = (
            connectivity[2 * linksConnect[1, i], links[1, i]],
            connectivity[2 * linksConnect[2, i], links[2, i]],
        )
        i != check[1] | i != check[2] ? error("Inconsistent link.
        links[:, $(i)] = $(links[:, i])
        linksConnect[:, $(i)] = $(linksConnect[:, i])
        connectivity[:, $(links[1, i])] = $(connectivity[:, links[1, i]])
        connectivity[:, $(links[2, i])] = $(connectivity[:, links[2, i]])") : nothing
    end

    return true
end


# Common backend methods.
#

# Generate single call, replace [xr]Type with typenames and assign letter.
macro blis_ccall_typed(funcname, ret, ctypename, rtypename, named_param_pairs...)
    # C and Julia function names.
    cfunc  = string("bli_", cblas_typechar[ eval(ctypename) ], funcname)
    jlfunc = Symbol("bli_", cblas_typechar[ eval(ctypename) ], funcname, "!")
    ccfunc = Expr(:string, cfunc)

    # Separate variable names and their types.
    names  = named_param_pairs[1:2:end]
    params = named_param_pairs[2:2:end]

    # For typed parameters and return value, replace [xr]Type with real typenames.
    typed_params = ( iType -> eval(:( ( (xType, rType) -> $iType )($ctypename, $rtypename) )) ).(params)
    typed_return =            eval(:( ( (xType, rType) -> $ret   )($ctypename, $rtypename) ))

    # Convert type specifiers to static symbol.
    typed_param_syms = Expr(:tuple, typed_params...)
    typed_return_sym = Symbol(typed_return)

    # Docstring.
    funcdoc = Expr(:string, """
        Autogenerated BLIS Typed-API function.
        For usage refer the
        [BLIS docs for `$cfunc`.](https://github.com/flame/blis/blob/0.8.0/docs/BLISTypedAPI.md#$funcname)
        """)

    # Define Julia function as ccall.
    return quote
        $(esc(jlfunc))( $(names...) ) = begin
            ccall(dlsym(libblis, $(esc(ccfunc))), $typed_return_sym, $typed_param_syms, $(names...))
        end

        @doc $funcdoc $jlfunc
    end
end

# Ganerate a famity of calls.
macro blis_ccall_group(funcname, ret, named_param_pairs...)
    return quote
        @blis_ccall_typed($funcname, $ret, Cfloat    , Cfloat , $(named_param_pairs...))
        @blis_ccall_typed($funcname, $ret, Cdouble   , Cdouble, $(named_param_pairs...))
        @blis_ccall_typed($funcname, $ret, ComplexF32, Cfloat , $(named_param_pairs...))
        @blis_ccall_typed($funcname, $ret, ComplexF64, Cdouble, $(named_param_pairs...))
    end
end


function get_dnl(::Type{UpperBounds}, arg0::RealFieldElement, arg1::RealFieldElement, arg2::jint, arg3::jint)
    return jcall(UpperBounds, "getDnl", RealFieldElement, (RealFieldElement, RealFieldElement, jint, jint), arg0, arg1, arg2, arg3)
end

function get_dnl(::Type{UpperBounds}, arg0::jdouble, arg1::jdouble, arg2::jint, arg3::jint)
    return jcall(UpperBounds, "getDnl", jdouble, (jdouble, jdouble, jint, jint), arg0, arg1, arg2, arg3)
end

function get_rnml(::Type{UpperBounds}, arg0::RealFieldElement, arg1::jint, arg2::jint, arg3::jint, arg4::jint, arg5::jint)
    return jcall(UpperBounds, "getRnml", RealFieldElement, (RealFieldElement, jint, jint, jint, jint, jint), arg0, arg1, arg2, arg3, arg4, arg5)
end

function get_rnml(::Type{UpperBounds}, arg0::jdouble, arg1::jint, arg2::jint, arg3::jint, arg4::jint, arg5::jint)
    return jcall(UpperBounds, "getRnml", jdouble, (jdouble, jint, jint, jint, jint, jint), arg0, arg1, arg2, arg3, arg4, arg5)
end