using GRIB

keylist = Vector{String}()
grib_file = "/home/neo/Projekte/Predly/weather/gfs_workflow/download/data/gfs.20210109/06/gfs.t06z.pgrb2.0p25.f002"

parameters = [7,11,586,422,423,424,426,428,430,433,434,435,437,438,439,440,441,442,443,451,453,455,456,457,458,472,473,474,476,479,480,481,482,483,495,496,498,588]

unneeded_identifiers = ["values", "codedValues", "latLonValues", "latitudes", "longitudes", "distinctLatitudes", "distinctLongitudes"]

start_time = time()

data_dict = Dict()

GribFile(grib_file) do f
    for (index, message) in enumerate(f)

        band_time = time()
        data_array = []
        parameter_name = message["parameterName"]
        parameter_unit = message["parameterUnits"]
        parameter_ecmf = message["nameECMF"]
        units_ecmf = message["unitsECMF"]
        name = message["name"]
        units = message["units"]
        level = message["level"]
        type_of_level = message["typeOfLevel"]
        short_name = message["shortName"]
        data_dict[index] = Dict()

        # calculate x/y raster from points where x == 2y - 2
        # x is 1440 here, y is 721, 1440x721=1038240
        number_points = message["numberOfDataPoints"]
        y =  round(sqrt(number_points/2), digits=0, RoundUp)
        x =  2 * y - 2

        lat0 = message["latitudeOfFirstGridPointInDegrees"]
        lon0 = message["longitudeOfFirstGridPointInDegrees"]
        lat1 = message["latitudeOfLastGridPointInDegrees"]
        lon1 = message["longitudeOfLastGridPointInDegrees"]

        if index in parameters

            data_dict[index]["latitudes"] = message["latitudes"]
            data_dict[index]["longitudes"] = message["longitudes"]
            data_dict[index]["values"] = message["values"]

            println("band ", index ," parameter name: ", parameter_name, " unit ", parameter_unit,
                    " name ecmf ", parameter_ecmf, " units ecmf ", units_ecmf,
                    " name ", name, " units ", units, " level ", level,
                    " type of level ", type_of_level, " short name ", short_name)

            println(" x ",x, " y ", y, " num points ", number_points, " ", lat0 , " ", lon0, " ", lat1, " ", lon1)

            Nearest(message) do near

                lons, lats, values, distances = find(near, message, 8, 50)
                println(lons, lats, values, distances)

            end

            for key in keys(message)
                if key ∉ keylist
                    try
                       push!(keylist, key)
                    catch
                        print(key)
                    end
                end
            end

            println(length(data_array), " took ", time() - band_time)
        end
    end
    println("Overall time: ", time() - start_time)
end

println(Base.summarysize(data_dict)/1000000000, " Gigabytes ")
print(keylist)


using ..Utils

"Map from literal to LogicalΔNode"
const LitCache = Dict{Lit, LogicalΔNode}

"Use literal to represent constraint (1 to X, -1 to not X), 0 to represent true"
const ⊤ = convert(Lit, 0)

"""
Learning from data a structured-decomposable circuit with several structure learning algorithms
"""
function learn_struct_prob_circuit(data::Union{XData, WXData};
        pseudocount = 1.0, algo = "chow-liu", algo_kwargs=(α=1.0, clt_root="graph_center"), vtree = "chow-liu", vtree_kwargs=(vtree_mode="balanced",))
    if algo == "chow-liu"
        clt = learn_chow_liu_tree(data; algo_kwargs...)
        vtree = learn_vtree_from_clt(clt; vtree_kwargs...);
        pc = compile_psdd_from_clt(clt, vtree);
        estimate_parameters(pc, convert(XBatches,data); pseudocount = pseudocount)
        pc, vtree
    else
        error("Cannot learn a structured-decomposable circuit with algorithm $algo")
    end
end

#############
# Learn PlainVtree from CLT
#############

"
Learn a vtree from clt,
with strategy (close to) `linear` or `balanced`
"
function learn_vtree_from_clt(clt::CLT; vtree_mode::String)::PlainVtree
    roots = [i for (i, x) in enumerate(parent_vector(clt)) if x == 0]
    rootnode = construct_children(Var.(roots), clt, vtree_mode)

    return node2dag(rootnode)
end

function construct_node(v::Var, clt::CLT, strategy::String)::PlainVtreeNode
    children = Var.(outneighbors(clt, v))
    if isempty(children) # leaf node
        return PlainVtreeLeafNode(v)
    else
        right = construct_children(children, clt, strategy)
        return add_parent(v, right)
    end
end

function construct_children(children::Vector{Var}, clt::CLT, strategy::String)::PlainVtreeNode
    sorted_vars = sort(collect(children))
    children_nodes = Vector{PlainVtreeNode}()
    foreach(x -> push!(children_nodes, construct_node(x, clt, strategy)), sorted_vars)

    if strategy == "linear"
        construct_children_linear(children_nodes, clt)
    elseif strategy == "balanced"
        construct_children_balanced(children_nodes, clt)
    else
        throw("Unknown type of strategy")
    end
end

function construct_children_linear(children_nodes::Vector{PlainVtreeNode}, clt::CLT)::PlainVtreeNode
    children_nodes = Iterators.Stateful(reverse(children_nodes))

    right = popfirst!(children_nodes)
    for left in children_nodes
        right = PlainVtreeInnerNode(left, right)
    end
    return right
end

function construct_children_balanced(children_nodes::Vector{PlainVtreeNode}, clt::CLT)::PlainVtreeNode
    if length(children_nodes) == 1
        return children_nodes[1]
    elseif length(children_nodes) == 2
        return PlainVtreeInnerNode(children_nodes[1], children_nodes[2])
    else
        len = trunc(Int64, length(children_nodes) / 2)
        left = construct_children_balanced(children_nodes[1 : len], clt)
        right = construct_children_balanced(children_nodes[len + 1 : end], clt)
        return PlainVtreeInnerNode(left, right)
    end
end

function add_parent(parent::Var, children::PlainVtreeNode)
    return PlainVtreeInnerNode(PlainVtreeLeafNode(parent), children)
end

#####################
# Compile PSDD from CLT and vtree
#####################

"Compile a psdd circuit from clt and vtree"
function compile_psdd_from_clt(clt::MetaDiGraph, vtree::PlainVtree)
    order = node2dag(vtree[end])
    parent_clt = Var.(parent_vector(clt))

    lin = Vector{ProbΔNode}()
    prob_cache = ProbCache()
    lit_cache = LitCache()
    v2p = Dict{PlainVtreeNode, ProbΔ}()

    get_params(cpt::Dict) = length(cpt) == 2 ? [cpt[1], cpt[0]] : [cpt[(1,1)], cpt[(0,1)], cpt[(1,0)], cpt[(0,0)]]
    function add_mapping!(v::PlainVtreeNode, circuits::ProbΔ)
        if !haskey(v2p, v); v2p[v] = Vector{ProbΔNode}(); end
        foreach(c -> if !(c in v2p[v]) push!(v2p[v], c);end, circuits)
    end

    # compile vtree leaf node to terminal/true node
    function compile_from_vtree_node(v::PlainVtreeLeafNode)
        var = v.var
        children = Var.(outneighbors(clt, var))
        cpt = get_prop(clt, var, :cpt)
        parent = parent_clt[var]
        if isequal(children, [])
            circuit = compile_true_nodes(var, v, get_params(cpt), lit_cache, prob_cache, lin)
        else
            circuit = compile_literal_nodes(var, v, get_params(cpt), lit_cache, prob_cache, lin)
        end
        add_mapping!(v, circuit)
    end

    # compile to decision node
    function compile_from_vtree_node(v::PlainVtreeInnerNode)
        left_var = left_most_child(v.left).var
        right_var = left_most_child(v.right).var
        left_circuit = v2p[v.left]
        right_circuit = v2p[v.right]

        if parent_clt[left_var] == parent_clt[right_var] # two nodes are independent, compile to seperate decision nodes
            circuit = [compile_decision_node([l], [r], v, [1.0], prob_cache, lin) for (l, r) in zip(left_circuit, right_circuit)]
        elseif left_var == parent_clt[right_var] # conditioned on left
            cpt = get_prop(clt, left_var, :cpt)
            circuit = compile_decision_nodes(left_circuit, right_circuit, v, get_params(cpt), prob_cache, lin)
        else
            throw("PlainVtree are not learned from the same CLT")
        end
        add_mapping!(v, circuit)
    end

    foreach(compile_from_vtree_node, vtree)
    return lin
end

#####################
# Construct probabilistic circuit node
#####################

prob_children(n, prob_cache) =
    copy_with_eltype(map(c -> prob_cache[c], n.children), ProbΔNode{<:StructLogicalΔNode})

"Add leaf nodes to circuit `lin`"
function add_prob_leaf_node(var::Var, vtree::PlainVtreeLeafNode, lit_cache::LitCache, prob_cache::ProbCache, lin)
    pos = StructLiteralNode{PlainVtreeNode}( var2lit(var), vtree)
    neg = StructLiteralNode{PlainVtreeNode}(-var2lit(var), vtree)
    lit_cache[var2lit(var)] = pos
    lit_cache[-var2lit(var)] = neg
    pos2 = ProbLiteral(pos)
    neg2 = ProbLiteral(neg)
    prob_cache[pos] = pos2
    prob_cache[neg] = neg2
    push!(lin, pos2)
    push!(lin, neg2)
    return (pos2, neg2)
end

"Add prob⋀ node to circuit `lin`"
function add_prob⋀_node(children::ProbΔ, vtree::PlainVtreeInnerNode, prob_cache::ProbCache, lin)::Prob⋀
    logic = Struct⋀Node{PlainVtreeNode}([c.origin for c in children], vtree)
    prob = Prob⋀(logic, prob_children(logic, prob_cache))
    prob_cache[logic] = prob
    push!(lin, prob)
    return prob
end

"Add prob⋁ node to circuit `lin`"
function add_prob⋁_node(children::ProbΔ, vtree::PlainVtreeNode, thetas::Vector{Float64}, prob_cache::ProbCache, lin)::Prob⋁
    logic = Struct⋁Node{PlainVtreeNode}([c.origin for c in children], vtree)
    prob = Prob⋁(logic, prob_children(logic, prob_cache))
    prob.log_thetas = log.(thetas)
    prob_cache[logic] = prob
    push!(lin, prob)
    return prob
end

"Construct decision nodes given `primes` and `subs`"
function compile_decision_node(primes::ProbΔ, subs::ProbΔ, vtree::PlainVtreeInnerNode, params::Vector{Float64}, prob_cache::ProbCache, lin)
    elements = [add_prob⋀_node([prime, sub], vtree, prob_cache, lin) for (prime, sub) in zip(primes, subs)]
    return add_prob⋁_node(elements, vtree, params, prob_cache, lin)
end

"Construct literal nodes given variable `var`"
function compile_literal_nodes(var::Var, vtree::PlainVtreeLeafNode, probs::Vector{Float64}, lit_cache::LitCache, prob_cache::ProbCache, lin)
    (pos, neg) = add_prob_leaf_node(var, vtree, lit_cache, prob_cache, lin)
    return [pos, neg]
end

"Construct true nodes given variable `var`"
function compile_true_nodes(var::Var, vtree::PlainVtreeLeafNode, probs::Vector{Float64}, lit_cache::LitCache, prob_cache::ProbCache, lin)
    (pos, neg) = add_prob_leaf_node(var, vtree, lit_cache, prob_cache, lin)
    return [add_prob⋁_node([pos, neg], vtree, probs[i:i+1], prob_cache, lin) for i in 1:2:length(probs)]
end

"Construct decision nodes conditiond on different distribution"
function compile_decision_nodes(primes::ProbΔ, subs::ProbΔ, vtree::PlainVtreeInnerNode, params::Vector{Float64}, prob_cache::ProbCache, lin)
    return [compile_decision_node(primes, subs, vtree, params[i:i+1], prob_cache, lin) for i in 1:2:length(params)]
end

#####################
# Map and cache constraints
#####################

function set_base(index, n::StructLiteralNode, bases)
    if positive(n)
        bases[n][variable(n)] = 1
    else
        bases[n][variable(n)] = -1
    end
end

function set_base(index, n::Struct⋁Node, bases)
    len = num_children(n)
    temp = sum([bases[c] for c in n.children])
    bases[n] = map(x-> if x == len 1; elseif -x == len; -1; else 0; end, temp)
end

function set_base(index, n::Struct⋀Node, bases)
    bases[n] = sum([bases[c] for c in n.children])
end

function calculate_all_bases(circuit::ProbΔ)::BaseCache
    num_var = num_variables(circuit[end].origin.vtree)
    bases = BaseCache()
    foreach(n -> bases[n.origin] = fill(⊤, num_var), circuit)
    foreach(n -> set_base(n[1], n[2].origin, bases), enumerate(circuit))
    @assert all(bases[circuit[end].origin] .== ⊤) "Base of root node should be true"
    return bases
end

#####################
# Compile fully factorized PSDD from vtree, all variables are independent initially
#####################

