struct Data
    mag::AbstractArray
    phase::AbstractArray
    header
    TEs::AbstractVector
    Data(mag, phase, header, TEs=1:size(mag,4)) = new(mag, phase, header, vec(TEs))
end

struct Options
    mag_combine
    mag_sens::Union{AbstractArray, Nothing}
    mag_softplus
    phase_unwrap::Symbol
    phase_hp_σ::AbstractArray
    phase_scaling_type::Symbol
    phase_scaling_strength::Real
    writesteps::Union{AbstractString, Nothing}
end

function Options(; mag_combine=:SNR, mag_sens=nothing, mag_softplus=true, phase_unwrap=:laplacian, phase_hp_σ=[4,4,0], phase_scaling_type=:tanh, phase_scaling_strength=4, writesteps=nothing)
    Options(mag_combine, mag_sens, mag_softplus, phase_unwrap, phase_hp_σ, phase_scaling_type, phase_scaling_strength, writesteps)
end

function saveconfiguration(options)
    open(joinpath(options.writesteps, "settings_swi.txt"), "w") do io
        for fname in fieldnames(typeof(options))
            val = getfield(options, fname)
            if !(val isa AbstractArray && !(val isa Vector))
                println(io, "$fname: " * string(val))
            end
        end
    end
end


function α(x) #clearing error function. To be replaced with solution to issue 1
      sum(x)
  end


using NaturalSort
using Test

# write your own tests here
STRS = [
    "",
    "a#",
    "a00",
    "a0a",
    "a1",
    "a1b1",
    "a1b2",
    "a01b3",
    "a1b11",
    "a2",
    "a11",
    "a11b1",
    "a11b2",
    "a11b11",
    "a12",
    "a12b1",
    "a22",
    "abc",
    "b"
]

for i = 1:length(STRS)
    for j = 1:i-1
        @test natural(STRS[j], STRS[i])
    end
    for j = i:length(STRS)
        @test !natural(STRS[j], STRS[i])
    end
end

EQ_STRS = [
    ("a01", "a1"),
    ("a01b2", "a1b2"),
    ("a0", "a00")
]
for (str, eqto) in EQ_STRS
    for i = 1:findfirst(x->x == eqto, STRS)-1
        @test natural(STRS[i], str)
        @test !natural(str, STRS[i])
    end
    for i = findfirst(x->x == eqto, STRS)+1:length(STRS)
        @test !natural(STRS[i], str)
        @test natural(str, STRS[i])
    end
end


export PhylogeneticTree, phylogenetic_tree, calculate_tree, newick_string, leaves, isleaf,
    set_property!, get_property, annotate!

struct PhylogeneticTree
    graph::MetaDiGraph
end

graph(tree::PhylogeneticTree) = getfield(tree, :graph)

function phylogenetic_tree(newick_string::String)
    graph = MetaDiGraph()
    add_vertex!(graph)
    set_prop!(graph, 1, :name, "root")

    components = newick_components(newick_string)
    for component in components
        add_to_graph!(graph, 1, component)
    end

    return PhylogeneticTree(graph)
end

function phylogenetic_tree(data::AbstractHLAData; model::String = "LG+G+I",
        verbose::Bool = false)
    !ismissing(data.tree) && return data.tree

    temp_prefix = tempname()
    temp_fasta = temp_prefix * ".fasta"

    write_numbered_fasta(data, temp_fasta)

    if verbose
        Base.run(`$(raxml_ng()) --msa $temp_fasta --model $model --threads auto
            --workers auto`)
    else
        @suppress Base.run(`$(raxml_ng()) --msa $temp_fasta --model $model --threads auto
        --workers auto`)
    end
    tree = phylogenetic_tree(readline(temp_fasta * ".raxml.bestTree"))

    raxml_files = ["bestModel", "bestTree", "log", "mlTrees", "rba", "startTree"]
    rm(temp_fasta)
    rm.(temp_fasta .* ".raxml." .* raxml_files)

    return tree
end

function write_numbered_fasta(data::AbstractHLAData, filepath::String)
    records = Escape.records(data)

    write_numbered_fasta(records, filepath)
end

function write_numbered_fasta(records::Vector{FASTX.FASTA.Record}, filepath::String)
    open(FASTA.Writer, filepath) do writer
        for (i, record) in enumerate(records)
            identifier = string(i)
            sequence = FASTA.sequence(record)
            new_record = FASTA.Record(identifier, sequence)

            write(writer, new_record)
        end
    end
end

function add_to_graph!(graph::AbstractMetaGraph, vertex::Int, s::String)
    name = extract_name(s)
    length = extract_length(s)

    add_vertex!(graph)
    set_prop!(graph, nv(graph), :name, name)

    add_edge!(graph, vertex, nv(graph))
    set_prop!(graph, Edge(vertex, nv(graph)), :branch_length, length)

    current_vertex = nv(graph)
    if occursin(')', s)
        components = newick_components(s)
        for c in components
            add_to_graph!(graph, current_vertex, c)
        end
    end
end

function extract_name(s::String)
    substring = occursin(')', s) ? s[findlast(isequal(')'), s) + 1 : end] : s
    substring != "" || return missing

    name = String(split(substring, ":")[1])
    return name != "" ? name : missing
end

function extract_length(s::String)
    substring = occursin(')', s) ? s[findlast(isequal(')'), s) + 1 : end] : s
    occursin(':', substring) || return missing

    length = parse(Float64, split(substring, ":")[end])
end

function newick_components(s::String)
    components = Vector{String}()
    component = ""

    start = occursin('(', s) ? findfirst(isequal('('), s) + 1 : 1
    stop = occursin(')', s) ? findlast(isequal(')'), s) - 1 : length(s)

    for c in s[start:stop]
        if c == ',' && count(x -> x == '(', component) == count(x -> x == ')', component)
            push!(components, strip(component))
            component = ""
        else
            component *= c
        end
    end

    push!(components, strip(component))
    return components
end

function newick_string(tree::PhylogeneticTree)
    graph = Escape.graph(tree)
    n = Graphs.neighbors(graph, 1)
    length(n) == 0 && return "();"
    length(n) == 1 && return "(" * newick_string(graph, n[1]) * ");"

    newick = "(" * newick_string(graph, Graphs.neighbors(graph, 1)[1])
    for n in Graphs.neighbors(graph, 1)[2:end]
        newick *= ","
        newick *= newick_string(graph, n)
    end
    newick *= ");"

    return newick
end

function newick_string(graph::AbstractMetaGraph, v::Int)
    name = get_prop(graph, v, :name)
    name = ismissing(name) ? "" : name
    branch_length = get_prop(graph, inneighbors(graph, v)[1], v, :branch_length)
    branch_length = ismissing(branch_length) ? "" : string(branch_length)

    if length(Graphs.neighbors(graph, v)) == 0
        return string(name, branch_length != "" ? ":" : "", branch_length)
    else
        newick = "(" * newick_string(graph, Graphs.neighbors(graph, v)[1])
        for n in Graphs.neighbors(graph, v)[2:end]
            newick *= ","
            newick *= newick_string(graph, n)
        end
        newick *= ")" * name
        newick = !ismissing(length) ? newick * ":$branch_length" : newick
    end

    return newick
end

function isleaf(v::Int, tree::PhylogeneticTree)
    v ∈ leaves(tree) && return true
    return false
end

function leaves(tree::PhylogeneticTree)
    graph = Escape.graph(tree)
    leaves = map(x -> x[1], attracting_components(graph))

    return leaves
end

function get_property(tree::PhylogeneticTree, v::Int, property::Symbol)
    graph = Escape.graph(tree)
    has_prop(graph, v, property) || return missing

    return get_prop(graph, v, property)
end

function set_property!(tree::PhylogeneticTree, v::Int, property::Symbol, value)
    graph = Escape.graph(tree)

    set_prop!(graph, v, property, value)
end

function annotate!(tree::PhylogeneticTree, data::AbstractHLAData, replacement::Replacement)
    name(data) == protein(replacement) || error("Replacement does not match HLA data.")
    records = Escape.records(data)

    for (i, record) in enumerate(records)
        sequence = FASTA.sequence(record)
        symbol = Char(sequence[position(replacement)])

        v = filter(x -> get_property(tree, x, :name) == string(i), leaves(tree))[1]

        if !negated(replacement)
            if symbol == Escape.replacement(replacement)
                set_property!(tree, v, :state, "1")
            elseif symbol in ('-', 'X')
                set_property!(tree, v, :state, missing)
            else
                set_property!(tree, v, :state, "0")
            end
        elseif negated(replacement)
            if symbol == Escape.replacement(replacement)
                set_property!(tree, v, :state, "0")
            elseif symbol in ('-', 'X')
                set_property!(tree, v, :state, missing)
            else
                set_property!(tree, v, :state, "1")
            end
        end
    end
end

function matching(tree::PhylogeneticTree, data::AbstractHLAData)
    records = Escape.records(data)
    matching = Escape.matching(tree, records)

    return matching
end

function matching(tree::PhylogeneticTree, records::Vector{FASTX.FASTA.Record})
    n_leaves = length(leaves(tree))

    msg = string(
        "Number of leaves ($n_leaves) does not match number of sequences ",
        "($(length(records)))"
    )
    n_leaves == length(records) || return DimensionMismatch(msg)

    msg = """:name property of leaves must be "1", "2", ..., "$n_leaves"."""
    leaf_names = Set(map(x -> get_property(tree, x, :name), leaves(tree)))
    leaf_names == Set(string.(1:n_leaves)) || return ErrorException(msg)

    return true
end

function states(tree::PhylogeneticTree)
    states = Vector{String}(undef, length(leaves(tree)))

    for leaf in leaves(tree)
        state = get_property(tree, leaf, :state)
        name = get_property(tree, leaf, :name)

        states[parse(Int, name)] = ismissing(state) ? "-" : state
    end

    return states
end


using Criterion

function test1_svd(k::Integer)
    for _ in k
        svd(zeros(20, 20))
    end
end

function test2_svd(k::Integer)
    for _ in k
        svd(ones(20, 20))
    end
end

function test3_svd(k::Integer)
    for _ in k
        svd(ones(10, 10))
    end
end

env = Criterion.measure_environment()

bench1 = Criterion.Benchmark("test1_svd", 1, test1_svd)
times, sample_analysis, outliers = Criterion.run_analyze_one(env, "test1_svd", bench1)
r1 = Criterion.Report(1, "test1/svd", times, sample_analysis, outliers)

bench2 = Criterion.Benchmark("test2_svd", 2, test2_svd)
times, sample_analysis, outliers = Criterion.run_analyze_one(env, "test2_svd", bench2)
r2 = Criterion.Report(2, "test2/svd", times, sample_analysis, outliers)

bench3 = Criterion.Benchmark("test3_svd", 2, test3_svd)
times, sample_analysis, outliers = Criterion.run_analyze_one(env, "test3_svd", bench3)
r3 = Criterion.Report(3, "test3/svd", times, sample_analysis, outliers)

Criterion.report([r1, r2, r3])

#using Winston
#p = FramedPlot()
#setattr(p, "title", "Benchmark Times (sec)")
#setattr(p, "ylabel", "counts")
#setattr(p, "xlabel", "times (sec)")
#add(p, Histogram(hist(times, 100)...))
#add(p, LineX(sa.mean.point, "color", "red"))
#add(p, LineX(sa.mean.lbound, "type", "dot"))
#add(p, LineX(sa.mean.ubound, "type", "dot"))
#Winston.display(p)


using Test

using POMDPs

mutable struct X <: POMDP{Float64,Bool,Int} end
abstract type Z <: POMDP{Float64,Bool,Int} end
mutable struct Y <: Z end

@test_throws ErrorException statetype(Int)
@test_throws ErrorException actiontype(Int)
@test_throws ErrorException obstype(Int)

@test statetype(X) == Float64
@test statetype(Y) == Float64
@test actiontype(X) == Bool
@test actiontype(Y) == Bool
@test obstype(X) == Int
@test obstype(Y) == Int


@noinline function dostuff(A; B = I)
    StrideArrays.@avx for i in eachindex(A)
        A[i] += 1
    end
    C = A * B
    s = zero(eltype(C))
    StrideArrays.@avx for i in eachindex(C)
        s += C[i]
    end
    s
end
function gc_preserve_test()
    A = @StrideArray rand(8, $(2<<2));
    B = @StrideArray rand($(1<<3), 8);
    @gc_preserve dostuff(A, B = B)
end

@testset "Miscellaneous" begin
    x = @StrideArray rand(127);
    @test maximum(abs, x) == maximum(abs, Array(x))
    y = @StrideArray rand(3);
    @test maximum(abs, y) == maximum(abs, Array(y))
    @test iszero(@allocated gc_preserve_test())

    A = @StrideArray rand($(5<<1),$(1<<3));
    @test StrideArrays.size(A) === (StaticInt(10), StaticInt(8))
    A_u = view(A, StaticInt(1):StaticInt(6), :)
    A_l = view(A, StaticInt(7):StaticInt(10), :)
    @test A == @inferred(vcat(A_u, A_l))
    @test StrideArrays.size(A) === StrideArrays.size(vcat(A_u, A_l))

    # On 1.5 tests fail if you don't do this first.
    @test pointer(view(A, 1:StaticInt(6), :)) == pointer(A)
    A_u = view(A, 1:StaticInt(6), :);
    A_l = view(A, StaticInt(7):StaticInt(10), :);
    @test A == @inferred(vcat(A_u, A_l))
    @test StrideArrays.size(A) == StrideArrays.size(vcat(A_u, A_l))
    @test StrideArrays.size(A) !== StrideArrays.size(vcat(A_u, A_l))

    Aa = Array(A)
    @test sum(A) ≈ sum(Aa)
    @test maximum(A) == maximum(Aa) == maximum(abs, @. -A)
    @test mapreduce(abs2, +, A) ≈ mapreduce(abs2, +, Aa)

    @test A[1,1] == A[1] == vec(A)[1]
    A[1,1] = 3
    @test A[1,1] == A[1] == vec(A)[1] == 3
    A[2] = 4
    @test A[2,1] == A[2] == vec(A)[2] == 4
    vec(A)[3] = 5
    @test A[3,1] == A[3] == vec(A)[3] == 5

    @test ArrayInterface.stride_rank(StrideArrays.similar_layout(A')) === ArrayInterface.stride_rank(StrideArrays.similar_layout(Aa')) === (StaticInt(2),StaticInt(1))
    @test ArrayInterface.contiguous_axis(StrideArrays.similar_layout(A')) === ArrayInterface.contiguous_axis(StrideArrays.similar_layout(Aa')) === StrideArrays.StaticInt{2}()

    let B = @StrideArray rand(10,10)
        @test exp(B) ≈ exp(Array(B))
    end