function compile_fully_factorized_psdd_from_vtree(vtree::PlainVtree)::ProbΔ

    function ful_factor_node(v::PlainVtreeLeafNode, lit_cache::LitCache, prob_cache::ProbCache, v2n, lin)
        var = variables(v)[1]
        pos, neg = add_prob_leaf_node(var, v, lit_cache, prob_cache, lin)
        prob_or = add_prob⋁_node([pos, neg], v, [0.5, 0.5], prob_cache, lin)
        v2n[v] = prob_or
        nothing
    end

    function ful_factor_node(v::PlainVtreeInnerNode, lit_cache::LitCache, prob_cache::ProbCache, v2n, lin)
        left = v2n[v.left]
        right = v2n[v.right]
        prob_and = add_prob⋀_node([left, right], v, prob_cache, lin)
        prob_or = add_prob⋁_node([prob_and], v, [1.0], prob_cache, lin)
        v2n[v] = prob_or
        nothing
    end

    lin = Vector{ProbΔNode}()
    prob_cache = ProbCache()
    lit_cache = LitCache()
    v2n = Dict{PlainVtreeNode, ProbΔNode}()

    for v in vtree
        ful_factor_node(v, lit_cache, prob_cache, v2n, lin)
    end

    lin
end
"""
Learning from data a structured-decomposable credal circuit with several structure learning algorithms
"""
function learn_struct_credal_circuit(data::Union{XData, WXData}, s_idm::Float64;
        pseudocount = 1.0, algo = "chow-liu", algo_kwargs=(α=1.0, clt_root="graph_center"), vtree = "chow-liu", vtree_kwargs=(vtree_mode="balanced",))
    if algo == "chow-liu"
        clt = learn_chow_liu_tree(data; algo_kwargs...)
        vtree = learn_vtree_from_clt(clt; vtree_kwargs...);
        pc = compile_csdd_from_clt(clt, vtree);
        estimate_credal_parameters(pc, convert(XBatches,data),s_idm ; pseudocount = pseudocount)
        pc, vtree
    else
        error("Cannot learn a structured-decomposable circuit with algorithm $algo")
    end
end

#####################
# Compile CSDD from CLT and vtree
#####################

"Compile a psdd circuit from clt and vtree"
function compile_csdd_from_clt(clt::MetaDiGraph, vtree::PlainVtree)
    order = node2dag(vtree[end])
    parent_clt = Var.(parent_vector(clt))

    lin = Vector{CredalΔNode}()
    cred_cache = CredalCache()
    lit_cache = LitCache()
    v2p = Dict{PlainVtreeNode, CredalΔ}()

    get_params_cred(cpt::Dict) = length(cpt) == 2 ? [cpt[1], cpt[0]] : [cpt[(1,1)], cpt[(0,1)], cpt[(1,0)], cpt[(0,0)]]
    function add_mapping_cred!(v::PlainVtreeNode, circuits::CredalΔ)
        if !haskey(v2p, v); v2p[v] = Vector{CredalΔNode}(); end
        foreach(c -> if !(c in v2p[v]) push!(v2p[v], c);end, circuits)
    end

    # compile vtree leaf node to terminal/true node
    function compile_from_vtree_node_cred(v::PlainVtreeLeafNode)
        var = v.var
        children = Var.(outneighbors(clt, var))
        cpt = get_prop(clt, var, :cpt)
        parent = parent_clt[var]
        if isequal(children, [])
            circuit = compile_true_nodes_cred(var, v, get_params_cred(cpt), lit_cache, cred_cache, lin)
        else
            circuit = compile_literal_nodes_cred(var, v, get_params_cred(cpt), lit_cache, cred_cache, lin)
        end
        add_mapping_cred!(v, circuit)
    end

    # compile to decision node
    function compile_from_vtree_node_cred(v::PlainVtreeInnerNode)
        left_var = left_most_child(v.left).var
        right_var = left_most_child(v.right).var
        left_circuit = v2p[v.left]
        right_circuit = v2p[v.right]

        if parent_clt[left_var] == parent_clt[right_var] # two nodes are independent, compile to seperate decision nodes
            circuit = [compile_decision_node_cred([l], [r], v, [1.0], cred_cache, lin) for (l, r) in zip(left_circuit, right_circuit)]
        elseif left_var == parent_clt[right_var] # conditioned on left
            cpt = get_prop(clt, left_var, :cpt)
            circuit = compile_decision_nodes_cred(left_circuit, right_circuit, v, get_params_cred(cpt), cred_cache, lin)
        else
            throw("PlainVtree are not learned from the same CLT")
        end
        add_mapping_cred!(v, circuit)
    end

    foreach(compile_from_vtree_node_cred, vtree)
    return lin
end

#####################
# Construct credal circuit node
#####################

cred_children(n, cred_cache) =
    copy_with_eltype(map(c -> cred_cache[c], n.children), CredalΔNode{<:StructLogicalΔNode})

"Add leaf nodes to circuit `lin`"
function add_cred_leaf_node(var::Var, vtree::PlainVtreeLeafNode, lit_cache::LitCache, cred_cache::CredalCache, lin)
    pos = StructLiteralNode{PlainVtreeNode}( var2lit(var), vtree)
    neg = StructLiteralNode{PlainVtreeNode}(-var2lit(var), vtree)
    lit_cache[var2lit(var)] = pos
    lit_cache[-var2lit(var)] = neg
    pos2 = CredalLiteral(pos)
    neg2 = CredalLiteral(neg)
    cred_cache[pos] = pos2
    cred_cache[neg] = neg2
    push!(lin, pos2)
    push!(lin, neg2)
    return (pos2, neg2)
end

"Add cred⋀ node to circuit `lin`"
function add_cred⋀_node(children::CredalΔ, vtree::PlainVtreeInnerNode, cred_cache::CredalCache, lin)::Credal⋀
    logic = Struct⋀Node{PlainVtreeNode}([c.origin for c in children], vtree)
    cred = Credal⋀(logic, cred_children(logic, cred_cache))
    cred_cache[logic] = cred
    push!(lin, cred)
    return cred
end

"Add cred⋁ node to circuit `lin`"
function add_cred⋁_node(children::CredalΔ, vtree::PlainVtreeNode, thetas::Vector{Float64}, cred_cache::CredalCache, lin)::Credal⋁
    logic = Struct⋁Node{PlainVtreeNode}([c.origin for c in children], vtree)
    cred = Credal⋁(logic, cred_children(logic, cred_cache))
    cred.log_thetas = log.(thetas)
    cred_cache[logic] = cred
    push!(lin, cred)
    return cred
end

"Construct decision nodes given `primes` and `subs`"
function compile_decision_node_cred(primes::CredalΔ, subs::CredalΔ, vtree::PlainVtreeInnerNode, params::Vector{Float64}, cred_cache::CredalCache, lin)
    elements = [add_cred⋀_node([prime, sub], vtree, cred_cache, lin) for (prime, sub) in zip(primes, subs)]
    return add_cred⋁_node(elements, vtree, params, cred_cache, lin)
end

"Construct literal nodes given variable `var`"
function compile_literal_nodes_cred(var::Var, vtree::PlainVtreeLeafNode, probs::Vector{Float64}, lit_cache::LitCache, cred_cache::CredalCache, lin)
    (pos, neg) = add_cred_leaf_node(var, vtree, lit_cache, cred_cache, lin)
    return [pos, neg]
end

"Construct true nodes given variable `var`"
function compile_true_nodes_cred(var::Var, vtree::PlainVtreeLeafNode, probs::Vector{Float64}, lit_cache::LitCache, cred_cache::CredalCache, lin)
    (pos, neg) = add_cred_leaf_node(var, vtree, lit_cache, cred_cache, lin)
    return [add_cred⋁_node([pos, neg], vtree, probs[i:i+1], cred_cache, lin) for i in 1:2:length(probs)]
end

"Construct decision nodes conditiond on different distribution"
function compile_decision_nodes_cred(primes::CredalΔ, subs::CredalΔ, vtree::PlainVtreeInnerNode, params::Vector{Float64}, cred_cache::CredalCache, lin)
    return [compile_decision_node_cred(primes, subs, vtree, params[i:i+1], cred_cache, lin) for i in 1:2:length(params)]
end


@testset "Demultiplexer" begin
    function randdna(n)
        return LongDNASeq(rand([DNA_A, DNA_C, DNA_G, DNA_T, DNA_N], n))
    end

    function make_errors(seq, p=0.03)
        seq = copy(seq)
        nucs = DNA['A', 'C', 'G', 'T', 'N']
        i = 1
        while i ≤ lastindex(seq)
            if rand() < p
                r = rand()
                if r < 0.1
                    # insertion
                    insert!(seq, i, rand(nucs))
                elseif r < 0.2
                    # deletion
                    deleteat!(seq, i)
                else
                    # substitution
                    seq[i] = rand(nucs)
                end
            end
            i += 1
        end
        return seq
    end

    @testset "Hamming distance" begin
        barcodes = [dna"ATGG", dna"CAGA", dna"GGAA", dna"TACG"]
        dplxr = Demultiplexer(barcodes, n_max_errors=1, distance=:hamming)

        for i in 1:lastindex(barcodes)
            @test dplxr[i] == barcodes[i]
        end

        @test demultiplex(dplxr, dna"ATGG") === (1, 0)
        @test demultiplex(dplxr, dna"CAGA") === (2, 0)
        @test demultiplex(dplxr, dna"GGAA") === (3, 0)
        @test demultiplex(dplxr, dna"TACG") === (4, 0)

        # every 1bp substitution is recoverable
        for (i, barcode) in enumerate(barcodes)
            # substitution
            for j in 1:lastindex(barcode)
                barcode′ = copy(barcode)
                for nt in [DNA_A, DNA_C, DNA_G, DNA_T, DNA_N]
                    barcode′[j] = nt
                    @test demultiplex(dplxr, barcode′) == (i, count(!=, barcode, barcode′))
                end
            end
        end

        n_ok = 0
        n_ok_with_fallback = 0
        for _ in 1:10_000
            i = rand(1:lastindex(barcodes))
            seq = make_errors(barcodes[i] * randdna(10))
            n_ok += demultiplex(dplxr, seq)[1] == i
            n_ok_with_fallback += demultiplex(dplxr, seq, true)[1] == i
        end
        # empirically, n_ok / 10_000 is ~0.985
        # @show n_ok
        @test n_ok / 10_000 > 0.98
        @test n_ok < n_ok_with_fallback
    end

    @testset "Levenshtein distance" begin
        barcodes = [dna"ATGG", dna"CAGA", dna"GGAA", dna"TACG"]
        dplxr = Demultiplexer(barcodes, n_max_errors=1, distance=:levenshtein)

        for i in 1:lastindex(barcodes)
            @test dplxr[i] == barcodes[i]
        end

        @test demultiplex(dplxr, dna"ATGG") === (1, 0)
        @test demultiplex(dplxr, dna"CAGA") === (2, 0)
        @test demultiplex(dplxr, dna"GGAA") === (3, 0)
        @test demultiplex(dplxr, dna"TACG") === (4, 0)

        # every 1bp substitution/insertion/deletion is recoverable
        for (i, barcode) in enumerate(barcodes)
            # substitution
            for j in 1:lastindex(barcode)
                barcode′ = copy(barcode)
                for nt in [DNA_A, DNA_C, DNA_G, DNA_T, DNA_N]
                    barcode′[j] = nt
                    i′, d = demultiplex(dplxr, barcode′)
                    @test i′ == i && 0 ≤ d ≤ 1
                end
            end
            # insertion
            for j in 1:lastindex(barcode)
                barcode′ = copy(barcode)
                insert!(barcode′, j, rand([DNA_A, DNA_C, DNA_G, DNA_T, DNA_N]))
                i′, d = demultiplex(dplxr, barcode′)
                @test i′ == i && 0 ≤ d ≤ 1
            end
            # deletion
            for j in 1:lastindex(barcode)
                barcode′ = copy(barcode)
                deleteat!(barcode′, j)
                i′, d = demultiplex(dplxr, barcode′)
                @test i′ == i && 0 ≤ d ≤ 1
            end
        end

        n_ok = 0
        n_ok_with_fallback = 0
        for _ in 1:10_000
            i = rand(1:lastindex(barcodes))
            seq = make_errors(barcodes[i] * randdna(10))
            n_ok += demultiplex(dplxr, seq)[1] == i
            n_ok_with_fallback += demultiplex(dplxr, seq, true)[1] == i
        end
        # empirically, n_ok / 10_000 is ~0.995
        # @show n_ok
        @test n_ok / 10_000 > 0.99
        @test n_ok < n_ok_with_fallback
    end
end


using ArgParse,FileIO,JLD,KUparser,Knet
macro msg(_x) :(if logging>0; join(STDERR,[Dates.format(now(),"HH:MM:SS"), $_x,'\n'],' '); end) end
macro log(_x) :(@msg($(string(_x))); $(esc(_x))) end

include("conll17.jl")

# for i in `cut -f2 languages.tsv | head -2`; do echo $i; julia gparser-mlp-static.jl --train $i-ud-train.conllu $i-ud-dev.conllu --vectors $i.vectors.xz --parse $i-ud-dev.txt --udpipe $i.udpipe > foo; ./eval.py $i-ud-dev.conllu foo; done

# Some example data
grctrn = "/mnt/ai/data/nlp/conll17/ud-treebanks-conll2017/UD_Ancient_Greek/grc-ud-train.conllu"
grcdev = "/mnt/ai/data/nlp/conll17/ud-treebanks-conll2017/UD_Ancient_Greek/grc-ud-dev.conllu"
grctxt = "/mnt/ai/data/nlp/conll17/ud-treebanks-conll2017/UD_Ancient_Greek/grc-ud-dev.txt"
grcvec = "/mnt/ai/data/nlp/conll17/word-embeddings-conll17/Ancient_Greek/grc.vectors.xz"
grcudp = "/mnt/ai/data/nlp/conll17/UDPipe/udpipe-ud-2.0-conll17-170315/models/ancient_greek-ud-2.0-conll17-170315.udpipe"

function main(args="")
    isa(args, AbstractString) && (args=split(args))
    s = ArgParseSettings()
    s.description="Using greedy parser with oracle (static) trained MLP model"
    s.exc_handler=ArgParse.debug_handler
    @add_arg_table s begin
        # Parsing options
        ("--parse"; help="file in plain text format to be parsed")
        ("--udpipe"; help="UDpipe model file")
        ("--load"; help="load model from jld file")
        ("--ptype"; default="ArcEagerR1"; help="Parser type")
        ("--flist"; default="Flist.zn11pv"; help="Feature list to use")

        # Training options
        ("--train"; nargs='*'; help="training file(s) in conllu format")
        ("--vectors"; help="word vectors")
        ("--save"; help="save model to jld file")
        ("--atype"; default=(gpu()>=0 ? "KnetArray{Float32}" : "Array{Float32}"); help="array type: Array for cpu, KnetArray for gpu")
        ("--otype"; default="Adam()"; help="Optimization algorithm and parameters.")
        ("--hidden"; nargs='*'; arg_type=Int; help="sizes of hidden layers, e.g. --hidden 128 64 for a net with two hidden layers")
        ("--epochs"; arg_type=Int; default=10; help="number of epochs for training")
        ("--batchsize"; arg_type=Int; default=100; help="minibatch size")
        ("--dropout"; arg_type=Float64; default=0.0; help="dropout probability")
        ("--l1"; arg_type=Float64; default=0.0; help="L1 regularization")
        ("--l2"; arg_type=Float64; default=0.0; help="L2 regularization")
        ("--fast"; action=:store_true; help="skip loss calculation for faster run")
        #TODO ("--gcheck"; arg_type=Int; default=0; help="check N random gradients per parameter")
        ("--seed"; arg_type=Int; default=-1; help="random number seed: use a nonnegative int for repeatable results")
        ("--logging"; arg_type=Int; default=1; help="Logging level: 0 to turn off")
    end

    # Process args
    o = parse_args(args, s; as_symbols=true)
    logging = o[:logging]
    @msg(s.description)
    @msg(string("opts=",[(k,v) for (k,v) in o]...))
    if o[:seed] > 0; setseed(o[:seed]); end
    atype = eval(parse(o[:atype]))
    ptype = eval(parse(o[:ptype]))
    flist = eval(parse(o[:flist]))

    # Load model
    if o[:load] != nothing
        @log a = load(o[:load])
        model = a["model"]
        optim = a["optim"]
        vocab = a["vocab"]
    else
        model = optim = nothing
        vocab = Vocab(); vocab.postags=UPOSTAG; vocab.deprels=UDEPREL
    end

    # Load raw text
    if o[:parse] != nothing
        # we need to test with predicted toks and tags from txt data.
        parse_conllu = tempname()
        @log run(pipeline(`udpipe --tokenize --tag $(o[:udpipe]) $(o[:parse])`, parse_conllu))
        @log parse_sentences = readconllu(parse_conllu, vocab)
    end

    # Load training data
    if !isempty(o[:train])
        readc(f)=readconllu(f,vocab)
        @log train_corpora = map(readc, o[:train])
    end

    # Load vectors: this has to be done after all conllu files have been loaded
    if o[:vectors] != nothing
        @log readvectors(o[:vectors], vocab)
    end

    # Fix words without vectors, this may happen if the model loaded does not have all the words in the corpora read
    if size(vocab.wvecs,2) < length(vocab.words)
        z = fill!(similar(vocab.wvecs, size(vocab.wvecs,1), length(vocab.words) - size(vocab.wvecs,2)), 0)
        @log vocab.wvecs = hcat(vocab.wvecs, z)
    end

    # Train model
    if !isempty(o[:train])
        data = map(train_corpora) do corpus
            @log (p,x,y) = oparse(ptype, corpus, flist)
            @log minibatch(x, y, o[:batchsize]; atype=atype)
        end
        if model == nothing
            xtrn,ytrn = data[1][1]
            model = initmodel(size(xtrn,1), o[:hidden]..., size(ytrn,1); atype=atype)
            optim = initoptim(model, o[:otype])
        end
        if o[:fast]
            report = (ep->nothing)
        else
            report = (ep->(@msg((ep,map(d->accuracy(model,d,mlp),data)...))))
            @msg((:epoch,:trnacc,:devacc))
        end
        report(0)
        for epoch=1:o[:epochs]
            for (x,y) in data[1]
                grads = softgrad(model,x,y,mlp;l1=o[:l1],l2=o[:l2],pdrop=o[:dropout])
                update!(model,grads,optim)
            end
            report(epoch)
        end
    end

    # Save model
    if o[:save] != nothing
        @log save(File(format"JLD",o[:save]), "model", model, "optim", optim, "vocab", vocab)
    end

    # Parse
    if o[:parse] != nothing
        pred(x,y)=copy!(y, mlp(model, atype(x)))
        @log parses = gparse(ptype, parse_sentences, flist, pred; nbatch=o[:batchsize])
        @log writeconllu(parse_sentences, parses, parse_conllu)
        rm(parse_conllu)
    end
end

# function savedata()
#     @date c = readconllu(corpus, wordvecs)
#     @date (p,x,y) = oparse(pt, c, ft)
#     @date save("foo.en.jld", "x", x, "y", y, "c", c, "p", p)
# end

function minibatch(x, y, batch; atype=Array{Float32})
    nx = size(x,2)
    data = Any[]
    for i=1:batch:nx
        j=min(nx,i+batch-1)
        push!(data, (atype(view(x,:,i:j)),atype(view(y,:,i:j))))
    end
    return data
end

function mlp(w,x;pdrop=0)
    for i=1:2:length(w)-2
        x = relu(w[i]*x .+ w[i+1])
        x = dropout(x, pdrop)
    end
    return w[end-1]*x .+ w[end]
end

# # For use with gparser until we change the interface:
# function predict(w,x,y)
#     x = KnetArray(x)
#     for i=1:2:length(w)-2
#         x = relu(w[i]*x .+ w[i+1])
#     end
#     # ccall(("cudaDeviceSynchronize","libcudart"),UInt32,()) # uncomment for profiling
#     copy!(y, w[end-1]*x .+ w[end])
# end

using Knet: @cuda, Cptr

# target may be a subarray, TODO: handle in Knet
function Base.copy!{T}(dest::SubArray{T}, src::KnetArray{T})
    if length(dest) < length(src); throw(BoundsError()); end
    @cuda(cudart,cudaMemcpy,(Cptr,Cptr,Csize_t,UInt32),
          pointer(dest), pointer(src), length(src)*sizeof(T), 2)
    return dest
end

# Calculate cross entropy loss of a model with weights w for one minibatch (x,p)
# Use non-zero l1 or l2 for regularization
function softloss(w,x,p,model;l1=0,l2=0,o...)
    y = model(w,x;o...)
    logphat = logp(y)
    J = -sum(p .* logphat) / size(x,2)  # dividing by number of instances for per-instance average
    if l1 != 0; J += l1 * sum(sumabs(wi)  for wi in w[1:2:end]); end
    if l2 != 0; J += l2 * sum(sumabs2(wi) for wi in w[1:2:end]); end
    return J
end

softgrad = grad(softloss)

function avgloss(w,data,model) # average loss for the whole dataset
    sum = cnt = 0
    for (x,y) in data
        sum += softloss(w,x,y,model)
        cnt += 1
    end
    return sum/cnt
end

function accuracy(w,data,model)
    corr = cnt = 0
    for (x,y) in data
        ypred = mlp(w,x)
        corr += sum(y .* (ypred .== maximum(ypred,1)))
        cnt += size(y,2)
    end
    return corr/cnt
end

function initmodel(d...; atype=Array{Float32})
    init(d...)=atype(xavier(Float32,d...))
    bias(d...)=atype(zeros(Float32,d...))
    w = Any[]
    for i=1:length(d)-1
        push!(w, init(d[i+1],d[i]))
        push!(w, bias(d[i+1],1))
    end
    return w
end

# This should work for any combination of tuple/array/dict
initoptim{T<:Number}(::KnetArray{T},otype)=eval(parse(otype))
initoptim{T<:Number}(::Array{T},otype)=eval(parse(otype))
initoptim(a::Associative,otype)=Dict(k=>initoptim(v,otype) for (k,v) in a)
initoptim(a,otype)=map(x->initoptim(x,otype), a)

# function train(dtrn,dtst;
#                hidden = [ 128 ],
#                optimizer = Adam,
#                epochs = 10,
#                l1 = 0, l2 = 0, pdrop = 0
#                )
#     (x,y)=dtrn[1]
#     w = initmodel(size(x,1), hidden..., size(y,1))
#     p = oparams(w, optimizer)
#     report(ep)=println((ep,accuracy(w,dtrn,mlp),accuracy(w,dtst,mlp)))
#     report(0)
#     for epoch=1:epochs
#         for (x,y) in dtrn
#             g = softgrad(w,x,y,mlp;l1=l1,l2=l2,pdrop=pdrop)
#             update!(w,g,p)
#         end
#         report(epoch)
#     end
#     return w
# end

# julia> dtrn,dtst = loaddata()
# julia> train(dtrn,dtst)
# (0,0.008304813f0,0.01f0)
# (1,0.9219519f0,0.9071f0)
# (2,0.9373984f0,0.9188f0)
# (3,0.9448262f0,0.9239f0)
# (4,0.9502219f0,0.923f0)
# (5,0.95345455f0,0.9229f0)
# (6,0.95614976f0,0.923f0)
# (7,0.95832354f0,0.9228f0)
# (8,0.9598797f0,0.9223f0)
# (9,0.962f0,0.9258f0)
# (10,0.9629759f0,0.9233f0)

# julia> train(dtrn,dtst;hidden=1000)
# (0,0.0073877005f0,0.0074f0)
# (1,0.9364385f0,0.9173f0)
# (2,0.9535187f0,0.9284f0)
# (3,0.9634011f0,0.9316f0)
# (4,0.97001874f0,0.9315f0)
# (5,0.97462296f0,0.9295f0)
# (6,0.9792781f0,0.9307f0)
# (7,0.982254f0,0.9316f0)
# (8,0.98386633f0,0.9325f0)
# (9,0.98568714f0,0.9291f0)
# (10,0.98813367f0,0.9312f0)

# julia> train(dtrn,dtst;hidden=10000)
# (0,0.0077700536f0,0.0077f0)
# (1,0.9415294f0,0.9223f0)
# (2,0.95814705f0,0.9313f0)
# (3,0.9689037f0,0.9358f0)
# (4,0.97555614f0,0.932f0)
# (5,0.98095185f0,0.9337f0)
# (6,0.9839626f0,0.934f0)
# (7,0.98727006f0,0.9358f0)
# (8,0.9889813f0,0.9356f0)
# (9,0.9904198f0,0.9358f0)
# (10,0.99142516f0,0.9338f0)

if basename(PROGRAM_FILE)==basename(@__FILE__); main(ARGS); end


# include packages -
include("Include.jl")