end


"""
    FerriteViz.update!(plotter::MakiePlotter, u::Vector)
Updates the Observable `plotter.u` and thereby, triggers the plot to update.
"""
function update!(plotter::MakiePlotter, u::Vector)
    @assert length(plotter.u[]) == length(u)
    plotter.u[] .= u
    Makie.notify(plotter.u)
end

function Makie.convert_arguments(P::Type{<:Makie.Mesh}, plotter::MakiePlotter)
    return Makie.convert_arguments(P,plotter.physical_coords,plotter.triangles)
end

"""
    solutionplot(plotter::MakiePlotter; kwargs...)
    solutionplot(dh::AbstractDofHandler, u::Vector; kwargs...)
    solutionplot!(plotter::MakiePlotter; kwargs...)
    solutionplot!(dh::AbstractDofHandler, u::Vector; kwargs...)
Solutionplot produces the classical contour plot onto the finite element mesh. Most important
keyword arguments are:

- `field::Symbol=:default` representing the field which gets plotted, defaults to the first field in the `dh`.
- `deformation_field::Symbol=:default` field that transforms the mesh by the given deformation, defaults to no deformation
- `process::Function=postprocess` function to construct nodal scalar values from a vector valued problem
- `colormap::Symbol=:cividis`
- `deformation_scale=1.0`
- `shading=false`
- `scale_plot=false`
- `transparent=false`
"""
@recipe(SolutionPlot) do scene
    Attributes(
    scale_plot=false,
    shading=false,
    field=:default,
    deformation_field=:default,
    process=postprocess,
    colormap=:cividis,
    colorrange=(0,1),
    transparent=false,
    deformation_scale = 1.0,
    )
end

function Makie.plot!(SP::SolutionPlot{<:Tuple{<:MakiePlotter}})
    plotter = SP[1][]
    solution = @lift($(SP[:field])===:default ? reshape(transfer_solution(plotter,$(plotter.u); field_idx=1, process=$(SP[:process])), num_vertices(plotter)) : reshape(transfer_solution(plotter,$(plotter.u); field_idx=Ferrite.find_field(plotter.dh,$(SP[:field])), process=$(SP[:process])), num_vertices(plotter)))
    u_matrix = @lift($(SP[:deformation_field])===:default ? zeros(0,3) : transfer_solution(plotter,$(plotter.u); field_idx=Ferrite.find_field(plotter.dh,$(SP[:deformation_field])), process=identity))
    coords = @lift($(SP[:deformation_field])===:default ? plotter.physical_coords : plotter.physical_coords .+ ($(SP[:deformation_scale]) .* $(u_matrix)))
    mins = @lift(minimum($solution))
    maxs = @lift(maximum($solution))
    SP[:colorrange] = @lift(isapprox($mins,$maxs) ? (0,1e-8) : ($mins,$maxs))
    return Makie.mesh!(SP, coords, plotter.triangles, color=solution, shading=SP[:shading], scale_plot=SP[:scale_plot], colormap=SP[:colormap], transparent=SP[:transparent])
end

"""
    cellplot(plotter::MakiePlotter,σ::Vector{T}; kwargs...) where T
    cellplot!(plotter::MakiePlotter,σ::Vector{T}; kwargs...) where T
`cellplot` plots constant scalar data on the cells of the finite element mesh. If `T` is not a number, the keyword argument `process`
can be passed in order to reduce the elements of `σ` to a scalar.

keyword arguments are:

- `deformation_field::Symbol=:default` field that transforms the mesh by the given deformation, defaults to no deformation
- `process::Function=identity` function to construct cell scalar values. Defaults to `identity`, i.e. scalar values.
- `colormap::Symbol=:cividis`
- `deformation_scale=1.0`
- `shading=false`
- `scale_plot=false`
- `transparent=false`
"""
@recipe(CellPlot) do scene
    Attributes(
    scale_plot=false,
    shading=false,
    deformation_field=:default,
    process=identity,
    colormap=:cividis,
    colorrange=(0,1),
    transparent=false,
    deformation_scale = 1.0,
    )
end

function Makie.plot!(CP::CellPlot{<:Tuple{<:MakiePlotter{dim},Vector}}) where dim
    plotter = CP[1][]
    qp_values = CP[2][]
    u_matrix = @lift($(CP[:deformation_field])===:default ? zeros(0,3) : transfer_solution(plotter,$(plotter.u); field_idx=Ferrite.find_field(plotter.dh,$(CP[:deformation_field])), process=identity))
    coords = @lift($(CP[:deformation_field])===:default ? plotter.physical_coords : plotter.physical_coords .+ ($(CP[:deformation_scale]) .* $(u_matrix)))
    mins = minimum(qp_values)
    maxs = maximum(qp_values)
    CP[:colorrange] = isapprox(mins,maxs) ? (0,1e-8) : (mins,maxs)
    solution =  @lift(reshape(transfer_scalar_celldata(plotter, qp_values; process=$(CP[:process])), num_vertices(plotter)))
    return Makie.mesh!(CP, coords, plotter.triangles, color=solution, shading=CP[:shading], scale_plot=CP[:scale_plot], colormap=CP[:colormap], transparent=CP[:transparent])
end

"""
    wireframe(plotter::MakiePlotter; kwargs...)
    wireframe(dh::AbstractDofHandler, u::Vector; kwargs...)
    wireframe(grid::AbstractGrid; kwargs...)
    wireframe!(plotter::MakiePlotter; kwargs...)
    wireframe!(dh::AbstractDofHandler, u::Vector; kwargs...)
    wireframe!(grid::AbstractGrid; kwargs...)
Plots the finite element mesh, optionally labels it and transforms it if a suitable `deformation_field` is given.

- `plotnodes::Bool=true` plots the nodes as circles/spheres
- `strokewidth::Int=2` how thick faces/edges are drawn
- `color::Symbol=theme(scene,:linecolor)` color of the faces/edges and nodes
- `markersize::Int=30` size of the nodes
- `deformation_field::Symbol=:default` field that transforms the mesh by the given deformation, defaults to no deformation
- `deformation_scale::Number=1.0` scaling of the deformation
- `cellsets=false` Color cells based on their cellset association. If no cellset is found for a cell, the cell is marked blue.
- `nodelables=false` global node id labels
- `nodelabelcolor=:darkblue`
- `celllabels=false` global cell id labels
- `celllabelcolor=:darkred`
- `textsize::Int=15` size of the label's text
- `visible=true`
"""
@recipe(Wireframe) do scene
    Attributes(
    plotnodes=true,
    strokewidth=2,
    color=theme(scene, :linecolor),
    markersize=30,
    deformation_field=:default,
    visible=true,
    deformation_scale=1,
    textsize=15,
    offset=(0.0,0.0),
    nodelabels=false,
    nodelabelcolor=:darkblue,
    celllabels=false,
    celllabelcolor=:darkred,
    cellsets=false,
    )
end

function Makie.plot!(WF::Wireframe{<:Tuple{<:MakiePlotter{dim}}}) where dim
    plotter = WF[1][]
    #triangle representation
    # can't use triangle representation, since we don't know by this information which edges to draw
    # further it would be in the incompressible example 2600 nodes vs 15000 in triangle representation
    # u_matrix = @lift($(WF[:deformation_field])===:default ? zeros(0,3) : transfer_solution(plotter; field_idx=Ferrite.find_field(plotter.dh,$(WF[:deformation_field])), process=identity))
    # coords = @lift($(WF[:deformation_field])===:default ? plotter.physical_coords : plotter.physical_coords .+ ($(WF[:scale]) .* $(u_matrix)))
    #original representation
    nodal_u_matrix = @lift($(WF[:deformation_field])===:default ? zeros(0,3) : dof_to_node(plotter.dh, $(WF[1][].u); field=Ferrite.find_field(plotter.dh,$(WF[:deformation_field])), process=identity))
    gridnodes = @lift($(WF[:deformation_field])===:default ? plotter.gridnodes : plotter.gridnodes .+ ($(WF[:deformation_scale]) .* $(nodal_u_matrix)))
    lines = @lift begin
        dim > 2 ? (lines = Point3f[]) : (lines = Point2f[])
        for cell in Ferrite.getcells(plotter.dh.grid)
            boundaryentities = dim < 3 ? Ferrite.faces(cell) : Ferrite.edges(cell)
            append!(lines, [$gridnodes[e,:] for boundary in boundaryentities for e in boundary])
        end
        lines
    end
    nodes = @lift($(WF[:plotnodes]) ? $(gridnodes) : zeros(Float32,0,3))
    #plot cellsets
    cellsets = plotter.dh.grid.cellsets
    cellset_to_value = Dict{String,Int}()
    for (cellsetidx,(cellsetname,cellset)) in enumerate(cellsets)
        cellset_to_value[cellsetname] = cellsetidx
    end
    cellset_u = zeros(Ferrite.getncells(plotter.dh.grid))
    for (cellidx,cell) in enumerate(Ferrite.getcells(plotter.dh.grid))
        for (cellsetname,cellsetvalue) in cellset_to_value
            if cellidx in cellsets[cellsetname]
                cellset_u[cellidx] = cellsetvalue
            end
        end
    end
    u_matrix = @lift($(WF[:deformation_field])===:default ? zeros(0,3) : transfer_solution(plotter,$(plotter.u); field_idx=Ferrite.find_field(plotter.dh,$(WF[:deformation_field])), process=identity))
    coords = @lift($(WF[:deformation_field])===:default ? plotter.physical_coords : plotter.physical_coords .+ ($(WF[:deformation_scale]) .* $(u_matrix)))
    colorrange = isempty(cellset_to_value) ? (0,1) : (0,maximum(values(cellset_to_value)))
    cellset_u =  reshape(transfer_scalar_celldata(plotter, cellset_u; process=identity), num_vertices(plotter))
    Makie.mesh!(WF, coords, plotter.triangles, color=cellset_u, shading=false, scale_plot=false, colormap=:darktest, visible=WF[:cellsets])
    #plot the nodes
    Makie.scatter!(WF,gridnodes,markersize=WF[:markersize], color=WF[:color], visible=WF[:visible])
    #set up nodelabels
    nodelabels = @lift $(WF[:nodelabels]) ? ["$i" for i in 1:size($gridnodes,1)] : [""]
    nodepositions = @lift $(WF[:nodelabels]) ? [dim < 3 ? Point2f(row) : Point3f(row) for row in eachrow($gridnodes)] : (dim < 3 ? [Point2f((0,0))] : [Point3f((0,0,0))])
    #set up celllabels
    celllabels = @lift $(WF[:celllabels]) ? ["$i" for i in 1:Ferrite.getncells(plotter.dh.grid)] : [""]
    cellpositions = @lift $(WF[:celllabels]) ? [midpoint(cell,$gridnodes) for cell in Ferrite.getcells(plotter.dh.grid)] : (dim < 3 ? [Point2f((0,0))] : [Point3f((0,0,0))])
    Makie.text!(WF,nodelabels, position=nodepositions, textsize=WF[:textsize], offset=WF[:offset],color=WF[:nodelabelcolor])
    Makie.text!(WF,celllabels, position=cellpositions, textsize=WF[:textsize], color=WF[:celllabelcolor], align=(:center,:center))
    #plot edges (3D) /faces (2D) of the mesh
    Makie.linesegments!(WF,lines,color=WF[:color], linewidth=WF[:strokewidth], visible=WF[:visible])
end

function Makie.plot!(WF::Wireframe{<:Tuple{<:Ferrite.AbstractGrid{dim}}}) where dim
    grid = WF[1][]
    coords = [Ferrite.getcoordinates(node)[i] for node in Ferrite.getnodes(grid), i in 1:dim]
    dim > 2 ? (lines = Point3f[]) : (lines = Point2f[])
    for cell in Ferrite.getcells(grid)
        boundaryentities = dim < 3 ? Ferrite.faces(cell) : Ferrite.edges(cell)
        append!(lines, [coords[e,:] for boundary in boundaryentities for e in boundary])
    end
    nodes = @lift($(WF[:plotnodes]) ? coords : zeros(Float32,0,3))
    Makie.scatter!(WF,nodes,markersize=WF[:markersize], color=WF[:color])
    nodelabels = @lift $(WF[:nodelabels]) ? ["$i" for i in 1:size(coords,1)] : [""]
    nodepositions = @lift $(WF[:nodelabels]) ? [dim < 3 ? Point2f(row) : Point3f(row) for row in eachrow(coords)] : (dim < 3 ? [Point2f((0,0))] : [Point3f((0,0,0))])
    celllabels = @lift $(WF[:celllabels]) ? ["$i" for i in 1:Ferrite.getncells(grid)] : [""]
    cellpositions = @lift $(WF[:celllabels]) ? [midpoint(cell,coords) for cell in Ferrite.getcells(grid)] : (dim < 3 ? [Point2f((0,0))] : [Point3f((0,0,0))])
    #cellsetsplot
    dh = Ferrite.DofHandler(grid)
    cellsets = grid.cellsets
    cellset_to_value = Dict{String,Int}()
    for (cellsetidx,(cellsetname,cellset)) in enumerate(cellsets)
        cellset_to_value[cellsetname] = cellsetidx
    end
    cellset_u = zeros(Ferrite.getncells(grid))
    for (cellidx,cell) in enumerate(Ferrite.getcells(grid))
        for (cellsetname,cellsetvalue) in cellset_to_value
            if cellidx in cellsets[cellsetname]
                cellset_u[cellidx] = cellsetvalue
            end
        end
    end
    plotter = MakiePlotter(dh,cellset_u)
    cellset_u =  reshape(transfer_scalar_celldata(plotter, cellset_u; process=identity), num_vertices(plotter))
    colorrange = isempty(cellset_to_value) ? (0,1) : (0,maximum(values(cellset_to_value)))
    Makie.mesh!(WF, plotter.physical_coords, plotter.triangles, color=cellset_u, shading=false, scale_plot=false, colormap=:darktest, visible=WF[:cellsets])
    Makie.text!(WF,nodelabels, position=nodepositions, textsize=WF[:textsize], offset=WF[:offset],color=WF[:nodelabelcolor])
    Makie.text!(WF,celllabels, position=cellpositions, textsize=WF[:textsize], color=WF[:celllabelcolor], align=(:center,:center))
    Makie.linesegments!(WF,lines,color=WF[:color], strokewidth=WF[:strokewidth])
end

"""
    surface(plotter::MakiePlotter; kwargs...)
    surface(dh::AbstractDofHandler, u::Vector; kwargs...)
    surface!(plotter::MakiePlotter; kwargs...)
    surface!(dh::AbstractDofHandler, u::Vector; kwargs...)
Uses the given `field` and plots the scalar values as a surface. If it's a vector valued problem, the nodal vector
values are transformed to a scalar based on `process` which defaults to the magnitude. Only availble in `dim=2`.

- `field = :default`
- `process = postprocess`
- `scale_plot = false`
- `shading = false`
- `colormap = :cividis`
"""
@recipe(Surface) do scene
    Attributes(
    field = :default,
    process = postprocess,
    scale_plot = false,
    shading = false,
    colormap = :cividis,
    )
end

function Makie.plot!(SF::Surface{<:Tuple{<:MakiePlotter{2}}})
    plotter = SF[1][]
    field = @lift($(SF[:field])===:default ? 1 : Ferrite.find_field(plotter.dh,$(SF[:field])))
    solution = @lift(reshape(transfer_solution(plotter, $(plotter.u); field_idx=$(field), process=$(SF[:process])), num_vertices(plotter)))
    points = @lift([Point3f(coord[1], coord[2], $(solution)[idx]) for (idx, coord) in enumerate(eachrow(plotter.physical_coords))])
    return Makie.mesh!(SF,points, plotter.triangles, color=solution, scale_plot=SF[:scale_plot], shading=SF[:shading], colormap=SF[:colormap])
end

"""
    arrows(plotter::MakiePlotter; kwargs...)
    arrows(dh::AbstractDofHandler, u::Vector; kwargs...)
    arrows!(plotter::MakiePlotter; kwargs...)
    arrows!(dh::AbstractDofHandler, u::Vector; kwargs...)
At every node position a arrows is drawn, where the arrow tip ends at the node. Only works in `dim >=2`. If a `color` is specified
the arrows are unicolored. Otherwise the color corresponds to the magnitude, or any other scalar value based on the `process` function.

- `arrowsize = 0.08`
- `normalize = true`
- `field = :default`
- `color = :default`
- `colormap = :cividis`
- `process=postprocess`
- `lengthscale = 1f0`
"""
@recipe(Arrows) do scene
    Attributes(
    arrowsize = 0.08,
    normalize = true,
    field = :default,
    color = :default,
    colormap = :cividis,
    process=postprocess,
    lengthscale = 1f0,
    )
end

function Makie.plot!(AR::Arrows{<:Tuple{<:MakiePlotter{dim}}}) where dim
    plotter = AR[1][]
    field = @lift($(AR[:field])===:default ? 1 : Ferrite.find_field(plotter.dh,$(AR[:field])))
    @assert Ferrite.getfielddim(plotter.dh,field[]) > 1
    solution = @lift(transfer_solution(plotter, $(plotter.u); field_idx=$(field), process=identity))
    if dim  == 2
        ps = [Point2f(i) for i in eachrow(plotter.physical_coords)]
        ns = @lift([Vec2f(i) for i in eachrow($(solution))])
        lengths = @lift($(AR[:color])===:default ? $(AR[:process]).($(ns)) : ones(length($(ns)))*$(AR[:color]))
    elseif dim  == 3
        ps = [Point3f(i) for i in eachrow(plotter.physical_coords)]
        ns = @lift([Vec3f(i) for i in eachrow($(solution))])
        lengths = @lift($(AR[:color])===:default ? $(AR[:process]).($(ns)) : ones(length($(ns)))*$(AR[:color]))
    else
        error("Arrows plots are only available in dim ≥ 2")
    end
    Makie.arrows!(AR, ps, ns, arrowsize=AR[:arrowsize], normalize=AR[:normalize], colormap=AR[:colormap], color=lengths, lengthscale=AR[:lengthscale])
end

"""
    ferriteviewer(plotter::MakiePlotter)
    ferriteviewer(plotter::MakiePlotter, u_history::Vector{Vector{T}}})
Constructs a viewer with a `solutionplot`, `Colorbar` as well as sliders,toggles and menus to change the current view.
If the second dispatch is called a timeslider is added, in order to step through a set of solutions obtained from a simulation.
"""
function ferriteviewer(plotter::MakiePlotter{dim}) where dim
    #set up figure and axis, axis either LScene in 3D or Axis if dim < 2
    fig = Figure()
    dim > 2 ? (ax = LScene(fig[1,1])) : (ax = Axis(fig[1,1]))
    #toggles and their labels for switching plot types on/off
    toggles = [Toggle(fig, active=active) for active in [true,true,false]]
    labels = [Label(fig,label) for label in ["mesh", "deformation", "labels"]]
    #setup the deformation_field as Node (Observable)
    deformation_field = Node(Ferrite.getfieldnames(plotter.dh)[1])
    #solutionplot main plot of the viewer
    solutionp = solutionplot!(plotter,colormap=:cividis,deformation_field=@lift $(toggles[2].active) ? $(deformation_field) : :default)

    #setting up various sliders
    markerslider = Slider(fig, range = 0:1:100, startvalue=5)
    strokewidthslider = Slider(fig, range = 0:1:10, startvalue=1)
    markersize = lift(x->x,markerslider.value)
    strokewidth = lift(x->x,strokewidthslider.value)

    #plot the fe-mesh
    wireframep = wireframe!(plotter,markersize=markersize,strokewidth=strokewidth,deformation_field= @lift $(toggles[2].active) ? $(deformation_field) : :default)
    #connect fe-mesh plot to the toggle
    connect!(wireframep.visible,toggles[1].active)
    connect!(wireframep.nodelabels,toggles[3].active)
    connect!(wireframep.celllabels,toggles[3].active)

    #set up dropdown menus for colormap, field, deformation field and processing function
    menu_cm = Menu(fig, options=["cividis", "inferno", "thermal"],label="colormap", direction=:up)
    menu_field = Menu(fig, options=Ferrite.getfieldnames(plotter.dh))
    menu_deformation_field = Menu(fig, options=Ferrite.getfieldnames(plotter.dh))
    menu_process = Menu(fig, options=[x₁,x₂,x₃,("magnitude",l2),("manhatten norm",l1),identity])
    #align all menus as a vgrid under each other
    fig[1,3] = vgrid!(grid!(hcat(toggles,labels), tellheight=false),
                      Label(fig,"nodesize",width=nothing), markerslider,
                      Label(fig,"strokewidth",width=nothing), strokewidthslider,
                      Label(fig,"processing function",width=nothing), menu_process,
                      Label(fig,"field",width=nothing), menu_field,
                      Label(fig, "deformation field",width=nothing),menu_deformation_field,
                      Label(fig, "colormap",width=nothing),menu_cm)
    #add colorbar
    cb = Colorbar(fig[1,2], solutionp)

    #event handling for selecting stuff from the menus
    on(menu_cm.selection) do s
        cb.colormap = s
        solutionp.colormap = s
    end

    on(menu_field.selection) do field
        solutionp.field = field
    end

    on(menu_deformation_field.selection) do field
        solutionp.deformation_field = field
        wireframep.deformation_field = field
    end

    on(menu_process.selection) do process_function
        solutionp.process=process_function
    end

    return fig
end

function ferriteviewer(plotter::MakiePlotter, data::Vector{Vector{T}}) where T
    fig = ferriteviewer(plotter)
    timeslider = labelslider!(fig, "timestep n:", 1:length(data); format = x->"$x", sliderkw = Dict(:snap=>false))
    fig[2,1] = timeslider.layout
    @lift(FerriteViz.update!(plotter,data[$(timeslider.slider.value)]))
    display(fig)
end

####### One Shot Methods #######
const FerriteVizPlots = Union{Type{<:Wireframe},Type{<:SolutionPlot},Type{<:Arrows}}

function Makie.convert_arguments(P::FerriteVizPlots, dh::Ferrite.AbstractDofHandler, u::Vector)
    return (MakiePlotter(dh,u),)
end


function vtk_grid(
        dtype::VTKImageData, filename::AbstractString, Ns::Dims{N};
        origin::NTuple{N} = ntuple(d -> 0.0, Val(N)),
        spacing::NTuple{N} = ntuple(d -> 1.0, Val(N)),
        extent = Ns, whole_extent = extent, kwargs...,
    ) where {N}
    Npts = prod(Ns)
    Ncls = num_cells_structured(Ns)

    xvtk = XMLDocument()
    vtk = DatasetFile(dtype, xvtk, filename, Npts, Ncls; kwargs...)

    # VTKFile node
    xroot = vtk_xml_write_header(vtk)

    # ImageData node
    xGrid = new_child(xroot, vtk.grid_type)
    let ext = extent_attribute(whole_extent)
        set_attribute(xGrid, "WholeExtent", ext)
    end

    origin_str = _tuple_to_str3(origin, zero(eltype(origin)))
    spacing_str = _tuple_to_str3(spacing, one(eltype(origin)))

    set_attribute(xGrid, "Origin", origin_str)
    set_attribute(xGrid, "Spacing", spacing_str)

    # Piece node
    xPiece = new_child(xGrid, "Piece")
    let ext = extent_attribute(extent)
        set_attribute(xPiece, "Extent", ext)
    end

    vtk
end

_tuple_to_str3(ts::NTuple{3, T}, default::T) where {T} = join(ts, " ")

function _tuple_to_str3(ts::NTuple{N, T}, default::T) where {N, T}
    @assert N < 3
    M = 3 - N
    tt = (ts..., ntuple(d -> default, Val(M))...)
    _tuple_to_str3(tt, default)
end

function vtk_grid(filename::AbstractString, Ns::Vararg{Integer}; kwargs...)
    # We put the origin at (0.0, 0.0, 0.0)
    vtk_grid(filename, map(N -> range(0.0; length = N), Ns)...; kwargs...)
end

"""
    vtk_grid(filename, x::AbstractRange{T}, y::AbstractRange{T}, [z::AbstractRange{T}];
             kwargs...)

Create image data (`.vti`) file.

Along each direction, the grid is specified in terms of an AbstractRange object.

# Examples

```jldoctest
julia> vtk = vtk_grid("abc", 1:0.2:5, 2:1.:3, 4:1.:5)  # 3D dataset
VTK file 'abc.vti' (ImageData file, open)

julia> vtk = vtk_grid("abc", 1:5, 2:1.:3, range(4, 5; length = 3))  # different kinds of ranges
VTK file 'abc.vti' (ImageData file, open)

julia> vtk = vtk_grid("abc", 1:0.2:5, 2:1.:3)  # 2D dataset
VTK file 'abc.vti' (ImageData file, open)

julia> vtk = vtk_grid("def",
                      LinRange(0., 5., 10),
                      LinRange(0., 2π, 16),
                      LinRange(1., 10., 12))
VTK file 'def.vti' (ImageData file, open)

```

"""
function vtk_grid(filename::AbstractString, xyz::Vararg{AbstractRange}; kwargs...)
    Nxyz = promote(length.(xyz)...)
    spacing = promote(step.(xyz)...)
    origin = promote(first.(xyz)...)

    if (extent = get(kwargs, :extent, nothing)) !== nothing
        # Shift origin accordingly
        length.(extent) == Nxyz ||
            throw(DimensionMismatch("`extent` argument doesn't match grid dimensions"))
        origin_new = @. origin + (1 - first(extent)) * spacing
        origin = oftype(origin, origin_new)
    end

    vtk_grid(VTKImageData(), filename, Nxyz;
             origin=origin, spacing=spacing, kwargs...)
end


module SpgpGpu

using LinearAlgebra

export predict
include("predict.jl")

export kernel
include("kernel.jl")

end


using Compat
using Compat.Libdl

const LIBMBHFMM2D = string(Base.Filesystem.dirname(Base.source_path()), "/../bin/libmbhfmm2d")

include("BoxFMMType.jl")
include("BoxFMMUtil.jl")

# Load OpenMP library
Libdl.dlopen("libgomp", Libdl.RTLD_GLOBAL)

mutable struct MBHFMM2DParams

    # storage of parameters which
    # describe the FMM tree and call

    lambda::Float64
    iprec::Int32
    ifalltarg::Bool
    src::Array{Float64,2}
    targ::Array{Float64,2}
    ifcharge::Bool
    ifdipole::Bool
    ifquad::Bool
    ifoct::Bool
    charge::Array{Float64}
    dipstr::Array{Float64}
    dipvec::Array{Float64,2}
    quadstr::Array{Float64}
    quadvec::Array{Float64,2}
    octstr::Array{Float64}
    octvec::Array{Float64,2}
    exrad::Array{Float64,1}
end

function MBHFMM2DParams(lambda::Float64,src::Array{Float64,2},
                        targ::Array{Float64,2},ifcharge::Bool,
                        ifdipole::Bool,ifquad::Bool,
                        ifoct::Bool,charge::Array{Float64,1},
                        dipstr::Array{Float64,1},
                        dipvec::Array{Float64,2},
                        quadstr::Array{Float64,1},
                        quadvec::Array{Float64,2},
                        octstr::Array{Float64,1},
                        octvec::Array{Float64,2};
                        iprec::Int=3,
                        ifalltarg::Bool=false,
                        exrad::Array{Float64,1}
                        =Array{Float64}(undef,0))

    iprec = convert(Int32,iprec)

    # constructor for MBHFMM2DParams, with a couple
    # default values

    return MBHFMM2DParams(lambda,iprec,ifalltarg,
                          src,targ,ifcharge,ifdipole,
                          ifquad,ifoct,
                          charge,dipstr,dipvec,quadstr,
                          quadvec,octstr,octvec,exrad)

end

mutable struct MBHFMM2DStorage

    # storage for FMM itself

    tree::BoxTree2D
    ifalltarg::Array{Int32,1}
    sorted_pts::SortedPts2D
    chargesort::Array{Float64}
    dipstrsort::Array{Float64}
    dipvecsort::Array{Float64,2}
    quadstrsort::Array{Float64}
    quadvecsort::Array{Float64,2}
    octstrsort::Array{Float64}
    octvecsort::Array{Float64,2}
    ifcharge::Array{Int32,1}
    ifdipole::Array{Int32,1}
    ifquad::Array{Int32,1}
    ifoct::Array{Int32,1}
    isave::Array{Int32,1}
    dsave::Array{Float64,1}
    csave::Array{Complex{Float64},1}
end