function objective_function(parameter_guess_array,time_start,time_step_size,time_stop,model_data_dictionary,exp_data_dictionary)

    # what is the host_type?
    host_type = :cell_free

    # Phase 1: parameter update =========================================================================== #
    # update the paramaters in the model data dictionary -
    # for now - lets only search over dG's -
    R = model_data_dictionary["R"]
    T_K = model_data_dictionary["T_K"]

    # compute W -
    tmp_W_array = Float64[]
    for index = 1:2
        parameter_guess = parameter_guess_array[index]
        value = exp(-1*parameter_guess/(R*T_K))
        push!(tmp_W_array,value)
    end

    # update the control W's -
    control_parameter_dictionary = model_data_dictionary["control_parameter_dictionary"]
	control_parameter_dictionary["W_deGFP_RNAP"] = tmp_W_array[1]       # 3
	control_parameter_dictionary["W_deGFP_sigma_70"] = tmp_W_array[2]   # 4
    model_data_dictionary["control_parameter_dictionary"] = control_parameter_dictionary

    binding_parameter_dictionary = model_data_dictionary["binding_parameter_dictionary"]
	binding_parameter_dictionary["n_deGFP_sigma_70"] = parameter_guess_array[3]
	binding_parameter_dictionary["K_deGFP_sigma_70"] = parameter_guess_array[4]
    model_data_dictionary["binding_parameter_dictionary"] = binding_parameter_dictionary

    # time constant modifier -
    # time constant modifiers -
	time_constant_modifier_array = [
		0.0	                        ;	# 1	deGFP
		0.0	                        ;	# 2	sigma_70
		parameter_guess_array[5] 	;	# 3	mRNA_deGFP
		1.0	                        ;	# 4	mRNA_sigma_70
		parameter_guess_array[6] 	;	# 5	protein_deGFP
		1.0	                        ;	# 6	protein_sigma_70
    ]
    model_data_dictionary["time_constant_modifier_array"] = time_constant_modifier_array

    # setup degradation_modifier_array -
    degradation_modifier_array = [
		0.0	;	# 1	deGFP
		0.0	;	# 2	sigma_70
		parameter_guess_array[7]	;	# 3	mRNA_deGFP
		1.0	;	# 4	mRNA_sigma_70
		parameter_guess_array[8]	;	# 5	protein_deGFP
		parameter_guess_array[9]	;	# 6	protein_sigma_70
    ]

    # update the translation time -
    model_data_dictionary["half_life_translation_capacity"] = parameter_guess_array[10]

    # lastly, update KL -
    biophysical_constants_dictionary = model_data_dictionary["biophysical_constants_dictionary"]
    biophysical_constants_dictionary["translation_saturation_constant"] = parameter_guess_array[11]
    model_data_dictionary["biophysical_constants_dictionary"] = biophysical_constants_dictionary

    # grab defaults -
    species_symbol_type_array = model_data_dictionary["species_symbol_type_array"]
    protein_coding_length_array = model_data_dictionary["protein_coding_length_array"]
    gene_coding_length_array = model_data_dictionary["gene_coding_length_array"]
    time_constant_modifier_array = model_data_dictionary["time_constant_modifier_array"]
    initial_condition_array = model_data_dictionary["initial_condition_array"]

    # # get gene IC -
    idx_gene = findall(x->x==:gene,species_symbol_type_array)
    gene_abundance_array = initial_condition_array[idx_gene]

    # Precompute the translation parameters -
	translation_parameter_array = precompute_translation_parameter_array(biophysical_constants_dictionary, protein_coding_length_array, time_constant_modifier_array,host_type)
    model_data_dictionary["translation_parameter_array"] = translation_parameter_array

	# Precompute the kinetic limit of transcription -
	transcription_kinetic_limit_array = precompute_transcription_kinetic_limit_array(biophysical_constants_dictionary, gene_coding_length_array, gene_abundance_array, time_constant_modifier_array, host_type)
    model_data_dictionary["transcription_kinetic_limit_array"] = transcription_kinetic_limit_array

    # Dilution degrdation matrix -
    dilution_degradation_matrix = build_dilution_degradation_matrix(biophysical_constants_dictionary, species_symbol_type_array, degradation_modifier_array)
    model_data_dictionary["dilution_degradation_matrix"] = dilution_degradation_matrix
    # ===================================================================================================== #

    # Phase 2:  solve model equations ===================================================================== #
    # solve the balance equations -
    (TSIM,XSIM) = SolveBalances(time_start,time_stop,time_step_size,model_data_dictionary)
    # ===================================================================================================== #

    # Phase 3:  compute simulation error ================================================================== #
    # compute the error - we need to do a bunch of interpolation -
    tsim_exp = exp_data_dictionary["mRNA_data_array"][:,1]

    # GFP mRNA -
    itp_gfp_mRNA =  LinearInterpolation(TSIM, (1000)*XSIM[:,3]);
    mRNA_GFP_sim = itp_gfp_mRNA[tsim_exp]  # convert to muM from nM

    # GFP protein -
    itp_gfp_protein =  LinearInterpolation(TSIM, XSIM[:,5]);
    protein_GFP_sim = itp_gfp_protein[tsim_exp]

    # get experimental data -
    mRNA_GFP_exp = exp_data_dictionary["mRNA_data_array"][:,2]          # mean is col 2 nM
    mRNA_GFP_std_exp = exp_data_dictionary["mRNA_data_array"][:,3]      # stdev is col 3 nM

    protein_GFP_exp = exp_data_dictionary["prot_data_array"][:,2]       # mean is col 2 muM
    protein_GFP_std_exp = exp_data_dictionary["prot_data_array"][:,3]   # stdev is col 3 muM

    # compute error terms -
    error_term_array = zeros(2)

    # mRNA GFP -
    #mRNA_GFP_std_exp[1] = 1.0   # we have 0 ic
    #tmp_arr = 1.0./((mRNA_GFP_std_exp).^2)
    #W_mRNA = diagm(tmp_arr)
    error_vector_1 = (mRNA_GFP_exp .- mRNA_GFP_sim)
    error_term_array[1] = transpose(error_vector_1)*error_vector_1

    # protein deGFP -
    #protein_GFP_std_exp[1] = 1.0
    #tmp_arr = 1.0./((protein_GFP_std_exp).^2)
    #W_prot = diagm(tmp_arr)
    error_vector_2 = (protein_GFP_exp .- protein_GFP_sim)
    error_term_array[2] = transpose(error_vector_2)*error_vector_2
    # ===================================================================================================== #

    # error total -
    error_total = sum(error_term_array)

    # return -
    return error_total
end

function main(path_to_data_dir::String; algorithm::Symbol = :LBFGS)

    # load the default data_dictionary -
    time_start = 0.0
    time_stop = 16.0
    time_step_size = 0.01

    # what is the host_type?
    host_type = :cell_free

    # path to parameters -
    path_to_biophysical_constants_file = "./CellFree.json"

    # Load the data dictionary (uses the default biophysical_constants file)
    default_data_dictionary = build_data_dictionary((time_start,time_stop,time_step_size), path_to_biophysical_constants_file, host_type)

    # Set some values for the weights, gene lengths etc -
    model_data_dictionary = customize_data_dictionary(default_data_dictionary,host_type)

    # load the experimental data -
    exp_data_dictionary = load_experimental_data_dictionary(path_to_data_dir)

    # setup the objective function -
    OF(P) = objective_function(P,time_start,time_step_size,time_stop,model_data_dictionary,exp_data_dictionary)

    # setup paramter bounds -
    pvec_bounds = [

        # dG's -
        20000.0  80000.0    ;   # 1     W_deGFP_RNAP
        -50000.0  -100.0    ;   # 2     W_deGFP_sigma_70

        # binding parameters -
        0.5 10.0            ;   # 3     n_deGFP_sigma_70
        0.001 100.0         ;   # 4     K_deGFP_sigma_70

        # time constants -
		0.001 100.0         ;	# 5	    mRNA_deGFP
		0.001 100.0         ;	# 6	    protein_deGFP

        # degradation mods -
		0.001 100.0 	    ;	# 7	    mRNA_deGFP
		0.001 100.0 	    ;	# 8	    protein_deGFP
        0.001 100.0 	    ;	# 9	    protein_sigma_70

         # w -
         4.0 10.0           ;   # 10    translation capacity half-life

        # KL value -
        100.0 500.0         ;   # 11    KL in muM
    ];

    # setup initial condition vector -
    pvec_initial = [

        # dG's -
        40000.0     ;   # 1     W_deGFP_RNAP
        -25000.0    ;   # 2     W_deGFP_sigma_70

        # binding parameters -
        1.0         ;   # 3     n_deGFP_sigma_70
        30.0        ;   # 4     K_deGFP_sigma_70

        # time constants -
		1.0         ;	# 5	    mRNA_deGFP
		1.0         ;	# 6	    protein_deGFP

        # degradation mods -
		1.0 	    ;	# 7	    mRNA_deGFP
		1.0 	    ;	# 8	    protein_deGFP
        1.0 	    ;	# 9	    protein_sigma_70

        # w -
        8.0         ;   # 10    translation capacity half-life

        # KL -
        250.0       ;   # 11    KL in muM
    ];

    # check bounds -
    number_of_parameters = length(pvec_initial)
    for parameter_index = 1:number_of_parameters

        # what is the parameter value?
        p_i = pvec_initial[parameter_index]

        # is p_i outside of the bounds?
        lb_value = pvec_bounds[parameter_index,1]
        ub_value = pvec_bounds[parameter_index,2]

        if (p_i<lb_value)
            pvec_initial[parameter_index,1] = lb_value
        end

        if (p_i>ub_value)
            pvec_initial[parameter_index,1] = ub_value
        end
    end

    # search -
    options = Optim.Options(show_trace=true,show_every=20,iterations=400)
    if algorithm == :SimulatedAnnealing
        result = optimize(OF,pvec_bounds[:,1],pvec_bounds[:,2],pvec_initial,Fminbox(SimulatedAnnealing()), options)
        return result
    elseif algorithm == :LBFGS
        options = Optim.Options(show_trace=true,show_every=10,iterations=100)
        result = optimize(OF,pvec_bounds[:,1],pvec_bounds[:,2],pvec_initial,Fminbox(LBFGS()),options)
        return result
    end

    result = optimize(OF,pvec_bounds[:,1],pvec_bounds[:,2],pvec_initial,Fminbox(SimulatedAnnealing()),options)
    return result
end

# setup paths -
path_to_data_dir = "$(pwd())/data"

# generate an ensemble of parameters -
number_of_p_sets = 100
number_of_parameters = 11+1 # end + 1 hold the error -
ensemble_array = zeros(number_of_parameters,number_of_p_sets)
algorithm_flag = :SimulatedAnnealing
for pset_index = 1:number_of_p_sets

    algorithm_flag = :SimulatedAnnealing
    if mod(pset_index,10) == 0
        algorithm_flag = :LBFGS
    end

    # call main -
    opt_result = main(path_to_data_dir; algorithm = algorithm_flag)

    # grab the best pset, and then dump to disk -
    pV = Optim.minimizer(opt_result)
    ensemble_array[end,pset_index] = Optim.minimum(opt_result)
    for (p_index,p_value) in enumerate(pV)
        ensemble_array[p_index,pset_index] = p_value
    end
end

# dump -
writedlm("Ensemble-T4.dat",ensemble_array)