function MBHFMM2DStorage_init(fmmpars::MBHFMM2DParams,
                              tree::BoxTree2D,
                              sorted_pts::SortedPts2D)

    # convert bools to fortran flags

    ifalltarg = boxfmm2d_booltoint32(fmmpars.ifalltarg)
    ifcharge = boxfmm2d_booltoint32(fmmpars.ifcharge)
    ifdipole = boxfmm2d_booltoint32(fmmpars.ifdipole)
    ifquad = boxfmm2d_booltoint32(fmmpars.ifquad)
    ifoct = boxfmm2d_booltoint32(fmmpars.ifoct)

    # if charge is used, sort it

    if fmmpars.ifcharge
        chargesort = fmmpars.charge[sorted_pts.isrcsort]
    else
        chargesort = zeros(Float64,1)
    end

    if fmmpars.ifdipole
        dipstrsort = fmmpars.dipstr[sorted_pts.isrcsort]
        dipvecsort = fmmpars.dipvec[:,sorted_pts.isrcsort]
    else
        dipstrsort = zeros(Float64,1)
        dipvecsort = zeros(Float64,2,1)
    end

    if fmmpars.ifquad
        quadstrsort = fmmpars.quadstr[sorted_pts.isrcsort]
        quadvecsort = fmmpars.quadvec[:,sorted_pts.isrcsort]
    else
        quadstrsort = zeros(Float64,1)
        quadvecsort = zeros(Float64,3,1)
    end

    if fmmpars.ifoct
        octstrsort = fmmpars.octstr[sorted_pts.isrcsort]
        octvecsort = fmmpars.octvec[:,sorted_pts.isrcsort]
    else
        octstrsort = zeros(Float64,1)
        octvecsort = zeros(Float64,4,1)
    end

    # initialize isave, dsave, csave to short arrays

    isave = Array{Int32}(undef, 1)
    dsave = Array{Float64}(undef, 1)
    csave = Array{Complex{Float64}}(undef, 1)

    return MBHFMM2DStorage(tree,ifalltarg,sorted_pts,
                           chargesort,dipstrsort,
                           dipvecsort,quadstrsort,
                           quadvecsort,octstrsort,
                           octvecsort,ifcharge,ifdipole,
                           ifquad,ifoct,isave,dsave,csave)

end

function mbhfmm2d_form(fmmpars::MBHFMM2DParams;maxnodes::Int=30,
                       maxboxes::Int=-1)
    # form tree
    tree, sorted_pts, ier = mbhfmm2d_tree(fmmpars, maxnodes=maxnodes, maxboxes=maxboxes,exrad=fmmpars.exrad)
    # do work
    mbhfmm2d_form(fmmpars, tree, sorted_pts)
end

function mbhfmm2d_tree(fmmpars::MBHFMM2DParams;maxnodes::Int=30,
                       maxboxes::Int=-1,
                       exrad::Array{Float64,1}
                       =Array{Float64}(undef,0))
    if maxboxes < 0
        return [], 1
    end
    tree, sorted_pts, ier = BoxTree2DMaxBoxesST(fmmpars.src, fmmpars.targ, maxboxes,
                                                maxnodes=maxnodes,
                                                ifverbose=false,
                                                ifalltarg=fmmpars.ifalltarg,
                                                exrad=exrad)
    return tree, sorted_pts, ier
end

function mbhfmm2d_form(fmmpars::MBHFMM2DParams,
                       tree::BoxTree2D,
                       sorted_pts::SortedPts2D)

    # initialize storage

    fmmstor = MBHFMM2DStorage_init(fmmpars,tree,sorted_pts)

    # query fmm for length of isave, dsave, csave

    lambda = fmmpars.lambda
    ier1 = zeros(Int32,1)
    iprec = fmmpars.iprec
    nlev = tree.nlev
    levelbox = tree.levelbox
    iparentbox = tree.iparentbox
    ichildbox = tree.ichildbox
    icolbox = tree.icolbox
    irowbox = tree.irowbox
    nboxes = tree.nboxes
    nblevel = tree.nblevel
    iboxlev = tree.iboxlev
    istartlev = tree.istartlev
    ifalltarg = fmmstor.ifalltarg
    localonoff = tree.localonoff
    zll = tree.zll
    blength = tree.blength
    ns = sorted_pts.ns
    srcsort = sorted_pts.srcsort
    isrcladder = sorted_pts.isrcladder
    ifcharge = fmmstor.ifcharge
    chargesort = fmmstor.chargesort
    ifdipole = fmmstor.ifdipole
    dipstrsort = fmmstor.dipstrsort
    dipvecsort = fmmstor.dipvecsort
    ifquad = fmmstor.ifquad
    quadstrsort = fmmstor.quadstrsort
    quadvecsort = fmmstor.quadvecsort
    ifoct = fmmstor.ifoct
    octstrsort = fmmstor.octstrsort
    octvecsort = fmmstor.octvecsort
    isave = fmmstor.isave
    dsave = fmmstor.dsave
    csave = fmmstor.csave

    lisave = [convert(Int32,-1)]
    ldsave = [convert(Int32,-1)]
    lcsave = [convert(Int32,-1)]

    ccall( (:mbhfmm2d_form_,LIBMBHFMM2D), Cvoid,
           (Ref{Float64},Ref{Int32},Ref{Int32},Ref{Int32},
            Ref{Int32},Ref{Int32},Ref{Int32},Ref{Int32},
            Ref{Int32},Ref{Int32},Ref{Int32},Ref{Int32},
            Ref{Int32},Ref{Int32},Ref{Int32},Ref{Float64},
            Ref{Float64},Ref{Int32},Ref{Float64},Ref{Int32},
            Ref{Int32},Ref{Float64},
            Ref{Int32},Ref{Float64},Ref{Float64},
            Ref{Int32},Ref{Float64},Ref{Float64},
            Ref{Int32},Ref{Float64},Ref{Float64},
            Ref{Int32},Ref{Int32},Ref{Float64},Ref{Int32},
            Ref{Complex{Float64}},Ref{Int32}),
           lambda,ier1,iprec,nlev,levelbox,iparentbox,
           ichildbox,icolbox,irowbox,nboxes,nblevel,
           iboxlev,istartlev,ifalltarg,localonoff,zll,
           blength,ns,srcsort,isrcladder,ifcharge,chargesort,
           ifdipole,dipstrsort,dipvecsort,ifquad,
           quadstrsort,quadvecsort,ifoct,octstrsort,
           octvecsort,isave,lisave,dsave,ldsave,csave,
           lcsave)

    # allocate

    fmmstor.isave = Array{Int32}(undef, lisave[1])
    fmmstor.dsave = Array{Float64}(undef, ldsave[1])
    fmmstor.csave = Array{Complex{Float64}}(undef, lcsave[1])

    isave = fmmstor.isave
    dsave = fmmstor.dsave
    csave = fmmstor.csave

    # form fmm

    ccall( (:mbhfmm2d_form_,LIBMBHFMM2D), Cvoid,
           (Ref{Float64},Ref{Int32},Ref{Int32},Ref{Int32},
            Ref{Int32},Ref{Int32},Ref{Int32},Ref{Int32},
            Ref{Int32},Ref{Int32},Ref{Int32},Ref{Int32},
            Ref{Int32},Ref{Int32},Ref{Int32},Ref{Float64},
            Ref{Float64},Ref{Int32},Ref{Float64},Ref{Int32},
            Ref{Int32},Ref{Float64},
            Ref{Int32},Ref{Float64},Ref{Float64},
            Ref{Int32},Ref{Float64},Ref{Float64},
            Ref{Int32},Ref{Float64},Ref{Float64},
            Ref{Int32},Ref{Int32},Ref{Float64},Ref{Int32},
            Ref{Complex{Float64}},Ref{Int32}),
           lambda,ier1,iprec,nlev,levelbox,iparentbox,
           ichildbox,icolbox,irowbox,nboxes,nblevel,
           iboxlev,istartlev,ifalltarg,localonoff,zll,
           blength,ns,srcsort,isrcladder,ifcharge,chargesort,
           ifdipole,dipstrsort,dipvecsort,ifquad,
           quadstrsort,quadvecsort,ifoct,octstrsort,
           octvecsort,isave,lisave,dsave,ldsave,csave,
           lcsave)

    return fmmstor

end

function mbhfmm2d_targ!(fmmpars::MBHFMM2DParams,
                        fmmstor::MBHFMM2DStorage,
                        targ::Array{Float64,2},
                        ifpot::Bool,pottarg::Array{Float64,1},
                        ifgrad::Bool,gradtarg::Array{Float64,2},
                        ifhess::Bool,hesstarg::Array{Float64,2},
                        ifignore_exrad::Bool=false)

    # determine if an exclusion radius should be used
    # if ifignore_exrad == true, then the exclusion
    # radius is ignored

    sorted_pts = fmmstor.sorted_pts
    ifexrad = sorted_pts.ifexrad

    if (!ifexrad || ifignore_exrad)
        mbhfmm2d_targ_all!(fmmpars,fmmstor,targ,ifpot,pottarg,
                           ifgrad,gradtarg,ifhess,hesstarg)
    else
        exradsort = sorted_pts.exradsort
        mbhfmm2d_targ_exrad!(fmmpars,fmmstor,exradsort,
                             targ,ifpot,pottarg,
                             ifgrad,gradtarg,ifhess,hesstarg)
    end

end

function mbhfmm2d_targ_all!(fmmpars::MBHFMM2DParams,
                            fmmstor::MBHFMM2DStorage,
                            targ::Array{Float64,2},
                         ifpot::Bool,pottarg::Array{Float64,1},
                         ifgrad::Bool,gradtarg::Array{Float64,2},
                         ifhess::Bool,hesstarg::Array{Float64,2})

    # unpack

    tree = fmmstor.tree
    sorted_pts = fmmstor.sorted_pts

    lambda = fmmpars.lambda
    ier1 = zeros(Int32,1)
    nlev = tree.nlev
    levelbox = tree.levelbox
    iparentbox = tree.iparentbox
    ichildbox = tree.ichildbox
    icolbox = tree.icolbox
    irowbox = tree.irowbox
    nboxes = tree.nboxes
    nblevel = tree.nblevel
    iboxlev = tree.iboxlev
    istartlev = tree.istartlev
    zll = tree.zll
    blength = tree.blength
    ns = sorted_pts.ns
    srcsort = sorted_pts.srcsort
    isrcladder = sorted_pts.isrcladder

    ifcharge = fmmstor.ifcharge
    chargesort = fmmstor.chargesort
    ifdipole = fmmstor.ifdipole
    dipstrsort = fmmstor.dipstrsort
    dipvecsort = fmmstor.dipvecsort
    ifquad = fmmstor.ifquad
    quadstrsort = fmmstor.quadstrsort
    quadvecsort = fmmstor.quadvecsort
    ifoct = fmmstor.ifoct
    octstrsort = fmmstor.octstrsort
    octvecsort = fmmstor.octvecsort
    isave = fmmstor.isave
    dsave = fmmstor.dsave
    csave = fmmstor.csave

    mtemp, nt = size(targ)
    nt = convert(Int32,nt)

    ifpot1 = boxfmm2d_booltoint32(ifpot)
    ifgrad1 = boxfmm2d_booltoint32(ifgrad)
    ifhess1 = boxfmm2d_booltoint32(ifhess)

    ifder31 = zeros(Int32,1)
    der3targ = zeros(Float64,4,1)

    ccall( (:mbhfmm2d3_targ_,LIBMBHFMM2D), Cvoid,
           (Ref{Float64},Ref{Int32},Ref{Int32},
            Ref{Int32},Ref{Int32},Ref{Int32},Ref{Int32},
            Ref{Int32},Ref{Int32},Ref{Int32},Ref{Int32},
            Ref{Int32},Ref{Float64},
            Ref{Float64},Ref{Int32},Ref{Float64},Ref{Int32},
            Ref{Int32},Ref{Float64},
            Ref{Int32},Ref{Float64},Ref{Float64},
            Ref{Int32},Ref{Float64},Ref{Float64},
            Ref{Int32},Ref{Float64},Ref{Float64},
            Ref{Int32},Ref{Float64},Ref{Complex{Float64}},
            Ref{Int32},Ref{Float64},Ref{Int32},
            Ref{Float64},Ref{Int32},Ref{Float64},
            Ref{Int32},Ref{Float64},Ref{Int32},
            Ref{Float64}),
           lambda,ier1,nlev,levelbox,iparentbox,
           ichildbox,icolbox,irowbox,nboxes,nblevel,
           iboxlev,istartlev,zll,
           blength,ns,srcsort,isrcladder,ifcharge,chargesort,
           ifdipole,dipstrsort,dipvecsort,ifquad,
           quadstrsort,quadvecsort,ifoct,octstrsort,
           octvecsort,isave,dsave,csave,nt,targ,ifpot1,
           pottarg,ifgrad1,gradtarg,ifhess1,hesstarg,
           ifder31,der3targ)

    return
end

function mbhfmm2d_targ_exrad!(fmmpars::MBHFMM2DParams,
                              fmmstor::MBHFMM2DStorage,
                              exradsort::Array{Float64,1},
                              targ::Array{Float64,2},
                           ifpot::Bool,pottarg::Array{Float64,1},
                           ifgrad::Bool,gradtarg::Array{Float64,2},
                           ifhess::Bool,hesstarg::Array{Float64,2})

    # unpack

    tree = fmmstor.tree
    sorted_pts = fmmstor.sorted_pts

    lambda = fmmpars.lambda
    ier1 = zeros(Int32,1)
    nlev = tree.nlev
    levelbox = tree.levelbox
    iparentbox = tree.iparentbox
    ichildbox = tree.ichildbox
    icolbox = tree.icolbox
    irowbox = tree.irowbox
    nboxes = tree.nboxes
    nblevel = tree.nblevel
    iboxlev = tree.iboxlev
    istartlev = tree.istartlev
    zll = tree.zll
    blength = tree.blength
    ns = sorted_pts.ns
    srcsort = sorted_pts.srcsort
    isrcladder = sorted_pts.isrcladder

    ifcharge = fmmstor.ifcharge
    chargesort = fmmstor.chargesort
    ifdipole = fmmstor.ifdipole
    dipstrsort = fmmstor.dipstrsort
    dipvecsort = fmmstor.dipvecsort
    ifquad = fmmstor.ifquad
    quadstrsort = fmmstor.quadstrsort
    quadvecsort = fmmstor.quadvecsort
    ifoct = fmmstor.ifoct
    octstrsort = fmmstor.octstrsort
    octvecsort = fmmstor.octvecsort
    isave = fmmstor.isave
    dsave = fmmstor.dsave
    csave = fmmstor.csave

    mtemp, nt = size(targ)
    nt = convert(Int32,nt)

    ifpot1 = boxfmm2d_booltoint32(ifpot)
    ifgrad1 = boxfmm2d_booltoint32(ifgrad)
    ifhess1 = boxfmm2d_booltoint32(ifhess)

    ifder31 = zeros(Int32,1)
    der3targ = zeros(Float64,4,1)

    ccall( (:mbhfmm2d3_targ_exrad_,LIBMBHFMM2D), Cvoid,
           (Ref{Float64},Ref{Int32},Ref{Int32},
            Ref{Int32},Ref{Int32},Ref{Int32},Ref{Int32},
            Ref{Int32},Ref{Int32},Ref{Int32},Ref{Int32},
            Ref{Int32},Ref{Float64},
            Ref{Float64},Ref{Int32},Ref{Float64},Ref{Int32},
            Ref{Float64},Ref{Int32},Ref{Float64},
            Ref{Int32},Ref{Float64},Ref{Float64},
            Ref{Int32},Ref{Float64},Ref{Float64},
            Ref{Int32},Ref{Float64},Ref{Float64},
            Ref{Int32},Ref{Float64},Ref{Complex{Float64}},
            Ref{Int32},Ref{Float64},Ref{Int32},
            Ref{Float64},Ref{Int32},Ref{Float64},
            Ref{Int32},Ref{Float64},Ref{Int32},
            Ref{Float64}),
           lambda,ier1,nlev,levelbox,iparentbox,
           ichildbox,icolbox,irowbox,nboxes,nblevel,
           iboxlev,istartlev,zll,
           blength,ns,srcsort,isrcladder,exradsort,
           ifcharge,chargesort,
           ifdipole,dipstrsort,dipvecsort,ifquad,
           quadstrsort,quadvecsort,ifoct,octstrsort,
           octvecsort,isave,dsave,csave,nt,targ,ifpot1,
           pottarg,ifgrad1,gradtarg,ifhess1,hesstarg,
           ifder31,der3targ)

    return