@testset "Field arithmetic" begin
    @field(A) do
        I₀ = 1.0
        T = 2.0
        σ = 3.0
        Tmax = 3.0
    end

    @field(A2) do
        I₀ = 1.0
        T = 2.0
        σ = 3.0
        Tmax = 3.0
        ξ = 0.6
    end

    @testset "Summed fields" begin
        @field(B) do
            I₀ = 0.5
            T = 1.0
            σ = 3.0
            Tmax = 3.0
        end

        C = A + B
        @test span(C) == span(A)
        @test span(B) ⊆ span(C)
        @test steps(C, 100) == steps(A, 100/austrip(period(B)))

        @test polarization(C) == LinearPolarization()
        @test dimensions(C) == 1

        @test vector_potential(C, 0.4) == vector_potential(A, 0.4) + vector_potential(B, 0.4)
        @test field_amplitude(C, 0.4) == field_amplitude(A, 0.4) + field_amplitude(B, 0.4)

        # These functions only make sense when adding two field of the same wavelength
        for fun in [wavelength, period, frequency, wavenumber, fundamental, photon_energy]
            @test fun(A+A) == fun(A)
        end

        @test max_frequency(C) == frequency(B)

        @test continuity(C) == Inf

        @test field_amplitude(phase_shift(C, π/3), 0.4) ≈ field_amplitude(phase_shift(A, π/3), 0.4) + field_amplitude(phase_shift(B, π/3), 0.4)

        @test field_amplitude(-A, 0.4) == -field_amplitude(A, 0.4)
        @test field_amplitude(A-B, 0.4) == field_amplitude(A, 0.4) - field_amplitude(B, 0.4)

        withenv("UNITFUL_FANCY_EXPONENTS" => true) do
            @test string(C) == """
                               ┌ Linearly polarized field with
                               │   - I₀ = 1.0000e+00 au = 3.5094452e16 W cm⁻² =>
                               │     - E₀ = 1.0000e+00 au = 514.2207 GV m⁻¹
                               │     - A₀ = 0.3183 au
                               │   – a Fixed carrier @ λ = 14.5033 nm (T = 48.3777 as, ω = 3.1416 Ha = 85.4871 eV)
                               │   – and a Gaussian envelope of duration 170.8811 as (intensity FWHM; ±2.00σ)
                               │   – Uₚ = 0.0253 Ha = 689.2724 meV => α = 0.1013 Bohr = 5.3617 pm
                               ⊕
                               │ Linearly polarized field with
                               │   - I₀ = 5.0000e-01 au = 1.7547226e16 W cm⁻² =>
                               │     - E₀ = 7.0711e-01 au = 363.6089 GV m⁻¹
                               │     - A₀ = 0.1125 au
                               │   – a Fixed carrier @ λ = 7.2516 nm (T = 24.1888 as, ω = 6.2832 Ha = 170.9742 eV)
                               │   – and a Gaussian envelope of duration 170.8811 as (intensity FWHM; ±1.00σ)
                               └   – Uₚ = 0.0032 Ha = 86.1591 meV => α = 0.0179 Bohr = 947.8211 fm
                               """
        end
        withenv("UNITFUL_FANCY_EXPONENTS" => false) do
            @test string(C) == """
                               ┌ Linearly polarized field with
                               │   - I₀ = 1.0000e+00 au = 3.5094452e16 W cm^-2 =>
                               │     - E₀ = 1.0000e+00 au = 514.2207 GV m^-1
                               │     - A₀ = 0.3183 au
                               │   – a Fixed carrier @ λ = 14.5033 nm (T = 48.3777 as, ω = 3.1416 Ha = 85.4871 eV)
                               │   – and a Gaussian envelope of duration 170.8811 as (intensity FWHM; ±2.00σ)
                               │   – Uₚ = 0.0253 Ha = 689.2724 meV => α = 0.1013 Bohr = 5.3617 pm
                               ⊕
                               │ Linearly polarized field with
                               │   - I₀ = 5.0000e-01 au = 1.7547226e16 W cm^-2 =>
                               │     - E₀ = 7.0711e-01 au = 363.6089 GV m^-1
                               │     - A₀ = 0.1125 au
                               │   – a Fixed carrier @ λ = 7.2516 nm (T = 24.1888 as, ω = 6.2832 Ha = 170.9742 eV)
                               │   – and a Gaussian envelope of duration 170.8811 as (intensity FWHM; ±1.00σ)
                               └   – Uₚ = 0.0032 Ha = 86.1591 meV => α = 0.0179 Bohr = 947.8211 fm
                               """
        end
    end

    @testset "Delayed fields" begin
        B = delay(A, 0.4)
        @test field_amplitude(delay(A,-0.4), 0.0) == field_amplitude(B, 0.8) == field_amplitude(A, 0.4)
        @test field_amplitude(delay(A, π*u"rad"), 0.0) == field_amplitude(A, -1.0)

        @test delay(B) == 0.4
        @test iszero(delay(A))

        @test span(B) == -5.6..6.4

        withenv("UNITFUL_FANCY_EXPONENTS" => true) do
            @test string(B) == """
                               Linearly polarized field with
                                 - I₀ = 1.0000e+00 au = 3.5094452e16 W cm⁻² =>
                                   - E₀ = 1.0000e+00 au = 514.2207 GV m⁻¹
                                   - A₀ = 0.3183 au
                                 – a Fixed carrier @ λ = 14.5033 nm (T = 48.3777 as, ω = 3.1416 Ha = 85.4871 eV)
                                 – and a Gaussian envelope of duration 170.8811 as (intensity FWHM; ±2.00σ)
                                 – Uₚ = 0.0253 Ha = 689.2724 meV => α = 0.1013 Bohr = 5.3617 pm
                                 – delayed by 0.4000 jiffies = 9.6755 as"""
        end
        withenv("UNITFUL_FANCY_EXPONENTS" => false) do
            @test string(B) == """
                               Linearly polarized field with
                                 - I₀ = 1.0000e+00 au = 3.5094452e16 W cm^-2 =>
                                   - E₀ = 1.0000e+00 au = 514.2207 GV m^-1
                                   - A₀ = 0.3183 au
                                 – a Fixed carrier @ λ = 14.5033 nm (T = 48.3777 as, ω = 3.1416 Ha = 85.4871 eV)
                                 – and a Gaussian envelope of duration 170.8811 as (intensity FWHM; ±2.00σ)
                                 – Uₚ = 0.0253 Ha = 689.2724 meV => α = 0.1013 Bohr = 5.3617 pm
                                 – delayed by 0.4000 jiffies = 9.6755 as"""
        end
    end

    @testset "Combined" begin
        @field(X) do
            I₀ = 1.0
            ω = 1.0
            cycles = 6.0
            env = :cos²
            ξ = 1.0
        end

        @field(Y) do
            I₀ = 1.0
            ω = 2.0
            cycles = 6.0
            env = :cos²
            ξ = -1.0
        end

        Z = X + delay(Y, 3/2π)

        @test field_amplitude(Z, 4.0) ≈ field_amplitude(X, 4.0) + field_amplitude(Y, 4.0 - 3/2π)

        W = X + delay(Y, 35.0)
        @test span(W) == span(X).left..(span(Y).right+35.0)
    end

    @testset "Padded fields" begin
        @test_throws ArgumentError PaddedField(A, -3, 1)
        @test_throws ArgumentError PaddedField(A, 3, -1)

        B = PaddedField(A, 3.0u"fs", 8.0u"fs")
        B2 = PaddedField(A2, 3.0u"fs", 8.0u"fs")

        @test parent(B) == A
        @test parent(B2) == A2

        for (F1,F2) in [(A,B), (A2,B2)]
            @test field_amplitude(F2, 0.4) == field_amplitude(F1, 0.4)
            @test field_amplitude(F2, 7.0) == zero(field_amplitude(F1, 0.4))
        end

        @test all(endpoints(span(B)) .≈ (-6-austrip(3.0u"fs"), 6+austrip(8.0u"fs")))

        @test time_integral(B) == time_integral(A)

        withenv("UNITFUL_FANCY_EXPONENTS" => true) do
            @test string(B) == """
                               Padding before 124.0241 jiffies = 3.0000 fs and after 330.7310 jiffies = 8.0000 fs of
                               Linearly polarized field with
                                 - I₀ = 1.0000e+00 au = 3.5094452e16 W cm⁻² =>
                                   - E₀ = 1.0000e+00 au = 514.2207 GV m⁻¹
                                   - A₀ = 0.3183 au
                                 – a Fixed carrier @ λ = 14.5033 nm (T = 48.3777 as, ω = 3.1416 Ha = 85.4871 eV)
                                 – and a Gaussian envelope of duration 170.8811 as (intensity FWHM; ±2.00σ)
                                 – Uₚ = 0.0253 Ha = 689.2724 meV => α = 0.1013 Bohr = 5.3617 pm"""
        end
        withenv("UNITFUL_FANCY_EXPONENTS" => false) do
            @test string(B) == """
                               Padding before 124.0241 jiffies = 3.0000 fs and after 330.7310 jiffies = 8.0000 fs of
                               Linearly polarized field with
                                 - I₀ = 1.0000e+00 au = 3.5094452e16 W cm^-2 =>
                                   - E₀ = 1.0000e+00 au = 514.2207 GV m^-1
                                   - A₀ = 0.3183 au
                                 – a Fixed carrier @ λ = 14.5033 nm (T = 48.3777 as, ω = 3.1416 Ha = 85.4871 eV)
                                 – and a Gaussian envelope of duration 170.8811 as (intensity FWHM; ±2.00σ)
                                 – Uₚ = 0.0253 Ha = 689.2724 meV => α = 0.1013 Bohr = 5.3617 pm"""
        end
    end

    @testset "Windowed fields" begin
        B = WindowedField(A, -0.1u"fs", 0.05u"fs")
        B2 = WindowedField(A2, -0.1u"fs", 0.05u"fs")

        @test parent(B) == A
        @test parent(B2) == A2

        for fun in [vector_potential, field_amplitude, intensity]
            for (F1,F2) in [(A,B), (A2,B2)]
                @test field_amplitude(F2, 0.4) == field_amplitude(F1, 0.4)
                @test field_amplitude(F2, -5.0) == zero(field_amplitude(F1, 0.4))
                @test field_amplitude(F2, 3.0) == zero(field_amplitude(F1, 0.4))
            end
        end

        @test all(endpoints(span(B)) .≈ (-austrip(0.1u"fs"), austrip(0.05u"fs")))

        @test field_amplitude(phase_shift(B, 2.0), 0.4) == field_amplitude(phase_shift(A, 2.0), 0.4)
        @test field_amplitude(phase_shift(B, 2.0), 3.0) == zero(field_amplitude(phase_shift(A, 2.0), 0.4))

        withenv("UNITFUL_FANCY_EXPONENTS" => true) do
            @test string(B) == """
                               Window from -4.1341 jiffies = -100.0000 as to 2.0671 jiffies = 50.0000 as of
                               Linearly polarized field with
                                 - I₀ = 1.0000e+00 au = 3.5094452e16 W cm⁻² =>
                                   - E₀ = 1.0000e+00 au = 514.2207 GV m⁻¹
                                   - A₀ = 0.3183 au
                                 – a Fixed carrier @ λ = 14.5033 nm (T = 48.3777 as, ω = 3.1416 Ha = 85.4871 eV)
                                 – and a Gaussian envelope of duration 170.8811 as (intensity FWHM; ±2.00σ)
                                 – Uₚ = 0.0253 Ha = 689.2724 meV => α = 0.1013 Bohr = 5.3617 pm"""
        end
        withenv("UNITFUL_FANCY_EXPONENTS" => false) do
            @test string(B) == """
                               Window from -4.1341 jiffies = -100.0000 as to 2.0671 jiffies = 50.0000 as of
                               Linearly polarized field with
                                 - I₀ = 1.0000e+00 au = 3.5094452e16 W cm^-2 =>
                                   - E₀ = 1.0000e+00 au = 514.2207 GV m^-1
                                   - A₀ = 0.3183 au
                                 – a Fixed carrier @ λ = 14.5033 nm (T = 48.3777 as, ω = 3.1416 Ha = 85.4871 eV)
                                 – and a Gaussian envelope of duration 170.8811 as (intensity FWHM; ±2.00σ)
                                 – Uₚ = 0.0253 Ha = 689.2724 meV => α = 0.1013 Bohr = 5.3617 pm"""
        end
    end
end


################################################################################
# Assign Model and Target Parameter
_modelProb = ModelWrapper(ProbModel(), val_dist)
_syms = [
    keys(_modelProb.val),
    keys(_modelProb.val)[1],
    keys(_modelProb.val)[end],
    (:d_a1, :d_a2, :d_a3, :d_a5, :d_a6, :d_a7, :d_a8),
]
_targets = [Tagged(_modelProb, _syms[iter]) for iter in eachindex(_syms)]
_params = [sample(_modelProb, _targets[iter]) for iter in eachindex(_syms)]

@testset "Tagged - Model parameter" begin
        ## Assign Sub Model
        for iter in eachindex(_syms)
        _sym = _syms[iter]
        _target = _targets[iter]
        _param = _params[iter]
        _model_temp = ModelWrapper(subset(val_dist, _sym))
        ## Compute logdensities and check length
        @test length(_model_temp) == length(_target)
        @test abs(log_prior(_model_temp, _target) - log_prior(_model_temp)) <= _TOL
        @test abs(log_abs_det_jac(_model_temp, _target) - log_abs_det_jac(_model_temp)) <=
              _TOL
        ## Type Check - Unflatten/Constrain
        theta_unconstrained_vec = randn(length(_target))
        theta_constrained = unflatten_constrain(_model_temp, theta_unconstrained_vec)
        theta_constrained2 = unflatten_constrain(
            _model_temp, _target, theta_unconstrained_vec
        )
        @test typeof(theta_constrained) == typeof(_model_temp.val)
        @test typeof(theta_constrained2) == typeof(_model_temp.val)
        _θ1  = flatten(_model_temp.info.reconstruct, theta_constrained)
        _θ2  = flatten(_target.info.reconstruct, theta_constrained2)
        @test sum(abs.(_θ1 - _θ2)) ≈ 0 atol = _TOL
        ## Utility functions
        log_prior(_target, _model_temp.val)
        θ_flat = flatten(_model_temp, _target)
        unflatten(_model_temp, _target, θ_flat)
        unflatten!(_model_temp, _target, θ_flat)
        θ_flat_unconstrained = unconstrain_flatten(_model_temp, _target)
        unflatten_constrain!(_model_temp, _target, θ_flat_unconstrained)
        log_prior_with_transform(_model_temp, _target)

        subset(_model_temp, _target)
        ModelWrappers.length(_target)
        ModelWrappers.paramnames(_target)
        fill(_model_temp, _target, _model_temp.val)
        fill!(_model_temp, _target, _model_temp.val)
        _model_temp.val
        sample(_RNG, _model_temp, _target)
        sample(_model_temp, _target)
        sample!(_RNG, _model_temp, _target)
        sample!(_model_temp, _target)
    end
end


using BioFeatures
using DataFrames
# Make some coordinates
region_lengths = rand(100)*1000
starts = int( 1:1000:100000 )
stops = int(  starts + region_lengths )

# Read these into a BioFeature
regions = DataFrame(start = starts, stop = stops)
F = Features( regions )

seq_ids = keys( F.features )
s = nothing
for seq_id in keys(F.features)
   s = F.features[seq_id]
end

for v = values(s)
  println(v[:start])
end


# Subsampling space can definitely be necessary
# Subsampling time may not be, maybe rather use DESolution methods.

# Supposing adaptive time, specify "saveat" -- simple.
# Optional, allow specifying "tstops" which forces non-interpolation.

import Base: to_indices, _maybetail, @_inline_meta, tail, getindex

abstract type AbstractSubsampler{D} end
function coordinate_indices(lattice::LATTICE,
        subsampler_arr::AbstractArray{<:AbstractSubsampler}) where {LATTICE<:AbstractLattice}
     idxs_arr = map(subsampler_arr) do subsampler
        coordinate_indices(lattice, subsampler)
    end
    all_keep = ones(Bool, size(lattice)...)
    for idxs in idxs_arr
        keep = zeros(Bool, size(lattice)...)
        keep[idxs] .= true
        all_keep .&= keep
    end
    all_indices = CartesianIndices(lattice)
    return all_indices[all_keep]
end
function subsample(lattice::LATTICE, sub) where {LATTICE<:AbstractLattice}
    idxs = coordinate_indices(lattice, sub)
    LATTICE(lattice.arr[idxs])
end
subsample(space::AbstractLattice, ::Nothing) = space
function getindex(lattice::LATTICE, sub) where {T,D,LATTICE<:AbstractLattice{T,D}}
    idxs = coordinate_indices(lattice, sub)
    LATTICE(lattice.arr[idxs])
end

@with_kw struct IndexSubsampler{D} <: AbstractSubsampler{D}
	strides::NTuple{D,Int}
end
@with_kw struct RightCutFromValue{D,T} <: AbstractSubsampler{D}
	cut::NTuple{D,T}
end
export RightCutFromValue

@with_kw struct RightCutProportionFromValue{D,T} <: AbstractSubsampler{D}
    cut::NTuple{D,T}
    proportion::NTuple{D,T}
end
export RightCutProportionFromValue

"IndexInfo specifies the index of 0 and the stride (`Δ`) between indices."
@with_kw struct IndexInfo{D,N}
	Δ::NTuple{N,D}
	origin_idx::CartesianIndex{N}
end

"""
	StrideToEnd(stride, start=1)

StrideToEnd is a custom index type that acts like `start:stride:end`, to circumvent the fact that you can't put `start:stride:end` into a variable.
"""
struct StrideToEnd{STOP}
    stride::Int
	start::Int
    stop::STOP
    StrideToEnd{STOP}(stride::Int, start::Int=1, stop::STOP=nothing) where {STOP<:Union{Nothing,Int}} = new(stride, start,stop)
end
StrideToEnd(stride, start, stop::S=nothing) where S = StrideToEnd{S}(stride, start, stop)
to_indices(A, inds, I::Tuple{StrideToEnd{Int}, Vararg{Any}})=
(@_inline_meta; (I[1].start:I[1].stride:I[1].stop, to_indices(A, _maybetail(inds), tail(I))...))
to_indices(A, inds, I::Tuple{StrideToEnd{Nothing}, Vararg{Any}})=
	(@_inline_meta; (I[1].start:I[1].stride:inds[1][end], to_indices(A, _maybetail(inds), tail(I))...))
to_indices(A, inds, I::Tuple{NTuple{N,StrideToEnd}, Vararg{Any}}) where N = to_indices(A, inds, (I[1]..., _maybetail(I)...))
getindex(A::AbstractArray, S::StrideToEnd) = getindex(A, to_indices(A, (S,))...)

function coordinate_indices(lattice::AbstractLattice{T}, subsampler::RightCutFromValue{D,T}) where {T,D}
    axis_dxs = map(zip(subsampler.cut, coordinate_axes(lattice))) do (value, axis)
        findfirst(axis .>= value)
    end
    stride = 1
    cuts = Tuple(axis_dxs)
    return StrideToEnd.(stride, cuts)
end
function coordinate_indices(lattice::AbstractLattice{T}, subsampler::RightCutProportionFromValue{D,T}) where {T,D}
    start_stop = map(zip(subsampler.cut, subsampler.proportion, coordinate_axes(lattice))) do (value, proportion, axis)
        axis_begin = findfirst(axis .>= value)
        remaining_span = length(axis) - axis_begin
        axis_end = floor(Int, axis_begin + remaining_span * proportion)
        (axis_begin, axis_end)
    end
    axis_starts, axis_stops = zip(start_stop...)
    return StrideToEnd.(1, axis_starts, axis_stops)
end

function coordinate_indices(::Any, subsampler::IndexSubsampler{D}) where D
    return StrideToEnd.(subsampler.strides, 1)
end

function coordinate_indices(lattice, sub_dxs::AbstractArray{<:CartesianIndex})
    #@warn "subsampling one pop"
    #FIXME don't subsample one pop
    idxs = population(sub_dxs, 1)
end


using Distributed
using Base.Filesystem
using DataFrames
using CSV
using Query
using Statistics
using UnicodePlots
using ClusterManagers
using Dates
using DelimitedFiles

## load the packages by covid19abm
using Parameters, Distributions, StatsBase, StaticArrays, Random, Match, DataFrames
using covid19abm
const cv=covid19abm

#addprocs(30, exeflags="--project=.")
#@everywhere using covid19abm

addprocs(SlurmManager(500), N=17, topology=:master_worker, exeflags="--project=.")
@everywhere using covid19abm

function run(myp::ModelParameters, nsims=500, folderprefix="./", writefiles = true)
    println("starting $nsims simulations...\nwriting files: $writefiles")
    #dump(myp)
    myp.calibration && error("can not run simulation, calibration is on.")
    # will return 6 dataframes. 1 total, 4 age-specific
    cdr = pmap(1:nsims) do x
            cv.runsim(x, myp)
    end

    println("simulations finished")
    println("total size of simulation dataframes: $(Base.summarysize(cdr))")
    ## write the infectors
    writefiles && writedlm("$(folderprefix)/infectors.dat", [cdr[i].infectors for i = 1:nsims])

    ## write contact numbers
    writefiles && writedlm("$(folderprefix)/ctnumbers.dat", [cdr[i].ct_numbers for i = 1:nsims])
    ## stack the sims together
    allag = vcat([cdr[i].a  for i = 1:nsims]...)
    ag1 = vcat([cdr[i].g1 for i = 1:nsims]...)
    ag2 = vcat([cdr[i].g2 for i = 1:nsims]...)
    ag3 = vcat([cdr[i].g3 for i = 1:nsims]...)
    ag4 = vcat([cdr[i].g4 for i = 1:nsims]...)
    ag5 = vcat([cdr[i].g5 for i = 1:nsims]...)
    #mydfs = Dict("all" => all, "ag1" => ag1, "ag2" => ag2, "ag3" => ag3, "ag4" => ag4, "ag5" => ag5)
    mydfs = Dict("all" => allag)

    ## save at the simulation and time level
    ## to ignore for now: miso, iiso, mild
    c1 = Symbol.((:LAT, :ASYMP, :MILD, :MISO, :INF, :IISO, :HOS, :ICU, :DED), :_INC)
    c2 = Symbol.((:LAT, :ASYMP, :MILD, :MISO, :INF, :IISO, :HOS, :ICU, :DED), :_PREV)
    for (k, df) in mydfs
        println("saving dataframe sim level: $k")
        # simulation level, save file per health status, per age group
        for c in vcat(c1..., c2...)
            udf = unstack(df, :time, :sim, c)
            fn = string("$(folderprefix)/simlevel_", lowercase(string(c)), "_", k, ".dat")
            writefiles && CSV.write(fn, udf)
        end
        println("saving dataframe time level: $k")
        # time level, save file per age group
        yaf = compute_yearly_average(df)
        fn = string("$(folderprefix)/timelevel_", k, ".dat")
        writefiles && CSV.write(fn, yaf)

        # write the secondary (Ro) values
        secvals = unstack(df, :time, :sim, :secondarys)
        fn = string("$(folderprefix)/secondarys_", k, ".dat")
        writefiles && CSV.write(fn, secvals)
    end
    return mydfs
end

function compute_yearly_average(df)
    ya = df |> @groupby(_.time) |> @map({time=key(_), cnt=length(_),
              sus_prev=mean(_.SUS_PREV),
              lat_prev=mean(_.LAT_PREV),
              pre_prev=mean(_.PRE_PREV),
              asymp_prev=mean(_.ASYMP_PREV),
              mild_prev=mean(_.MILD_PREV),
              miso_prev=mean(_.MISO_PREV),
              inf_prev=mean(_.INF_PREV),
              iiso_prev=mean(_.IISO_PREV),
              hos_prev=mean(_.HOS_PREV),
              icu_prev=mean(_.ICU_PREV),
              rec_prev=mean(_.REC_PREV),
              ded_prev=mean(_.DED_PREV),
              sus_inc=mean(_.SUS_INC),
              lat_inc=mean(_.LAT_INC),
              pre_inc=mean(_.PRE_INC),
              asymp_inc=mean(_.ASYMP_INC),
              mild_inc=mean(_.MILD_INC),
              miso_inc=mean(_.MISO_INC),
              inf_inc=mean(_.INF_INC),
              iiso_inc=mean(_.IISO_INC),
              hos_inc=mean(_.HOS_INC),
              icu_inc=mean(_.ICU_INC),
              rec_inc=mean(_.REC_INC),
              ded_inc=mean(_.DED_INC)
              }) |> DataFrame
    return ya
end

function savestr(p::ModelParameters, custominsert="simresults", customstart="")
    ## format: /data/covid19abm/$(custominsert)/$(customstart)_$(params)
    datestr = (Dates.format(Dates.now(), dateformat"mmdd_HHMM"))
    ## setup folder name based on model parameters
    taustr = replace(string(p.τmild), "." => "")
    fstr = replace(string(p.fmild), "." => "")
    rstr = replace(string(p.β), "." => "")
    prov = replace(string(p.prov), "." => "")
    eldr = replace(string(p.eldq), "." => "")
    eldqag = replace(string(p.eldqag), "." => "")
    fpreiso = replace(string(p.fpreiso), "." => "")
    tpreiso = replace(string(p.tpreiso), "." => "")
    fsev = replace(string(p.fsevere), "." => "")
    frelasymp = replace(string(p.frelasymp), "." => "")

    if p.ctstrat == 1
        strat = "strat1sc"
    elseif p.ctstrat == 3
        strat = "strat3s" * string(p.strat3qdays)
    else
        strat = "strat0"
    end
    pct = replace(string(p.fctcapture), "." => "")
    cct = replace(string(p.fcontactst), "." => "")
    idt = replace(string(p.cidtime), "." => "")
    tback = replace(string(p.cdaysback), "." => "")
    fldrname = "/data/covid19abm/$(custominsert)/$(customstart)_$(prov)_$(strat)_pct$(pct)_cct$(cct)_idt$(idt)_tback$(tback)_fsev$(fsev)_tau$(taustr)_fmild$(fstr)_q$(eldr)_qag$(eldqag)_relasymp$(frelasymp)_tpreiso$(tpreiso)_preiso$(fpreiso)/"
    mkpath(fldrname)
end

function _calibrate(nsims, myp::ModelParameters)
    myp.calibration != true && error("calibration parameter not turned on")
    vals = zeros(Int64, nsims)
    println("calibrating with beta: $(myp.β), total sims: $nsims, province: $(myp.prov)")
    println("calibration parameters:")
    #dump(myp)
    cdr = pmap(1:nsims) do i
        h = main(myp) ## gets the entire model.
        val = sum(_get_column_incidence(h, covid19abm.LAT))
        return val
    end
    return mean(cdr), std(cdr)
end

function calibrate(beta, nsims=1000, init_inf=1, size=10000, prov=:newyork)
    myp = ModelParameters() # set up default parameters
    myp.β = beta
    myp.prov = prov
    myp.popsize = size
    myp.modeltime = 50
    myp.calibration = true
    myp.initialinf = init_inf
    m, sd = _calibrate(nsims, myp)
    println("mean R0: $(m) with std: $(sd)")
    myp.calibration = false
    return m
end

function calibrate_regression()
    ## runs through a range of beta values and fits a linear regression line
    betas = 0.0225:0.0005:0.085
    ros = similar(betas)
    for i = 1:length(betas)
        ros[i] = calibrate(betas[i])
    end
    # run regression, myb is a dataframe of ros and betas
    #linearRegressor = lm(@formula(ros ~ betas), df)  ## formula expression requires dataframe
    regressed = lm(hcat(ones(length(betas)), betas), ros)
    return regressed
    # intercept = -0.00285 + 46.1761(beta) = ro
    # (ro + 0.00285)/46.1761
end

function calibrate_robustness(beta, reps, prov=:ontario)
    #[:ontario, :alberta, :bc, :manitoba, :newbruns, :newfdland, :nwterrito, :novasco, :nunavut, :pei, :quebec, :saskat, :yukon]
    # once a beta is found based on nsims simulations,
    # see how robust it is. run calibration with same beta 100 times
    # to see the variation in R0 produced.
    nsims = [500, 1000]
    means = zeros(Float64, reps, length(nsims))
    for (i, ns) in enumerate(nsims)
        cd = map(1:reps) do x
            println("iter: $x, sims: $ns")
            mval = calibrate(beta, ns, prov)
            return mval
        end
        means[:, i] = cd
    end
    # for i in 2:nworkers()
    #     ## mf defined as: @everywhere mg() = covid19abm.p.β
    #     rpr = remotecall_fetch(mf,  i+1).prov
    #     rpr != prov && error("province didn't get set in the remote workers")
    # end
    return means
end
#provs = []
#for x in [:ontario, :alberta, :bc, :manitoba, :newbruns]
#    m = calibrate_robustness(0.101, x)
#    push!(provs, m)
#end

#@everywhere mg() = covid19abm.p.β
#@everywhere mf(b) = covid19abm.p.β = b
# function hmatrix_analysis()
#     ## function takes the output of the model
#     ## and splits it into an array of 2d matrices for plotting purposes.
#     myp = cv.ModelParameters()
#     myp.β = 0.038
#     reset_params($myp)
#     cd, ags = main(1)

#     #mats = Array{Int64, 3}(undef, 500)
#     mats = Array{Int64, 3}(undef, 100,100,500)
#     for i = 1:500
#         col = cd[:, i]
#         mats[:, :, i] = reshape(col, (100, 100))
#     end