end

function mbhfmm2d_srcsrc!(fmmpars::MBHFMM2DParams,
                        fmmstor::MBHFMM2DStorage,
                        ifpot::Bool,pottarg::Array{Float64,1},
                        ifgrad::Bool,gradtarg::Array{Float64,2},
                        ifhess::Bool,hesstarg::Array{Float64,2})

    # temporary storage (you get the sorted versions from
    # the fmm)

    potsort = copy(pottarg)
    gradsort = copy(gradtarg)
    hesssort = copy(hesstarg)

    # unpack

    tree = fmmstor.tree
    sorted_pts = fmmstor.sorted_pts

    lambda = fmmpars.lambda
    ier1 = zeros(Int32,1)
    nlev = tree.nlev
    levelbox = tree.levelbox
    iparentbox = tree.iparentbox
    ichildbox = tree.ichildbox
    icolbox = tree.icolbox
    irowbox = tree.irowbox
    nboxes = tree.nboxes
    nblevel = tree.nblevel
    iboxlev = tree.iboxlev
    istartlev = tree.istartlev
    zll = tree.zll
    blength = tree.blength
    ns = sorted_pts.ns
    srcsort = sorted_pts.srcsort
    isrcladder = sorted_pts.isrcladder
    isrcsort = sorted_pts.isrcsort

    ifcharge = fmmstor.ifcharge
    chargesort = fmmstor.chargesort
    ifdipole = fmmstor.ifdipole
    dipstrsort = fmmstor.dipstrsort
    dipvecsort = fmmstor.dipvecsort
    ifquad = fmmstor.ifquad
    quadstrsort = fmmstor.quadstrsort
    quadvecsort = fmmstor.quadvecsort
    ifoct = fmmstor.ifoct
    octstrsort = fmmstor.octstrsort
    octvecsort = fmmstor.octvecsort
    isave = fmmstor.isave
    dsave = fmmstor.dsave
    csave = fmmstor.csave

    ifpot1 = boxfmm2d_booltoint32(ifpot)
    ifgrad1 = boxfmm2d_booltoint32(ifgrad)
    ifhess1 = boxfmm2d_booltoint32(ifhess)

    ifder31 = zeros(Int32,1)
    der3targ = zeros(Float64,4,1)

    ccall( (:mbhfmm2d3_srcsrc_,LIBMBHFMM2D), Cvoid,
           (Ref{Float64},Ref{Int32},Ref{Int32},
            Ref{Int32},Ref{Int32},Ref{Int32},Ref{Int32},
            Ref{Int32},Ref{Int32},Ref{Int32},Ref{Int32},
            Ref{Int32},Ref{Float64},
            Ref{Float64},Ref{Int32},Ref{Float64},Ref{Int32},
            Ref{Int32},Ref{Float64},
            Ref{Int32},Ref{Float64},Ref{Float64},
            Ref{Int32},Ref{Float64},Ref{Float64},
            Ref{Int32},Ref{Float64},Ref{Float64},
            Ref{Int32},Ref{Float64},Ref{Complex{Float64}},
            Ref{Int32},Ref{Float64},Ref{Int32},
            Ref{Float64},Ref{Int32},Ref{Float64},
            Ref{Int32},Ref{Float64}),
           lambda,ier1,nlev,levelbox,iparentbox,
           ichildbox,icolbox,irowbox,nboxes,nblevel,
           iboxlev,istartlev,zll,
           blength,ns,srcsort,isrcladder,ifcharge,chargesort,
           ifdipole,dipstrsort,dipvecsort,ifquad,
           quadstrsort,quadvecsort,ifoct,octstrsort,
           octvecsort,isave,dsave,csave,ifpot1,
           potsort,ifgrad1,gradsort,ifhess1,hesssort,
           ifder31,der3targ)

    if ifpot
        for i = 1:ns
            pottarg[isrcsort[i]] = potsort[i]
        end
    end
    if ifgrad
        for i = 1:ns
            gradtarg[:,isrcsort[i]] = gradsort[:,i]
        end
    end
    if ifhess
        for i = 1:ns
            hesstarg[:,isrcsort[i]] = hesssort[:,i]
        end
    end

    return
end

function mbhfmm2d_direct!(fmmpars::MBHFMM2DParams,
                          targ::Array{Float64,2},
                          ifpot::Bool,pottarg::Array{Float64,1},
                          ifgrad::Bool,gradtarg::Array{Float64,2},
                          ifhess::Bool,hesstarg::Array{Float64,2};
                          exrad::Array{Float64,1}=Array{Float64}(undef,0))

    src = fmmpars.src
    mtemp, ns = size(src)
    ns = convert(Int32,ns)
    mtemp, nt = size(targ)
    nt = convert(Int32,nt)

    ifexrad = false
    if length(exrad) > 0
        @assert length(exrad) == ns "length of exrad should equal number of sources"
        ifexrad = true
    end

    lambda = fmmpars.lambda
    ifcharge = boxfmm2d_booltoint32(fmmpars.ifcharge)
    ifdipole = boxfmm2d_booltoint32(fmmpars.ifdipole)
    ifquad = boxfmm2d_booltoint32(fmmpars.ifquad)
    ifoct = boxfmm2d_booltoint32(fmmpars.ifoct)

    charge = fmmpars.charge
    dipstr = fmmpars.dipstr
    dipvec = fmmpars.dipvec
    quadstr = fmmpars.quadstr
    quadvec = fmmpars.quadvec
    octstr = fmmpars.octstr
    octvec = fmmpars.octvec

    ifpot1 = boxfmm2d_booltoint32(ifpot)
    ifgrad1 = boxfmm2d_booltoint32(ifgrad)
    ifhess1 = boxfmm2d_booltoint32(ifhess)

    pottemp = zeros(Float64,1)
    gradtemp = zeros(Float64,2)
    hesstemp = zeros(Float64,3)
    der3temp = zeros(Float64,4)

    if ifexrad
        for i = 1:nt
            pottemp = zeros(Float64,1)
            gradtemp = zeros(Float64,2)
            hesstemp = zeros(Float64,3)
            der3temp = zeros(Float64,4)

            ccall((:mbhpotgrad2dall_cdqo3_add_exrad_,LIBMBHFMM2D),Cvoid,
                  (Ref{Float64},Ref{Float64},Ref{Float64},
                   Ref{Int32},Ref{Int32},
                   Ref{Float64},Ref{Int32},Ref{Float64},
                   Ref{Float64},
                   Ref{Int32},Ref{Float64},Ref{Float64},
                   Ref{Int32},Ref{Float64},Ref{Float64},
                   Ref{Float64},Ref{Int32},Ref{Float64},
                   Ref{Int32},Ref{Float64},Ref{Int32},
                   Ref{Float64}),
                  lambda,src,exrad,ns,ifcharge,charge,
                  ifdipole,dipstr,dipvec,
                  ifquad,quadstr,quadvec,ifoct,octstr,octvec,
                  targ[:,i],ifpot1,pottemp,ifgrad1,gradtemp,
                  ifhess1,hesstemp)

            pottarg[i] = pottemp[1]
            gradtarg[:,i] = gradtemp
            hesstarg[:,i] = hesstemp

        end
    else
        for i = 1:nt

            ccall((:mbhpotgrad2dall_cdqo_,LIBMBHFMM2D),Cvoid,
                  (Ref{Float64},Ref{Float64},Ref{Int32},Ref{Int32},
                   Ref{Float64},Ref{Int32},Ref{Float64},
                   Ref{Float64},
                   Ref{Int32},Ref{Float64},Ref{Float64},
                   Ref{Int32},Ref{Float64},Ref{Float64},
                   Ref{Float64},Ref{Int32},Ref{Float64},
                   Ref{Int32},Ref{Float64},Ref{Int32},
                   Ref{Float64}),
                  lambda,src,ns,ifcharge,charge,ifdipole,
                  dipstr,dipvec,
                  ifquad,quadstr,quadvec,ifoct,octstr,octvec,
                  targ[:,i],ifpot1,pottemp,ifgrad1,gradtemp,
                  ifhess1,hesstemp)

            pottarg[i] = pottemp[1]
            gradtarg[:,i] = gradtemp
            hesstarg[:,i] = hesstemp

        end
    end
    return
end

function mbhfmm2d_direct_self!(fmmpars::MBHFMM2DParams,
                          ifpot::Bool,pottarg::Array{Float64,1},
                          ifgrad::Bool,gradtarg::Array{Float64,2},
                          ifhess::Bool,hesstarg::Array{Float64,2})

    src = fmmpars.src
    mtemp, ns = size(src)
    ns = convert(Int32,ns)

    lambda = fmmpars.lambda
    ifcharge = boxfmm2d_booltoint32(fmmpars.ifcharge)
    ifdipole = boxfmm2d_booltoint32(fmmpars.ifdipole)
    ifquad = boxfmm2d_booltoint32(fmmpars.ifquad)
    ifoct = boxfmm2d_booltoint32(fmmpars.ifoct)

    charge = fmmpars.charge
    dipstr = fmmpars.dipstr
    dipvec = fmmpars.dipvec
    quadstr = fmmpars.quadstr
    quadvec = fmmpars.quadvec
    octstr = fmmpars.octstr
    octvec = fmmpars.octvec

    ifpot1 = boxfmm2d_booltoint32(ifpot)
    ifgrad1 = boxfmm2d_booltoint32(ifgrad)
    ifhess1 = boxfmm2d_booltoint32(ifhess)

    ifder31 = zeros(Int32,1)
    der3targ = zeros(Float64,4,1)

    ccall((:mbhpotgrad2dall_cdqo3_self_,LIBMBHFMM2D),Cvoid,
          (Ref{Float64},Ref{Float64},Ref{Int32},Ref{Int32},
           Ref{Float64},Ref{Int32},Ref{Float64},Ref{Float64},
           Ref{Int32},Ref{Float64},Ref{Float64},
           Ref{Int32},Ref{Float64},Ref{Float64},
           Ref{Int32},Ref{Float64},Ref{Int32},Ref{Float64},
           Ref{Int32},Ref{Float64},Ref{Int32},Ref{Float64}),
          lambda,src,ns,ifcharge,charge,ifdipole,dipstr,dipvec,
          ifquad,quadstr,quadvec,ifoct,octstr,octvec,
          ifpot1,pottarg,ifgrad1,gradtarg,
          ifhess1,hesstarg,ifder31,der3targ)

    return
end


import LPDM: bounds, best_lb_action, best_ub_action

mutable struct LDBounds1d{S,A,O}
    lb_    ::Float64
    ub_    ::Float64
    best_lb_action_::Float64
    best_ub_action_::Float64
    rng::RNGVector # dummy RNG, not used

    function LDBounds1d{S,A,O}(::POMDP{S,A,O}) where {S,A,O}
        this = new()
        this.lb_ = +Inf
        this.ub_ = -Inf
        this.best_lb_action_ = NaN
        this.best_ub_action_ = NaN
        this.rng = RNGVector(1, UInt32(42))
        return this
    end
end

function LPDM.bounds(b::LDBounds1d{S,A,O},
                     pomdp::AbstractLD1,
                     particles::Vector{LPDMParticle{S}},
                     config::LPDMConfig) where {S,A,O}

    # reset every time bounds are about to be recomputed
    b.lb_ = +Inf
    b.ub_ = -Inf
    b.best_lb_action_ = NaN
    b.best_ub_action_ = NaN

    tmp_lb = NaN
    tmp_ub = NaN
    tmp_lb_action = NaN
    tmp_ub_action = NaN

    for p in particles
        tmp_ub, tmp_ub_action = upper_bound(pomdp,p)
        if p.state >= pomdp.min_noise_loc # assume min_noise_loc > 0
            # both will have to do roughly the same thing (+/- discretization differences),
            # so make them the same
            #TODO: review for lb > ub
            tmp_lb, tmp_lb_action = tmp_ub, tmp_ub_action
        else
            tmp_lb, tmp_lb_action = lower_bound(pomdp,p)
        end

        if tmp_lb < b.lb_
            b.lb_             = tmp_lb
            b.best_lb_action_ = tmp_lb_action
            # if abs(tmp_lb + 30.0) < 0.1 # DEBUG: remove
            #     println("particles: $particles")
            #     println("s: $(p.state), b.lb_=$(b.lb_), tmp_lb_action = $(tmp_lb_action), b.lba_ = $(b.best_lb_action_)")
            # end
        end
        if tmp_ub > b.ub_
            b.ub_             = tmp_ub
            b.best_ub_action_ = tmp_ub_action
        end
    end

    # Sanity check
    if b.ub_ < b.lb_
        # show(particles); println("")
        # error("BOUNDS: ub=$(b.ub_) < lb=$(b.lb_), lba = $(b.best_lb_action_), uba = $(b.best_ub_action_), tmp_lb = $(tmp_lb), tmp_ub = $(tmp_ub), tmp_lb_action = $(tmp_lb_action), tmp_ub_action = $(tmp_ub_action)")
        error("BOUNDS: ub=$(b.ub_) < lb=$(b.lb_)")
    end
    return b.lb_, b.ub_
end

LPDM.best_lb_action(b::LDBounds1d) = isnan(b.best_lb_action_) ? error("best_lb_action undefined. Call bounds() first") : b.best_lb_action_
LPDM.best_ub_action(b::LDBounds1d) = isnan(b.best_ub_action_) ? error("best_ub_action undefined. Call bounds() first") : b.best_ub_action_

function move(p::AbstractLD1, x1::Float64, x2::Float64)#::(Float64,Float64)
    direction = x2 > x1 ? 1 : -1
    all_actions = POMDPs.actions(p)
    # I = findall(!iszero, all_actions)
    pos_actions = all_actions[all_actions .> 0]
    # nz_actions = all_actions[I]
    min_a = minimum(pos_actions)

    r = 0.0
    x = x1
    first_a = NaN
    a_dir = NaN

    if abs(x1) <= p.term_radius
        return reward(p,x1,0.0), 0.0
    end

    if abs(x2-x1) < min_a #too close to take any action
        return reward(p,x1,0.0), 0.0
    end

    while (abs(x2-x) > min_a) && (abs(x) >= p.term_radius)
        a = maximum(pos_actions[pos_actions .<= abs(x2-x)]) # maximum action not exceeding Δx
        a_dir = direction * a
        if isnan(first_a)
            first_a = a_dir # assign first action (directional)
        end
        r += reward(p, x, a_dir) # use current state for computing the reward
        # println("BOUNDS: x=$x, x1=$x1, x2=$x2, min_a = $min_a, all_actions=$all_actions, pos_actions=$pos_actions, av_actions=$(all_actions[all_actions .< abs(x2-x)]), a_dir=$a_dir,  r=$r ")
        x += a_dir # take the step
    end
    # error("review below, might be a bug with assigning first_a = 0.0")
    if abs(x2) <= p.term_radius
        # x1 < 0.7 && println("TERMINATION REWARD #2 FOR x2=$x")
        r += reward(p,x2,0.0) #termination reward
        first_a = 0.0
    end
    # x1 < 0.7 && println("exiting move $x1 -> $x2")
    # if isnan(first_a) #DEBUG: remove
    #     error("move: first_a = $first_a, x1=$x1, x2=$x2")
    # end
    return r, first_a #action sign depends on the direction
end

# NOTE: not for direct calling
# computes the cost of traveling to the low noise region and only then towards target. i.e. a slow approach
function lower_bound(p::AbstractLD1, particle::LPDMParticle{Float64})

    if abs(particle.state) < p.term_radius
        return reward(p, particle.state, 0.0), 0.0
    else
        r1,a1 = move(p, particle.state, p.min_noise_loc)         # estimate cost of moving x to low noise x
        # estimate cost of moving x to the target from the low noise region (will terminate earlier)
        r2,a2 = move(p, p.min_noise_loc, 0.0)
    end
    return r1+r2, a1
end
# lower_bound(p::LightDarkLPDM.LightDark1DDespot, particle::LPDM.LPDMParticle{Float64}) = lower_bound(p, POMDPToolbox.Particle{Float64}(particle.state, particle.weight))

# NOTE: not for direct calling
# computes the reward for the straight-line path to target
function upper_bound(p::AbstractLD1, particle::LPDMParticle{Float64})
    if abs(particle.state) < p.term_radius
        return reward(p, particle.state, 0.0), 0.0
    else
        return move(p, particle.state, 0.0)
    end
end


"""
    function addIntegratorMediumLax_!(p::Program_,abm::Agent,space::SimulationSpace,platform::String)

Generate the functions related with Medium Update.
"""
function addIntegratorMediumLax_!(p::Program_,abm::Agent,space::SimulationSpace,platform::String)

    if "UpdateMedium" in keys(abm.declaredUpdates)

        #Check updated
        up = symbols_(abm,abm.declaredUpdates["UpdateMedium"])
        up = up[Bool.((up[:,"placeDeclaration"].==:Model) .* Bool.((up[:,"assigned"].==true) .+ (up[:,"updated"].==true))),:]

        #Construct functions

            #Make function to compute everything inside the volume

        f = abm.declaredUpdates["UpdateMedium"]
        for (i,j) in enumerate(abm.declaredSymbols["Medium"]) # Change symbol for update

            if abm.dims == 1
                add = :((mediumV[ic1_+1,$i] + mediumV[ic1_+1,$i])/2)
            elseif abm.dims == 2
                add = :((mediumV[ic1_-1,ic2_,$i] + mediumV[ic1_+1,ic2_,$i]
                        +mediumV[ic1_,ic2_-1,$i] + mediumV[ic1_,ic2_+1,$i])/4)
            elseif abm.dims == 3
                add = :((mediumV[ic1_-1,ic2_,ic3_,$i] + mediumV[ic1_+1,ic2_,ic3_,$i]
                        +mediumV[ic1_,ic2_-1,ic3_,$i] + mediumV[ic1_,ic2_+1,ic3_,$i]
                        +mediumV[ic1_,ic2_,ic3_-1,$i] + mediumV[ic1_,ic2_,ic3_+1,$i])/6)
            end

            s = Meta.parse(string(MEDIUMSYMBOL,j))
            f = postwalk(x -> @capture(x,ss_=v_) && ss == s ? :($j = $add + $v*dt) : x, f)
        end
        f = postwalk(x->@capture(x,s_(v_)) ? :($(adaptOperatorsMediumLax_(v,s,abm,space,p))) : x, f) # Adapt
        f = postwalk(x->@capture(x,s_) ? adaptSymbolsMediumLax_(s,abm,space,p) : x, f) # Adapt

        f = vectorizeMedium_(abm,f,p)

        f = simpleGridLoopWrapInFunction_(platform,:mediumInnerStep_!, f, abm.dims)

        #Remove count over boundaries
        if platform == "cpu"
            f = postwalk(x->@capture(x,1:Nx_) && Nx == :Nx_ ? :(2:Nx_-1) : x, f)
            f = postwalk(x->@capture(x,1:Ny_) && Ny == :Ny_ ? :(2:Ny_-1) : x, f)
            f = postwalk(x->@capture(x,1:Nz_) && Nz == :Nz_ ? :(2:Nz_-1) : x, f)
        elseif platform == "gpu"
            f = postwalk(x->@capture(x,indexX_:strideX_:Nx_) && Nx == :Nx_ ? :(indexX_+1:strideX_:Nx_-1) : x, f)
            f = postwalk(x->@capture(x,indexY_:strideY_:Ny_) && Ny == :Ny_ ? :(indexY_+1:strideY_:Ny_-1) : x, f)
            f = postwalk(x->@capture(x,indexZ_:strideZ_:Nz_) && Nz == :Nz_ ? :(indexZ_+1:strideZ_:Nz_-1) : x, f)
        end

            #Make function to compute the boundaries
        code = quote end

        lUpdate = length(abm.declaredSymbols["Medium"])
        for i in 1:abm.dims #Adapt boundary updates depending on the boudary type
            ic = Meta.parse(string("ic",i,"_"))
            nm = Meta.parse(string("N",["x","y","z"][i],"_"))
            subcode = quote end

            if abm.dims == 1
                v = :(mediumVCopy[ic1_,ic4_])
            elseif abm.dims == 2
                v = :(mediumVCopy[ic1_,ic2_,ic4_])
            elseif abm.dims == 3
                v = :(mediumVCopy[ic1_,ic2_,ic3_,ic4_])
            end

            if space.medium[i].minBoundaryType == "Periodic"

                #Lower boundary
                vUp = postwalk(x->@capture(x, t_) && t == ic ? 1 : x, v)
                vAss = postwalk(x->@capture(x, t_) && t == ic ? :($nm-1) : x, v)

                push!(subcode.args, :($vUp=$vAss))

                #Upper boundary
                vUp = postwalk(x->@capture(x, t_) && t == ic ? :($nm) : x, v)
                vAss = postwalk(x->@capture(x, t_) && t == ic ? 2 : x, v)

                push!(subcode.args, :($vUp=$vAss))
            end

            if space.medium[i].minBoundaryType == "Newmann"

                #Lower boundary
                vUp = postwalk(x->@capture(x, t_) && t == ic ? 1 : x, v)
                vAss = postwalk(x->@capture(x, t_) && t == ic ? :(2) : x, v)

                push!(subcode.args, :($vUp=$vAss))
            end

            if space.medium[i].maxBoundaryType == "Newmann"

                #Lower boundary
                vUp = postwalk(x->@capture(x, t_) && t == ic ? nm : x, v)
                vAss = postwalk(x->@capture(x, t_) && t == ic ? :($nm-1) : x, v)

                push!(subcode.args, :($vUp=$vAss))
            end

            if space.medium[i].minBoundaryType == "Dirichlet"

                #Lower boundary
                vUp = postwalk(x->@capture(x, t_) && t == ic ? 1 : x, v)

                push!(subcode.args, :($vUp=0.))
            end

            if space.medium[i].maxBoundaryType == "Dirichlet"

                #Lower boundary
                vUp = postwalk(x->@capture(x, t_) && t == ic ? nm : x, v)

                push!(subcode.args, :($vUp=0.))
            end

            subcode = :(for ic4_ in 1:$lUpdate
                            $subcode
                        end
                    )

            subcode = simpleGridLoop_(platform,subcode,abm.dims-1, indexes = [1,2,3][findall([1,2,3].!=i)])

            push!(code.args,subcode)
        end

        f2 = wrapInFunction_(:mediumBoundaryStep_!,code)

        if "UpdateMediumInteraction" in keys(p.agent.declaredUpdates)
            fWrap = wrapInFunction_(:mediumStep_!,
                                    :(begin
                                        @platformAdapt mediumInnerStep_!(ARGS_)
                                        @platformAdapt mediumBoundaryStep_!(ARGS_)
                                        @platformAdapt mediumInteractionStep_!(ARGS_)
                                    end)
                                    )
        else
            fWrap = wrapInFunction_(:mediumStep_!,
                                    :(begin
                                        @platformAdapt mediumInnerStep_!(ARGS_)
                                        @platformAdapt mediumBoundaryStep_!(ARGS_)
                                    end)
                                    )
        end

        # println(prettify(code))

        push!(p.declareF.args,f,f2,fWrap)

        push!(p.execInloop.args,
                :(mediumStep_!(ARGS_))
            )
    end

    return nothing
end

function adaptSymbolsMediumLax_(s,abm,space,p)

    if s == :xₘ

        return :(ic1_*dxₘ_+$(space.box[1].min))

    elseif s == :yₘ

        return :(ic2_*dyₘ_+$(space.box[2].min))

    elseif s == :zₘ

        return :(ic1_*dzₘ_+$(space.box[3].min))

    end

    return s

end