#     colorarray = IndirectArray(mats, [colorant"green", colorant"teal", colorant"blue", colorant"azure",
#     colorant"red", colorant"lightsalmon", colorant"magenta", colorant"magenta", colorant"grey60", colorant"black"])
#     save("myanimation.gif", colorarray; fps=5)

#     idxarray = rand(1:3, 100, 100, 50);
#     colorarray = IndirectArray(idxarray, [colorant"red", colorant"purple", colorant"indigo"])

#     return mats
# end

# function makeheatmaps(bufs)
#     heatmaps = map(bufs) do buf
#         heatmap(
#             buf, scale_plot = false, show_axis = false)

#     end
#     scene = hbox(map(i-> vbox(heatmaps[i, :]), 1:size(bufs, 1))...)
#     scene, last.(heatmaps)
# end
#datarows = 500; datacols = 500
#plotrows = 4; plotcols = 4
#bufs = [fill(0.0f0, datarows, datacols) for _ in 1:plotrows, _ in 1:plotcols]

# http://juliaplots.org/MakieReferenceImages/gallery//lots_of_heatmaps/index.html
# http://juliaplots.org/MakieReferenceImages/gallery//chess_game/index.html


"""
linearStateTolerances!(n)
# the purpose of this function is to taper the tolerances on the constant state constraints
# the idea is that when doing MPC, the final states are well within the bounds so that the next optimization is not initalized at an infeasible point
# if you want a constant bond, set the slope to zero
# default is a positive slope on the lower bound and a negative slope on the upper bound
# this functionality in not needed for states like position, so you do not need to add a linearStateTol for all states

--------------------------------------------------------------------------------------\n
Author: Huckleberry Febbo, Graduate Student, University of Michigan
Date Create: 3/23/2017, Last Modified: 4/13/2018 \n
-------------------------------------------------------------------------------------\n
"""
function linearStateTolerances!(n::NLOpt;
                                mXL::Array{Any,1}=falses(n.ocp.state.num),
                                mXU::Array{Any,1}=falses(n.ocp.state.num))
  n.ocp.mXL=mXL;  n.ocp.mXU=mXU;
  for st in 1:n.ocp.state.num
    # lower state constraint
    if n.ocp.mXL[st]!=false
      if !isnan(n.ocp.XL[st])
        for j in 1:n.ocp.state.pts
        n.ocp.XL_var[st,j]=n.ocp.XL[st] + n.ocp.mXL[st]*(j/n.ocp.state.pts); # lower
        end
      end
    end

    # upper state constraint
    if n.ocp.XU[st]!=false
      if !isnan(n.ocp.XU[st])
        for j in 1:n.ocp.state.pts
        n.ocp.XU_var[st,j]=n.ocp.XU[st] + n.ocp.mXU[st]*(j/n.ocp.state.pts); # upper
        end
      end
    end
  end
  return nothing
end

"""
defineTolerances!(n)
--------------------------------------------------------------------------------------\n
Author: Huckleberry Febbo, Graduate Student, University of Michigan
Date Create: 2/8/2017, Last Modified: 4/13/2018 \n
-------------------------------------------------------------------------------------\n
"""
function defineTolerances!(n::NLOpt;
                          X0_tol::Array{Float64,1}=0.05*ones(Float64,n.ocp.state.num,),
                          XF_tol::Array{Float64,1}=0.05*ones(Float64,n.ocp.state.num,))
  # TODO error checking, if the user does not pass tolerances etc.
  n.ocp.X0_tol = X0_tol
  n.ocp.XF_tol = XF_tol
  return nothing
end

"""
create_tV!(n)
# define a time vector (n.ocp.tV) for use with time varying constraints when (finalTimeDV=>true)
--------------------------------------------------------------------------------------\n
Author: Huckleberry Febbo, Graduate Student, University of Michigan
Date Create: 3/17/2017, Last Modified: 4/13/2018 \n
--------------------------------------------------------------------------------------\n
"""
function create_tV!(n::NLOpt)

  if n.s.ocp.integrationMethod==:ps
    # create mesh points, interval size = tf_var/Ni
    tm = @NLexpression(n.ocp.mdl, [idx=1:n.ocp.Ni+1], (idx-1)*n.ocp.tf/n.ocp.Ni)
    # go through each mesh interval creating time intervals; [t(i-1),t(i)] --> [-1,1]
    ts = [Array{Any}(n.ocp.Nck[int]+1,) for int in 1:n.ocp.Ni]
    for int in 1:n.ocp.Ni
      ts[int][1:end-1] = @NLexpression(n.ocp.mdl,[j=1:n.ocp.Nck[int]], (tm[int+1]-tm[int])/2*n.ocp.tau[int][j] +  (tm[int+1]+tm[int])/2);
      ts[int][end] = @NLexpression(n.ocp.mdl, n.ocp.tf/n.ocp.Ni*int) # append +1 at end of each interval
    end
    tt1 = [idx for tempM in ts for idx = tempM[1:end-1]]
    tmp = [tt1;ts[end][end]]
    n.ocp.tV = @NLexpression(n.ocp.mdl,[j=1:n.ocp.state.pts], n.ocp.t0 + tmp[j])
  else
    # create vector with the design variable in it
    t = Array{Any}(n.ocp.N+1,1)
    tm = @NLexpression(n.ocp.mdl, [idx=1:n.ocp.N], n.ocp.tf/n.ocp.N*idx)
    tmp = [0;tm]
    n.ocp.tV = @NLexpression(n.ocp.mdl,[j=1:n.ocp.state.pts], n.ocp.t0 + tmp[j])
  end
  return nothing
end

"""
obj=integrate!(n,:(u1))
--------------------------------------------------------------------------------------\n
Author: Huckleberry Febbo, Graduate Student, University of Michigan
Date Create: 1/2/2017, Last Modified: 4/13/2018 \n
--------------------------------------------------------------------------------------\n
"""
function integrate!(n::NLOpt,V::Expr)
  if n.s.ocp.integrationMethod==:ps
    integral_expr = [Array{Any}(n.ocp.Nck[int]) for int in 1:n.ocp.Ni]
    for int in 1:n.ocp.Ni
      x_int,u_int = intervals(n,int,n.r.ocp.x,n.r.ocp.u)
      L = size(x_int)[1]-1
      integral_expr[int][:] = NLExpr(n,V,x_int,u_int,L)
    end
    @NLexpression(n.ocp.mdl, temp[int=1:n.ocp.Ni], (n.ocp.tf-n.ocp.t0)/2*sum(n.ocp.ws[int][j]*integral_expr[int][j] for j = 1:n.ocp.Nck[int]) )
    expression = @NLexpression(n.ocp.mdl, sum(temp[int] for int = 1:n.ocp.Ni))
  elseif n.s.ocp.integrationMethod==:tm
    L = size(n.r.ocp.x)[1]
    temp = NLExpr(n,V,n.r.ocp.x,n.r.ocp.u,L);
    if n.s.ocp.integrationScheme==:bkwEuler
      # NOTE integration this way does not penalize the first control
      expression = @NLexpression(n.ocp.mdl, sum(temp[j+1]*n.ocp.tf/n.ocp.N for j = 1:n.ocp.N) )
      #expression = @NLexpression(n.ocp.mdl, sum(temp[j]*n.ocp.tf/n.ocp.N for j = 1:n.ocp.N) )
    elseif n.s.ocp.integrationScheme==:trapezoidal
      expression = @NLexpression(n.ocp.mdl, sum(0.5*(temp[j]+temp[j+1])*n.ocp.tf/n.ocp.N for j = 1:n.ocp.N) )
    else
      error("\n $(n.s.ocp.integrationScheme) not defined in integrationSchemes\n")
    end
  else
    error("\n $(n.s.ocp.integrationMethod) not defined in integrationMethods \n")
  end
  return expression
end

########################################################################################
# constraint data functions
########################################################################################
"""
# funtionality to save constraint data
--------------------------------------------------------------------------------------\n
Author: Huckleberry Febbo, Graduate Student, University of Michigan
Date Create: 2/7/2017, Last Modified: 4/13/2018 \n
--------------------------------------------------------------------------------------\n
"""
function initConstraint!(n)
  if n.r.ocp.constraint==nothing
    n.r.ocp.constraint = Constraint()
  end
  return nothing
end

function newConstraint!(n,handle,name::Symbol)
  initConstraint!(n)
  n.r.ocp.constraint::Constraint = n.r.ocp.constraint
  push!(n.r.ocp.constraint.handle,handle)
  push!(n.r.ocp.constraint.name,name)
  return nothing
end

"""
maxDualInf = evalMaxDualInf(n)
# funtionality to evaluate maximum dual infeasibility of problem
--------------------------------------------------------------------------------------\n
Author: Huckleberry Febbo, Graduate Student, University of Michigan
Date Create: 2/13/2017, Last Modified: 4/13/2018 \n
--------------------------------------------------------------------------------------\n
"""
function evalMaxDualInf(n::NLOpt)
  num=length(n.r.ocp.constraint.handle); dual_con_temp=zeros(num);
  for i = 1:length(n.r.ocp.constraint.handle)
    if !isempty(n.r.ocp.constraint.handle[i])
      if n.r.ocp.constraint.name[i]==:dyn_con  # state constraits
        temp1=zeros(n.ocp.state.num);
        for st in 1:n.ocp.state.num
          if n.s.ocp.integrationMethod==:ps
            temp=[getdual(n.r.ocp.constraint.handle[i][int][:,st]) for int in 1:n.ocp.Ni];
            vals=[idx for tempM in temp for idx=tempM];
            temp1[st]=maximum(vals);
          else
            temp1[st] = maximum(getdual(n.r.ocp.constraint.handle[i][:,st]));
          end
        end
        dual_con_temp[i]=maximum(temp1);
      else
        S=JuMP.size(n.r.ocp.constraint.handle[i])
        if length(S)==1
          dual_con_temp[i]=maximum(getdual(n.r.ocp.constraint.handle[i][:]));
        elseif length(S)==2
          temp1=zeros(S[1]);
          for idx in 1:S[1]
            temp1[idx] = maximum(getdual(n.r.ocp.constraint.handle[i][idx,:]));
          end
          dual_con_temp[i]=maximum(temp1);
        end
      end
    end
  end
  return maximum(maximum(maximum(dual_con_temp)))
end

########################################################################################
# state data functions
########################################################################################
"""
# initialize states
--------------------------------------------------------------------------------------\n
Author: Huckleberry Febbo, Graduate Student, University of Michigan
Date Create: 2/7/2017, Last Modified: 4/13/2018 \n
--------------------------------------------------------------------------------------\n
"""
function initState(numStates)
  s = State()
  s.num = numStates
  s.name = [Symbol("x$i") for i in 1:numStates]
  s.description = [String("x$i") for i in 1:numStates]
  return s
end

function states!(n::NLOpt,names;descriptions=[])
  if !n.f.ocp.defined
    error("\n call define() before calling states!() \n")
  end
  if length(names)!=n.ocp.state.num
    error("\n Check size of names \n")
  end
  if !isempty(descriptions) && length(descriptions)!=n.ocp.state.num
    error("\n Check size of descriptions \n")
  end

  for i in 1:n.ocp.state.num
    if names[i]==:xxx
      error("xxx is OFF limits for a state name; please choose something else. \n")
    end
  end

   n.ocp.state.name = names
   if !isempty(descriptions)
     n.ocp.state.description = descriptions
   end
  return nothing
end

########################################################################################
# control data functions
########################################################################################
"""
# initialize control
--------------------------------------------------------------------------------------\n
Author: Huckleberry Febbo, Graduate Student, University of Michigan
Date Create: 2/7/2017, Last Modified: 4/13/2018 \n
--------------------------------------------------------------------------------------\n
"""
function initControl(numControls)
  c = Control()
  c.num = numControls
  c.name = [Symbol("u$i") for i in 1:numControls]
  c.description = [String("u$i") for i in 1:numControls]
  return c
end

function controls!(n::NLOpt,names;descriptions=[])
  if !n.f.ocp.defined
    error("\n call define() before calling controls!() \n")
  end
  if length(names)!=n.ocp.control.num
    error("\n Check sizes of names \n")
  end
  if !isempty(descriptions) && length(descriptions)!=n.ocp.control.num
    error("\n Check size of descriptions \n")
  end

  for i in 1:n.ocp.control.num
    if names[i]==:uuu
      error("uuu is OFF limits for a control name; please choose something else. \n")
    end
   end

  n.ocp.control.name = names
  if !isempty(descriptions)
    n.ocp.control.description = descriptions
  end
  return nothing
end

########################################################################################
# data related functions
########################################################################################