function adaptOperatorsMediumLax_(f,op,abm,space,p)

    f = vectorizeMedium_(abm,f,p)

    if op == :δx

        if space.medium[1].minBoundaryType == "Periodic"
            f = :(δMedium_(round(Int,($f-$(space.box[1].min))/dxₘ_)+1,ic1_))
        else
            f = :(δMedium_(round(Int,($f-$(space.box[1].min))/dxₘ_),ic1_))
        end

    elseif op == :δy

        if abm.dims < 2
            error("Operator δy cannot be used in Agent with dimension 1.")
        end

        if space.medium[2].minBoundaryType == "Periodic"
            f = :(δMedium_(round(Int,($f-$(space.box[2].min))/dyₘ_)+1,ic2_))
        else
            f = :(δMedium_(round(Int,($f-$(space.box[2].min))/dyₘ_),ic2_))
        end

    elseif op == :δz

        if abm.dims < 3
            error("Operator δz cannot be used in Agent with dimension 1 or 2.")
        end

        if space.medium[3].minBoundaryType == "Periodic"
            f = :(δMedium_(round(Int,($f-$(space.box[3].min))/dzₘ_)+1,ic3_))
        else
            f = :(δMedium_(round(Int,($f-$(space.box[3].min))/dzₘ_),ic3_))
        end

    elseif op == :Δ
        if abm.dims == 1
            add = postwalk(x->@capture(x,s_[ic1_,j_]) && s == :mediumV ? :(mediumV[ic1_+1,$j]) : x,f)
            subs = postwalk(x->@capture(x,s_[ic1_,j_]) && s == :mediumV ? :(mediumV[ic1_-1,$j]) : x,f)

            f = :(($add+$subs-2*$f)/dxₘ_^2)
        elseif abm.dims == 2
            add = postwalk(x->@capture(x,s_[ic1_,ic2_,j_]) && s == :mediumV ? :(mediumV[ic1_+1,ic2_,$j]) : x,f)
            subs = postwalk(x->@capture(x,s_[ic1_,ic2_,j_]) && s == :mediumV ? :(mediumV[ic1_-1,ic2_,$j]) : x,f)
            addy = postwalk(x->@capture(x,s_[ic1_,ic2_,j_]) && s == :mediumV ? :(mediumV[ic1_,ic2_+1,$j]) : x,f)
            subsy = postwalk(x->@capture(x,s_[ic1_,ic2_,j_]) && s == :mediumV ? :(mediumV[ic1_,ic2_-1,$j]) : x,f)

            f = :(($add+$subs-2*$f)/dxₘ_^2+($addy+$subsy-2*$f)/dyₘ_^2)
        elseif abm.dims == 3
            add = postwalk(x->@capture(x,s_[ic1_,ic2_,ic3_,j_]) && s == :mediumV ? :(mediumV[ic1_+1,ic2_,ic3_,$j]) : x,f)
            subs = postwalk(x->@capture(x,s_[ic1_,ic2_,ic3_,j_]) && s == :mediumV ? :(mediumV[ic1_-1,ic2_,ic3_,$j]) : x,f)
            addy = postwalk(x->@capture(x,s_[ic1_,ic2_,ic3_,j_]) && s == :mediumV ? :(mediumV[ic1_,ic2_+1,ic3_,$j]) : x,f)
            subsy = postwalk(x->@capture(x,s_[ic1_,ic2_,ic3_,j_]) && s == :mediumV ? :(mediumV[ic1_,ic2_-1,ic3_,$j]) : x,f)
            addz = postwalk(x->@capture(x,s_[ic1_,ic2_,ic3_,j_]) && s == :mediumV ? :(mediumV[ic1_,ic2_,ic3_+1,$j]) : x,f)
            subsz = postwalk(x->@capture(x,s_[ic1_,ic2_,ic3_,j_]) && s == :mediumV ? :(mediumV[ic1_,ic2_,ic3_-1,$j]) : x,f)

            f = :(($add+$subs-2*$f)/dxₘ_^2+($addy+$subsy-2*$f)/dyₘ_^2+($addz+$subsz-2*$f)/dzₘ_^2)

        end
    elseif op == :Δx
        if abm.dims == 1
            add = postwalk(x->@capture(x,s_[ic1_,j_]) && s == :mediumV ? :(mediumV[ic1_+1,$j]) : x,f)
            subs = postwalk(x->@capture(x,s_[ic1_,j_]) && s == :mediumV ? :(mediumV[ic1_-1,$j]) : x,f)

            f = :(($add+$subs-2*$f)/dxₘ_^2)
        elseif abm.dims == 2
            add = postwalk(x->@capture(x,s_[ic1_,ic2_,j_]) && s == :mediumV ? :(mediumV[ic1_+1,ic2_,$j]) : x,f)
            subs = postwalk(x->@capture(x,s_[ic1_,ic2_,j_]) && s == :mediumV ? :(mediumV[ic1_-1,ic2_,$j]) : x,f)

            f = :(($add+$subs-2*$f)/dxₘ_^2)
        elseif abm.dims == 3
            add = postwalk(x->@capture(x,s_[ic1_,ic2_,ic3_,j_]) && s == :mediumV ? :(mediumV[ic1_+1,ic2_,ic3_,$j]) : x,f)
            subs = postwalk(x->@capture(x,s_[ic1_,ic2_,ic3_,j_]) && s == :mediumV ? :(mediumV[ic1_-1,ic2_,ic3_,$j]) : x,f)

            f = :(($add+$subs-2*$f)/dxₘ_^2)
        end
    elseif op == :Δy
        if abm.dims == 1
            error("Operator Δy cannot exist in Agent with dimension 1.")
        elseif abm.dims == 2
            addy = postwalk(x->@capture(x,s_[ic1_,ic2_,j_]) && s == :mediumV ? :(mediumV[ic1_,ic2_+1,$j]) : x,f)
            subsy = postwalk(x->@capture(x,s_[ic1_,ic2_,j_]) && s == :mediumV ? :(mediumV[ic1_,ic2_-1,$j]) : x,f)

            f = :(($addy+$subsy-2*$f)/dyₘ_^2)
        elseif abm.dims == 3
            addy = postwalk(x->@capture(x,s_[ic1_,ic2_,ic3_,j_]) && s == :mediumV ? :(mediumV[ic1_,ic2_+1,ic3_,$j]) : x,f)
            subsy = postwalk(x->@capture(x,s_[ic1_,ic2_,ic3_,j_]) && s == :mediumV ? :(mediumV[ic1_,ic2_-1,ic3_,$j]) : x,f)

            f = :(($addy+$subsy-2*$f)/dyₘ_^2)
        end
    elseif op == :Δz
        if abm.dims == 1
            error("Operator Δz cannot exist in Agent with dimension 1.")
        elseif abm.dims == 2
            error("Operator Δy cannot exist in Agent with dimension 2.")
        elseif abm.dims == 3
            addz = postwalk(x->@capture(x,s_[ic1_,ic2_,ic3_,j_]) && s == :mediumV ? :(mediumV[ic1_,ic2_,ic3_+1,$j]) : x,f)
            subsz = postwalk(x->@capture(x,s_[ic1_,ic2_,ic3_,j_]) && s == :mediumV ? :(mediumV[ic1_,ic2_,ic3_-1,$j]) : x,f)

            f = :(($addz+$subsz-2*$f)/dzₘ_^2)
        end
    elseif op == :∇
        if abm.dims == 1
            add = postwalk(x->@capture(x,s_[ic1_,j_]) && s == :mediumV ? :(mediumV[ic1_+1,$j]) : x,f)
            subs = postwalk(x->@capture(x,s_[ic1_,j_]) && s == :mediumV ? :(mediumV[ic1_-1,$j]) : x,f)

            f = :(($add-$subs)/(2*dxₘ_))
        elseif abm.dims == 2
            add = postwalk(x->@capture(x,s_[ic1_,ic2_,j_]) && s == :mediumV ? :(mediumV[ic1_+1,ic2_,$j]) : x,f)
            subs = postwalk(x->@capture(x,s_[ic1_,ic2_,j_]) && s == :mediumV ? :(mediumV[ic1_-1,ic2_,$j]) : x,f)
            addy = postwalk(x->@capture(x,s_[ic1_,ic2_,j_]) && s == :mediumV ? :(mediumV[ic1_,ic2_+1,$j]) : x,f)
            subsy = postwalk(x->@capture(x,s_[ic1_,ic2_,j_]) && s == :mediumV ? :(mediumV[ic1_,ic2_-1,$j]) : x,f)

            f = :(($add-$subs)/(2*dxₘ_)+($addy-$subsy)/(2*dyₘ_))
        elseif abm.dims == 3
            add = postwalk(x->@capture(x,s_[ic1_,ic2_,ic3_,j_]) && s == :mediumV ? :(mediumV[ic1_+1,ic2_,ic3_,$j]) : x,f)
            subs = postwalk(x->@capture(x,s_[ic1_,ic2_,ic3_,j_]) && s == :mediumV ? :(mediumV[ic1_-1,ic2_,ic3_,$j]) : x,f)
            addy = postwalk(x->@capture(x,s_[ic1_,ic2_,ic3_,j_]) && s == :mediumV ? :(mediumV[ic1_,ic2_+1,ic3_,$j]) : x,f)
            subsy = postwalk(x->@capture(x,s_[ic1_,ic2_,ic3_,j_]) && s == :mediumV ? :(mediumV[ic1_,ic2_-1,ic3_,$j]) : x,f)
            addz = postwalk(x->@capture(x,s_[ic1_,ic2_,ic3_,j_]) && s == :mediumV ? :(mediumV[ic1_,ic2_,ic3_+1,$j]) : x,f)
            subsz = postwalk(x->@capture(x,s_[ic1_,ic2_,ic3_,j_]) && s == :mediumV ? :(mediumV[ic1_,ic2_,ic3_-1,$j]) : x,f)

            f = :(($add-$subs)/(2*dxₘ_)+($addy-$subsy)/(2*dyₘ_)+($addz-$subsz)/(2*dzₘ_))
        end
    elseif op == :∇x
        if abm.dims == 1
            add = postwalk(x->@capture(x,s_[ic1_,j_]) && s == :mediumV ? :(mediumV[ic1_+1,$j]) : x,f)
            subs = postwalk(x->@capture(x,s_[ic1_,j_]) && s == :mediumV ? :(mediumV[ic1_-1,$j]) : x,f)

            f = :(($add-$subs)/(2*dxₘ_))
        elseif abm.dims == 2
            add = postwalk(x->@capture(x,s_[ic1_,ic2_,j_]) && s == :mediumV ? :(mediumV[ic1_+1,ic2_,$j]) : x,f)
            subs = postwalk(x->@capture(x,s_[ic1_,ic2_,j_]) && s == :mediumV ? :(mediumV[ic1_-1,ic2_,$j]) : x,f)

            f = :(($add-$subs)/(2*dxₘ_))
        elseif abm.dims == 3
            add = postwalk(x->@capture(x,s_[ic1_,ic2_,ic3_,j_]) && s == :mediumV ? :(mediumV[ic1_+1,ic2_,ic3_,$j]) : x,f)
            subs = postwalk(x->@capture(x,s_[ic1_,ic2_,ic3_,j_]) && s == :mediumV ? :(mediumV[ic1_-1,ic2_,ic3_,$j]) : x,f)

            f = :(($add-$subs)/(2*dxₘ_))
        end
    elseif op == :∇y
        if abm.dims == 1
            error("Operator Δy cannot exist in Agent with dimension 1.")
        elseif abm.dims == 2
            addy = postwalk(x->@capture(x,s_[ic1_,ic2_,j_]) && s == :mediumV ? :(mediumV[ic1_,ic2_+1,$j]) : x,f)
            subsy = postwalk(x->@capture(x,s_[ic1_,ic2_,j_]) && s == :mediumV ? :(mediumV[ic1_,ic2_-1,$j]) : x,f)

            f = :(($addy-$subsy)/(2*dyₘ_))
        elseif abm.dims == 3
            addy = postwalk(x->@capture(x,s_[ic1_,ic2_,ic3_,j_]) && s == :mediumV ? :(mediumV[ic1_,ic2_+1,ic3_,$j]) : x,f)
            subsy = postwalk(x->@capture(x,s_[ic1_,ic2_,ic3_,j_]) && s == :mediumV ? :(mediumV[ic1_,ic2_-1,ic3_,$j]) : x,f)

            f = :(($addy-$subsy)/(2*dyₘ_))
        end
    elseif op == :∇z
        if abm.dims == 1
            error("Operator Δz cannot exist in Agent with dimension 1.")
        elseif abm.dims == 2
            error("Operator Δy cannot exist in Agent with dimension 2.")
        elseif abm.dims == 3
            addz = postwalk(x->@capture(x,s_[ic1_,ic2_,ic3_,j_]) && s == :mediumV ? :(mediumV[ic1_,ic2_,ic3_+1,$j]) : x,f)
            subsz = postwalk(x->@capture(x,s_[ic1_,ic2_,ic3_,j_]) && s == :mediumV ? :(mediumV[ic1_,ic2_,ic3_-1,$j]) : x,f)

            f = :(($addz-$subsz)/(2*dzₘ_))
        end
    end

    return f

end

δMedium_(i1,i2) = if i1==i2 1. else 0. end


module CURANDkernel
using ..Codegen

export curandState_t, curandStateSobol32_t
export curand_init, curand, curand_uniform, curand_normal, curand_log_normal, curand_poisson

const __current_compiler__ = Codegen.__current_compiler__

cxx"""
#include <curand_kernel.h>
"""

"""
    curandState_t

"""
typealias curandState_t cxxt"curandState_t"{48}

export new, delete
new(::Type{curandState_t}) = icxx"return (curandState_t*) malloc(sizeof(curandState_t));"
delete(x) = icxx"free($x);"

"""
   curand_init(see, sequence, offset, state)

Initialisation for `curandState_t`

# Arguments
- `seed`
- `sequence`
- `offset`
- `state` Cxx pointer to a state
"""
function curand_init(seed, sequence, offset, state)
  icxx"""
    curand_init($seed, $sequence, $offset, $state);
  """
end

"""
   curand(state)

"""
curand(state) = icxx"curand($state);"

curand_uniform(::Type{Float32}, state) = icxx"curand_uniform($state);"
curand_uniform(::Type{Float64}, state) = icxx"curand_uniform_double($state);"

curand_normal(::Type{Float32}, state) = icxx"curand_normal($state);"
curand_normal(::Type{Float64}, state) = icxx"curand_normal_double($state);"

curand_log_normal(::Type{Float32}, state, mean, stddev) =
  icxx"curand_log_normal($state, $mean, $stddev);"
curand_log_normal(::Type{Float64}, state, mean, stddev) =
  icxx"curand_log_normal_double($state, $mean, $stddev);"

curand_poisson(state, lambda) = icxx"curand_poisson($state, $lambda);"

end


@testset "Test 32 9-Bus Machine Only System" begin
    path = (joinpath(pwd(), "test-9Bus-system"))
    !isdir(path) && mkdir(path)
    try
        sys = System(joinpath(TEST_FILES_DIR, "data_tests/9BusSystem.json"))
        for l in get_components(PSY.PowerLoad, sys)
            PSY.set_model!(l, PSY.LoadModels.ConstantImpedance)
        end

        gen_static = get_component(ThermalStandard, sys, "generator-2-1")
        gen_dynamic = get_dynamic_injector(gen_static)
        sim = Simulation(
            ResidualModel,
            sys,
            path,
            (0.0, 2.0),
            ControlReferenceChange(1.0, gen_dynamic, :P_ref, 0.78),
        )
        small_sig = small_signal_analysis(sim)
        @test execute!(sim, IDA()) == PSID.SIMULATION_FINALIZED

        sim = Simulation(
            MassMatrixModel,
            sys,
            path,
            (0.0, 2.0),
            ControlReferenceChange(1.0, gen_dynamic, :P_ref, 0.78),
        )
        small_sig = small_signal_analysis(sim)
        @test execute!(sim, Rodas4()) == PSID.SIMULATION_FINALIZED

    finally
        @info("removing test files")
        rm(path, force = true, recursive = true)
    end
end


function get_target_pv(obj::GroundPointing, arg0::FieldPVCoordinatesProvider, arg1::FieldAbsoluteDate, arg2::Frame)
    return jcall(obj, "getTargetPV", TimeStampedFieldPVCoordinates, (FieldPVCoordinatesProvider, FieldAbsoluteDate, Frame), arg0, arg1, arg2)
end

function get_target_pv(obj::GroundPointing, arg0::PVCoordinatesProvider, arg1::AbsoluteDate, arg2::Frame)
    return jcall(obj, "getTargetPV", TimeStampedPVCoordinates, (PVCoordinatesProvider, AbsoluteDate, Frame), arg0, arg1, arg2)
end


insts = OrderedDict()
insts["minimal"] = [
    ((5, 2, 2, true, true, true),),
    ((5, 1, 2, false, true, true),),
    ((5, 1, 2, false, true, false),),
    ((5, 1, 2, true, false, false), :ExpPSD),
    ((5, 1, 2, true, false, false), :SOCExpPSD),
    ((:iris, 2, true, false, true),),
    ]
insts["fast"] = [
    ((10, 1, 10, true, false, false), :ExpPSD),
    ((10, 1, 10, true, false, false), :SOCExpPSD),
    ((100, 1, 250, true, true, true),),
    ((100, 2, 5, false, true, true),),
    ((200, 2, 20, true, true, true),),
    ((50, 3, 2, true, true, true),),
    ((50, 3, 2, false, false, true),),
    ((50, 3, 2, true, false, true), :ExpPSD),
    ((50, 3, 2, true, false, true), :SOCExpPSD),
    ((50, 3, 4, true, true, true),),
    ((500, 3, 14, true, true, true),),
    ((100, 8, 2, true, true, true),),
    ((250, 4, 4, true, true, true),),
    ((250, 4, 4, false, false, true),),
    ((200, 32, 2, true, false, true),),
    ((:iris, 4, false, false, false),),
    ((:iris, 5, true, false, false),),
    ((:iris, 6, true, true, true),),
    ((:cancer, 4, true, true, true),),
    ]
insts["various"] = [
    ((50, 2, 16, true, false, false),),
    ((50, 2, 16, true, false, false), :ExpPSD),
    ((50, 2, 16, true, false, false), :SOCExpPSD),
    ((50, 2, 16, false, false, false),),
    ((50, 2, 16, false, false, true),),
    ((50, 8, 4, true, false, false),),
    ((50, 8, 4, true, false, false), :ExpPSD),
    ((50, 8, 4, true, false, false), :SOCExpPSD),
    ((50, 8, 4, false, false, false),),
    ((50, 8, 4, false, false, true),),
    ((50, 32, 2, true, false, false),),
    ((50, 32, 2, true, false, false), :ExpPSD),
    ((50, 32, 2, true, false, false), :SOCExpPSD),
    ((:iris, 6, true, true, true),),
    ((:cancer, 4, true, true, true),),
    ]
return (DensityEstJuMP, insts)


using Oceananigans.Simulations:
    stop, iteration_limit_exceeded, stop_time_exceeded, wall_time_limit_exceeded, TimeStepWizard

@testset "Simulations" begin
    @info "Testing simulations..."

    for arch in archs

        grid = RegularCartesianGrid(size=(16, 16, 16), extent=(1, 1, 1))
        model = IncompressibleModel(architecture=arch, grid=grid)

        for Δt in (3, TimeStepWizard(Δt=5.0))
            model.clock.time = 0.0
            model.clock.iteration = 0

            simulation = Simulation(model, Δt=Δt, stop_iteration=1)

            # Just make sure we can construct a simulation without any errors.
            @test simulation isa Simulation

            @test iteration_limit_exceeded(simulation) == false
            @test stop(simulation) == false

            run!(simulation)

            # Just make sure run! executes without any errors.
            @test simulation isa Simulation

            # Some basic tests
            @test iteration_limit_exceeded(simulation) == true
            @test stop(simulation) == true

            t = Δt isa Number ? 3 : 5
            @test model.clock.time ≈ t
            @test model.clock.iteration == 1
            @test simulation.run_time > 0

            @test stop_time_exceeded(simulation) == false
            simulation.stop_time = 1e-12
            @test stop_time_exceeded(simulation) == true

            @test wall_time_limit_exceeded(simulation) == false
            simulation.wall_time_limit = 1e-12
            @test wall_time_limit_exceeded(simulation) == true
        end
    end
end


module OneDTmazeERExperiment

import Flux: Descent, ADAM
import Random

using GVFHordes
using Curiosity
using MinimalRLCore
using SparseArrays
using ProgressMeter

const BU = Curiosity.BaselineUtils
const ODTMU = Curiosity.OneDTMazeUtils
const SRCU = Curiosity.SRCreationUtils

default_args() =
    Dict(
        "logger_interval" => 100,

        # Behaviour Items
        # "behaviour_eta" => 0.1/8,
        "behaviour_gamma" => 0.9,
        "behaviour_learner" => "RoundRobin",
        "behaviour_update" => "TB",
        "behaviour_reward_projector" => "maze",
        "behaviour_rp_tilings" => 1,
        "behaviour_rp_tiles" => 16,
        "behaviour_trace" => "ReplacingTraces",
        "behaviour_opt" => "Auto",
        "behaviour_lambda" => 0.95,
        "behaviour_alpha_init" => 0.1,
        "exploration_param" => 1.0,
        "exploration_strategy" => "epsilon_greedy",
        "ϵ_range" => (0.4,0.1),
        "decay_period" => 5000,
        "warmup_steps" => 12000,
        "behaviour_w_init" => 4,

        # Demon Attributes
        # "demon_eta" => 0.1/8,
        "demon_discounts" => 0.9,
        "demon_learner" => "SR",
        "demon_update" => "TB",
        "demon_interest_set" => "oned_tmaze",
        "demon_policy_type" => "greedy_to_cumulant",
        "demon_opt" => "Auto+ADAM",
        "demon_lambda" => 0.9,
        "demon_trace"=> "AccumulatingTraces",

        "demon_eta" => 0.001,
        "demon_alpha_init"=>0.1,
        "demon_beta" => (0.99, 0.999),

        #shared
        # "num_tiles" => 2,
        # "num_tilings" =>8,
        "tiling_structure" => [2,8],
        "demon_rep" => "ideal",
        # "demon_rep" => "tilecoding",
        "demon_num_tiles" => 8,
        "demon_num_tilings" => 1,
        "eta" => 0.05,
        "alpha_init" => 0.1,
        "replay_size" => 10000,
        "batch_size" => 4,

        # Environment Config
        "constant_target"=> [-10.0,10.0],
        "cumulant_schedule" => "DrifterDistractor",
        "cumulant" => 1,
        "distractor" => (1.0, 5.0),
        "drifter" => (1.0, sqrt(0.01)),
        "exploring_starts"=>"beg",
        "env_step_penalty" => -0.01,

        # Agent and Logger
        "horde_type" => "regular",
        "intrinsic_reward" => "weight_change",
        # "logger_keys" => [LoggerKey.TTMAZE_ERROR],
        "save_dir" => "OneDTMazeExperiment",
        "seed" => 1,
        "steps" => 20000,
        "use_external_reward" => true,
        "random_first_action" => false,
        "logger_keys" => [LoggerKey.ONEDTMAZEERROR,LoggerKey.ONEDTMAZEERROR_DPI,LoggerKey.ONEDTMAZEERROR_UNIFORM, LoggerKey.ONED_GOAL_VISITATION, LoggerKey.EPISODE_LENGTH, LoggerKey.INTRINSIC_REWARD]
    )

function construct_agent(parsed)
    action_space = 4
    obs_size = 6

    intrinsic_reward_type = parsed["intrinsic_reward"]
    use_external_reward = parsed["use_external_reward"]

    # Unpack Tiling Config
    if "tiling_structure" ∈ keys(parsed)
        parsed["num_tilings"] = parsed["tiling_structure"][1]
        parsed["num_tiles"] = parsed["tiling_structure"][2]
    end

    # Unpack joint eta param
    if "eta" in keys(parsed)
        prefixes = ["behaviour","demon"]
        for prefix in prefixes
            parsed[join([prefix, "eta"], "_")] = parsed["eta"]
        end
    end

    if "alpha_init" in keys(parsed)
        prefixes = ["behaviour","demon"]
        for prefix in prefixes
            parsed[join([prefix, "alpha_init"], "_")] = parsed["alpha_init"]
        end
    end

    if parsed["demon_opt"] == "Auto+Descent"
        parsed["demon_auto_alpha_init"] = parsed["demon_auto_alpha_init"] / parsed["num_tilings"]
    end

    # Unpack Auto Init
    if parsed["demon_opt"] == "Auto"
        parsed["demon_alpha_init"] =
            parsed["demon_alpha_init"] / parsed["num_tilings"]
    end

    if "behaviour_opt" ∈ keys(parsed) && parsed["behaviour_opt"] == "Auto"
        parsed["behaviour_alpha_init"] =
            parsed["behaviour_alpha_init"] / parsed["num_tilings"]
    end
    if parsed["behaviour_opt"] == "Descent"
        if "behaviour_eta" in keys(parsed)
            parsed["behaviour_eta"] = parsed["behaviour_eta"] / parsed["num_tilings"]
        end
    end

    if parsed["behaviour_w_init"] != 0.0
        parsed["behaviour_w_init"] = parsed["behaviour_w_init"] / parsed["num_tilings"]
    end

    fc = Curiosity.FeatureSubset(
        Curiosity.SparseTileCoder(parsed["num_tilings"], parsed["num_tiles"], 2),
        1:2)
    feat_size = size(fc)

    demon_projected_fc = if "demon_rep" ∉ keys(parsed)
        nothing
    elseif parsed["demon_rep"] == "tilecoding"
        Curiosity.FeatureProjector(Curiosity.FeatureSubset(
                Curiosity.SparseTileCoder(parsed["demon_num_tilings"], parsed["demon_num_tiles"], 2),
            1:2), false)
    elseif parsed["demon_rep"] == "ideal"
        Curiosity.FeatureSubset(ODTMU.IdealDemonFeatures(), 1:2)
    elseif parsed["demon_rep"] == "ideal_state_action"
        t = ODTMU.IdealStateActionDemonFeatures(4)
        Curiosity.FeatureSubset(t, 1:2)
    elseif parsed["demon_rep"] == "ideal_martha"
        Curiosity.FeatureProjector(Curiosity.FeatureSubset(ODTMU.MarthaIdealDemonFeatures(4), 1:2), false)
    else
        throw(ArgumentError("Not a valid demon projection rep for SR"))
    end

    demons = ODTMU.create_demons(parsed, demon_projected_fc)

    demon_learner = Curiosity.get_linear_learner(parsed,
                                feat_size,
                                action_space,
                                demons,
                                "demon",
                                demon_projected_fc)

    exploration_strategy = Curiosity.get_exploration_strategy(parsed, 1:action_space)

    behaviour_reward_projector = if "behaviour_reward_projector" ∉ keys(parsed)
        nothing
    elseif parsed["behaviour_reward_projector"] == "tilecoding"
        Curiosity.FeatureProjector(Curiosity.FeatureSubset(
            Curiosity.SparseTileCoder(parsed["behaviour_rp_tilings"], parsed["behaviour_rp_tiles"], 2),
            1:2), false)
    elseif parsed["behaviour_reward_projector"] == "base"
        Curiosity.FeatureProjector(fc, false)
    elseif parsed["behaviour_reward_projector"] == "identity"
        Curiosity.FeatureProjector(Curiosity.FeatureSubset(identity, 1:2), false)
    elseif parsed["behaviour_reward_projector"] == "ideal"
        Curiosity.FeatureProjector(Curiosity.FeatureSubset(
            ODTMU.IdealDemonFeatures(), 1:2), true)
    elseif parsed["behaviour_reward_projector"] == "ideal_martha"
        Curiosity.FeatureProjector(Curiosity.FeatureSubset(
            ODTMU.MarthaIdealDemonFeatures(), 1:2), false)
    elseif parsed["behaviour_reward_projector"] == "maze"
        Curiosity.FeatureSubset(ODTMU.TMazeEncoding(),1:2)
    else
        throw(ArgumentError("Not a valid demon projection rep for SR"))
    end

    behaviour_learner, behaviour_demons, behaviour_discount = if parsed["behaviour_learner"] == "RoundRobin"
        ODTMU.RoundRobinPolicy(), nothing, 0.0
    elseif parsed["behaviour_learner"] == "RandomDemons"
        action_set = 1:action_space
        blearner = if parsed["demon_learner"] == "SR"
            #We want the demons specifying the cumulants and not predicting the SF
            BU.RandomDemons(demons.PredHorde, action_set)
        else
            BU.RandomDemons(demons, action_set)
        end
        blearner, nothing, 0.0
    else
        behaviour_num_tasks = 1
        num_SFs = 4
        num_demons = if parsed["behaviour_learner"] ∈ ["GPI"]
            # num_SFs * size(behaviour_reward_projector) * action_space + behaviour_num_tasks
            num_SFs * size(behaviour_reward_projector) + behaviour_num_tasks

        elseif parsed["behaviour_learner"] ∈ ["Q"]
            behaviour_num_tasks
        elseif parsed["behaviour_learner"] == "RoundRobin"
            0
        else
            throw(ArgumentError("Hacky thing not working, yay!"))
        end
        behaviour_learner = Curiosity.get_linear_learner(parsed,
                                                         size(fc),
                                                         action_space,
                                                         num_demons,
                                                         behaviour_num_tasks,
                                                         "behaviour",
                                                         behaviour_reward_projector)

        behaviour_demons = if behaviour_learner isa GPI
            @assert !(behaviour_reward_projector isa Nothing)
            bh_gvf = ODTMU.make_behaviour_gvf(behaviour_learner,
                                              0.0,
                                              fc,
                                              exploration_strategy)
            pred_horde = GVFHordes.Horde([bh_gvf])
            SF_policies = [ODTMU.GoalPolicy(i) for i in 1:4]
            SF_discounts = [ODTMU.GoalTermination(parsed["behaviour_gamma"]) for i in 1:4]
            num_SFs = length(SF_policies)
            SF_horde = SRCU.create_SF_horde(SF_policies, SF_discounts, behaviour_reward_projector, 1:action_space)
            Curiosity.GVFSRHordes.SRHorde(pred_horde, SF_horde, num_SFs, behaviour_reward_projector)
        elseif behaviour_learner isa QLearner
            bh_gvf = ODTMU.make_behaviour_gvf(behaviour_learner,
                                              parsed["behaviour_gamma"],
                                              fc,
                                              exploration_strategy)
            GVFHordes.Horde([bh_gvf])
        elseif behaviour_learner isa ODTMU.RoundRobinPolicy
            nothing
        else
            throw(ArgumentError("goes with which horde? " ))
        end

        behaviour_learner, behaviour_demons, parsed["behaviour_gamma"]
    end #end behaviour_learner, behaviour_demons, behaviour_discount = if parsed["behaviour_learner"] == "RoundRobin"

    random_first_action = parsed["random_first_action"]
    if behaviour_learner isa ODTMU.RoundRobinPolicy
        if random_first_action == false
            @warn "Round Robin with random first action is recommended"
        end
    elseif random_first_action == true
        @warn "Random first action is enabled"
    end

    AgentER(demons,
            feat_size,
            behaviour_learner,
            behaviour_demons,
            behaviour_discount,
            demon_learner,
            obs_size,
            action_space,
            intrinsic_reward_type,
            fc,
            use_external_reward,
            exploration_strategy,
            random_first_action,
            parsed["batch_size"],
            parsed["replay_size"])
end

function main_experiment(parsed=default_args(); progress=false, working=false)

    num_steps = parsed["steps"]
    Random.seed!(parsed["seed"])

    logger_init_dict = Dict(
        LoggerInitKey.TOTAL_STEPS => num_steps,
        LoggerInitKey.INTERVAL => parsed["logger_interval"],
        # LoggerInitKey.ENV => "tabular_tmaze"
    )

    Curiosity.experiment_wrapper(parsed, logger_init_dict, working) do parsed, logger

        cumulant_schedule = ODTMU.get_cumulant_schedule(parsed)
        env = OneDTMaze(cumulant_schedule, parsed["exploring_starts"], parsed["env_step_penalty"])

        agent = construct_agent(parsed)
        goal_visitations = zeros(4)

        eps = 1
        max_num_steps = num_steps
        steps = Int[]

        logger_start!(logger, env, agent)

        prg_bar = ProgressMeter.Progress(num_steps, "Step: ")
        while sum(steps) < max_num_steps
            cur_step = 0
            max_episode_steps = max_num_steps - sum(steps)
            tr, stp =
                run_episode!(env, agent, max_episode_steps) do (s, a, s_next, r, t)
                    #This is a callback for every timestep where logger can go
                    # agent is accesible in this scope
                    if t == true && working==true
                        goals = s_next[3:end]
                        f = findfirst(!iszero, goals)
                        goal_visitations[f] += 1
                    end
                    if progress
                        next!(prg_bar)
                    end

                    logger_step!(logger, env, agent, s, a, s_next, r, t)
                    cur_step+=1
                end
                logger_episode_end!(logger)
            push!(steps, stp)
            eps += 1
        end
        if working == true
            println("Goal Visitation: ", goal_visitations)
        end
        # agent
    end

end

end