"""
------------------------------------------------------------------\n
Author: Huckleberry Febbo, Graduate Student, University of Michigan
Date Create: 3/26/2017, Last Modified: 2/6/2018 \n
--------------------------------------------------------------------------------------\n
"""
function savePlantData!(n)
  if isempty(n.r.ip.plant)
    warn("No plant data to save.\n
          Make sure that n.s.save = true ")
    return nothing
  end
  cd(n.r.resultsDir)
    CSV.write("plant.csv", n.r.ip.plant; quotechar = ' ');
  cd(n.r.mainDir)
  return nothing
end

"""
------------------------------------------------------------------\n
Author: Huckleberry Febbo, Graduate Student, University of Michigan
Date Create: 6/7/2018, Last Modified: 6/7/2018 \n
--------------------------------------------------------------------------------------\n
"""
function saveOptData(n)
  if isempty(n.r.ocp.dfsOpt)
    warn("No optimization data to save.\n
          Make sure that n.s.save = true ")
    return nothing
  end
  cd(n.r.resultsDir)
    CSV.write("opt.csv", n.r.ocp.dfsOpt; quotechar = ' ');
  cd(n.r.mainDir)
  return nothing
end

"""
------------------------------------------------------------------\n
Author: Huckleberry Febbo, Graduate Student, University of Michigan
Date Create: 10/3/2017, Last Modified: 11/11/2017 \n
--------------------------------------------------------------------------------------\n
"""
function saveData(n)
  # all polynomial data
  dfs=DataFrame();

  dfs[:t] = n.r.ocp.tpts
  for st in 1:n.ocp.state.num # state
    dfs[n.ocp.state.name[st]]=n.r.ocp.Xpts[:,st];
  end

  for ctr in 1:n.ocp.control.num # control
    dfs[n.ocp.control.name[ctr]]=n.r.ocp.Upts[:,ctr];
  end

  if n.s.ocp.evalCostates && n.s.ocp.integrationMethod == :ps && n.s.ocp.evalConstraints
    for st in 1:n.ocp.state.num # state
      dfs[Symbol(n.ocp.state.name[st],:cs)]=n.r.ocp.CSpts[:,st];
    end
  end

  cd(n.r.resultsDir)
    CSV.write("st_ctr.csv",n.r.ocp.dfs[end]; quotechar = ' '); # assuming only want the last one is needed
    CSV.write("st_ctr_poly.csv",dfs; quotechar = ' ');
  cd(n.r.mainDir)
  return nothing
end
"""
------------------------------------------------------------------\n
Author: Huckleberry Febbo, Graduate Student, University of Michigan
Date Create: 9/14/2017, Last Modified: 9/19/2017 \n
--------------------------------------------------------------------------------------\n
"""
function saveBenchMarkData!(n)
  first=2
  dfs=DataFrame();
  temp = [n.r.ocp.dfs[jj][:t][1:end-1,:] for jj in first:length(n.r.ocp.dfs)]; # time
  U=[idx for tempM in temp for idx=tempM]; dfs[:t]=U;

  for st in 1:n.ocp.state.num # state
    temp = [n.r.ocp.dfs[jj][n.ocp.state.name[st]][1:end-1,:] for jj in first:length(n.r.ocp.dfs)];
    U=[idx for tempM in temp for idx=tempM];
    dfs[n.ocp.state.name[st]]=U;
  end

  for ctr in 1:n.ocp.control.num # control
    temp = [n.r.ocp.dfs[jj][n.ocp.control.name[ctr]][1:end-1,:] for jj in first:length(n.r.ocp.dfs)];
    U=[idx for tempM in temp for idx=tempM];
    dfs[n.ocp.control.name[ctr]]=U;
  end

  # save optimization times
  temp = [n.r.ocp.dfsOpt[jj][n.ocp.control.name[ctr]][1:end-1,:] for jj in first:length(n.r.ocp.dfs)];

  cd(n.r.resultsDir)
    CSV.write("bench_data.csv",dfs; quotechar = ' ');
  cd(n.r.mainDir)
  return nothing
end

"""
# maximum(x->maximum(x[:A]), dfs) -> consider
# maximum(x->maximum(filter(y->y !== nothing, x[:A])), dfs)
# minimum(x->maximum(filter(y->y !== nothing, x[:t])), dfs)
--------------------------------------------------------------------------------------\n
Author: Huckleberry Febbo, Graduate Student, University of Michigan
Date Create: 12/8/2016, Last Modified: 3/29/2017 \n
--------------------------------------------------------------------------------------\n
"""
function minDF(dfs,varb)
  k=length(dfs); tmp=1e10*ones(k,1);
  for i in 1:k
    if dfs[i]!=nothing
      tmp[i]=minimum(dfs[i][varb])
    end
  end
  minimum(tmp)  # find the minimum
end

"""
--------------------------------------------------------------------------------------\n
Author: Huckleberry Febbo, Graduate Student, University of Michigan
Date Create: 12/8/2016, Last Modified: 3/29/2017 \n
--------------------------------------------------------------------------------------\n
"""
function maxDF(dfs,varb)
  k=length(dfs); tmp=1e-10*ones(k,1);
  for i in 1:k
    if dfs[i]!=nothing
      tmp[i]=maximum(dfs[i][varb])
    end
  end
  maximum(tmp)  # find the minimum
end

function try_import(name::Symbol)
 try
      @eval using $name
      return true
  catch e
      return false
  end
end

"""
# t must be the time vector
# V is any vector that you are interpolating
--------------------------------------------------------------------------------------\n
Author: Huckleberry Febbo, Graduate Student, University of Michigan
Date Create: 2/06/2017, Last Modified: 2/06/2018 \n
--------------------------------------------------------------------------------------\n
"""
function linearSpline(t::Vector,V::Vector)
  if !isequal(length(t),length(V))
    error("!isequal(length(t),length(V))")
  end
  # remove any repeating values
  M = Array{Bool}(length(t)); M[1:length(t)] = false;
  for i in 1:length(t)-1
      if t[i]==t[i+1]
          M[i]=true
      else
          M[i]=false
      end
  end
  rm_idx = find(M)
  if (length(t)==length(rm_idx))
    error("No time has elapsed and there will be an issue with interpolation. \n
          Cannot simulate the vehicle.")
  end

  # initialize vetors
  t_new = Array{Float64}(length(t)-length(rm_idx))
  V_new = Array{Float64}(length(t)-length(rm_idx))
  q = 1
  for i in 1:length(V) #TODO put an error message here if V and t are different sizes
      if !M[i]
          t_new[q] = t[i]
          V_new[q] = V[i]
          q = q + 1
      end
  end

  # make interpolant using Dierckx.jl
  #Spline1D(t_new,V_new,k=1)    # linear spline

  # make interpolant using Interpolations.jl
  knots = (t_new,)
  it = interpolate(knots,V_new,Gridded(Linear()))
  et = extrapolate(it, Flat())
  return et
end


function SubLine(arg0::Line, arg1::IntervalsSet)
    return SubLine((Line, IntervalsSet), arg0, arg1)
end

function SubLine(arg0::Segment)
    return SubLine((Segment,), arg0)
end

function SubLine(arg0::Vector3D, arg1::Vector3D, arg2::jdouble)
    return SubLine((Vector3D, Vector3D, jdouble), arg0, arg1, arg2)
end

function get_segments(obj::SubLine)
    return jcall(obj, "getSegments", List, ())
end

function intersection(obj::SubLine, arg0::SubLine, arg1::jboolean)
    return jcall(obj, "intersection", Vector3D, (SubLine, jboolean), arg0, arg1)
end


abstract type AbstractServiceFormulation end

mutable struct ServiceModel{D<:PSY.Service,
                            B<:AbstractServiceFormulation}
    service::Type{D}
    formulation::Type{B}
end

function construct_service!(canonical::Canonical,
                           service_model::ServiceModel,
                           system_formulation::Type{S},
                           sys::PSY.System;
                           kwargs...) where {S<:PM.AbstractPowerModel}

    construct_service!(canonical,
                      service_model.service,
                      service_model.formulation,
                      system_formulation,
                      sys;
                      kwargs...)

    return

end


function layout_graph_vega(case::Dict{String,Any}, spring_const;
    node_types::Array{String,1} = ["bus","gen","storage"],
    edge_types::Array{String,1} = ["switch","branch","dcline","transformer"],
    )

    data = deepcopy(case)
    node_comp_map = Dict()
    for node_type in node_types
        temp_node = get(data, node_type, Dict())
        temp_map = Dict(string(comp["source_id"][1],"_",comp["source_id"][2]) => comp  for (comp_id, comp) in temp_node)
        merge!(node_comp_map,temp_map)
    end

    edge_comp_map = Dict()
    for edge_type in edge_types
        temp_edge = get(data, edge_type, Dict())
        for (id,edge) in temp_edge
            edge["src"] = "bus_$(edge["f_bus"])"
            edge["dst"] = "bus_$(edge["t_bus"])"
        end
        temp_map = Dict(string(comp["source_id"][1],"_",comp["source_id"][2]) => comp for (comp_id, comp) in temp_edge)
        merge!(edge_comp_map,temp_map)
    end
    # connectors
    connector_map = Dict()
    for node_type in node_types
        if node_type!="bus"
            nodes = get(data, node_type, Dict())
            for (id,node) in nodes
                temp_connector = Dict()
                temp_connector["src"] = "$(node_type)_$(id)"
                temp_connector["dst"] = "bus_$(node["$(node_type)_bus"])"

                temp_map = Dict(string("connector_",(length(connector_map)+1)) => temp_connector)
                merge!(connector_map,temp_map)
            end
        end
    end

    # find fixed positions of nodes
    pos = Dict()
    for (id,node) in node_comp_map
        if haskey(node, "xcoord_1") && haskey(node, "ycoord_1")
            pos[id] = [node["xcoord_1"], node["ycoord_1"]]
        else
            pos[id] = missing
        end
    end
    fixed = [node for (node, p) in pos if !ismissing(p)]

    G = nx.Graph()
    for (id,node) in node_comp_map
        G.add_node(id)
    end
    for (id,edge) in edge_comp_map
        G.add_edge(edge["src"], edge["dst"], weight=1.0)
    end
    for (id,edge) in connector_map
        G.add_edge(edge["src"], edge["dst"], weight=1.0)
    end

    if isempty(fixed)
        positions = nx.kamada_kawai_layout(G, dist=nothing, pos=nothing, weight="weight", scale=1.0, center=nothing, dim=2)
    else
        avg_x, avg_y = mean(hcat(skipmissing([v for v in values(pos)])...), dims=2)
        std_x, std_y = std(hcat(skipmissing([v for v in values(pos)])...), dims=2)
        for (v, p) in pos
            if ismissing(p)
                #get parent bus coord, or center of figure
                comp_type, comp_id = split(v, "_")
                x1 = get(get(data["bus"],string(get(case[comp_type][comp_id],"$(comp_type)_bus", NaN)),Dict()),"xcoord_1", avg_x)
                y1 = get(get(data["bus"],string(get(case[comp_type][comp_id],"$(comp_type)_bus", NaN)),Dict()),"ycoord_1", avg_x)
                pos[v] = [x1,y1] + [std_x*(rand()-0.5), std_y*(rand()-0.5)]*300
            end
        end
        # spring_const = 1e-2
        k=spring_const*minimum(std([p for p in values(pos)]))
        positions = nx.spring_layout(G; pos,  fixed, k,  iterations=100)
        # positions = pos
    end

    # Set Node Positions
    for (node, (x, y)) in positions
        (comp_type,comp_id) = split(node, "_")
        data[comp_type][comp_id]["xcoord_1"] = x
        data[comp_type][comp_id]["ycoord_1"] = y
    end
    # Set Edge positions
    for (edge, val) in (edge_comp_map)
        (x,y) = positions[val["src"]]
        (x2,y2) = positions[val["dst"]]
        (comp_type,comp_id) = split(edge, "_")
        data[comp_type][comp_id]["xcoord_1"] = x
        data[comp_type][comp_id]["ycoord_1"] = y
        data[comp_type][comp_id]["xcoord_2"] = x2
        data[comp_type][comp_id]["ycoord_2"] = y2
    end

    # Create connector dictionary
    data["connector"] = Dict()
    for (edge, con) in connector_map
        _,id = split(edge, "_")
        data["connector"][id]=  Dict(
            "src" => con["src"],
            "dst" => con["dst"],
            "xcoord_1" => 0.0,
            "ycoord_1" => 0.0,
            "xcoord_2" => 0.0,
            "ycoord_2" => 0.0,
        )
    end
    # Set Connector positions
    for (edge, val) in (connector_map)
        (x,y) = positions[val["src"]]
        (x2,y2) = positions[val["dst"]]
        (comp_type,comp_id) = split(edge, "_")
        data[comp_type][comp_id]["xcoord_1"] = x
        data[comp_type][comp_id]["ycoord_1"] = y
        data[comp_type][comp_id]["xcoord_2"] = x2
        data[comp_type][comp_id]["ycoord_2"] = y2
    end

    return data
end


function differentiate(::Type{Differentiation}, arg0::ParameterFunction, arg1::jint, arg2::jdouble)
    return jcall(Differentiation, "differentiate", ParameterFunction, (ParameterFunction, jint, jdouble), arg0, arg1, arg2)
end

function differentiate(::Type{Differentiation}, arg0::StateFunction, arg1::jint, arg2::AttitudeProvider, arg3::OrbitType, arg4::PositionAngle, arg5::jdouble, arg6::jint)
    return jcall(Differentiation, "differentiate", StateJacobian, (StateFunction, jint, AttitudeProvider, OrbitType, PositionAngle, jdouble, jint), arg0, arg1, arg2, arg3, arg4, arg5, arg6)
end